"""Tier 2: acceptance check for the filename-derived session/modality parser.

Produces three things:
  1. A stratified sample of >=200 images with every parse field, written for
     human review. Human sign-off is NOT performed here and remains outstanding.
  2. An independent-reimplementation agreement rate. A second parser written
     with different logic (separator splitting, not the production regexes) must
     agree on session identity.
  3. A per-image behavioural check: is an image's nearest same-identity
     neighbour in embedding space also in its parsed session? Real capture
     sessions should dominate the nearest-neighbour relation; this gives a
     falsifiable error rate rather than an assertion.
"""
from __future__ import annotations
import json, re, sys
from pathlib import Path
import numpy as np, pandas as pd

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT)); sys.path.insert(0, str(ROOT / 'tools'))
from ext_sessions import COPY_PREFIX, load as load_sessions
from face_research.prepare import load_prepared

OUT = ROOT / 'artifacts' / 'extensions' / 'parser_acceptance_v1'
A = ROOT / 'artifacts' / 'thesis_study_v1'
PER_STRATUM = 40


def independent_session(filename: str, collection: str) -> str:
    """Deliberately different implementation: split on separators, no shared regex."""
    name = COPY_PREFIX.sub('', filename)
    stem = name.rsplit('.', 1)[0]
    if collection == 'HUST_Face_data':
        if '_frame' in stem:
            return 'clip' + stem.split('_frame')[0]
        parts = stem.split('_')
        if parts[0] == 'IMG' and len(parts) >= 3 and len(parts[1]) == 8 and parts[1].isdigit():
            return 'ts' + parts[1] + parts[2][:4]
        if len(parts) >= 2 and len(parts[0]) == 8 and parts[0].isdigit() and parts[1].isdigit():
            return 'ts' + parts[0] + parts[1][:4]
        if parts[0] == 'IMG' and len(parts) == 2 and parts[1].isdigit():
            return 'img' + parts[1][:4]
        return 'unparsed'
    if collection == 'ProcessedData2324':
        parts = stem.split('_')
        return '_'.join(parts[:2]) if len(parts) >= 3 and parts[0] == 'IMG' else 'unparsed'
    if collection == 'ProcessedData20241':
        parts = stem.split('-')
        return '-'.join(parts[:5]) if len(parts) >= 6 else 'unparsed'
    if collection == 'ProcessedData20242':
        parts = stem.rsplit('_', 1)
        return parts[0] if len(parts) == 2 and parts[1].isdigit() else 'unparsed'
    return 'unparsed'


def main():
    OUT.mkdir(parents=True, exist_ok=True)
    df = load_sessions(ROOT / 'artifacts' / 'data_audit_v1')
    df['bare_session'] = df.session.str.split('|').str[1:].str.join('|')

    # (2) agreement of an independently written parser, on the FULL inventory
    indep = [independent_session(f, c) for f, c in zip(df.filename, df.collection)]
    prod = [s if r != 'none' else 'unparsed' for s, r in
            zip(df.bare_session.str.replace(r'\|.*$', '', regex=True), df.session_rule)]
    agree = np.mean([a == b for a, b in zip(indep, prod)])
    per_coll = {}
    for coll in sorted(df.collection.unique()):
        m = df.collection == coll
        per_coll[coll] = float(np.mean([a == b for a, b, k in zip(indep, prod, m) if k]))
    print(f"independent-reimplementation agreement: {agree:.4f} overall")
    for k, v in per_coll.items():
        print(f"   {k:<20} {v:.4f}")

    # (1) stratified reviewable sample
    rng = np.random.default_rng(20260908)
    parts = []
    for (coll, rule), g in df.groupby(['collection', 'session_rule']):
        take = min(PER_STRATUM, len(g))
        parts.append(g.iloc[rng.choice(len(g), take, replace=False)])
    sample = pd.concat(parts)[['path', 'filename', 'collection', 'session_rule',
                               'modality', 'bare_session', 'person_id']]
    sample['independent_session'] = [independent_session(f, c) for f, c in
                                     zip(sample.filename, sample.collection)]
    sample.to_csv(OUT / 'review_sample.csv', index=False)
    print(f"\nstratified review sample: {len(sample)} images across "
          f"{sample.groupby(['collection','session_rule']).ngroups} strata -> review_sample.csv")

    # (3) behavioural check on the sample, using untouched M0 embeddings
    prepared = load_prepared(A)
    emb_by_id, order = {}, {}
    for fold in range(3):
        meta = json.loads((A / 'runs' / f'fold_{fold}' / 'M0' / 'seed_42' / 'scores_test.json').read_text())
        cache = A / 'cache' / 'embeddings' / meta['cache_key']
        e = np.load(cache / 'embeddings.npy', allow_pickle=False)
        ids = pd.read_csv(cache / 'images.csv', dtype=str).image_id.tolist()
        for i, k in enumerate(ids):
            emb_by_id[k] = e[i]
    sess_of = df.set_index('image_id').session.to_dict()
    person_of = df.set_index('image_id').person_id.to_dict()
    by_person = {}
    for k in emb_by_id:
        by_person.setdefault(person_of[k], []).append(k)
    hits = tot = 0
    rows = []
    for k in df[df.image_id.isin(emb_by_id)].image_id:
        peers = [q for q in by_person[person_of[k]] if q != k]
        if len(peers) < 2 or 'unparsed' in sess_of[k]:
            continue
        same = [q for q in peers if sess_of[q] == sess_of[k]]
        if not same or len(same) == len(peers):
            continue
        v = emb_by_id[k]
        sims = np.array([float(v @ emb_by_id[q]) for q in peers])
        nn = peers[int(np.argmax(sims))]
        ok = sess_of[nn] == sess_of[k]
        hits += ok; tot += 1
        rows.append({'image_id': k, 'nearest_in_same_session': bool(ok)})
    print(f"\nbehavioural check: nearest same-identity neighbour lies in the parsed session "
          f"for {hits}/{tot} images ({100*hits/tot:.1f}%)")
    pd.DataFrame(rows).to_csv(OUT / 'nearest_neighbour_check.csv', index=False)
    json.dump({'independent_agreement_overall': float(agree),
               'independent_agreement_by_collection': per_coll,
               'review_sample_size': int(len(sample)),
               'nearest_neighbour_same_session_rate': hits / tot,
               'nearest_neighbour_evaluated_images': int(tot),
               'human_signoff': 'OUTSTANDING - sample written for manual review, not yet reviewed'},
              open(OUT / 'acceptance.json', 'w'), indent=2)
    print('wrote', OUT)


if __name__ == '__main__':
    main()