#!/usr/bin/env python3
"""Phase 4N checkpoint-specific whole-brain/credit audit and conservative sealing."""
from __future__ import annotations

import argparse
import gc
import json
import resource
import shutil
import subprocess
import time
from pathlib import Path

import numpy as np
import pandas as pd
import torch

from phase4h_audit import sha256, verify_checksum, write_json
from phase4i import Validator, phase4c_regression
from phase4j import effect
from phase4k import dataset_checks, seal
from phase4l import OdorSimulator, pn_drive
from phase4m import frozen_checks, frozen_hashes, load_odor_inputs
from phase4n import GROUPS, load_config


def state(exp, seed, pair):
    if pair == 0:
        return np.ones(14232, np.float32)
    with np.load(exp / "acquisition-trajectory" / f"E-seed-{seed}-pairings-{pair}.npz") as f:
        return f["E_bounded"].copy()


def curves(ncfg, exp, csi, replay, identity):
    csi = csi[csi.mechanism == "bounded"]
    peak = json.loads((exp / "specificity-curve/peak.json").read_text())["pairings"]
    onset = None
    if peak is not None:
        p = csi[csi.pairings == peak]
        baseline_csi = p.CSI.mean()
        baseline_other = max(p.shared.mean(), p["B-specific"].mean())
        for pair in ncfg["checkpoints"]:
            if pair <= peak:
                continue
            q = csi[csi.pairings == pair]
            other = max(q.shared.mean(), q["B-specific"].mean())
            if q.CSI.mean() <= .8 * baseline_csi and other >= 1.1 * baseline_other and baseline_csi > 0 and baseline_other > 0:
                onset = pair
                break
    write_json(exp / "specificity-curve/generalization-onset.json", {"status": "GENERALIZATION-ONSET-IDENTIFIED" if onset is not None else "UNIDENTIFIED",
        "pairings": onset, "rule": ncfg["generalization_onset_rule"]})
    groups = pd.read_parquet(exp / "acquisition-trajectory/group-plasticity.parquet")
    groups = groups[groups.mechanism == "bounded"]
    # Dependency-free vector plot; values remain in the accompanying Parquet files.
    parts = ['<svg xmlns="http://www.w3.org/2000/svg" width="900" height="640" viewBox="0 0 900 640">',
             '<rect width="900" height="640" fill="white"/>', '<text x="55" y="30" font-size="17">Phase 4N acquisition and specificity</text>']
    def panel(series, y0, height, title):
        allv = np.concatenate([np.asarray(v, np.float64) for _, v, _ in series]); lo, hi = float(np.min(allv)), float(np.max(allv))
        if hi <= lo: hi = lo + 1.0
        parts.append(f'<text x="55" y="{y0-10}" font-size="13">{title}</text>')
        parts.append(f'<path d="M60 {y0} V{y0+height} H850" fill="none" stroke="black"/>')
        for label, vals, color in series:
            xs = np.asarray(ncfg["checkpoints"], np.float64)
            ys = np.asarray(vals, np.float64)
            coords = ' '.join(f'{60+790*x/100:.1f},{y0+height-20-(height-35)*(y-lo)/(hi-lo):.1f}' for x,y in zip(xs,ys))
            parts.append(f'<polyline points="{coords}" fill="none" stroke="{color}" stroke-width="2"/>')
            parts.append(f'<text x="{625}" y="{y0+18+17*len([s for s in series if s[0] < label])}" font-size="11" fill="{color}">{label}</text>')
        for p,color in ((peak,"green"),(onset,"red")):
            if p is not None:
                x = 60 + 790*p/100
                parts.append(f'<path d="M{x:.1f} {y0} V{y0+height}" stroke="{color}" stroke-dasharray="4,4"/>')
    first = []
    for name,color in zip(GROUPS,("#1261a0","#d4871e","#b72230","#777777")):
        g = groups[groups.group == name].groupby("pairings").weighted_transmission_reduction.mean().reindex(ncfg["checkpoints"]).fillna(0.0)
        first.append((name,g.to_numpy(),color))
    panel(first, 80, 210, "Mean anatomical-weighted E reduction")
    panel([("A-B spike L1",replay.mean_spike_L1_separation.to_numpy(),"#2059a6"),
           ("A-B latency ms",replay.mean_latency_separation_ms.to_numpy(),"#b72230")], 370, 210, "Functional MBON effect difference")
    parts.append('</svg>')
    (exp / "specificity-curve/trajectory.svg").write_text('\n'.join(parts), encoding="utf-8")
    return peak, onset


def shared_decomposition(ncfg, exp, identity, sim, v, peak):
    if peak is None:
        write_json(exp / "shared-edge/summary.json", {"status": "NOT-RUN-NO-SPECIFICITY-PEAK"})
        return None
    from phase4m import ReplayEngine, replay_effect
    comp = identity.compartment.astype(str).to_numpy() == str(ncfg["matched_compartment_verbatim"])
    targets = np.unique(identity.loc[comp, "target_neuron_idx"].to_numpy(np.int32))
    engine = ReplayEngine(sim, targets)
    w = identity.anatomical_weight_mv.to_numpy(np.float32)
    groups = identity.identity_group.astype(str).to_numpy()
    rows = []
    from phase4n import load_config
    _, mcfg, _, _ = load_config(Path(__file__).with_name("phase4n.yaml"))
    for seed, es in zip(ncfg["train_seeds"], ncfg["evaluation_seeds"]):
        E = state(exp, seed, peak)
        for odor in ("A", "B"):
            spikes = pd.read_parquet(Path(mcfg["experiment_root"]) / "isolated-replay" / f"KC-spikes-ODOR_{odor}-seed-{es}.parquet")
            baseline = engine.run(spikes, np.zeros(len(E), np.float32), float(mcfg["protocol"]["duration_ms"]))
            for label, included in (("A-only", ("A-specific",)), ("A+shared", ("A-specific", "shared")),
                                    ("A+shared+B", ("A-specific", "shared", "B-specific")), ("all-updated", GROUPS)):
                delta = np.where(np.isin(groups, included), w * (E - 1), 0).astype(np.float32)
                post = engine.run(spikes, delta, float(mcfg["protocol"]["duration_ms"]))
                metrics = replay_effect(baseline, post, sim.dt)
                rows.append({"train_seed": seed, "eval_seed": es, "pairings": peak, "odor": odor, "decomposition": label,
                    **{k: val for k, val in metrics.items() if k != "isolated_input_scope"}})
    pd.DataFrame(rows).to_parquet(exp / "shared-edge/isolated-decomposition.parquet", index=False)
    return pd.DataFrame(rows)


def wholebrain(ncfg, mcfg, cfg, exp, identity, sim, v, peak):
    wanted = sorted(set([10, 20, 50, 75, 100] + ([peak] if peak is not None else [])))
    panel, vectors, selection, mapping = load_odor_inputs(cfg)
    gloms = selection["glomeruli_order"]
    matched = float(np.median(panel.mapped_response_sum))
    pn = {}
    for label in ("A", "B"):
        oid = selection[f"ODOR_{label}"]["odor_id"]
        pn[label] = pn_drive(oid, vectors[oid], mapping, gloms, str(mcfg["protocol"]["input_mode"]), matched, cfg)[:2]
    dan = np.unique(v.confirmed.source_neuron_idx.to_numpy(np.int32))
    group_sources = {name: np.unique(identity.loc[identity.identity_group == name, "source_neuron_idx"].to_numpy(np.int32)) for name in GROUPS}
    w = identity.anatomical_weight_mv.to_numpy(np.float32)
    rows, calls = [], []
    for seed, es in zip(ncfg["train_seeds"], ncfg["evaluation_seeds"]):
        baseline = {label: sim.run_rates(*pn[label], es, float(mcfg["protocol"]["duration_ms"]), collect_kc_spikes=True) for label in pn}
        for pair in wanted:
            delta = w * (state(exp, seed, pair) - 1)
            for label in ("A", "B"):
                post = sim.run_rates(*pn[label], es, float(mcfg["protocol"]["duration_ms"]), delta=delta, collect_kc_spikes=True)
                calls.append(post["wall_seconds"])
                eff = effect(baseline[label], post)
                row = {"train_seed": seed, "eval_seed": es, "pairings": pair, "odor": label,
                    "MBON_spike_L1": eff["MBON_spike_l1"], "MBON_g_max_abs_mv": eff["MBON_g_max_abs"],
                    "MBON_v_max_abs_mv": eff["MBON_v_max_abs"], "MBON_latency_delta_ms": eff["MBON_first_latency_delta_ms"],
                    "downstream_spike_L1": eff["downstream_spike_l1"], "KC_spike_L1": int(np.abs(post["kc_counts"]-baseline[label]["kc_counts"]).sum()),
                    "endogenous_DAN_spikes": int(post["counts"][dan].sum()), "baseline_endogenous_DAN_spikes": int(baseline[label]["counts"][dan].sum()),
                    "finite": bool(post["finite"]), "peak_vram_bytes": int(post["peak_vram_bytes"]), "wall_seconds": post["wall_seconds"]}
                for group, nodes in group_sources.items():
                    row[f"{group}_KC_spikes"] = int(post["counts"][nodes].sum())
                    row[f"{group}_KC_newly_active_vs_E1"] = int(np.count_nonzero((post["counts"][nodes] > 0) & (baseline[label]["counts"][nodes] == 0)))
                row["MBON_spikes"] = int(post["mbon_counts"].sum())
                row["other_feedback_spikes"] = int(post["downstream_counts"].sum())
                rows.append(row)
    frame = pd.DataFrame(rows)
    frame.to_parquet(exp / "whole-brain/checkpoint-responses.parquet", index=False)
    write_json(exp / "performance/whole-brain.json", {"calls": len(calls), "mean_wall_seconds": float(np.mean(calls)),
        "sum_wall_seconds": float(np.sum(calls)), "max_peak_vram_bytes": int(frame.peak_vram_bytes.max())})
    return frame


def audit_updates(ncfg, mcfg, exp, identity, history, csi, peak):
    group = identity.set_index("edge_id").identity_group
    hist = history.copy()
    hist["identity_group"] = group.loc[hist.edge_id.to_numpy(np.int64)].to_numpy()
    hist["delta_E_bounded"] = hist.E_bounded_after - hist.E_bounded_before
    hist["DAN_source"] = "external-US-associated"
    hist["matched_compartment"] = True
    hist["endogenous_DAN_used_as_gate"] = False
    hist["local_KC_trace_nonzero"] = hist.eligibility > 0
    hist.to_parquet(exp / "local-eligibility/phase4m-updates-classified.parquet", index=False)
    first = hist[(hist.identity_group == "B-specific") & (hist.delta_E_bounded < -1e-7)].groupby("train_seed").pairing.min().reset_index(name="first_B_specific_update_pairing")
    first.to_parquet(exp / "b-contamination/first-B-updates.parquet", index=False)
    causes = {"direct_A_driven_KC_activity": None, "recurrent_KC_reactivation": None, "MBON_DAN_feedback": None,
        "olfactory_feedback": None, "long_eligibility_trace": None, "endogenous_DAN_teaching_interpretation": 0.0,
        "unresolved_runtime_KC_recruitment_or_trace": 1.0}
    write_json(exp / "b-contamination/provenance.json", {"classification": causes,
        "reason": "Frozen Phase 4M update history records trace at forced matched-US but no event-level causal predecessor or KC spike times. Assigning a recurrent/feedback percentage would be invented.",
        "earliest_B_update_by_seed": first.to_dict("records"), "all_updates_gated_by_external_matched_US": True})
    write_json(exp / "dan-source/phase4m-gate-audit.json", {"external_US_gates": int(history[["train_seed", "pairing"]].drop_duplicates().shape[0]),
        "endogenous_DAN_as_teaching_gate": 0, "policy": "Phase 4M implementation opens plasticity once at imposed 90 ms US, not at endogenous DAN spike times",
        "observed_endogenous_DAN_activity_is_not_equivalent_to_dopamine_concentration": True})
    score = csi[csi.mechanism == "bounded"].groupby("pairings").CSI.mean()
    write_json(exp / "stopping-analysis/summary.json", {"status": "POST-HOC-ONLY-NO-ADAPTIVE-POLICY",
        "peak_CSI_checkpoint": int(score.idxmax()), "peak_CSI": float(score.max()), "functional_peak_checkpoint": peak,
        "online_validated_criterion": False, "heldout_training_validation_required": True})


def controls(exp, identity, ncfg, mcfg):
    rows = []
    for seed in ncfg["train_seeds"]:
        for pair in ncfg["checkpoints"]:
            E = state(exp, seed, pair)
            rows += [{"seed": seed, "pairings": pair, "condition": "A+matched-US", "E_changed": int(np.count_nonzero(E < 1 - 1e-7)), "status": "RUN-RECONSTRUCTED-FROZEN-PHASE4M"},
                     {"seed": seed, "pairings": pair, "condition": "A+D0", "E_changed": 0, "status": "EXACT-ZERO-GATE-IDENTITY"},
                     {"seed": seed, "pairings": pair, "condition": "frozen-E1", "E_changed": 0, "status": "EXACT-FROZEN-IDENTITY"},
                     {"seed": seed, "pairings": pair, "condition": "external-US-only-gate", "E_changed": int(np.count_nonzero(E < 1 - 1e-7)), "status": "IDENTICAL-TO-FROZEN-PHASE4M-RULE"}]
    pd.DataFrame(rows).to_parquet(exp / "controls/trajectory-status.parquet", index=False)
    report = {"A_matched_US": "RUN-PHASE4M-TRAJECTORY-RECONSTRUCTED", "A_D0": "ANALYTIC-EXACT-IDENTITY", "frozen_E1": "ANALYTIC-EXACT-IDENTITY",
        "external_US_only": "RUN-IDENTICAL-TO-PHASE4M", "B_matched_US": "NOT-RUN", "random_US_timing": "NOT-RUN",
        "unrelated_compartment_US": "NOT-RUN", "all_DAN_gate": "NOT-RUN", "full_control_matrix_pass": False}
    write_json(exp / "controls/summary.json", report)
    return report


def integrity(ncfg, mcfg, cfg, base_path, exp, frozen_before):
    gc.collect(); torch.cuda.empty_cache()
    checks = frozen_checks(cfg)
    mroot, mm = Path(mcfg["experiment_root"]), Path(mcfg["model_root"])
    mchecks = {"phase4m_exp": verify_checksum(mroot, mroot / "manifests/checksums.sha256")["passed"],
        "phase4m_model": verify_checksum(mm, mm / "checksums.sha256")["passed"]}
    dataset = dataset_checks(cfg)
    regression = phase4c_regression(cfg, base_path, exp)
    after = frozen_hashes(cfg)
    report = {"all_passed": bool(checks["all_passed"] and all(mchecks.values()) and dataset["all_passed"] and regression["all_passed"] and frozen_before == after),
        "phase4c_to_4l": checks, "phase4m": mchecks, "phase1_to_3b": dataset, "phase4c_reference_regression": regression,
        "before_after_identical": frozen_before == after, "base_csr_unmodified": True, "APL_v2_encoder_mask_unmodified": True}
    write_json(exp / "audit/frozen-integrity.json", report)
    return report


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--config", type=Path, required=True); args = ap.parse_args()
    ncfg, mcfg, cfg, base_path = load_config(args.config)
    exp = Path(ncfg["experiment_root"])
    identity = pd.read_parquet(exp / "edge-identity/frozen-edge-identity.parquet")
    csi = pd.read_parquet(exp / "specificity-curve/CSI.parquet")
    replay = pd.read_parquet(exp / "specificity-curve/memory-specificity.parquet")
    history = pd.read_parquet(Path(mcfg["experiment_root"]) / "efficacy-overlay/update-history.parquet")
    before = frozen_hashes(cfg)
    peak, onset = curves(ncfg, exp, csi, replay, identity)
    print(json.dumps({"stage": "specificity", "peak": peak, "generalization_onset": onset}), flush=True)
    v = Validator(cfg, base_path, exp); sim = OdorSimulator(v, cfg, exp)
    shared = shared_decomposition(ncfg, exp, identity, sim, v, peak)
    print(json.dumps({"stage": "shared_decomposition", "rows": 0 if shared is None else len(shared)}), flush=True)
    whole = wholebrain(ncfg, mcfg, cfg, exp, identity, sim, v, peak)
    print(json.dumps({"stage": "wholebrain", "rows": len(whole)}), flush=True)
    audit_updates(ncfg, mcfg, exp, identity, history, csi, peak)
    control = controls(exp, identity, ncfg, mcfg)
    del sim, v
    integ = integrity(ncfg, mcfg, cfg, base_path, exp, before)
    ru = resource.getrusage(resource.RUSAGE_SELF)
    perf = {"peak_RAM_KiB": int(ru.ru_maxrss), "edge_audit_rows": int(len(history)),
        "edge_audit_memory_bytes": int(history.memory_usage(deep=True).sum()), "shadow_trace_overhead": "NOT-MEASURED-NOT-RUN",
        "wholebrain": json.loads((exp / "performance/whole-brain.json").read_text()),
        "isolated": json.loads((exp / "performance/isolated.json").read_text())}
    write_json(exp / "performance/summary.json", perf)
    summary = {"phase": "4N", "status": "PARTIAL-MECHANISM-AUDIT-NO-CANDIDATE", "peak": peak, "generalization_onset": onset,
        "edge_groups": identity.identity_group.value_counts().to_dict(), "controls": control, "frozen_integrity": integ["all_passed"],
        "candidate_generated": False, "limitations": ["Endogenous DAN event-level gate overlay not run", "Full matched control trajectories not run", "Runtime causal predecessor and local shadow traces not captured; contamination source percentages unresolved"]}
    write_json(exp / "audit/final-summary.json", summary)
    model = Path(ncfg["model_root"]); model.mkdir(parents=True, exist_ok=True)
    shutil.copy2(args.config, model / "phase4n.yaml"); shutil.copy2(Path(__file__).resolve(), model / "phase4n_finish.py")
    shutil.copy2(Path(__file__).with_name("phase4n.py"), model / "phase4n.py")
    shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4n_finish.py")
    write_json(model / "VERSION.json", {"phase": "4N", "status": summary["status"], "candidate_generated": False,
        "identity_sha256": sha256(exp / "edge-identity/frozen-edge-identity.parquet"), "config_sha256": sha256(args.config)})
    seal(model, model / "MANIFEST.json", model / "checksums.sha256")
    seal(exp, exp / "manifests/manifest.json", exp / "manifests/checksums.sha256")
    print(json.dumps(summary), flush=True)


if __name__ == "__main__": main()