#!/usr/bin/env python
from __future__ import annotations
import json, os
import numpy as np, torch
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, bridge
from scoobdoob.validation import mb_ground_truth as gt
from scoobdoob.validation import mb_data as V
import make_mb_panels as MP

OUT = "figures/mb_unconditional"
DATA_DIR = "data/mb_unconditional_data"
HORIZONS = [16, 32, 64, 128, 256, 512]


_CACHE_KEYS = [
    "gx",
    "gy",
    "gv",
    "S1",
    "S2",
    "Axy",
    "Bxy",
    "xlo",
    "xhi",
    "ylo",
    "yhi",
    "xy",
    "allxy",
    "paths0",
    "reach0",
    "succ",
    "N0",
]
CACHE_NPZ = f"{DATA_DIR}/uncond_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():
        out[k] = (
            np.array([np.asarray(x) for x in v], dtype=object)
            if k in ("paths0",)
            else 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_unconditional.py --cache-only"
        )
    z = np.load(path, allow_pickle=True)
    D = {k: z[k] for k in z.files}
    D["paths0"] = list(D["paths0"])
    for k in ("xlo", "xhi", "ylo", "yhi", "N0"):
        D[k] = float(D[k]) if k != "N0" else int(D[k])
    print(f"[cache] loaded {path}")
    return D


def compute():
    """Everything numeric; no matplotlib."""
    trajs = V.load_trajs(str(V.DEFAULT_DATA))
    model = mb.build_full_msm(
        trajs,
        trajs,
        lag=MP.LAG_OPT,
        tica_lag=MP.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
    sad = gt.saddle_points()
    S1 = sad[np.argmin(sad[:, 0])]
    S2 = sad[np.argmax(sad[:, 0])]
    Axy = model.active_to_xy(start).mean(0)
    Bxy = model.active_to_xy(end).mean(0)
    endset = set(end)

    def reach_frac(paths):
        return np.mean([any(s in endset for s in p) for p in paths]) * 100

    N0 = MP.N_OPT
    paths0 = bridge.sample_unconditional_paths(model.P, start, N0, n_paths=3000, seed=1)
    reach0 = reach_frac(paths0)

    succ = []
    for N in HORIZONS:
        ps = bridge.sample_unconditional_paths(model.P, start, N, n_paths=3000, seed=1)
        succ.append(reach_frac(ps))

    allxy = model.active_xy
    xlo, xhi = allxy[:, 0].min() - 0.2, allxy[:, 0].max() + 0.2
    ylo, yhi = allxy[:, 1].min() - 0.2, allxy[:, 1].max() + 0.2
    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; recomputes nothing."""
    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"]
    xy = D["xy"]
    allxy = D["allxy"]
    paths0 = D["paths0"]
    reach0 = D["reach0"]
    succ = D["succ"]
    N0 = D["N0"]

    fig = plt.figure(figsize=(10.6, 5.0))
    gs = GridSpec(1, 2, width_ratios=[1.25, 1], wspace=0.30)

    ax = fig.add_subplot(gs[0])
    ax.contour(gx, gy, gv, levels=25, colors="0.82", linewidths=0.5, zorder=0)
    ax.scatter(
        allxy[:, 0], allxy[:, 1], s=9, color="0.72", zorder=1, label="microstates"
    )
    for p in paths0[:250]:
        pp = xy[list(p)]
        ax.plot(pp[:, 0], pp[:, 1], "-", color="#4477aa", lw=0.5, alpha=0.06, zorder=2)
    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.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",
    )
    ax.text(
        Axy[0],
        Axy[1] + 0.13,
        "A",
        ha="center",
        fontsize=11,
        color="#228833",
        weight="bold",
    )
    ax.text(
        Bxy[0],
        Bxy[1] - 0.17,
        "B",
        ha="center",
        fontsize=11,
        color="#ee6677",
        weight="bold",
    )
    ax.set_xlim(xlo, xhi)
    ax.set_ylim(ylo, yhi)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.text(
        0.03,
        0.03,
        f"{reach0:.0f}% reach B",
        transform=ax.transAxes,
        fontsize=10,
        fontweight="bold",
        bbox=dict(boxstyle="round", fc="white", ec="#4477aa", alpha=0.9),
    )
    ax.legend(loc="upper right", fontsize=8, frameon=False)
    ax.set_title(f"A  Unconditional MSM rollouts (N={N0})", fontsize=10.5, loc="left")

    ax = fig.add_subplot(gs[1])
    ax.axhline(100, ls=":", color="0.6", lw=1)
    ax.text(
        HORIZONS[-1],
        100,
        "conditioned = 100%\nat N=16",
        va="top",
        ha="right",
        fontsize=8,
        color="0.45",
    )
    ax.plot(HORIZONS, succ, "-o", color="#4477aa", lw=1.8, ms=5)
    ax.set_xscale("log", base=2)
    ax.set_xticks(HORIZONS)
    ax.set_xticklabels(HORIZONS)
    ax.set_xlabel("horizon N (lag steps)")
    ax.set_ylabel("% of rollouts that reach B")
    ax.set_ylim(-4, 108)
    ax.set_title("B  Unconditional reaction vs horizon", fontsize=10.5, loc="left")
    ax.spines[["top", "right"]].set_visible(False)
    for xN, yv in zip(HORIZONS, succ):
        ax.annotate(
            f"{yv:.0f}%",
            (xN, yv),
            xytext=(0, 7),
            textcoords="offset points",
            fontsize=7.5,
            ha="center",
            color="#22447a",
        )

    fig.suptitle(
        "Naive unconditional MSM rollouts almost never react; conditioning is "
        "required.  $\\tau$=5, TICA lag 20",
        fontsize=11,
        y=1.0,
    )
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    for ext in ("svg", "png", "pdf"):
        fig.savefig(f"{OUT}.{ext}", dpi=200, bbox_inches="tight")

    print(
        f"[stats] N={N0}: {reach0:.1f}% reach B; "
        f"horizon {list(zip(HORIZONS, [round(s, 1) for s in succ]))}"
    )
    print(f"[done] -> {OUT}.svg/.png/.pdf")


def main():
    import argparse

    ap = argparse.ArgumentParser(description="MB unconditional-rollout 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 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()