"""Recompute the frozen exploratory case for transparent plotted probabilities.

No new selection, no threshold changes, no full-network intervention after failure.
"""
import argparse
import json
from pathlib import Path
import numpy as np
from functional.run import save_csv,write_json
from .core import QKVModel,distribution_metrics
from .conditional import tv_metrics


def ranks_and_mass(a,b):
    result=[]
    for bi in range(a.shape[0]):
        d={'target_self_mass':[],'donor_self_mass':[],'argmax_agreement':[],'shared_top5_mass_target':[],'shared_top5_mass_donor':[]}
        for t in range(7,32):
            x,y=a[bi,0,t,:t+1],b[bi,0,t,:t+1]
            common=np.intersect1d(np.argsort(x)[-5:],np.argsort(y)[-5:])
            d['target_self_mass'].append(float(x[t]));d['donor_self_mass'].append(float(y[t]))
            d['argmax_agreement'].append(float(np.argmax(x)==np.argmax(y)))
            d['shared_top5_mass_target'].append(float(x[common].sum()));d['shared_top5_mass_donor'].append(float(y[common].sum()))
        result.append({k:float(np.mean(v)) for k,v in d.items()})
    return result


def main():
    p=argparse.ArgumentParser();p.add_argument('--source',type=Path,required=True);p.add_argument('--root',type=Path,required=True);args=p.parse_args();root=args.root;source=args.source
    status=json.loads((root/'conditional-status.json').read_text());assert status['state']=='completed'
    data=np.load(root/'cache'/'conditional-fresh-inputs.npz')
    samples=json.loads((root/'conditional-fresh-samples.json').read_text())['samples']
    lm=QKVModel(source/'weights',json.loads((source/'model_config.json').read_text()),seconds=300)
    own=lm.parts(data['attn_23'],23);same=lm.parts(data['attn_23'],3);actual=lm.parts(data['attn_3'],3)
    target=own['p'][:,6:7];donor=same['p'][:,7:8];cached=actual['p'][:,7:8]
    np.savez(root/'conditional-fresh-probabilities.npz',target_layer23_head6=target,donor_layer3_head7_on_H23=donor,donor_layer3_head7_on_H3=cached)
    rows=[]
    from .run import outrows
    for mode,profile,prob in [('same_H23',same,donor),('actual_H3_to_H23',actual,cached)]:
        stats={**distribution_metrics(target,prob),**tv_metrics(target,prob)}
        mass=ranks_and_mass(target,prob)
        changed=lm.head_transplant_output(own,profile,23,6,7)
        out=outrows(own['heads'][:,:,6],changed,samples,{},1e-6,1e-6)
        for n,s in enumerate(samples):rows.append({'mode':mode,'sample':s['id'],**mass[n],**{k:(None if v[n] is None else float(v[n])) for k,v in stats.items()},'head_output_error':max(out[n]['error_a'],out[n]['error_b'])})
    save_csv(root/'conditional-fresh-source-diagnostic.csv',rows)
    with (root/'conditional-function-metrics.csv').open() as f:import csv;allrows=list(csv.DictReader(f))
    failures=[]
    for r in allrows:
        if r['split']!='fresh_test':continue
        for field,limit in [('js',.01),('js_p95',.05),('total_variation',.10),('total_variation_p95',.20)]:
            if float(r[field])>limit:failures.append({'sample':r['sample'],'field':field,'value':float(r[field]),'limit':limit})
    same_rows=[r for r in rows if r['mode']=='same_H23']
    classification='self_token_dominated' if min(min(r['target_self_mass'],r['donor_self_mass']) for r in same_rows)>.85 else 'inspect detailed source diagnostic'
    write_json(root/'conditional-case-audit.json',{'passed_recomputation':True,'strict_failures':failures,'pattern_classification':classification,
                'minimum_same_input_target_self_mass':min(r['target_self_mass'] for r in same_rows),'minimum_same_input_donor_self_mass':min(r['donor_self_mass'] for r in same_rows),
                'minimum_same_input_argmax_agreement':min(r['argmax_agreement'] for r in same_rows),
                'minimum_shared_top5_probability_mass':min(min(r['shared_top5_mass_target'],r['shared_top5_mass_donor']) for r in same_rows),
                'maximum_same_input_head_output_error':max(r['head_output_error'] for r in same_rows),
                'no_success_relabeling':True,'network_intervention_not_run_due_to_fresh_gate_failure':True})


if __name__=='__main__':main()