#!/usr/bin/env python
from __future__ import annotations
import sys
from pathlib import Path
import numpy as np
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
from matplotlib.collections import LineCollection
from matplotlib.lines import Line2D
from matplotlib.colors import LinearSegmentedColormap
from scipy.ndimage import gaussian_filter, gaussian_filter1d

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
from scoobdoob.validation import aib9, bridge, msm_build as mb, tpt

LAG, K, NC, ST, H, NP = 5, 100, 4, 5, 320, 600
KAPPA_RHO, TSg, W = 0.20, 0.30, 1e-12
MACRO_C = {"R": "#1F3A93", "L": "#B22222", "I_N": "#1E7A46", "I_C": "#C4720A"}
MACRO_LBL = {"R": "R", "L": "L", "I_N": "I$_N$", "I_C": "I$_C$"}

ORCHID_P = LinearSegmentedColormap.from_list(
    "orchid_blue",
    ["#12225E", "#28409B", "#4470CB", "#7CA9DF", "#AED2EE", "#D8EAF8", "#F2F8FD"],
    N=256,
)

COMMIT = LinearSegmentedColormap.from_list(
    "purple_red",
    ["#3B0F70", "#7B2D8E", "#B47ACB", "#EBDDF0", "#F6C3AE", "#DA5432", "#7E1008"],
    N=256,
)
GATE_C = "#FF2D95"
PAGE = "#FFFFFF"


def main():
    trajs, _ = aib9.load_openmm_colvar(stride=ST)
    feat, hand = aib9.featurize_sincos(trajs), aib9.handedness_indicator(trajs)
    model = mb.build_full_msm(
        feat, hand, lag=LAG, k=K, n_comp=NC, seed=0, use_tica=True
    )
    labels = aib9.classify_microstates(model.active_xy)
    sets = {
        s: np.where(labels == s)[0].tolist() for s in ["R", "L", "I_N", "I_C", "other"]
    }
    q = tpt.committor_agreement(model, sets["R"], sets["L"])[1]
    pi = model.stationary
    P = model.P
    NA = model.n_active
    proj = model.centers[model.active_symbols][:, :2].copy()
    INT = sorted(set(sets["I_N"]) | set(sets["I_C"]))
    qIL = tpt.committor_agreement(model, INT, sets["L"])[1]
    Nn = np.asarray(tpt.reactive_flux(model, sets["R"], sets["L"]).net_flux)
    netf = (Nn.sum(1) + Nn.sum(0)) / 2
    INcen, ICcen = proj[sets["I_N"]].mean(0), proj[sets["I_C"]].mean(0)
    exitg = sorted(
        [
            i
            for i in range(NA)
            if i not in set(INT) and i not in set(sets["L"]) and abs(qIL[i] - 0.5) < 0.2
        ],
        key=lambda i: -netf[i],
    )[:4]
    egch = [
        "I_N" if np.hypot(*(proj[i] - INcen)) < np.hypot(*(proj[i] - ICcen)) else "I_C"
        for i in exitg
    ]
    egset = set(exitg)
    half = np.where(np.abs(q - 0.5) < 0.08)[0]

    SL = model.active_xy
    n_half, c_half = SL[:, 1:4].mean(axis=1), SL[:, 5:8].mean(axis=1)
    Nchan = set(sets["I_N"]) | {i for i in sets["other"] if n_half[i] > c_half[i]}
    Cchan = set(sets["I_C"]) | {i for i in sets["other"] if n_half[i] <= c_half[i]}

    chan_state = np.array(
        [
            "I_N"
            if np.hypot(*(proj[i] - INcen)) < np.hypot(*(proj[i] - ICcen))
            else "I_C"
            for i in range(NA)
        ],
        dtype=object,
    )
    rf = tpt.reactive_flux(model, sets["R"], sets["L"])
    pf = tpt.pathway_fractions(rf, fraction=0.95)

    part4 = [sets["R"], sorted(Nchan), sorted(Cchan), sets["L"]]
    _, cg4 = rf.coarse_grain(part4)
    F4 = np.asarray(cg4.net_flux)
    tot4 = float(cg4.total_flux)
    fRN4, fRC4 = F4[0, 1] / tot4, F4[0, 2] / tot4
    print(
        f"[reference MSM] merged-partition flux out of R:  N-channel {fRN4 * 100:.1f}%  "
        f"C-channel {fRC4 * 100:.1f}%"
    )

    def chroute(tgt):
        idx = [
            i
            for i, p_ in enumerate(pf["paths"])
            if tpt.path_passes_through(p_, set(tgt))
        ]
        return (
            (
                np.array(pf["paths"][max(idx, key=lambda i: pf["frac_of_total"][i])]),
                float(sum(pf["frac_of_total"][i] for i in idx)),
            )
            if idx
            else (None, 0.0)
        )

    rIN, fIN = chroute(sets["I_N"])
    rIC, fIC = chroute(sets["I_C"])

    Cc = np.asarray(model.model.count_model.count_matrix, float)
    rho = np.clip(Cc.sum(1), 1.0, None) / Cc.sum(1).mean()
    wb = np.maximum(np.exp(-KAPPA_RHO * rho), W)
    tw = np.exp(-0.5 * ((np.arange(H) - (H - 1) / 2) / (TSg * H)) ** 2)
    gz = np.log(netf / np.maximum(pi, 1e-15) + 1e-9)
    gz = (gz - gz.mean()) / gz.std()

    def kern(wl, Pm=None):
        Pm = P if Pm is None else Pm
        return bridge.doob_kernels(Pm, sets["L"], H, w=wl)

    def samp(Ph, seed, n_=NP):
        return bridge.sample_bridge_paths(Ph, sets["R"], n_paths=n_, seed=seed)

    naive = [
        np.asarray(p, int)
        for p in bridge.sample_unconditional_paths(P, sets["R"], H, n_paths=NP, seed=1)
    ]
    doob = samp(bridge.doob_kernels(P, sets["L"], H), 1)
    wfk = [np.maximum(wb * (np.exp(0.5 * gz) ** tw[t]), W) for t in range(H)]
    dfk = samp(kern(wfk), 1)

    def confined(forbid):
        Pm = P.copy()
        Pm[:, list(forbid)] = 0.0
        sums = Pm.sum(1, keepdims=True)
        return np.divide(Pm, sums, out=np.zeros_like(Pm), where=sums > 0)

    cfk = samp(kern(wfk, confined(Cchan)), 11, NP // 2) + samp(
        kern(wfk, confined(Nchan)), 12, NP // 2
    )
    Rs, Ls = set(sets["R"]), set(sets["L"])

    def reactive(paths):
        segs = []
        for p in paths:
            inL = np.array([int(x) in Ls for x in p])
            if not inL.any():
                continue
            fL = int(np.argmax(inL))
            pre = np.array([int(x) in Rs for x in p[:fL]])
            if not pre.any():
                continue
            segs.append(p[int(np.where(pre)[0][-1]) : fL + 1])
        return segs

    S = {
        "naive": reactive(naive),
        "doob": reactive(doob),
        "dfk": reactive(dfk),
        "cfk": reactive(cfk),
    }
    M = {}
    for k_, segs in S.items():
        reach = len(segs) / NP
        via = (
            float(np.mean([bool(set(int(x) for x in s_) & egset) for s_ in segs]))
            if segs
            else 0.0
        )

        def seg_chan(s_):
            fN = sum(int(x) in Nchan for x in s_)
            fC = sum(int(x) in Cchan for x in s_)
            return "I_N" if fN > fC else ("I_C" if fC > fN else None)

        cs = [seg_chan(s_) for s_ in segs]
        nN = sum(1 for c in cs if c == "I_N")
        nC = sum(1 for c in cs if c == "I_C")
        M[k_] = (reach, via, len(segs), nN, nC)
        print(
            f"{k_:>6}: reach-L {reach * 100:3.0f}%  n_react {len(segs):3d}  "
            f"via gateway {via * 100:3.0f}%   channel I_N {nN:3d} / I_C {nC:3d} "
            f"(I_C share {100 * nC / max(nN + nC, 1):.0f}%)"
        )

    if proj[sets["R"], 0].mean() > proj[sets["L"], 0].mean():
        proj[:, 0] *= -1
        flip = True
    else:
        flip = False
    Y = np.vstack(
        [model.tica.transform(np.asarray(f, dtype=np.float64)) for f in feat]
    )[:, :2]
    if flip:
        Y = Y.copy()
        Y[:, 0] *= -1
    x, y = Y[:, 0], Y[:, 1]
    rr = [
        [x.min() - 0.04 * np.ptp(x), x.max() + 0.04 * np.ptp(x)],
        [y.min() - 0.04 * np.ptp(y), y.max() + 0.04 * np.ptp(y)],
    ]
    Hc, xe, ye = np.histogram2d(x, y, bins=140, range=rr)
    dd = gaussian_filter(Hc, 2.0)
    with np.errstate(divide="ignore"):
        F = -np.log(np.where(dd > 0, dd / dd.sum(), np.nan))
    F -= np.nanmin(F[np.isfinite(F)])
    F = np.ma.masked_where((dd <= 0) | ~np.isfinite(F) | (F > 6.5), F)
    xc, yc = 0.5 * (xe[:-1] + xe[1:]), 0.5 * (ye[:-1] + ye[1:])
    setsA = {s: np.array(sets[s]) for s in sets}

    def fes(ax, a=1.0):
        cf = ax.contourf(
            xc,
            yc,
            F.T,
            levels=np.arange(0, 6.75, 0.25),
            cmap=ORCHID_P,
            zorder=0,
            alpha=a,
        )
        ax.set_xlim(xc[0], xc[-1])
        ax.set_ylim(yc[0], yc[-1])
        ax.set_xlabel("TICA 1", fontsize=10.5)
        for s in ("top", "right"):
            ax.spines[s].set_visible(False)
        return cf

    def rings(ax, lw=1.6, s=140):
        for lab in ["R", "L", "I_N", "I_C"]:
            idx = setsA[lab]
            if len(idx):
                ax.scatter(
                    proj[idx, 0],
                    proj[idx, 1],
                    marker="o",
                    s=s,
                    facecolors="none",
                    edgecolors=MACRO_C[lab],
                    linewidths=lw,
                    zorder=5,
                )

    def gates(ax, s=190):
        for i, ch in zip(exitg, egch):
            ax.scatter(
                proj[i, 0],
                proj[i, 1],
                marker="X",
                s=s,
                c=GATE_C,
                edgecolors=MACRO_C[ch],
                linewidths=2.4,
                zorder=9,
            )

    def sm(seq):
        d = seq[np.concatenate([[True], seq[1:] != seq[:-1]])]
        xy = proj[d]
        return (
            xy
            if len(xy) < 3
            else np.column_stack(
                [
                    gaussian_filter1d(xy[:, 0], 1.0, mode="nearest"),
                    gaussian_filter1d(xy[:, 1], 1.0, mode="nearest"),
                ]
            )
        )

    def crossings(ax, segs, color):
        """Color crossings by channel and emphasize gateway visitation."""
        buckets = {
            ("I_N", True): [],
            ("I_N", False): [],
            ("I_C", True): [],
            ("I_C", False): [],
        }
        for s_ in segs:
            xy = sm(s_)
            if len(xy) < 2:
                continue
            fN = sum(int(x) in Nchan for x in s_)
            fC = sum(int(x) in Cchan for x in s_)
            ch = "I_N" if fN >= fC else "I_C"
            buckets[(ch, bool(set(int(t) for t in s_) & egset))].append(xy)
        for (ch, viaG), L_ in buckets.items():
            if not L_:
                continue
            ax.add_collection(
                LineCollection(
                    [[xy[i], xy[i + 1]] for xy in L_ for i in range(len(xy) - 1)],
                    colors=MACRO_C[ch],
                    linewidths=0.85 if viaG else 0.5,
                    alpha=0.20 if viaG else 0.11,
                    zorder=4 if viaG else 3,
                    capstyle="round",
                )
            )

    fig = plt.figure(figsize=(23.5, 11.4))
    fig.patch.set_facecolor(PAGE)
    gs = fig.add_gridspec(2, 4, height_ratios=[1.12, 1.0], hspace=0.24, wspace=0.15)

    axa = fig.add_subplot(gs[0, 0:2])
    axa.set_facecolor(PAGE)
    fes(axa)
    sc = axa.scatter(
        proj[:, 0],
        proj[:, 1],
        c=q,
        cmap=COMMIT,
        s=54,
        vmin=0,
        vmax=1,
        edgecolors=PAGE,
        linewidths=0.9,
        zorder=6,
    )
    rings(axa)
    xs = axa.scatter(
        proj[half, 0],
        proj[half, 1],
        marker="x",
        c="#101010",
        s=70,
        linewidths=2,
        zorder=7,
    )
    xs.set_path_effects([pe.withStroke(linewidth=3.4, foreground=PAGE)])
    gates(axa)
    axa.set_ylabel("TICA 2", fontsize=10.5)
    axa.set_title("a  Committor + macrostates", fontsize=12, fontweight="bold", pad=8)
    cb = fig.colorbar(sc, ax=axa, pad=0.012, fraction=0.045)
    cb.set_label("committor $q^+$ (R$\\rightarrow$L)", fontsize=9)
    cb.outline.set_visible(False)
    cb.ax.tick_params(length=0, labelsize=8)
    lh = [
        Line2D(
            [0], [0], marker="o", color="w", mfc="none", mec=MACRO_C[m], mew=1.8, ms=10
        )
        for m in ["R", "I_N", "I_C", "L"]
    ]
    lh += [
        Line2D([0], [0], marker="x", color="#101010", lw=0, mew=2, ms=9),
        Line2D(
            [0],
            [0],
            marker="X",
            color="w",
            mfc=GATE_C,
            mec=MACRO_C["I_N"],
            mew=1.6,
            ms=11,
        ),
    ]
    axa.legend(
        lh,
        [
            "R state",
            "I$_N$ state",
            "I$_C$ state",
            "L state",
            "committor-½",
            "rate-limiting exit gate",
        ],
        loc="lower left",
        fontsize=7.3,
        framealpha=0.95,
        ncol=2,
    ).get_frame().set_edgecolor("#C9D3DE")

    axb = fig.add_subplot(gs[0, 2:4])
    axb.set_facecolor(PAGE)
    cfb = fes(axb)
    rings(axb)
    for r, cc, lab, fr in (
        (rIN, MACRO_C["I_N"], "I$_N$", fIN),
        (rIC, MACRO_C["I_C"], "I$_C$", fIC),
    ):
        if r is None:
            continue
        xy = proj[r]
        axb.plot(
            xy[:, 0],
            xy[:, 1],
            "-o",
            lw=3.0,
            ms=6,
            color=cc,
            mec=PAGE,
            mew=1.0,
            zorder=8,
            path_effects=[pe.withStroke(linewidth=5.4, foreground=PAGE)],
            label=f"channel via {lab}: {fr * 100:.0f}% of flux",
        )
    gates(axb)
    axb.set_title(
        "b  The two reactive channels (TPT)", fontsize=12, fontweight="bold", pad=8
    )
    axb.legend(
        loc="upper left", fontsize=8.2, framealpha=0.95
    ).get_frame().set_edgecolor("#C9D3DE")
    cbF = fig.colorbar(cfb, ax=axb, pad=0.012, fraction=0.045, ticks=range(0, 7))
    cbF.set_label("free energy $F/k_\\mathrm{B}T$ (background)", fontsize=9)
    cbF.outline.set_visible(False)
    cbF.ax.tick_params(length=0, labelsize=8)

    panels = [
        ("c  naive MSM (unconditioned)", "naive", "#5C6570"),
        ("d  Doob bridge (conditioned)", "doob", "#1F4EA1"),
        ("e  Doob + FK  (global tilt)", "dfk", "#A3161C"),
        ("f  Doob + FK  (channel-dependent)", "cfk", "#1E7A46"),
    ]
    samp_axes = []
    for col, (title, key, color) in enumerate(panels):
        ax = fig.add_subplot(gs[1, col])
        samp_axes.append(ax)
        ax.set_facecolor(PAGE)
        fes(ax, 0.9)
        crossings(ax, S[key], color)
        rings(ax, lw=1.2, s=90)
        gates(ax, s=150)
        reach, via, ns, nN, nC = M[key]
        ax.set_title(title, fontsize=12, fontweight="bold", pad=8)
        if col == 0:
            ax.set_ylabel("TICA 2", fontsize=10.5)
        ax.text(
            0.5,
            0.03,
            f"reach L: {reach * 100:.0f}%   ({ns} of {NP} crossings)\n"
            f"via rate-limiting gateway: {via * 100:.0f}%\n"
            f"N-channel {nN} / C-channel {nC}",
            transform=ax.transAxes,
            ha="center",
            va="bottom",
            fontsize=11,
            fontweight="bold",
            color="#111",
            linespacing=1.5,
            bbox=dict(boxstyle="round,pad=0.5", fc=PAGE, ec=color, lw=1.8, alpha=0.96),
            zorder=13,
        )
    axc0 = samp_axes[0]
    axc0.legend(
        [
            Line2D([0], [0], color=MACRO_C["I_N"], lw=2.2),
            Line2D([0], [0], color=MACRO_C["I_C"], lw=2.2),
            Line2D([0], [0], color="#555", lw=2.2),
            Line2D([0], [0], color="#555", lw=0.9),
        ],
        [
            "crossing via N-channel",
            "crossing via C-channel",
            "uses rate-limiting gateway",
            "does not",
        ],
        loc="upper left",
        fontsize=7.6,
        framealpha=0.95,
    ).get_frame().set_edgecolor("#C9D3DE")

    fig.suptitle(
        "Aib9 R→L:  Doob fixes ARRIVAL (40%→100%);  global FK routes crossings through the "
        "rate-limiting exit gateway (40%→82%) but COLLAPSES the C-channel (47%→2%);\n"
        "channel-dependent FK keeps the gateway emphasis (68%) AND both channels",
        fontsize=13,
        fontweight="bold",
        y=1.02,
    )
    fig.text(
        0.006,
        0.006,
        f"Aib9 OpenMM MSM (lag {LAG}, K={K}); {NP} rollouts/sampler, horizon {H}.  Panels c–f draw the "
        f"REACTIVE SEGMENT of each rollout (last frame in R → first in L), so only genuine crossings are\n"
        f"shown — naive produces {M['naive'][2]}, the conditioned samplers {NP}.  Rate-limiting exit gateway = "
        f"the 4 highest net-flux committor-½ states of the trap→L barrier (magenta X, edge = channel).\n"
        f"Channel = MERGED handedness partition (each 'other' microstate joins the channel whose half leads "
        f"the flip): counting only the labelled I$_N$/I$_C$ helices under-counts a channel, because most\n"
        f"reactive flux runs R→other→L.  The two channels are intrinsically BALANCED — merged-partition flux "
        f"out of R is {fRN4 * 100:.0f}% N / {fRC4 * 100:.0f}% C (Aib9 is a homopolymer, so N↔C is a symmetry) —\n"
        f"so f's equal 300/300 allocation matches the true split; the measured results there are that both "
        f"confined bridges reach L 100% independently and keep 68% gateway use.",
        fontsize=8.2,
        color="#6B7785",
        ha="left",
        va="bottom",
        linespacing=1.5,
    )
    fig.subplots_adjust(left=0.045, right=0.988, top=0.90, bottom=0.115)
    out = ROOT / "figures" / "aib9_method"
    out.parent.mkdir(parents=True, exist_ok=True)
    for ext, kw in (("png", dict(dpi=180)), ("pdf", {})):
        fig.savefig(
            out.with_suffix("." + ext), facecolor=PAGE, bbox_inches="tight", **kw
        )
        print(f"[fig] {out.with_suffix('.' + ext)}")
    plt.close(fig)


if __name__ == "__main__":
    main()