#!/usr/bin/env python
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import matplotlib

matplotlib.use("Agg")

from scoobdoob.validation import msm_build as mb
from scoobdoob.validation import tpt, bridge, aib9

ROOT = Path(__file__).resolve().parent
DEFAULT_DATA = aib9.OPENMM_DIR
OUTDIR = ROOT / "data" / "validation_aib9"

LAG = 5
K = 100
N_COMP = 4
HORIZON = 60
TOL_Q = 0.10


def gate_msm(feat, hand, temps):
    print("\n" + "=" * 70)
    print("PART 2.0/2.1  Build Aib9 MSM (36D sin/cos -> TICA -> connected -> rev MLE)")
    print("=" * 70)
    model = mb.build_full_msm(
        feat, hand, lag=LAG, k=K, n_comp=N_COMP, seed=0, use_tica=True
    )
    n_off = K - len(model.active_symbols)
    lam1 = float(np.real(model.model.eigenvalues(k=1)[0]))
    print(
        f"  microstates K={K}, connected set={model.n_active} "
        f"({n_off} dropped), lambda_1={lam1:.5f}"
    )
    lags = [2, 5, 10, 20, 40]
    lg, its = mb.implied_timescales(model.dtrajs, lags, k=K, n_its=3)
    ck = mb.ck_test(model.dtrajs, lag=LAG, k=K, factors=(2, 3))
    print(
        f"  ITS(tau*={LAG}, frames) slowest 3 = "
        f"{np.array2string(its[lags.index(LAG)], precision=1)}"
    )
    print(f"  CK error: 2tau={ck[2]:.3f}, 3tau={ck[3]:.3f}")
    return model, {
        "n_active": int(model.n_active),
        "n_dropped": int(n_off),
        "lambda1": lam1,
        "ck": {int(k): v for k, v in ck.items()},
    }


def gate_endpoints(model):
    print("\n" + "=" * 70)
    print("GATE 2.2  Endpoints R / L from helical handedness")
    print("=" * 70)
    labels = aib9.classify_microstates(model.active_xy)
    sets = {
        k: np.where(labels == k)[0].tolist() for k in ["R", "L", "I_N", "I_C", "other"]
    }
    print(
        f"  microstate labels: R={len(sets['R'])} L={len(sets['L'])} "
        f"I_N={len(sets['I_N'])} I_C={len(sets['I_C'])} other={len(sets['other'])}"
    )
    pL_R = model.active_xy[sets["R"]][:, aib9.CORE].mean()
    pL_L = model.active_xy[sets["L"]][:, aib9.CORE].mean()
    print(
        f"  R: core P(left)={pL_R:.2f} (right helix);  "
        f"L: core P(left)={pL_L:.2f} (left helix)"
    )

    A2, B2, _ = mb.q2_sign_endpoints(model, k_each=6)
    a2_isR = np.mean([labels[i] in ("R",) for i in A2])
    b2_isL = np.mean([labels[i] in ("L",) for i in B2])
    print(
        f"  q2-sign split agreement with R/L handedness: "
        f"{max(a2_isR, 1 - a2_isR) * 100:.0f}% / {max(b2_isL, 1 - b2_isL) * 100:.0f}%"
    )
    return labels, sets


def gate_committor(model, sets, labels):
    print("\n" + "=" * 70)
    print("GATE 2.3  Committor-1/2 vs KNOWN intermediates I_N, I_C  (THE GATE)")
    print("=" * 70)
    R, L = sets["R"], sets["L"]
    agree, q, _ = tpt.committor_agreement(model, R, L)
    print(f"  committor solver cross-check (deeptime vs linsolve): max|dq|={agree:.2e}")
    print(f"  q on R ~ {q[R].mean():.3f},  q on L ~ {q[L].mean():.3f}")

    half = tpt.committor_half_states(q, tol=TOL_Q)
    inter = set(sets["I_N"]) | set(sets["I_C"])
    half_in_inter = [h for h in half if h in inter]
    frac_states = len(half_in_inter) / max(len(half), 1)
    pi = model.stationary
    mass_frac = (
        pi[half_in_inter].sum() / max(pi[list(half)].sum(), 1e-12) if len(half) else 0.0
    )

    in_basin = [h for h in half if labels[h] in ("R", "L")]
    qIN = q[sets["I_N"]].mean() if sets["I_N"] else np.nan
    qIC = q[sets["I_C"]].mean() if sets["I_C"] else np.nan
    n_half_IN = sum(h in set(sets["I_N"]) for h in half)
    n_half_IC = sum(h in set(sets["I_C"]) for h in half)
    print(
        f"  committor-1/2 ensemble: {len(half)} states "
        f"({n_half_IN} in I_N, {n_half_IC} in I_C)"
    )
    print(
        f"  overlap with intermediates: {frac_states * 100:.0f}% of states, "
        f"{mass_frac * 100:.0f}% of committor-1/2 mass"
    )
    print(f"  committor-1/2 states inside an endpoint (R/L) basin: {len(in_basin)}")
    sym = "symmetric" if abs(qIN - qIC) < 0.1 else "asymmetric"
    print(
        f"  mean committor of I_N={qIN:.2f}, I_C={qIC:.2f}  "
        f"-> the two channels are {sym}. Aib is achiral, so a symmetric split "
        f"is the physically expected result."
    )
    passed = (agree < 1e-6) and (mass_frac >= 0.5) and (len(in_basin) == 0)
    print(
        f"  -> {'PASS' if passed else 'FAIL'}  "
        f"(committor-1/2 coincides with the known intermediates)"
    )
    return (
        passed,
        {
            "committor_agreement": agree,
            "n_half": int(len(half)),
            "overlap_frac_states": frac_states,
            "overlap_frac_mass": float(mass_frac),
            "n_half_in_basin": len(in_basin),
            "q_IN": float(qIN),
            "q_IC": float(qIC),
            "pass": bool(passed),
        },
        q,
        half,
    )


def gate_flux(model, sets, q, half):
    print("\n" + "=" * 70)
    print("GATE 2.4  Reactive flux + bridge rollouts route through intermediates")
    print("=" * 70)
    R, L = sets["R"], sets["L"]
    IN, IC = set(sets["I_N"]), set(sets["I_C"])

    rf = tpt.reactive_flux(model, R, L)
    order = ["R", "I_N", "I_C", "other", "L"]
    _, cg = rf.coarse_grain([sets[k] for k in order])
    F = np.asarray(cg.net_flux)
    tot = float(cg.total_flux)
    f_IN = float(F[0, 1] / tot)
    f_IC = float(F[0, 2] / tot)
    f_other = float(F[0, 3] / tot)
    f_direct = float(F[0, 4] / tot)
    f_thru = f_IN + f_IC + f_other
    print(f"  total reactive flux = {tot:.3e}")
    print(
        f"  flux routed through resolved intermediates: {f_thru * 100:.0f}%  "
        f"(direct R->L: {f_direct * 100:.0f}% -- small once swap artifacts removed)"
    )
    print(
        f"  TWO CHANNELS: flip-from-N (I_N) {f_IN * 100:.0f}%  |  "
        f"flip-from-C (I_C) {f_IC * 100:.0f}%  |  middle/other {f_other * 100:.0f}%"
    )

    band = set(np.where((q > 0.2) & (q < 0.8))[0].tolist())
    Ph = bridge.doob_kernels(model.P, L, HORIZON)
    paths = bridge.sample_bridge_paths(Ph, R, n_paths=300, seed=0)
    succ = bridge.success_with_ci(paths, L, seed=0)
    visit_band = np.mean([bool(set(map(int, p)) & band) for p in paths])
    visit_IN = np.mean([bool(set(map(int, p)) & IN) for p in paths])
    visit_IC = np.mean([bool(set(map(int, p)) & IC) for p in paths])
    print(
        f"  bridge rollouts R->L: success {succ['rate'] * 100:.0f}% "
        f"CI[{succ['ci'][0] * 100:.0f},{succ['ci'][1] * 100:.0f}], "
        f"cross committor band {visit_band * 100:.0f}%"
    )
    print(
        f"  rollouts touching I_N {visit_IN * 100:.0f}%, I_C {visit_IC * 100:.0f}% "
        f"(both channels sampled)"
    )

    passed = (f_IN > 0 and f_IC > 0) and (succ["rate"] > 0.9)
    print(
        f"  -> {'PASS' if passed else 'FAIL'}  "
        f"(multi-channel reactive ensemble through both known intermediates)"
    )
    return (
        passed,
        {
            "total_flux": tot,
            "flux_through_resolved_inter": f_thru,
            "flux_direct": f_direct,
            "flux_IN": f_IN,
            "flux_IC": f_IC,
            "flux_other": f_other,
            "rollout_success": succ["rate"],
            "rollout_cross_band": float(visit_band),
            "rollout_visit_IN": float(visit_IN),
            "rollout_visit_IC": float(visit_IC),
            "pass": bool(passed),
        },
        paths,
    )


def run_robustness(feat, hand):
    """Opt-in K x lag + ITS-convergence robustness report (gate stability)."""
    from scoobdoob.validation import robustness as rob

    print("\n" + "=" * 70)
    print("ROBUSTNESS  Are the conclusions stable to K and lag?")
    print("=" * 70)
    Ks, lags = [75, 100, 150, 200], [2, 5, 10, 20]
    results = rob.sweep(feat, hand, Ks, lags, n_comp=N_COMP)
    print(
        f"  {'K':>4} {'lag':>4} {'active':>6} {'lam1':>5} {'CK2t':>5} "
        f"{'chalf%':>7} {'direct%':>7} {'I_N%':>5} {'I_C%':>5}"
    )
    for r in results:
        if r.get("ok"):
            print(
                f"  {r['K']:4d} {r['lag']:4d} {r['n_active']:6d} {r['lam1']:5.2f} "
                f"{r['ck2']:5.2f} {r['committor_half_mass'] * 100:6.0f}% "
                f"{r['direct'] * 100:6.0f}% {r['flux_IN'] * 100:5.0f} {r['flux_IC'] * 100:5.0f}"
            )
    summ = rob.summarize(results)
    print(
        f"\n  committor-1/2 overlap across all {summ['n_points']} (K,lag): "
        f"{summ['committor_half_mass_min'] * 100:.0f}-{summ['committor_half_mass_max'] * 100:.0f}%"
    )
    print(
        f"  spurious direct flux across all (K,lag): "
        f"{summ['direct_min'] * 100:.0f}-{summ['direct_max'] * 100:.0f}%"
    )
    print(
        f"  both channels carry flux at every (K,lag): {summ['both_channels_always']}"
    )
    its_lags = [1, 2, 3, 5, 8, 12, 20, 30]
    lg, its = rob.its_convergence(feat, hand, K=K, lags=its_lags, n_comp=N_COMP)
    OUTDIR.mkdir(parents=True, exist_ok=True)
    rob.make_figure(
        results,
        lg,
        its,
        K_default=K,
        lag_default=LAG,
        dt_ps=1.0,
        outpath=ROOT / "figures" / "validation_aib9" / "aib9_robustness.png",
    )
    print(f"  robustness figure -> {OUTDIR / 'aib9_robustness.png'}")
    with open(OUTDIR / "robustness.json", "w") as f:
        json.dump(
            {
                "grid": results,
                "summary": summ,
                "its_lags": [int(x) for x in lg],
                "its": np.where(np.isfinite(its), its, None).tolist(),
            },
            f,
            indent=2,
            default=float,
        )
    print(f"  robustness report -> {OUTDIR / 'robustness.json'}")
    return summ


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=DEFAULT_DATA)
    ap.add_argument(
        "--robustness",
        action="store_true",
        help="also run the K x lag + ITS-convergence stability report",
    )
    ap.add_argument("--colvar_stride", type=int, default=5)
    args = ap.parse_args()

    trajs, temps = aib9.load_openmm_colvar(args.data, stride=args.colvar_stride)
    print(f"loaded {sum(len(t) for t in trajs)} frames")
    feat = aib9.featurize_sincos(trajs)
    hand = aib9.handedness_indicator(trajs)

    model, qc = gate_msm(feat, hand, temps)
    labels, sets = gate_endpoints(model)
    p23, r23, q, half = gate_committor(model, sets, labels)
    p24, r24, paths = gate_flux(model, sets, q, half)

    rob_summary = run_robustness(feat, hand) if args.robustness else None

    print("\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    for name, ok in [
        ("2.3 committor-1/2 <-> known intermediates", p23),
        ("2.4 reactive flux / rollouts through intermediates", p24),
    ]:
        print(f"  [{'PASS' if ok else 'FAIL'}]  {name}")
    gate2 = p23
    print(
        f"\n  GATE 2 (2.3) -> {'PASS' if gate2 else 'FAIL'}  "
        f"(structural-overlap criterion {'met' if gate2 else 'not met'})"
    )

    OUTDIR.mkdir(parents=True, exist_ok=True)
    report = {
        "config": {
            "lag": LAG,
            "K": K,
            "n_comp": N_COMP,
            "horizon": HORIZON,
            "tol_q": TOL_Q,
            "source": "openmm",
        },
        "qc": qc,
        "gate_2_3": r23,
        "gate_2_4": r24,
        "gate_2_pass": bool(gate2),
        "robustness": rob_summary,
    }
    with open(OUTDIR / "report.json", "w") as f:
        json.dump(report, f, indent=2, default=float)
    print(f"  report -> {OUTDIR / 'report.json'}")
    return 0 if gate2 else 1


if __name__ == "__main__":
    raise SystemExit(main())