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

import argparse
import json
import pickle
from collections import Counter
from pathlib import Path
from types import SimpleNamespace

import numpy as np
from scipy.stats import wilcoxon

from scoobdoob.validation import msm_build as mb
from scoobdoob.validation.selectors import (
    SEED_RELAX,
    STRIDE,
    build_selectors,
    enrich,
    state_means,
)

HERE = Path(__file__).resolve().parent
CACHE = HERE / "data/p53/pocket_sc_feats_s10.pkl"
FRAMES = HERE / "data/p53/frame_drug/frames.json"
DEFAULT_OUT = HERE / "data/p53/trajectory_cv_tilt.json"


def assign_folds(meta, nfold=10):
    """Keep every fold heterogeneous while holding out independent trajectories."""
    folds = [[] for _ in range(nfold)]
    folded = [i for i, m in enumerate(meta) if m["cls"] == "folded"]
    regular = [
        i
        for i, m in enumerate(meta)
        if m["cls"] == "seed" and not m["name"].startswith("barrier_")
    ]
    barrier = [
        i
        for i, m in enumerate(meta)
        if m["cls"] == "seed" and m["name"].startswith("barrier_")
    ]
    for group, offset in ((folded, 0), (regular, 0), (barrier, 5)):
        for j, i in enumerate(group):
            folds[(j + offset) % nfold].append(i)
    return folds


def fit_model(data, train, k, lag, ncomp, seed):
    feats = [data["feats"][i] for i in train]
    cvs = [data["cvs"][i] for i in train]

    proj, tica = mb.tica_reduce(feats, lag=lag, n_comp=ncomp, scaling="kinetic_map")
    dtrajs, centers, km = mb.cluster_microstates(proj, k=k, seed=seed)
    state_xy, counts = mb.state_mean_features(dtrajs, cvs, k)
    model, active, sym2active, _ = mb.build_msm(dtrajs, lag=lag, k=k, reversible=True)
    msm = mb.MSM(
        model=model,
        dtrajs=dtrajs,
        centers=centers,
        state_xy=state_xy,
        counts=counts,
        active_symbols=active,
        lag=lag,
        tica=tica,
        _sym2active=sym2active,
    )

    held_dtrajs = {}
    for i in range(len(data["feats"])):
        if i in train:
            continue
        x = tica.transform(np.asarray(data["feats"][i], dtype=np.float64))
        held_dtrajs[i] = km.predict(x).astype(np.int64)
    return msm, held_dtrajs


def training_context(msm, data, train, lag):
    """Quantities needed by ``build_selectors``, calculated on training only."""
    active = np.asarray(msm.active_symbols)
    s2a = {int(s): i for i, s in enumerate(active)}
    n = len(active)
    cvsum, cvcnt = np.zeros((n, 2)), np.zeros(n)
    dsum, dcnt = np.zeros(n), np.zeros(n)
    C = np.zeros((n, n))
    starts = []

    for local, orig in enumerate(train):
        a = np.array([s2a.get(int(s), -1) for s in msm.dtrajs[local]])
        cv = np.asarray(data["cvs"][orig])
        feat = np.asarray(data["feats"][orig])
        nt = min(len(a), len(cv), len(feat))
        a, cv, feat = a[:nt], cv[:nt], feat[:nt]
        good = a >= 0
        np.add.at(cvsum, a[good], cv[good])
        np.add.at(cvcnt, a[good], 1)
        np.add.at(dsum, a[good], feat[good].mean(axis=1))
        np.add.at(dcnt, a[good], 1)
        if data["meta"][orig]["cls"] == "folded" and good[0]:
            starts.append(int(a[0]))
        pairs = (a[:-lag] >= 0) & (a[lag:] >= 0)
        np.add.at(C, (a[:-lag][pairs], a[lag:][pairs]), 1)

    if not starts:
        raise RuntimeError("training split has no active folded start state")
    cvm = cvsum / np.maximum(cvcnt[:, None], 1)
    dilation = np.full(n, np.nan)
    nz = dcnt > 0
    dilation[nz] = dsum[nz] / dcnt[nz]
    a0 = Counter(starts).most_common(1)[0][0]
    return cvm, C, a0, dilation


def frame_table(data, held_dtrajs, held, args):
    """Map held-out fpocket frame records into the held-out model labels."""
    by_orig = {int(m.get("orig_traj", i)): i for i, m in enumerate(data["meta"])}
    held = set(held)
    rows = json.loads(FRAMES.read_text())
    out_state, out_drug, out_traj = [], [], []

    for r in rows:
        ci = by_orig.get(int(r["traj"]))
        if ci not in held:
            continue
        value = r.get(args.observable)
        if value is None:
            continue
        core = r.get("core_rmsd_nm")
        zn = r.get("zn_n_coord", -1)
        if core is None or core > args.max_core_rmsd:
            continue
        if zn >= 0 and zn < args.min_zn:
            continue
        off = SEED_RELAX if data["meta"][ci]["cls"] == "seed" else 0
        pos = int(r["frame"]) // STRIDE - off
        dt = held_dtrajs[ci]
        if not 0 <= pos < len(dt):
            continue
        out_state.append(int(dt[pos]))
        out_drug.append(float(value))
        out_traj.append(ci)
    return (np.asarray(out_state), np.asarray(out_drug), np.asarray(out_traj))


def percentile_ci(x):
    if len(x) < 2:
        return [float("nan"), float("nan")]
    return np.percentile(x, [2.5, 97.5]).tolist()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--k", type=int, default=100)
    ap.add_argument("--lag", type=int, default=1)
    ap.add_argument("--ncomp", type=int, default=6)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--target-q", type=float, default=0.15)
    ap.add_argument(
        "--observable",
        default="site_drug",
        choices=["site_drug", "site_volume", "site_apolar", "site_sasa"],
    )
    ap.add_argument("--drug-thresh", type=float, default=0.5)
    ap.add_argument(
        "--alpha-rho-sweep",
        type=float,
        nargs="*",
        default=[0.0, 0.1, 0.25, 0.36, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0, 8.0, 12.0, 20.0],
        help="Algorithm-1 tilt strengths to evaluate (w=exp(-alpha_rho*rho), no "
        "committor shaping). alpha_rho=0 is the no-density-tilt control. The "
        "tilt strength is a knob; report the sweep, not one value.",
    )
    ap.add_argument("--max-core-rmsd", type=float, default=0.30)
    ap.add_argument("--min-zn", type=int, default=3)
    ap.add_argument("--out", type=Path, default=DEFAULT_OUT)
    args = ap.parse_args()

    data = pickle.load(open(CACHE, "rb"))
    folds = assign_folds(data["meta"])
    all_idx = set(range(len(data["meta"])))
    result = {"config": vars(args) | {"out": str(args.out)}, "folds": []}

    selector_names = (
        "equilibrium",
        "uniform_states",
        "tpt_flux",
        "doob_untilted",
        "fk_doob",
    ) + tuple(f"alg1[a={a:g}]" for a in args.alpha_rho_sweep)
    for fi, held in enumerate(folds):
        train = sorted(all_idx - set(held))
        print(f"[fold {fi}] train={len(train)} held={len(held)}")
        msm, held_dt = fit_model(data, train, args.k, args.lag, args.ncomp, args.seed)
        cvm, C, a0, dilation = training_context(msm, data, train, args.lag)

        fit_args = SimpleNamespace(
            target="dilation",
            target_q=args.target_q,
            cleft_thr=0.45,
            horizon=None,
            fk_variant="rho",
            comm_l_sweep=[],
            horizon_sweep=[],
            alpha_rho_sweep=args.alpha_rho_sweep,
        )
        selectors, q, B, a0, horizon = build_selectors(
            msm, cvm, C, a0, fit_args, dil=dilation
        )

        orig_state, drug, traj = frame_table(data, held_dt, held, args)
        s2a = {int(s): i for i, s in enumerate(msm.active_symbols)}
        state = np.array([s2a.get(int(s), -1) for s in orig_state])
        keep = state >= 0
        state, drug, traj = state[keep], drug[keep], traj[keep]
        n = msm.n_active
        mu, cnt = state_means(state, drug, n)
        po, _ = state_means(state, (drug > args.drug_thresh).astype(float), n)

        eq = enrich(
            selectors["equilibrium"],
            mu,
            cnt,
            po,
            drug.mean(),
            (drug > args.drug_thresh).mean(),
        )
        fold = {
            "fold": fi,
            "held_trajectories": held,
            "held_names": [data["meta"][i]["name"] for i in held],
            "n_frames": int(len(drug)),
            "n_active": int(n),
            "state_coverage": int((cnt > 0).sum()),
            "B_size": int(len(B)),
            "a0": int(a0),
            "horizon": int(horizon),
            "selectors": {},
        }
        for name in selector_names:
            if name not in selectors:
                continue
            e = enrich(
                selectors[name],
                mu,
                cnt,
                po,
                drug.mean(),
                (drug > args.drug_thresh).mean(),
            )
            if e is None:
                continue
            e["ratio_to_equilibrium"] = e["E_drug"] / eq["E_drug"]
            e["open_ratio_to_equilibrium"] = e["P_open"] / eq["P_open"]
            fold["selectors"][name] = e
        result["folds"].append(fold)
        fk = fold["selectors"].get("fk_doob", {})
        print(
            f"  n={len(drug)} coverage={fk.get('coverage', np.nan):.2f} "
            f"FK/equilibrium={fk.get('ratio_to_equilibrium', np.nan):.3f}"
        )

    summary = {}
    for name in selector_names:
        vals = [
            f["selectors"][name]["ratio_to_equilibrium"]
            for f in result["folds"]
            if name in f["selectors"]
        ]
        opens = [
            f["selectors"][name]["open_ratio_to_equilibrium"]
            for f in result["folds"]
            if name in f["selectors"]
        ]
        if not vals:
            continue
        arr = np.asarray(vals)
        if name == "equilibrium":
            p = 1.0
        else:
            try:
                p = float(wilcoxon(arr - 1, alternative="two-sided").pvalue)
            except ValueError:
                p = 1.0
        summary[name] = {
            "n_folds": len(vals),
            "mean_ratio_to_equilibrium": float(arr.mean()),
            "median_ratio_to_equilibrium": float(np.median(arr)),
            "fold_percentile_interval": percentile_ci(arr),
            "folds_above_1": int((arr > 1).sum()),
            "wilcoxon_p_vs_1": p,
            "mean_open_ratio_to_equilibrium": float(np.mean(opens)),
            "values": vals,
        }
    result["summary"] = summary
    args.out.write_text(json.dumps(result, indent=2, default=float))

    print("\nTRAJECTORY-HELD-OUT RESULT")
    for name, s in summary.items():
        print(
            f"  {name:<18} {s['mean_ratio_to_equilibrium']:.3f}x "
            f"(median {s['median_ratio_to_equilibrium']:.3f}; "
            f"{s['folds_above_1']}/{s['n_folds']} folds >1; "
            f"p={s['wilcoxon_p_vs_1']:.3g})"
        )
    print(f"-> {args.out}")


if __name__ == "__main__":
    main()