import bootstrap
from bootstrap import BASELINE
from collections import Counter
from dataclasses import replace
import hashlib
import json
import os
import statistics
import time
from inputs import OUT
from integrated import Integrated
from evaluate import stats
from features import Settings
from state import Episode


def group_states(view,result,t,holds):
    by={p['token']:p for p in result['pairs']}
    for p in result['pairs']:
        if p['reset']:holds.pop(p['token'],None)
        if p['hit']:holds[p['token']]=t
    for key,ts in list(holds.items()):
        if t-ts>1.0+1e-8:del holds[key]
    graph={}
    for p in view['pairs']:
        a,b=p['orders'];graph.setdefault(a,set()).add(b);graph.setdefault(b,set()).add(a)
    seen=set();groups=[]
    for start in sorted(graph):
        if start in seen:continue
        todo=[start];members=set()
        while todo:
            x=todo.pop()
            if x in members:continue
            members.add(x);seen.add(x);todo.extend(graph[x]-members)
        ps=[p for p in result['pairs'] if set(p['orders'])<=members]
        current=[p['token'] for p in ps if p['hit']]
        held=[p['token'] for p in ps if p['token'] in holds and not p['hit'] and not p['reset']]
        groups.append(dict(id='G'+str(len(groups)+1),orders=sorted(members),
            state='RED_CURRENT' if current else 'YELLOW_HOLD' if held else 'BLUE_CANDIDATE',
            contributing_pairs=current,held_pairs=held,
            hold_last_hit={k:holds[k] for k in held}))
    return groups


def main():
    clip=os.environ['CLIP'];dest=OUT/'download'/clip/'data';dest.mkdir(parents=True,exist_ok=True)
    candidate=json.loads((OUT/'candidates.json').read_text());cfg=candidate['candidates'][0]['config']
    assert candidate['frozen'] and len(candidate['candidates'])==1
    c_hash=hashlib.sha256((OUT/'candidates.json').read_bytes()).hexdigest()
    raw=(OUT/'work'/clip/'enriched.jsonl').read_bytes();f_hash=hashlib.sha256(raw).hexdigest()
    assert f_hash==json.loads((OUT/'work'/clip/'enriched_summary.json').read_text())['sha256']
    complete=dest/'summary.json'
    if complete.exists():
        old=json.loads(complete.read_text());assert old['candidate_sha256']==c_hash and old['feature_sha256']==f_hash
        assert hashlib.sha256((dest/old['trace']).read_bytes()).hexdigest()==old['sha256']
        print('REUSED',clip);return
    engine=Integrated(cfg);ep=Episode(replace(Settings(),display_hold=1.0));holds={};outputs=[];times=[]
    rule_count=Counter();veto_count=Counter();reasons=Counter()
    for line in raw.splitlines():
        r=json.loads(line);start=time.perf_counter();v=engine.step(r)
        episode=ep.step(r['seconds'],v['hit'],v['quiet'])
        groups=group_states(r,v,r['seconds'],holds);times.append((time.perf_counter()-start)*1000)
        source_pairs={p['token']:p for p in r['pairs']}
        for p in v['pairs']:
            p['insufficiency_reasons']=source_pairs[p['token']]['feature_rejections']
            p['relation']=source_pairs[p['token']]['relation']
        out=dict(sample=r['sample'],source_frame=r['source_frame'],seconds=r['seconds'],displayed=r['displayed'],
            reference=r['reference'],case_ids=r['case_ids'],dedupe_key=r['dedupe_key'],review_frame=r['review_frame'],
            groups=groups,episode=episode,valid_people=len(r['bodies']),rejected=r['rejected'],
            ambiguous=r['ambiguous'],pairs_over_cap=r['pairs_over_cap'],retained_absent=r['retained_absent'],**v)
        outputs.append(out)
        if r['displayed']:
            rule_count.update({x for p in v['pairs'] for f in p['passing'] for x in f['rules']})
            veto_count.update({x for p in v['pairs'] for f in p['diagnostics'] if f['suppression_applied'] for x in f['suppression_evidence']})
            if not v['observable']:reasons['NO_OBSERVABLE_PAIR']+=1
            elif not v['hit']:reasons['NO_FINAL_HIT_WITH_OBSERVED_FEATURES']+=1
    trace=dest/'B_Integrated_diagnostics.jsonl'
    with trace.open('w') as f:
        for r in outputs:f.write(json.dumps(r,separators=(',',':'))+'\n')
    baseline=json.loads((BASELINE/'download'/clip/'data/summary.json').read_text())
    for entry in baseline['candidates'].values():
        path=BASELINE/'download'/clip/'data'/entry['trace']
        assert hashlib.sha256(path.read_bytes()).hexdigest()==entry['sha256']
        old=[json.loads(s) for s in path.read_text().splitlines()]
        assert [(r['sample'],r['source_frame'],r['seconds'],r['displayed']) for r in old]==[(r['sample'],r['source_frame'],r['seconds'],r['displayed']) for r in outputs]
    result=dict(clip=clip,trace=trace.name,sha256=hashlib.sha256(trace.read_bytes()).hexdigest(),
        candidate_sha256=c_hash,feature_sha256=f_hash,stats=stats(outputs),
        rule_observation_counts=dict(rule_count),applied_exclusion_observation_counts=dict(veto_count),
        insufficiency=dict(reasons),baseline={k:v for k,v in baseline['candidates'].items()},
        decision_p95_ms=statistics.quantiles(times,n=100,method='inclusive')[94],decision_max_ms=max(times))
    complete.write_text(json.dumps(result,ensure_ascii=False,indent=2))
    print('EVAL_PASS',clip,result['stats']['condition_frames'],dict(rule_count),flush=True)


if __name__=='__main__':main()