#!/usr/bin/env python3
"""Targeted GPU replays of the two Phase 4N/4M eligibility discrepancies."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from phase4n import load_config
from phase4i import Validator
from phase4j import apply_v2
from phase4l import OdorSimulator,pn_drive
from phase4m import load_odor_inputs
from phase4o_a import EXP,ROOT,save


def main():
    ncfg,mcfg,cfg,cfgpath=load_config(ROOT/'models/male-cns-lif/credit-assignment-v1/phase4n.yaml')
    history=pd.read_parquet(Path(mcfg['experiment_root'])/'efficacy-overlay/update-history.parquet')
    by={(int(s),int(p)):g for (s,p),g in history.groupby(['train_seed','pairing'])}
    v=Validator(cfg,cfgpath,EXP);sim=OdorSimulator(v,cfg,EXP)
    panel,vectors,selection,mapping=load_odor_inputs(cfg)
    oid=selection['ODOR_A']['odor_id']
    pn=pn_drive(oid,vectors[oid],mapping,selection['glomeruli_order'],str(mcfg['protocol']['input_mode']),float(np.median(panel.mapped_response_sum)),cfg)[:2]
    target=str(mcfg['protocol']['target_compartment_verbatim'])
    dan=v.dan_targets(str(mcfg['protocol']['matched_dan_type']),target)
    lookup=np.full(v.n,-1,np.int32);lookup[v.kc]=np.arange(len(v.kc),dtype=np.int32)
    loc=lookup[v.cand.source_neuron_idx.to_numpy(np.int32)]
    eid_to_local=pd.Series(np.arange(len(v.edge_ids),dtype=np.int32),index=v.edge_ids.astype(np.int64))
    rows=[]
    for seed,pairing in [(51202,40),(51203,39)]:
        delta=np.zeros(len(v.edge_ids),np.float32)
        for p in range(1,pairing):
            rec=by[(seed,p)]
            local=eid_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
            frozen_trace=np.zeros(len(v.kc),np.float32)
            frozen_trace[loc[local]]=rec.eligibility.to_numpy(np.float32)
            delta,_=apply_v2(v,frozen_trace,target,1.,delta,cfg)
        rec=by[(seed,pairing)]
        local=eid_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
        old_trace=rec.eligibility.to_numpy(np.float32)
        for repeat in range(15):
            run=sim.run_rates(*pn,seed+pairing,float(mcfg['protocol']['duration_ms']),delta=delta,dan_targets=dan,
                              dan_pulses=[float(mcfg['protocol']['dan_time_ms'])],collect_kc_spikes=True)
            observed=run['kc_trace_at_dan'][loc[local]]
            diff=np.abs(observed-old_trace)
            j=int(np.argmax(diff)) if len(diff) else None
            events=run['kc_spikes'][['step','neuron_idx']].to_numpy(np.int32)
            row={'seed':seed,'pairing':pairing,'repeat':repeat,'event_count':len(events),
                 'event_sha256':hashlib.sha256(events.tobytes()).hexdigest(),
                 'trace_sha256':hashlib.sha256(observed.tobytes()).hexdigest(),
                 'max_abs_error_vs_frozen':float(diff.max()) if len(diff) else 0.,
                 'exact_to_frozen':bool(len(diff)==0 or float(diff.max())<=2e-5),
                 'first_largest_error_edge_id':int(rec.edge_id.iloc[j]) if j is not None else None,
                 'first_largest_error_KC_neuron_idx':int(v.cand.source_neuron_idx.iloc[local[j]]) if j is not None else None,
                 'observed_eligibility':float(observed[j]) if j is not None else None,
                 'frozen_eligibility':float(old_trace[j]) if j is not None else None,
                 'wall_seconds':run['wall_seconds']}
            rows.append(row)
            print(json.dumps({'seed':seed,'pairing':pairing,'repeat':repeat,'exact':row['exact_to_frozen'],
                              'event_sha256':row['event_sha256'][:12]}),flush=True)
            save(EXP/'eligibility-replay/targeted-gpu-repeats.json',rows)
    save(EXP/'eligibility-replay/targeted-summary.json',{'rows':len(rows),
         'exact_count':sum(r['exact_to_frozen'] for r in rows),
         'distinct_event_hashes_by_trial':{f'{s}/{p}':len({r['event_sha256'] for r in rows if r['seed']==s and r['pairing']==p}) for s,p in [(51202,40),(51203,39)]},
         'same_seed_config_and_sealed_additive_prefix':True,
         'note':'This is an isolated 15-repeat targeted diagnostic of the two formerly discrepant trials, not a 5x500 complete brain rerun.'})


if __name__=='__main__':main()