from __future__ import annotations

import numpy as np

from . import msm_build as mb
from . import tpt, aib9


def analyze_point(feat, hand, K, lag, n_comp=4, seed=0):
    """One (K, lag) build -> dict of the key stability metrics."""
    out = {"K": int(K), "lag": int(lag), "ok": False}
    try:
        m = mb.build_full_msm(
            feat, hand, lag=lag, k=K, n_comp=n_comp, use_tica=True, seed=seed
        )
    except Exception as e:
        out["error"] = str(e)
        return out
    out["n_active"] = int(m.n_active)
    out["lam1"] = float(np.real(m.model.eigenvalues(k=1)[0]))
    labels = aib9.classify_microstates(m.active_xy)
    sets = {
        k: np.where(labels == k)[0].tolist() for k in ["R", "L", "I_N", "I_C", "other"]
    }
    out["n_R"], out["n_L"] = len(sets["R"]), len(sets["L"])
    if not sets["R"] or not sets["L"]:
        return out
    q = tpt.forward_committor(m, sets["R"], sets["L"])
    half = tpt.committor_half_states(q, tol=0.1)
    inter = set(sets["I_N"]) | set(sets["I_C"])
    pi = m.stationary
    out["committor_half_mass"] = (
        float(
            pi[[h for h in half if h in inter]].sum() / max(pi[list(half)].sum(), 1e-12)
        )
        if len(half)
        else 0.0
    )
    rf = tpt.reactive_flux(m, sets["R"], sets["L"])
    _, cg = rf.coarse_grain([sets[k] for k in ["R", "I_N", "I_C", "other", "L"]])
    F = np.asarray(cg.net_flux)
    tot = float(cg.total_flux)
    out["direct"] = float(F[0, 4] / tot)
    out["flux_IN"] = float(F[0, 1] / tot)
    out["flux_IC"] = float(F[0, 2] / tot)
    try:
        out["ck2"] = mb.ck_test(m.dtrajs, lag=lag, k=K, factors=(2,))[2]
    except Exception:
        out["ck2"] = float("nan")
    out["ok"] = True
    return out


def sweep(feat, hand, Ks, lags, n_comp=4, seed=0):
    return [analyze_point(feat, hand, K, lag, n_comp, seed) for K in Ks for lag in lags]


def its_convergence(feat, hand, K, lags, n_comp=4, n_its=3, seed=0):
    """ITS vs lag at a FIXED discretization (vary only the count lag)."""
    ref_lag = lags[len(lags) // 2]
    m = mb.build_full_msm(
        feat, hand, lag=ref_lag, k=K, n_comp=n_comp, use_tica=True, seed=seed
    )
    lg, its = mb.implied_timescales(m.dtrajs, lags, k=K, n_its=n_its)
    return lg, its


def summarize(results):
    """Range of the two decision metrics across all valid sweep points."""
    ok = [r for r in results if r.get("ok")]
    if not ok:
        return {}
    ch = [r["committor_half_mass"] for r in ok]
    di = [r["direct"] for r in ok]
    return {
        "n_points": len(ok),
        "committor_half_mass_min": float(min(ch)),
        "committor_half_mass_max": float(max(ch)),
        "direct_min": float(min(di)),
        "direct_max": float(max(di)),
        "both_channels_always": bool(
            all(r["flux_IN"] > 0 and r["flux_IC"] > 0 for r in ok)
        ),
    }


def make_figure(results, its_lags, its, K_default, lag_default, dt_ps, outpath):
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    plt.rcParams.update(
        {
            "font.family": "DejaVu Sans",
            "mathtext.fontset": "custom",
            "mathtext.rm": "DejaVu Sans",
            "mathtext.it": "DejaVu Sans:italic",
            "mathtext.bf": "DejaVu Sans:bold",
            "mathtext.sf": "DejaVu Sans",
            "mathtext.cal": "DejaVu Sans",
            "pdf.fonttype": 42,
            "ps.fonttype": 42,
        }
    )

    Ks = sorted({r["K"] for r in results})
    lags = sorted({r["lag"] for r in results})
    grid = {(r["K"], r["lag"]): r for r in results}

    def heat(metric):
        M = np.full((len(Ks), len(lags)), np.nan)
        for i, K in enumerate(Ks):
            for j, lg in enumerate(lags):
                r = grid.get((K, lg))
                if r and r.get("ok"):
                    M[i, j] = r[metric]
        return M

    fig, ax = plt.subplots(1, 3, figsize=(13.2, 4.1), layout="constrained")

    for c in range(its.shape[1]):
        ax[0].plot(
            np.array(its_lags) * dt_ps,
            its[:, c] * dt_ps,
            "-o",
            label=rf"$t_{c + 1}$",
            markersize=4,
        )
    ax[0].axvline(
        lag_default * dt_ps,
        color="gray",
        ls="--",
        label=f"{lag_default * dt_ps:.0f} ps lag",
    )
    ax[0].set_xlabel("Lag time (ps)")
    ax[0].set_ylabel("Implied timescale (ps)")
    ax[0].set_yscale("log")
    ax[0].set_title("Implied timescales")
    ax[0].legend(fontsize=10, frameon=False)

    for a, metric, title, vmax, decimals, ticks in [
        (
            ax[1],
            "committor_half_mass",
            "Intermediate overlap",
            100.0,
            1,
            [0, 25, 50, 75, 100],
        ),
        (
            ax[2],
            "direct",
            r"Direct $R$-to-$L$ flux",
            1.1,
            2,
            [0, 0.25, 0.5, 0.75, 1.0, 1.1],
        ),
    ]:
        M = 100.0 * heat(metric)
        finite = M[np.isfinite(M)]
        if finite.size and finite.max() > vmax:
            vmax = float(finite.max())
            ticks = np.linspace(0, vmax, 5)
        im = a.imshow(
            M,
            aspect="auto",
            cmap="viridis" if metric != "direct" else "magma_r",
            vmin=0,
            vmax=vmax,
            origin="lower",
        )
        a.set_xticks(range(len(lags)))
        a.set_xticklabels(lags)
        a.set_yticks(range(len(Ks)))
        a.set_yticklabels(Ks)
        a.set_xlabel("Lag time (ps)")
        a.set_ylabel("Microstates ($K$)")
        for i in range(len(Ks)):
            for j in range(len(lags)):
                if np.isfinite(M[i, j]):
                    rgb = im.cmap(im.norm(M[i, j]))[:3]
                    lightness = np.dot(rgb, [0.2126, 0.7152, 0.0722])
                    a.text(
                        j,
                        i,
                        f"{M[i, j]:.{decimals}f}",
                        ha="center",
                        va="center",
                        color="black" if lightness > 0.5 else "white",
                        fontsize=11,
                    )

        a.scatter(
            [lags.index(lag_default) + 0.30],
            [Ks.index(K_default) + 0.29],
            marker="*",
            c="red",
            s=105,
            edgecolors="white",
            linewidths=0.6,
            zorder=3,
        )
        cbar = fig.colorbar(im, ax=a, ticks=ticks, fraction=0.048, pad=0.03)
        cbar.ax.set_title("%", fontsize=11, pad=6)
        cbar.ax.tick_params(labelsize=10)
        a.set_title(title)

    for letter, a in zip("ABC", ax):
        a.tick_params(labelsize=10)
        a.xaxis.label.set_size(11)
        a.yaxis.label.set_size(11)
        a.title.set_fontsize(12)
        a.text(
            -0.16,
            1.07,
            letter,
            transform=a.transAxes,
            fontsize=15,
            fontweight="bold",
            va="bottom",
        )

    from pathlib import Path as _P

    _o = _P(outpath)
    _o.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(_o.with_suffix(".png"), dpi=220)
    fig.savefig(_o.with_suffix(".pdf"), bbox_inches="tight")
    plt.close(fig)