#!/usr/bin/env python3
"""Five read-only repeats of sealed Phase 4M's 500-trial eligibility trajectory."""
from __future__ import annotations

import argparse
import hashlib
import json
import time
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

ROOT = Path("/data/digital-drosophila")
EXP = ROOT / "experiments/phase4o-runtime-credit"


def save(path: Path, x: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(x, indent=2, sort_keys=True, default=str) + "\n")


def arrhash(x: np.ndarray) -> str:
    return hashlib.sha256(x.tobytes()).hexdigest()


def first_diff(a: np.ndarray, b: np.ndarray):
    n = min(len(a), len(b))
    mismatched = np.flatnonzero(a[:n] != b[:n])
    j = int(mismatched[0]) if len(mismatched) else (n if len(a) != len(b) else None)
    if j is None:
        return None
    return {"event_position": j, "reference": None if j >= len(a) else int(a[j]), "observed": None if j >= len(b) else int(b[j])}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repeats", type=int, default=5)
    args = ap.parse_args()
    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")
    frozen = {(int(s), int(p)): f for (s, p), f in history.groupby(["train_seed", "pairing"])}
    old = pd.read_parquet(ROOT / "experiments/phase4n-credit-assignment/local-eligibility/KC-spike-history.parquet")
    old_events = {(int(s), int(p)): f[["step", "neuron_idx"]].to_numpy(np.int32)
                  for (s, p), f in old.groupby(["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)
    src = v.cand.source_neuron_idx.to_numpy(np.int32)
    lookup = np.full(v.n, -1, np.int32)
    lookup[v.kc] = np.arange(len(v.kc), dtype=np.int32)
    loc = lookup[src]
    comp = v.cand.target_compartment_verbatim.astype(str).to_numpy() == target
    ids = v.edge_ids.astype(np.int64)
    id_to_local = pd.Series(np.arange(len(ids), dtype=np.int32), index=ids)
    seen = {}
    rows = []
    first_failures = []
    for repeat in range(args.repeats):
        for seed in ncfg["train_seeds"]:
            delta = np.zeros(len(ids), np.float32)
            for pairing in range(1, 101):
                start = time.perf_counter()
                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)
                trace = run["kc_trace_at_dan"]
                events = run["kc_spikes"][["step", "neuron_idx"]].to_numpy(np.int32)
                old_ev = old_events.get((seed, pairing), np.empty((0, 2), np.int32))
                rec = frozen[(seed, pairing)]
                local = id_to_local.loc[rec.edge_id.to_numpy(np.int64)].to_numpy(np.int32)
                value = trace[loc[local]]
                recorded = rec.eligibility.to_numpy(np.float32)
                max_err = float(np.max(np.abs(value-recorded))) if len(value) else 0.0
                nonzero_now = np.flatnonzero(comp & (trace[loc] > 0))
                frozen_active = set(map(int, local))
                active_symmetric_diff = len(set(map(int, nonzero_now)) ^ frozen_active)
                key = (seed,pairing)
                h = arrhash(events)
                if key not in seen:
                    seen[key] = h
                row = {"repeat": repeat, "seed": seed, "pairing": pairing,
                       "KC_event_count": len(events), "KC_event_sha256": h,
                       "KC_events_same_as_first_repeat": h == seen[key],
                       "KC_events_same_as_phase4n_audit": bool(np.array_equal(events, old_ev)),
                       "first_KC_event_difference_from_phase4n": first_diff(old_ev[:,0],events[:,0]) if not np.array_equal(events,old_ev) else None,
                       "eligibility_max_abs_error_vs_frozen_phase4m": max_err,
                       "eligibility_active_edge_symmetric_difference": active_symmetric_diff,
                       "exact_eligibility": bool(max_err <= 2e-5 and active_symmetric_diff == 0),
                       "trace_sha256": arrhash(trace), "runtime_wall_seconds": run["wall_seconds"],
                       "audit_wall_seconds": time.perf_counter()-start}
                if not row["exact_eligibility"]:
                    err = np.abs(value-recorded)
                    j = int(np.argmax(err)) if len(err) else None
                    row["first_edge_mismatch"] = None if j is None else {
                        "edge_id": int(rec.edge_id.iloc[j]), "KC_neuron_idx": int(src[local[j]]),
                        "DAN_pulse_ms": float(mcfg["protocol"]["dan_time_ms"]),
                        "observed_trace": float(value[j]), "frozen_trace": float(recorded[j]),
                        "delta_trace": float(value[j]-recorded[j])}
                    if len(first_failures) < 20:
                        first_failures.append(row)
                        path = EXP / "eligibility-replay" / f"mismatch-repeat-{repeat}-seed-{seed}-pairing-{pairing}-KC-events.parquet"
                        path.parent.mkdir(parents=True,exist_ok=True)
                        run["kc_spikes"].to_parquet(path,index=False)
                rows.append(row)
                # Keep subsequent brain input equal to the sealed historical additive trajectory,
                # even when a diagnostic trial diverged; no frozen state is ever overwritten.
                sealed_trace = np.zeros(len(v.kc), np.float32)
                sealed_trace[loc[local]] = recorded
                delta, _ = apply_v2(v, sealed_trace, target, 1.0, delta, cfg)
                if pairing % 20 == 0:
                    save(EXP / "eligibility-replay/reproducibility-matrix.json", rows)
                    print(json.dumps({"repeat":repeat,"seed":seed,"pairing":pairing,
                                      "exact_in_block":sum(r["exact_eligibility"] for r in rows[-20:]),
                                      "event_repeat_exact_in_block":sum(r["KC_events_same_as_first_repeat"] for r in rows[-20:])}),flush=True)
    save(EXP / "eligibility-replay/reproducibility-matrix.json", rows)
    save(EXP / "eligibility-replay/summary.json", {"complete_repeats": args.repeats, "trials_per_repeat":500,
         "exact_frozen_eligibility_count":sum(r["exact_eligibility"] for r in rows),
         "exact_repeat_event_count":sum(r["KC_events_same_as_first_repeat"] for r in rows),
         "all_exact":all(r["exact_eligibility"] and r["KC_events_same_as_first_repeat"] for r in rows),
         "first_failures":first_failures,
         "historical_training_KC_event_sequence_available":False,
         "phase4n_event_sequence_used_as_comparator":True,
         "note":"Phase4M sealed only eligibility/update history, not all training KC events; Phase4N KC events are a later audit comparator."})


if __name__ == "__main__": main()