import numpy as np

from scoobdoob import _conditioning
from scoobdoob.validation import tpt

STRIDE, SEED_RELAX = 10, 50
ALPHA_V, TAU = 3.0, 0.12
TIME_SIGMA, COMM_L, COMM_S, W_FLOOR = 0.30, 3.0, 0.35, 1e-12


def fk_weights(
    C, q, horizon, comm_l=COMM_L, comm_s=COMM_S, variant="rho", alpha_rho=None
):
    """(untilted, committor-tilted) FK weight lists, one entry per horizon step."""
    row = np.asarray(C, dtype=np.float64).sum(axis=1)
    a_eff = (TAU * ALPHA_V) if alpha_rho is None else float(alpha_rho)
    if variant == "rho":
        rho = np.clip(row, 1.0, None) / row.mean()
        w_base = np.exp(-a_eff * rho)
    else:
        w_base = np.exp(-a_eff * (1.0 / (row + 1.0)))
    q_pen = np.exp(comm_l * (-((q - 0.5) ** 2)) / comm_s**2)
    t = np.arange(horizon)
    time_win = np.exp(-0.5 * ((t - (horizon - 1) / 2) / (TIME_SIGMA * horizon)) ** 2)
    wA = [np.maximum(w_base, W_FLOOR) for _ in range(horizon)]
    wB = [np.maximum(w_base * (q_pen ** time_win[i]), W_FLOOR) for i in range(horizon)]
    return wA, wB


def crossing_horizon(m, a0, B, cap=20000):
    """Return the bounded horizon and reference MFPT in lag steps."""
    P = np.asarray(m.P, dtype=np.float64)
    n = P.shape[0]
    Bs = set(int(b) for b in B)
    T = np.array([i for i in range(n) if i not in Bs])
    if not len(T):
        return cap
    try:
        mfpt = np.zeros(n)
        mfpt[T] = np.linalg.solve(np.eye(len(T)) - P[np.ix_(T, T)], np.ones(len(T)))
        h = int(np.clip(round(mfpt[a0]), 5, cap))
        return h, float(mfpt[a0])
    except Exception:
        return cap, float("nan")


def state_means(st, dg, n_states):
    """E[drug | state] and per-state counts, on whatever frame subset is passed in."""
    mu = np.full(n_states, np.nan)
    cnt = np.zeros(n_states, dtype=int)
    for s in range(n_states):
        m = st == s
        cnt[s] = m.sum()
        if cnt[s]:
            mu[s] = dg[m].mean()
    return mu, cnt


def target_set(cvm, dil, a0, n, a, dil_override=None):
    """The B set the tilt aims at."""
    if a.target == "cleft":
        B = np.array([b for b in np.where(cvm[:, 1] > a.cleft_thr)[0] if b != a0])
        if not len(B):
            raise SystemExit(
                f"empty target set B at cleft_thr={a.cleft_thr}. Under a pocket-feature MSM the "
                f"per-state mean cleft CV never reaches this threshold because pocket states mix "
                f"open and closed frames of the GLOBAL coordinate. Use --target dilation."
            )
        return B
    d = dil if dil_override is None else dil_override
    ok = np.isfinite(d)
    thr = np.quantile(d[ok], 1.0 - a.target_q)
    B = np.array([b for b in np.where(ok & (d >= thr))[0] if b != a0])
    if not len(B):
        raise SystemExit("empty dilation target set")
    return B


def build_selectors(m, cvm, C, a0, a, dil=None, B=None):
    """Weight vectors over ACTIVE states for every selector, each summing to 1."""
    P = np.asarray(m.P, dtype=np.float64)
    pi = np.asarray(m.stationary, dtype=np.float64)
    n = len(pi)

    if B is None:
        B = target_set(cvm, dil, a0, n, a)
    q = np.asarray(tpt.forward_committor(m, [a0], list(B)))

    sel = {}
    sel["equilibrium"] = pi / pi.sum()
    sel["uniform_states"] = np.full(n, 1.0 / n)

    half = np.abs(q - 0.5) <= 0.1
    if half.sum():
        w = pi * half
        sel["committor_half"] = w / w.sum()

    try:
        rf = tpt.reactive_flux(m, [a0], list(B))
        gross = np.asarray(rf.gross_flux)
        carried = gross.sum(axis=1) + gross.sum(axis=0)
        if carried.sum() > 0:
            sel["tpt_flux"] = carried / carried.sum()
    except Exception as e:
        print(f"  [warn] TPT flux unavailable: {e}")

    horizon = a.horizon
    mfpt = float("nan")
    if horizon is None:
        horizon, mfpt = crossing_horizon(m, a0, B)
    print(f"  [fk] horizon = {horizon} lag steps (MFPT a0->B = {mfpt:.1f})")

    def occupancy(wl, H):
        Ph, _ = _conditioning.build(P, wl, list(B), H, starts=[a0])
        mu = np.zeros(n)
        mu[a0] = 1.0
        acc = np.zeros(n)
        for Pn in Ph:
            mu = mu @ np.asarray(Pn)
            s = mu.sum()
            if not np.isfinite(s) or s <= 0:
                raise FloatingPointError(
                    "Conditioned occupancy lost all probability mass"
                )
            mu = mu / s
            acc += mu
        return acc / acc.sum()

    wA, wB = fk_weights(C, q, horizon, variant=a.fk_variant)
    for name, wl in (("doob_untilted", wA), ("fk_doob", wB)):
        try:
            sel[name] = occupancy(wl, horizon)
        except Exception as e:
            print(f"  [warn] {name} failed: {e}")

    for a_rho in getattr(a, "alpha_rho_sweep", []) or []:
        try:
            wa, _ = fk_weights(C, q, horizon, variant=a.fk_variant, alpha_rho=a_rho)
            sel[f"alg1[a={a_rho:g}]"] = occupancy(wa, horizon)
        except Exception as e:
            print(f"  [warn] alg1[a={a_rho:g}] failed: {e}")

    for cl in a.comm_l_sweep:
        for hf in a.horizon_sweep:
            H = max(5, int(round(horizon * hf)))
            try:
                _, wb = fk_weights(C, q, H, comm_l=cl, variant=a.fk_variant)
                sel[f"fk_doob[L={cl},H={H}]"] = occupancy(wb, H)
            except Exception:
                pass

    return sel, q, B, a0, horizon


def enrich(weights, mu_state, cnt, p_open_state, base_mean, base_open):
    """Expected druggability and open-probability under a state-weight vector."""
    ok = np.isfinite(mu_state) & (cnt > 0)
    w = np.where(ok, weights, 0.0)
    if w.sum() <= 0:
        return None
    w = w / w.sum()
    em = float((w * np.nan_to_num(mu_state)).sum())
    eo = float((w * np.nan_to_num(p_open_state)).sum())
    return {
        "E_drug": em,
        "P_open": eo,
        "enrich_mean": em / base_mean if base_mean > 0 else float("nan"),
        "enrich_open": eo / base_open if base_open > 0 else float("nan"),
        "coverage": float(weights[ok].sum() / max(weights.sum(), 1e-300)),
    }