import numpy as np, collections
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
import make_mb_panels as MP

Ac = np.array([-200, -100, -170, 15.0])
a = np.array([-1, -1, -6.5, 0.7])
b = np.array([0, 0, 11, 0.6])
c = np.array([-10, -10, -6.5, 0.7])
X = np.array([1, 0, -0.5, -1.0])
Y = np.array([0, 0.5, 1.5, 1.0])


def E(p):
    x, y = p
    return float(
        (Ac * np.exp(a * (x - X) ** 2 + b * (x - X) * (y - Y) + c * (y - Y) ** 2)).sum()
    )


LAG, N, TL = MP.LAG_OPT, MP.N_OPT, 20
sad = gt.saddle_points()
S1 = sad[np.argmin(sad[:, 0])]
S2 = sad[np.argmax(sad[:, 0])]
trajs = V.load_trajs(str(V.DEFAULT_DATA))
AGG = collections.defaultdict(list)
DIV = collections.defaultdict(list)
for seed in [0, 1, 2, 3, 4]:
    model = mb.build_full_msm(
        trajs, trajs, lag=LAG, tica_lag=TL, k=V.K, n_comp=2, seed=seed, 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)
    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
    on = on & (
        (xy[:, 0] >= lo[0])
        & (xy[:, 0] <= hi[0])
        & (xy[:, 1] >= lo[1])
        & (xy[:, 1] <= hi[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)
    flr = np.percentile(rho[react], 25)
    n1 = np.linalg.norm(xy - S1, axis=1) < 0.25
    qs2 = q[np.argmin(np.linalg.norm(xy - S2, axis=1))]
    for al in [0.0, 0.5, 1.0, 1.5, 2.0, 3.0]:
        w = np.ones_like(rho) if al == 0 else np.exp(-LAG * al * np.maximum(rho, flr))
        w = w.copy()
        w[~on] = 0.0
        ps = bridge.sample_bridge_paths(
            bridge.doob_kernels(model.P, end, N, w=w), start, n_paths=3000, seed=0
        )
        rec = np.mean([any(n1[s] for s in p) for p in ps]) * 100

        DIV[al].append((len({s for p in ps for s in p}), len({tuple(p) for p in ps})))
        r = MP.channel(ps, q, xy, on, thru, start, end)
        qr = [q[np.argmin(np.linalg.norm(xy - p, axis=1))] for p in r]
        cr = [i for i in range(len(qr) - 1) if (qr[i] < qs2) != (qr[i + 1] < qs2)]
        d2 = (E((r[cr[0]] + r[cr[0] + 1]) / 2) - E(S2)) / 16.7 if cr else np.nan
        AGG[al].append((rec, d2))
print(f"TICA lag {TL}, count lag {LAG}, N={N}; 5 clustering seeds; kT=16.7")
print(f"{'alpha':>6} {'S1 recov':>16} {'S2 dE (kT)':>16} {'states':>10} {'routes':>10}")
for al, v in sorted(AGG.items()):
    rc = np.array([x[0] for x in v])
    d2 = np.array([x[1] for x in v])
    st = np.array([x[0] for x in DIV[al]])
    rt = np.array([x[1] for x in DIV[al]])
    print(
        f"{al:6.1f} {rc.mean():8.0f}%+-{rc.std():3.0f} {np.nanmean(d2):11.1f}+-{np.nanstd(d2):3.1f} "
        f"{st.mean():6.0f}+-{st.std():2.0f} {rt.mean():6.0f}+-{rt.std():2.0f}"
    )