#!/usr/bin/env python
import json
import sys
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parents[2]
FIG = ROOT / "data" / "p53"
CACHE = FIG / "fig_data"
sys.path.insert(0, str(ROOT))


HIGHLIGHT = [56, 0]
RARE_THRESH = 0.98


def main():
    from scoobdoob.validation.pocket import site_drug

    CACHE.mkdir(parents=True, exist_ok=True)

    M = np.load(FIG / "rmsd_matrix.npy")
    mids = sorted(
        int(p.stem.replace("micro", ""))
        for p in (FIG / "pocket_reps").glob("micro*.pdb")
    )
    if len(mids) != M.shape[0]:
        raise SystemExit(
            f"{len(mids)} representatives but rmsd_matrix has {M.shape[0]} rows"
        )

    drug = np.array([site_drug(m) for m in mids], dtype=float)
    truth = np.nanmedian(M, axis=1)
    closest = np.nanmin(M, axis=1)
    order = np.argsort(truth)
    rank = {mids[j]: r + 1 for r, j in enumerate(order)}

    pij = json.loads((FIG / "pocket_msm_pi.json").read_text())
    pi = np.array([float(pij[str(m)]) if str(m) in pij else np.nan for m in mids])
    if not np.isfinite(pi).all():
        missing = [m for m, p in zip(mids, pi) if not np.isfinite(p)]
        raise SystemExit(
            f"no stationary probability for microstates {missing}; "
            f"run upstream/dump_pocket_msm_pi.py first"
        )

    rows = json.loads((FIG / "frame_drug" / "frames.json").read_text())
    v = np.array(
        [
            r.get("site_drug") if r.get("site_drug") is not None else np.nan
            for r in rows
        ],
        dtype=float,
    )

    import pickle

    meta = pickle.load(open(FIG / "pocket_feats_s10_cb_ALL.pkl", "rb"))["meta"]
    cls = np.array(
        [meta[r["traj"]]["cls"] if r["traj"] < len(meta) else "?" for r in rows]
    )
    fin = np.isfinite(v)
    v_ref, v_seed = v[fin & (cls == "folded")], v[fin & (cls == "seed")]
    v = v[fin]

    cov = json.loads((FIG / "coverage_all_y220c.json").read_text())
    apo_profile = np.sort(np.array([r["apo_to_holo"] for r in cov], dtype=float))

    hi_idx = [mids.index(m) for m in HIGHLIGHT]
    n_rare = int((v >= RARE_THRESH).sum())
    share = {
        f"ge_{t}": {
            "n": int((v >= t).sum()),
            "reference": int((v_ref >= t).sum()),
            "seeded": int((v_seed >= t).sum()),
        }
        for t in (0.8, 0.9, 0.95, RARE_THRESH)
    }

    np.savez_compressed(
        CACHE / "discovery.npz",
        mids=np.array(mids),
        drug=drug,
        truth=truth,
        closest=closest,
        pi=pi,
        order=order,
        frame_scores=v,
        frame_scores_ref=v_ref,
        frame_scores_seed=v_seed,
        highlight=np.array(HIGHLIGHT),
        highlight_idx=np.array(hi_idx),
        highlight_rmsd_profiles=np.sort(M[hi_idx], axis=1),
        apo_rmsd_profile=apo_profile,
        n_complexes=np.array(M.shape[1]),
    )

    stats = {
        "n_microstates": len(mids),
        "n_complexes": int(M.shape[1]),
        "n_frames": int(len(v)),
        "frame_median": float(np.median(v)),
        "rare_threshold": RARE_THRESH,
        "n_frames_ge_thresh": n_rare,
        "n_frames_reference": int(len(v_ref)),
        "n_frames_seeded": int(len(v_seed)),
        "high_druggability_share": share,
        "frame_max": float(v.max()),
        "top_rep_drug": float(drug.max()),
        "apo_to_holo_median": float(np.median(apo_profile)),
        "n_ensemble_closer_than_apo": int(
            sum(1 for r in cov if r["best_ens"] < r["apo_to_holo"])
        ),
        "highlight": {
            str(m): {
                "rank": int(rank[m]),
                "median_rmsd": float(truth[mids.index(m)]),
                "closest_rmsd": float(closest[mids.index(m)]),
                "rep_drug": float(drug[mids.index(m)]),
                "pi": float(pi[mids.index(m)]),
            }
            for m in HIGHLIGHT
        },
        "table_top10": [
            {
                "rank": r + 1,
                "micro": int(mids[j]),
                "median_rmsd": float(truth[j]),
                "closest_rmsd": float(closest[j]),
                "rep_drug": float(drug[j]),
                "pi": float(pi[j]),
            }
            for r, j in enumerate(order[:10])
        ],
    }
    (CACHE / "discovery.json").write_text(json.dumps(stats, indent=1, default=float))

    print(f"  {len(mids)} microstates, {M.shape[1]} complexes, {len(v)} scored frames")
    print(
        f"  most holo-like: microstate {mids[order[0]]} "
        f"(median {truth[order[0]]:.2f} A, closest {closest[order[0]]:.2f} A, "
        f"druggability {drug[order[0]]:.3f}, pi {pi[order[0]]:.4f})"
    )
    print(
        f"  frames >= {RARE_THRESH}: {n_rare}/{len(v)} "
        f"({100 * n_rare / len(v):.3f}%), ensemble median {np.median(v):.3f}"
    )
    print(f"-> {CACHE / 'discovery.npz'}\n-> {CACHE / 'discovery.json'}")


if __name__ == "__main__":
    main()