"""Portfolio adaptation study: live small-model NER and cached redaction policies.
Run offline: HF_DATASETS_OFFLINE=1 HF_HUB_OFFLINE=1 legal-anon-env/bin/python scripts/evaluate_portfolio_adaptations.py
No training, threshold tuning, deployment changes or external requests.
"""
import csv
import hashlib
import json
import re
import sys
from pathlib import Path
import numpy as np
import spacy

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / 'src'))
from anonymisation.data import load_tab
from anonymisation.evaluation import score_spans, extract_gold_entities
from anonymisation.ensemble import EnsemblePredictor

cache_path = ROOT / 'results/predictions_cache.json'
cache = json.loads(cache_path.read_text())
docs = list(load_tab()['test'])
assert len(cache) == len(docs) == 555
for c, d in zip(cache, docs):
    assert c['doc_id'] == d['doc_id']
    assert c['gold'] == [list(s) for s in extract_gold_entities(d)]

# Run actual inference once per unique text, then score every annotation row.
nlp = spacy.load('en_core_web_sm', disable=['parser', 'lemmatizer', 'attribute_ruler'])
texts = dict((d['doc_id'], d['text']) for d in docs)
small = {}
for (key, text), parsed in zip(texts.items(), nlp.pipe(texts.values(), batch_size=8)):
    small[key] = [(e.start_char,e.end_char,e.label_,e.text) for e in parsed.ents if e.label_ in ('PERSON','ORG')]
print('Small-model inference complete:', len(small), 'unique judgments', flush=True)

ner_counts = {(k,g): np.zeros(3, dtype=int) for k in ('spacy_small','roberta_names_only') for g in ('masked','all')}
policies = ('roberta_spans','majority_spans','union_spans','union_sentences','suppress_all')
totals = {(p,g): dict(mentions=0,fully_removed=0,rows=0,clean_rows=0,removed_chars=0,total_chars=0,non_target_removed=0,non_target_chars=0,person_n=0,person_full=0,org_n=0,org_full=0) for p in policies for g in ('masked','all')}
for index,(c,d) in enumerate(zip(cache,docs)):
    text=d['text']; size=len(text)
    gold=[s for s in c['gold'] if s[2] in ('PERSON','ORG')]
    for key,preds in [('spacy_small',small[d['doc_id']]),('roberta_names_only',[s for s in c['roberta'] if s[2] in ('PERSON','ORG')])]:
        for g in ('masked','all'):
            target_gold = gold if g == 'masked' else [(m['start_offset'],m['end_offset'],m['entity_type'],m['span_text']) for m in d['entity_mentions'] if m['entity_type'] in ('PERSON','ORG')]
            r=score_spans(preds,target_gold,'partial')['_ALL'];ner_counts[key,g] += [r.tp,r.fp,r.fn]
    members={m:(lambda _,v=c[m]:v) for m in ('roberta','legalbert','longformer')}
    majority=EnsemblePredictor(members,min_votes=2,tie_break='longest').predict('x')
    union=[s for m in members for s in c[m]]
    spans={'roberta_spans':c['roberta'],'majority_spans':majority,'union_spans':union}
    masks={}
    for name,items in spans.items():
        mask=np.zeros(size,dtype=bool)
        for start,end,*_ in items: mask[start:end]=True
        masks[name]=mask
    # Deliberately broad, deterministic sentence-like segments. Abbreviations
    # can split segments; this is a policy simulation, not linguistic parsing.
    mask=np.zeros(size,dtype=bool)
    for match in re.finditer(r'[^.!?\n]+(?:[.!?]+|\n|$)|[.!?\n]+',text):
        start,end=match.span()
        if masks['union_spans'][start:end].any(): mask[start:end]=True
    masks['union_sentences']=mask
    masks['suppress_all']=np.ones(size,dtype=bool)
    # Independent set-based spot check of union character coverage.
    if index < 5:
        covered=set(k for start,end,*_ in union for k in range(start,end))
        assert int(masks['union_spans'].sum()) == len(covered)
    for gold_set in ('masked','all'):
        mentions=[m for m in d['entity_mentions'] if m['entity_type'] in ('PERSON','ORG') and (gold_set=='all' or m['identifier_type'] in ('DIRECT','QUASI'))]
        target=np.zeros(size,dtype=bool)
        for m in mentions: target[m['start_offset']:m['end_offset']]=True
        for policy,mask in masks.items():
            t=totals[policy,gold_set]; full=0
            for m in mentions:
                ok=bool(mask[m['start_offset']:m['end_offset']].all());full+=ok
                kind=m['entity_type'].lower();t[kind+'_n']+=1;t[kind+'_full']+=ok
            t['mentions']+=len(mentions);t['fully_removed']+=full
            t['rows']+=bool(mentions);t['clean_rows']+=bool(mentions) and full==len(mentions)
            t['removed_chars']+=int(mask.sum());t['total_chars']+=size
            t['non_target_removed']+=int((mask & ~target).sum());t['non_target_chars']+=int((~target).sum())
rows=[]
for (policy,gold_set),t in totals.items():
    rows.append(dict(policy=policy,gold_set=gold_set,**t,full_removal=t['fully_removed']/t['mentions'],clean_rate=t['clean_rows']/t['rows'],text_removed=t['removed_chars']/t['total_chars'],non_target_removal=t['non_target_removed']/t['non_target_chars']))
with (ROOT/'results/portfolio_adaptations.csv').open('w') as f:
    w=csv.DictWriter(f,fieldnames=rows[0]);w.writeheader();w.writerows(rows)
ners=[]
for (name,gold_set),(tp,fp,fn) in ner_counts.items():
    ners.append(dict(model=name,gold_set=gold_set,tp=int(tp),fp=int(fp),fn=int(fn),precision=float(tp/(tp+fp)),recall=float(tp/(tp+fn)),f1=float(2*tp/(2*tp+fp+fn))))
result=dict(rows=len(docs),unique_documents=len(texts),cache_sha256=hashlib.sha256(cache_path.read_bytes()).hexdigest(),small_model_version=nlp.meta['version'],ner=ners,policies=rows,notes=['NER: PERSON/ORG only, DIRECT+QUASI gold, typed partial matching; small model run live, RoBERTa replayed at 512 tokens.','Redaction uses predictions of ALL entity types; union combines character coverage without dropping overlaps.','Full removal means every gold mention character removed; clean rate is over annotation rows, not unique judgments.','Text removed measures all characters, including whitespace; non-target means outside the selected PERSON/ORG gold, not necessarily harmless text.','Sentence policy removes punctuation/newline-delimited segments containing any prediction. Suppress-all is a zero-utility bound.','Exploratory held-out replay, no new external validation, speed benchmarks or confidence-threshold sweep.'])
(ROOT/'results/portfolio_adaptations.json').write_text(json.dumps(result,indent=2)+'\n')
print(json.dumps(dict(ner=ners,policies=[{k:r[k] for k in ('policy','gold_set','full_removal','clean_rate','text_removed')} for r in rows]),indent=2))