#!/usr/bin/env python
from __future__ import annotations
import json, os
import numpy as np, torch
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec

from scoobdoob.mb_potential import mb_potential_xy
from scoobdoob.validation import msm_build as mb, tpt, bridge
from scoobdoob.validation import mb_ground_truth as gt
from scoobdoob.validation import mb_data as V


plt.rcParams["svg.fonttype"] = "none"
plt.rcParams["pdf.fonttype"] = 42

DATA_DIR = "data/mb_panels_data"

LAG_OPT, N_OPT = 5, 16
TICA_LAG = 20


def channel(paths, q, xy, on, thru, start, end, nbins=18):
    """Dominant route = the maximum-bottleneck (widest) path through the ensemble's OWN empirical transition graph, from the start set to the target set."""
    import heapq
    from collections import defaultdict

    W = defaultdict(float)
    for p in paths:
        for a, b in zip(p[:-1], p[1:]):
            if a != b and on[a] and on[b]:
                W[(int(a), int(b))] += 1.0
    if not W:
        return np.zeros((0, 2))
    adj = defaultdict(list)
    for (a, b), c in W.items():
        adj[a].append((b, c))
    S, T = set(map(int, start)), set(map(int, end))

    best = defaultdict(float)
    prev = {}
    heap = []
    for a in S:
        best[a] = float("inf")
        heapq.heappush(heap, (-best[a], a))
    seen = set()
    while heap:
        nb, u = heapq.heappop(heap)
        nb = -nb
        if u in seen:
            continue
        seen.add(u)
        if u in T:
            path = [u]
            while path[-1] not in S:
                path.append(prev[path[-1]])
            return xy[path[::-1]]
        for v, c in adj[u]:
            w = min(nb, c)
            if w > best[v]:
                best[v] = w
                prev[v] = u
                heapq.heappush(heap, (-w, v))
    return np.zeros((0, 2))


_CACHE_KEYS = [
    "gx",
    "gy",
    "gv",
    "S1",
    "S2",
    "Axy",
    "Bxy",
    "xlo",
    "xhi",
    "ylo",
    "yhi",
    "cx",
    "cy",
    "q",
    "keep",
    "xy",
    "mep",
    "arc",
    "Emep",
    "ch_raw",
    "ch_fk",
    "praw",
    "pfk",
    "s1_raw",
    "s1_fk",
    "counts",
    "thru",
    "on",
]
CACHE_NPZ = "data/mb_panels_data/panels_cache.npz"


def save_cache(D, path=CACHE_NPZ):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    out = {}
    for k, v in D.items():
        if k in ("praw", "pfk"):
            out[k] = np.array([np.asarray(p) for p in v], dtype=object)
        else:
            out[k] = np.asarray(v)
    np.savez_compressed(path, **out)
    print(f"[cache] -> {path}")


def load_cache(path=CACHE_NPZ):
    if not os.path.exists(path):
        raise SystemExit(f"missing {path}\nrun: python make_mb_panels.py --cache-only")
    z = np.load(path, allow_pickle=True)
    D = {k: z[k] for k in z.files}
    for k in ("praw", "pfk"):
        D[k] = list(D[k])
    for k in ("xlo", "xhi", "ylo", "yhi", "s1_raw", "s1_fk"):
        D[k] = float(D[k])
    print(f"[cache] loaded {path}")
    return D


def compute():
    """Everything numeric. Returns the dict the plot stage needs; touches no matplotlib."""
    trajs = V.load_trajs(str(V.DEFAULT_DATA))
    model = mb.build_full_msm(
        trajs,
        trajs,
        lag=LAG_OPT,
        tica_lag=TICA_LAG,
        k=V.K,
        n_comp=2,
        seed=0,
        use_tica=True,
    )
    A, B, _ = mb.q2_sign_endpoints(model)
    start, end = V.orient_endpoints(model, A, B)
    xy = model.state_xy
    _, q, _ = tpt.committor_agreement(model, start, end)
    rf = tpt.reactive_flux(model, start, end)
    thru = np.asarray(rf.gross_flux).sum(1)
    sad = gt.saddle_points()
    S1 = sad[np.argmin(sad[:, 0])]
    S2 = sad[np.argmax(sad[:, 0])]
    q1 = q[np.argmin(np.linalg.norm(xy - S1, axis=1))]
    q2 = q[np.argmin(np.linalg.norm(xy - S2, axis=1))]

    share = thru / thru.sum()
    on = share > 0.5 / model.n_active
    on[list(start)] = True
    on[list(end)] = True
    Axy = model.active_to_xy(start).mean(0)
    Bxy = model.active_to_xy(end).mean(0)
    land = np.vstack([Axy, Bxy, sad])
    lo = land.min(0) - 0.35
    hi = land.max(0) + 0.35
    inbox = (
        (xy[:, 0] >= lo[0])
        & (xy[:, 0] <= hi[0])
        & (xy[:, 1] >= lo[1])
        & (xy[:, 1] <= hi[1])
    )
    on = on & inbox
    keep = np.where(on)[0]
    cx, cy = xy[keep, 0], xy[keep, 1]

    labels = np.concatenate(model.dtrajs)
    counts = np.bincount(labels, minlength=model.n_active)[: model.n_active].astype(
        float
    )
    rho = np.maximum(counts, 1) / counts.mean()
    react = (q > 0.05) & (q < 0.95)
    rho_floor = np.percentile(rho[react], 25)
    w_fk = np.exp(-LAG_OPT * 1.5 * np.maximum(rho, rho_floor))
    w_fk[~on] = 0.0
    w_raw = np.ones_like(w_fk)
    w_raw[~on] = 0.0

    def rollout(w):
        Ph = bridge.doob_kernels(model.P, end, N_OPT, w=w)
        return bridge.sample_bridge_paths(Ph, start, n_paths=3000, seed=0)

    praw = rollout(w_raw)
    pfk = rollout(w_fk)
    ch_raw = channel(praw, q, xy, on, thru, start, end)
    ch_fk = channel(pfk, q, xy, on, thru, start, end)
    n1 = np.linalg.norm(xy - S1, axis=1) < 0.25
    s1_raw = np.mean([any(n1[s] for s in p) for p in praw]) * 100
    s1_fk = np.mean([any(n1[s] for s in p) for p in pfk]) * 100

    M = np.load("data/mb_mep.npz")
    mep = M["mep"]
    Emep = M["Emep"]
    arc = M["arc"]
    _keep_mep = mep[:, 1] >= -0.05
    mep = mep[_keep_mep]
    xlo, xhi = cx.min() - 0.12, cx.max() + 0.12
    ylo, yhi = cy.min() - 0.12, cy.max() + 0.12
    gx, gy = np.meshgrid(np.linspace(xlo, xhi, 300), np.linspace(ylo, yhi, 300))
    with torch.no_grad():
        gv = np.clip(
            mb_potential_xy(torch.tensor(gx), torch.tensor(gy)).numpy(), -150, 100
        )

    return {k: v for k, v in locals().items() if k in _CACHE_KEYS}


def render(D):
    os.makedirs("figures", exist_ok=True)
    """Plotting only. Reads the cache dict; recomputes nothing. Safe to re-run while restyling."""
    gx = D["gx"]
    gy = D["gy"]
    gv = D["gv"]
    S1 = D["S1"]
    S2 = D["S2"]
    Axy = D["Axy"]
    Bxy = D["Bxy"]
    xlo = D["xlo"]
    xhi = D["xhi"]
    ylo = D["ylo"]
    yhi = D["yhi"]
    cx = D["cx"]
    cy = D["cy"]
    q = D["q"]
    keep = D["keep"]
    xy = D["xy"]
    mep = D["mep"]
    arc = D["arc"]
    Emep = D["Emep"]
    ch_raw = D["ch_raw"]
    ch_fk = D["ch_fk"]
    praw = D["praw"]
    pfk = D["pfk"]
    s1_raw = D["s1_raw"]
    s1_fk = D["s1_fk"]
    counts = D["counts"]
    thru = D["thru"]
    on = D["on"]

    fig = plt.figure(figsize=(15, 5.4))
    gs = GridSpec(1, 4, width_ratios=[1, 1, 1, 0.9], wspace=0.28)

    def base(ax):
        ax.contour(gx, gy, gv, levels=22, colors="0.82", linewidths=0.5, zorder=0)
        ax.scatter(*S1, marker="X", s=150, color="black", zorder=6)
        ax.scatter(*S2, marker="X", s=150, color="0.5", zorder=6)
        ax.scatter(*Axy, marker="*", s=260, color="#228833", ec="white", zorder=7)
        ax.scatter(*Bxy, marker="*", s=260, color="#ee6677", ec="white", zorder=7)
        ax.set_xlim(xlo, xhi)
        ax.set_ylim(ylo, yhi)
        ax.set_xlabel("x")

    ax = fig.add_subplot(gs[0])
    base(ax)
    sc = ax.scatter(cx, cy, c=q[keep], cmap="coolwarm", s=22, zorder=3, vmin=0, vmax=1)
    ax.plot(
        mep[:, 0],
        mep[:, 1],
        "--",
        color="0.25",
        lw=1.6,
        zorder=4,
        label="min-energy path",
    )
    ax.set_ylabel("y")
    ax.set_title("a  MSM microstates + committor", fontsize=10.5, loc="left")
    ax.legend(fontsize=7.5, loc="lower right", frameon=False)
    ax.annotate(
        "S1", S1, xytext=(S1[0] - 0.13, S1[1] + 0.10), fontsize=8.5, weight="bold"
    )
    ax.annotate(
        "S2",
        S2,
        xytext=(S2[0] + 0.06, S2[1] + 0.10),
        fontsize=8.5,
        color="0.4",
        weight="bold",
    )
    cax = ax.inset_axes([0.60, 0.90, 0.36, 0.045])
    cb = fig.colorbar(sc, cax=cax, orientation="horizontal")
    cb.set_label("committor $q^+$", fontsize=7.5)
    cb.ax.tick_params(labelsize=6.5)

    ax = fig.add_subplot(gs[1])
    base(ax)
    ax.scatter(cx, cy, s=9, color="0.72", zorder=1)
    for p in praw[:150]:
        pp = xy[list(p)]
        ax.plot(pp[:, 0], pp[:, 1], "-", color="#4477aa", lw=0.5, alpha=0.06, zorder=2)
    ax.plot(mep[:, 0], mep[:, 1], "--", color="0.35", lw=1.3, zorder=4)
    ax.plot(
        ch_raw[:, 0],
        ch_raw[:, 1],
        "-o",
        color="#22447a",
        lw=2.6,
        ms=4,
        zorder=5,
        label="dominant route (max-flux)",
    )
    ax.set_title(f"b  Without FK ({s1_raw:.0f}% via S1)", fontsize=10.5, loc="left")
    ax.legend(fontsize=7.5, loc="lower right", frameon=False)

    ax = fig.add_subplot(gs[2])
    base(ax)
    ax.scatter(cx, cy, s=9, color="0.72", zorder=1)
    for p in pfk[:150]:
        pp = xy[list(p)]
        ax.plot(pp[:, 0], pp[:, 1], "-", color="#cc6677", lw=0.5, alpha=0.06, zorder=2)
    ax.plot(mep[:, 0], mep[:, 1], "--", color="0.35", lw=1.3, zorder=4)
    ax.plot(
        ch_fk[:, 0],
        ch_fk[:, 1],
        "-o",
        color="#a03050",
        lw=2.6,
        ms=4,
        zorder=5,
        label="dominant route (max-flux)",
    )
    ax.set_title(f"c  With floored FK ({s1_fk:.0f}% via S1)", fontsize=10.5, loc="left")
    ax.legend(fontsize=7.5, loc="lower right", frameon=False)

    ax = fig.add_subplot(gs[3])
    ax.plot(arc, Emep, "-", color="0.2", lw=2)
    ipk = np.argmax(Emep)
    ax.scatter([arc[ipk]], [Emep[ipk]], s=90, color="black", zorder=5)
    ax.annotate(
        "S1 (rate-limiting)\nbarrier 106 = 6.4 $k_BT$",
        (arc[ipk], Emep[ipk]),
        xytext=(arc[ipk] - 0.05, Emep[ipk] - 30),
        fontsize=8,
        ha="center",
        arrowprops=dict(arrowstyle="->"),
    )

    iS2 = np.argmin(np.linalg.norm(mep - S2, axis=1))
    ax.scatter([arc[iS2]], [Emep[iS2]], s=70, color="0.5", zorder=5)
    ax.annotate(
        "S2\nbarrier 9 = 0.5 $k_BT$",
        (arc[iS2], Emep[iS2]),
        xytext=(arc[iS2] - 0.05, Emep[iS2] - 34),
        fontsize=8,
        color="0.4",
        ha="center",
        arrowprops=dict(arrowstyle="->", color="0.4"),
    )
    ax.set_xlabel("path length A$\\to$B")
    ax.set_ylabel("potential energy (MB units)")
    ax.set_title("d  barrier heights along the MEP", fontsize=10.5, loc="left")
    ax.spines[["top", "right"]].set_visible(False)

    fig.suptitle(
        "The transition is obligated through S1 (d); the raw MSM under-samples it (b), "
        "floored FK recovers it (c).  $\\tau$=5, N=16, $\\alpha$=1.5, TICA lag 20",
        fontsize=11.5,
        y=1.0,
    )
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig("figures/mb_panels.svg", bbox_inches="tight")
    fig.savefig("figures/mb_panels.png", dpi=200, bbox_inches="tight")
    fig.savefig("figures/mb_panels.pdf", bbox_inches="tight")

    print(
        f"[stats] S1 via raw {s1_raw:.0f}% -> FK {s1_fk:.0f}%; MEP peak at S1 (+106 vs S2 +9)"
    )
    print("[done] -> figures/mb_panels.svg / .png / .pdf")


def main():
    import argparse

    ap = argparse.ArgumentParser(description="Muller-Brown panel figure.")
    ap.add_argument(
        "--cache-only", action="store_true", help="compute and cache, do not plot"
    )
    ap.add_argument(
        "--from-cache",
        action="store_true",
        help="plot from the cache without recomputing (fast; for restyling)",
    )
    ap.add_argument("--cache", default=CACHE_NPZ)
    a = ap.parse_args()

    if a.from_cache:
        render(load_cache(a.cache))
        return
    D = compute()
    save_cache(D, a.cache)
    if a.cache_only:
        return
    render(D)


if __name__ == "__main__":
    main()