#!/usr/bin/env python3
"""Exact-order replay of all frozen Phase 4N KC events through Phase 4M eligibility."""
from __future__ import annotations

import hashlib
import json
import math
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import yaml

ROOT=Path('/data/digital-drosophila')
EXP=ROOT/'experiments/phase4o-runtime-credit'
MODEL=ROOT/'models/male-cns-lif'


def save(p,x):
    p.parent.mkdir(parents=True,exist_ok=True)
    p.write_text(json.dumps(x,indent=2,sort_keys=True,default=str)+'\n')


def h(x): return hashlib.sha256(x.tobytes()).hexdigest()


def main():
    ncfg=yaml.safe_load((MODEL/'credit-assignment-v1/phase4n.yaml').read_text())
    mcfg=yaml.safe_load((MODEL/'kcmbon-efficacy-v1/phase4m.yaml').read_text())
    lcfg=yaml.safe_load(Path(mcfg['frozen_phase4l_config']).read_text())
    cfg=lcfg
    identity=pd.read_parquet(ROOT/'experiments/phase4n-credit-assignment/edge-identity/frozen-edge-identity.parquet')
    historical=pd.read_parquet(Path(mcfg['experiment_root'])/'efficacy-overlay/update-history.parquet')
    kc_events=pd.read_parquet(ROOT/'experiments/phase4n-credit-assignment/local-eligibility/KC-spike-history.parquet')
    graph=Path(cfg['paths']['connectome_root'])
    meta=pq.read_table(graph/'neuron-index/unified-metadata.parquet',columns=['neuron_idx','annotation_class']).to_pandas()
    kc=meta.loc[meta.annotation_class=='Kenyon_Cell','neuron_idx'].to_numpy(np.int32)
    kc.sort()
    lookup=np.full(len(meta),-1,np.int32);lookup[kc]=np.arange(len(kc),dtype=np.int32)
    src=identity.source_neuron_idx.to_numpy(np.int32);local=lookup[src]
    if (local<0).any(): raise RuntimeError('non-KC edge source')
    comp=identity.compartment.astype(str).to_numpy()==str(ncfg['matched_compartment_verbatim'])
    edgeid=identity.edge_id.to_numpy(np.int64)
    id_to_local=pd.Series(np.arange(len(identity),dtype=np.int32),index=edgeid)
    events={(int(seed),int(pair)):group[['step','neuron_idx']].to_numpy(np.int32)
            for (seed,pair),group in kc_events.groupby(['seed','pairing'])}
    histories={(int(seed),int(pair)):group for (seed,pair),group in historical.groupby(['train_seed','pairing'])}
    dt=float(cfg['model']['dt_ms']);tau=float(cfg['eligibility_v2']['kc_trace_tau_ms'])
    decay=np.float32(math.exp(-dt/tau));step_gate=int(round(float(mcfg['protocol']['dan_time_ms'])/dt))
    k=np.float32(cfg['eligibility_v2']['eta_effective_mv']/cfg['model']['w_syn_mv_per_anatomical_synapse'])
    bound=np.float32(1.0-cfg['eligibility_v2']['max_fractional_depression'])
    results=[];first_signatures={};mismatch_examples=[]
    for repeat in range(5):
        for seed in ncfg['train_seeds']:
            E=np.ones(len(identity),np.float32)
            for pairing in range(1,101):
                arr=events.get((seed,pairing),np.empty((0,2),np.int32))
                # Phase4N emitted events in simulation-time, then neuron_idx order.
                if len(arr) and (np.any(arr[:,0]<0) or np.any(arr[:,0]>=int(round(mcfg['protocol']['duration_ms']/dt)))): raise RuntimeError('invalid event timestamp')
                if len(arr)>1 and (np.any(arr[1:,0]<arr[:-1,0]) or np.any((arr[1:,0]==arr[:-1,0])&(arr[1:,1]<=arr[:-1,1]))): raise RuntimeError('event tie/duplicate/order ambiguity')
                relevant=arr[arr[:,0]<=step_gate]
                if len(relevant) and (lookup[relevant[:,1]]<0).any(): raise RuntimeError('non-KC event')
                by_step={int(step):lookup[relevant[relevant[:,0]==step,1]] for step in np.unique(relevant[:,0])}
                trace=np.zeros(len(kc),np.float32)
                for step in range(step_gate+1):
                    trace*=decay
                    firing=by_step.get(step)
                    if firing is not None: trace[firing]+=np.float32(1)
                    np.minimum(trace,np.float32(1),out=trace)
                rec=histories[(seed,pairing)]
                rid=id_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
                observed=trace[local[rid]]
                expected=rec.eligibility.to_numpy(np.float32)
                err=np.abs(observed-expected)
                actual=np.flatnonzero(comp&(trace[local]>0))
                missing=np.setdiff1d(rid,actual,assume_unique=True)
                extra=np.setdiff1d(actual,rid,assume_unique=True)
                predicted=np.maximum(bound,E[rid]*(np.float32(1)-k*observed))
                frozen_after=rec.E_bounded_after.to_numpy(np.float32)
                update_error=float(np.max(np.abs(predicted-frozen_after))) if len(predicted) else 0.0
                trace_error=float(err.max()) if len(err) else 0.0
                signature=h(np.concatenate((arr.flatten(),observed.view(np.int32),predicted.view(np.int32))))
                key=(seed,pairing)
                first_signatures.setdefault(key,signature)
                row={'repeat':repeat,'seed':seed,'pairing':pairing,'event_count':len(arr),'event_sha256':h(arr),
                     'gate_step':step_gate,'eligibility_sha256':h(observed),'update_sha256':h(predicted),
                     'repeat_exact':signature==first_signatures[key],
                     'frozen_trace_max_abs_error':trace_error,'frozen_update_max_abs_error':update_error,
                     'missing_frozen_active_edges':len(missing),'extra_active_edges':len(extra),
                     'exact_to_phase4m':trace_error<=2e-5 and update_error<=2e-5 and len(missing)==0 and len(extra)==0}
                if not row['exact_to_phase4m'] and len(mismatch_examples)<30:
                    j=int(np.argmax(err)) if len(err) else None
                    row['example']={'edge_id':int(rec.edge_id.iloc[j]) if j is not None else None,
                                    'KC_neuron_idx':int(src[rid[j]]) if j is not None else None,
                                    'timestamp_ms':float(mcfg['protocol']['dan_time_ms']),
                                    'event_trace':float(observed[j]) if j is not None else None,
                                    'frozen_trace':float(expected[j]) if j is not None else None,
                                    'frozen_delta_E':float(frozen_after[j]-rec.E_bounded_before.iloc[j]) if j is not None else None,
                                    'missing_edge_ids':edgeid[missing[:20]].tolist(),'extra_edge_ids':edgeid[extra[:20]].tolist()}
                    mismatch_examples.append(row.copy())
                results.append(row)
                # Each trial is run against the sealed Phase4M before/after trajectory.
                # A divergent diagnostic event trace never feeds the next trial.
                E[rid]=frozen_after
            print(json.dumps({'repeat':repeat,'seed':seed,'total_so_far':len(results),
                              'exact_to_phase4m_so_far':sum(r['exact_to_phase4m'] for r in results)}),flush=True)
            save(EXP/'eligibility-replay/event-reproducibility-matrix.json',results)
    save(EXP/'eligibility-replay/event-reproducibility-matrix.json',results)
    save(EXP/'eligibility-replay/event-replay-summary.json',{'complete_repeats':5,'events_per_repeat':500,
         'same_event_and_update_sequence_5_of_5':all(r['repeat_exact'] for r in results),
         'all_exact_to_frozen_phase4m':all(r['exact_to_phase4m'] for r in results),
         'frozen_exact_count_by_repeat':{str(rep):sum(r['exact_to_phase4m'] for r in results if r['repeat']==rep) for rep in range(5)},
         'mismatch_examples':mismatch_examples,
         'event_order_contract_observed':'simulation_step ascending, neuron_idx ascending, no duplicate KC/step; trace updated after spike; external gate sampled after step 900 trace update',
         'limitation':'Phase4M did not seal its training KC event sequence. Phase4N sealed the later event replay; exact repeat of those events does not prove the original GPU trajectory.'})


if __name__=='__main__':main()