"""Predeclare a broad source-grounded holdout, without reading retrieval outputs.

All profiles have already been exposed in v1-v3. This is a NEW QUERY-GROUP
holdout over the same corpus, explicitly not an unseen-person experiment.
Historic cases remain DEV; paraphrases of the same logical task stay together.
"""
from __future__ import annotations

import argparse
import copy
import json
import random
import re
from collections import Counter
from pathlib import Path

from .benchmark_v2 import BASE, academic_binding, academics, audit_dataset, fingerprint, independent_pool, make_case, validate_case, write_json
from .source_oracle import load_source_oracle, norm

SEED = 20261025
QUOTAS = {
    'shared_keyword_pair': 40, 'same_academic_year_pair': 40,
    'same_start_year_pair': 32, 'person_title_year': 24,
    'person_keyword_title': 32, 'document_author_category_multi_hop': 24,
    'person_multiple_categories': 24, 'candidate_set_aggregation': 32,
    'conjunctive_disambiguation': 40, 'negative_lookup': 32,
}


def logical_group(case: dict) -> str:
    """Normalize task/subjects/constraints; never make paraphrases independent."""
    graph = case['graph']; category = case['category']
    names = sorted(norm(row['value']) for row in graph['start_entities'] if row['type'] == 'person')
    bindings = graph.get('expected_bindings', [])
    keywords = sorted({norm(row['keyword']) for row in bindings if row.get('keyword')})
    years = sorted({int(value) for value in re.findall(r'\b20\d{2}\b', case['question'])})
    # A shared-year query has no requested year; its complete answer is derived.
    selectors = {'keywords': keywords, 'requested_years': years}
    if category == 'document_author_category_multi_hop':
        selectors['categories'] = sorted({norm(row['category']) for row in bindings if row.get('category')})
    if graph.get('expected_no_answer'):
        selectors['negative_constraints'] = {
            key: norm(value) if isinstance(value, str) else value
            for key, value in graph.get('negative_constraints', {}).items()}
    return fingerprint({'category': category, 'subjects': names, 'selectors': selectors})


def paraphrase(case: dict, variant: int) -> str:
    names = [e['value'] for e in case['graph']['start_entities']]
    category = case['category']; rows = case['graph']['expected_bindings']
    year = next(iter(re.findall(r'\b20\d{2}\b', case['question'])), '')
    tag = next((r['keyword'] for r in rows if r.get('keyword')), '')
    if variant == 0 or category == 'negative_lookup':
        return case['question']
    if category == 'shared_keyword_pair':
        return (f'Majú osoby {names[0]} a {names[1]} spoločné kľúčové slovo {tag}?'
                if variant == 1 else f'Pre osoby {names[0]} a {names[1]} over rovnaké kľúčové slovo {tag}.')
    if category == 'same_academic_year_pair':
        return (f'Ktoré spoločné roky prác majú osoby {names[0]} a {names[1]}?'
                if variant == 1 else f'Porovnaj rovnaké roky práce pri osobách {names[0]} a {names[1]}.')
    if category == 'same_start_year_pair':
        return f'Aký spoločný rok začiatku štúdia majú osoby {names[0]} a {names[1]}?'
    if category == 'person_multiple_categories':
        return f'Aké všetky kategórie majú dokumenty osoby {names[0]}?'
    if category == 'document_author_category_multi_hop':
        return f'Ktorý autor dokumentu je pri osobe {names[0]} spojený s kategóriou {rows[0]["category"]}?'
    if category == 'person_title_year':
        # Preserve an explicit project selector, if present in the source query.
        label = 'diplomového projektu' if 'diplomoveho projektu' in norm(case['question']) else 'práce'
        return f'Uveď názov {label} osoby {names[0]} a rok práce pre rok {year}.'
    if category == 'person_keyword_title':
        return f'Uveď názov práce osoby {names[0]} v roku {year} s kľúčovým slovom {tag}.'
    if category == 'candidate_set_aggregation':
        return f'Z kandidátov {", ".join(names)} vyber osoby s prácou v roku {year}. Aký je ich počet?'
    if category == 'conjunctive_disambiguation':
        return f'Ktorá z osôb {names[0]} a {names[1]} spĺňa kľúčové slovo {tag} a rok práce {year}?'
    return case['question']


def build_pool(oracle: dict) -> dict[str, list[dict]]:
    names = sorted(oracle['people'], key=norm)
    pool = independent_pool(oracle, names)
    negatives = []
    # Realistic absent combinations as well as unknown entities. Positive cases
    # make up 90% of TEST; no easy-negative inflation of the headline score.
    for index, name in enumerate(names):
        other_tags = sorted({tag for p in oracle['people'].values() for tag in p['tags']
                             if norm(tag) not in {norm(t) for t in oracle['people'][name]['tags']}}, key=norm)
        if other_tags:
            tag = other_tags[index % len(other_tags)]
            negatives.append(make_case('', 'test', 'negative_lookup',
                f'Aký názov práce má osoba {name} s kľúčovým slovom {tag}?', [name], oracle, [],
                no_answer=True, reason='Requested literal tag is absent from this source document family.',
                negative_constraints={'keyword': tag}))
        observed = {r['year'] for r in academics(oracle, name)}
        absent_year = next(y for y in range(2018, 2027) if y not in observed)
        negatives.append(make_case('', 'test', 'negative_lookup',
            f'Aký názov práce má osoba {name} v roku {absent_year}?', [name], oracle, [],
            no_answer=True, reason='Requested academic year is absent from the original profile.',
            negative_constraints={'year': absent_year}))
    for first, last in [('Jozef', 'Neznámy Overovací'), ('Petra', 'Neznáma Overovacia'),
                        ('Juraj', 'Nevyskytujúci Skúšobný'), ('Lucia', 'Chýbajúca Vzorová'),
                        ('Adam', 'Nedostupný Kontrolný'), ('Eva', 'Neprítomná Kontrolná')]:
        name = f'{first} {last}'
        negatives.append(make_case('', 'test', 'negative_lookup', f'Aký názov práce má osoba {name}?',
            [name], oracle, [], no_answer=True, reason='Full name absent from all original profiles.',
            negative_constraints={'unknown_person': name}))
    pool['negative_lookup'] = negatives
    return pool


def refresh_titles(case: dict, oracle: dict) -> dict:
    """Adjudicate title gold from raw sources before looking at TEST predictions."""
    if case['category'] not in {'person_title_year', 'person_keyword_title'} or not case['should_answer']:
        return copy.deepcopy(case)
    name = case['graph']['start_entities'][0]['value']
    years = [int(v) for v in re.findall(r'\b20\d{2}\b', case['question'])]
    tag = next((b.get('keyword') for b in case['graph']['expected_bindings'] if b.get('keyword')), None)
    text = norm(case['question'])
    kind = next((kind for phrase, kind in [('diplomoveho projektu', 'diploma_project'),
        ('diplomovej prace', 'diploma_thesis'), ('bakalarskej prace', 'bachelor_thesis')]
        if phrase in text), None)
    records = [r for r in academics(oracle, name) if r.get('title')
               and (not years or r['year'] in years) and (kind is None or r['academic_type'] == kind)]
    bindings = [academic_binding(name, r, title=True, keyword=tag) for r in records]
    return make_case(case['id'], case['split'], case['category'], case['question'], [name], oracle, bindings,
        no_answer=not bindings, reason='No explicit title for the requested academic record in original Markdown.',
        negative_constraints={'title_required': True, **({'year': years[-1]} if years else {}),
                              **({'keyword': tag} if tag else {}), **({'academic_type': kind} if kind else {})})


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=BASE / 'graphrag_questions_v4.json')
    parser.add_argument('--artifacts', type=Path, default=BASE / 'results/v4_preparation')
    parser.add_argument('--seed', type=int, default=SEED)
    parser.add_argument('--extra-exposed', type=Path, action='append', default=[])
    args = parser.parse_args()
    if args.output.exists() or args.artifacts.exists():
        parser.error('Refusing to overwrite a declared benchmark or its preparation artifacts.')
    oracle = load_source_oracle()
    historic = json.loads((BASE / 'graphrag_questions_v3.json').read_text())
    for path in args.extra_exposed:
        historic.extend(json.loads(path.read_text()))
    old_questions = {norm(q['question']) for q in historic}
    exposed = {logical_group(q) for q in historic}
    adjudications = []
    for index, original in enumerate(historic):
        corrected = refresh_titles(original, oracle)
        if corrected['graph']['expected_bindings'] != original['graph']['expected_bindings']:
            adjudications.append({'id': original['id'], 'split': 'dev',
                'old_bindings': original['graph']['expected_bindings'], 'new_bindings': corrected['graph']['expected_bindings'],
                'source_evidence': corrected['source_evidence']})
        historic[index] = corrected
    exposed.update(logical_group(q) for q in historic)
    dev = []
    for index, original in enumerate(historic):
        case = copy.deepcopy(original); case['split'] = 'dev'
        case['id'] = f'regression-{index + 1:04d}'
        case['graph']['partition_group'] = logical_group(case)
        dev.append(case)
    rng = random.Random(args.seed); holdout = []; capacities = {}
    for category, cases in build_pool(oracle).items():
        unique = {logical_group(c): c for c in cases
                  if logical_group(c) not in exposed and norm(c['question']) not in old_questions}
        candidates = sorted(unique.values(), key=lambda q: logical_group(q)); rng.shuffle(candidates)
        quota = QUOTAS[category]; capacities[category] = len(candidates)
        if len(candidates) < quota:
            raise RuntimeError(f'Insufficient unseen logical groups for {category}: {len(candidates)} < {quota}')
        # Random selection is committed before running either split. Never
        # select or remove a TEST question based on a retrieval outcome.
        for index, case in enumerate(candidates[:quota]):
            case = refresh_titles(case, oracle); case['split'] = 'test'
            case['id'] = f'wide-{len(holdout) + 1:03d}'
            case['question'] = paraphrase(case, index % 3)
            case['graph']['partition_group'] = logical_group(case)
            holdout.append(case)
        for index, case in enumerate(candidates[quota:quota + min(12, max(0, len(candidates) - quota))]):
            case = refresh_titles(case, oracle); case['split'] = 'dev'
            case['id'] = f'wide-dev-{len(dev) + 1:04d}'
            case['question'] = paraphrase(case, index % 3)
            case['graph']['partition_group'] = logical_group(case)
            dev.append(case)
    questions = dev + holdout
    problems = {q['id']: validate_case(q, oracle) for q in questions if validate_case(q, oracle)}
    if problems:
        raise RuntimeError(f'Source oracle rejected cases: {problems}')
    groups = {split: {q['graph']['partition_group'] for q in questions if q['split'] == split}
              for split in ['dev', 'test']}
    assert not groups['dev'] & groups['test']
    assert not {norm(q['question']) for q in dev} & {norm(q['question']) for q in holdout}
    write_json(args.output, questions)
    audit = audit_dataset(questions, oracle, args.output)
    audit['partition_policy'] = {'method': 'disjoint logical query groups',
        'shared_source_profiles': True, 'group_function': 'benchmark_v4.logical_group',
        'seed': args.seed, 'previously_exposed_groups': sorted(exposed),
        'limitations': ['Known source corpus, persons, vocabulary and template families.',
                        'Fresh questions/fact combinations; not an unseen-person or LLM-answer benchmark.']}
    write_json(args.artifacts / 'source_audit.json', audit)
    write_json(args.artifacts / 'source_adjudications.json', {
        'reason': 'Raw-source review before any new TEST predictions; preserve historic datasets unchanged.',
        'questions': adjudications})
    write_json(args.artifacts / 'partition.json', {
        **audit['partition_policy'], 'dataset_sha256': audit['dataset_sha256'],
        'dev_count': len(dev), 'test_count': len(holdout), 'available_unseen_groups': capacities,
        'test_categories': dict(Counter(q['category'] for q in holdout)),
        'test_positive_count': sum(q['should_answer'] for q in holdout),
        'test_negative_count': sum(not q['should_answer'] for q in holdout),
        'dev_groups': sorted(groups['dev']), 'test_groups': sorted(groups['test']),
        'test_profile_count': len({p for q in holdout for p in q['graph']['source_profiles']}),
        'source_corpus': oracle['source_corpus']})
    print(json.dumps({'dev': len(dev), 'test': len(holdout), 'source_issues': 0,
                      'available_unseen_groups': capacities}, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()