import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from scoobdoob.validation import msm_build as mb
from scoobdoob.validation import tpt, bridge


def _toy_dtrajs():

    rng = np.random.default_rng(0)
    P = np.array(
        [
            [0.80, 0.20, 0.00, 0.00, 0.00],
            [0.15, 0.70, 0.15, 0.00, 0.00],
            [0.00, 0.15, 0.70, 0.15, 0.00],
            [0.00, 0.00, 0.15, 0.70, 0.15],
            [0.00, 0.00, 0.00, 0.20, 0.80],
        ]
    )
    traj = [0]
    for _ in range(40000):
        traj.append(int(rng.choice(5, p=P[traj[-1]])))
    return [np.array(traj)], P


def test_connected_and_reversible():
    dtrajs, _ = _toy_dtrajs()
    model, syms, sym2a, _ = mb.build_msm(dtrajs, lag=1, k=5, reversible=True)
    assert model.n_states == 5
    assert abs(model.stationary_distribution.sum() - 1) < 1e-10

    pi = model.stationary_distribution
    P = model.transition_matrix
    assert np.allclose(pi[:, None] * P, (pi[:, None] * P).T, atol=1e-8)


def test_committor_solvers_agree():
    dtrajs, _ = _toy_dtrajs()
    model, *_ = mb.build_msm(dtrajs, lag=1, k=5, reversible=True)
    q_dt = np.asarray(model.committor_forward([0], [4]))
    q_ls = tpt.forward_committor_linsolve(model.transition_matrix, [0], [4])
    assert np.max(np.abs(q_dt - q_ls)) < 1e-9
    assert q_dt[0] == 0.0 and q_dt[4] == 1.0
    assert np.all(np.diff(q_dt) >= -1e-9)


def test_committor_half_localizes_in_middle():
    dtrajs, _ = _toy_dtrajs()
    model, *_ = mb.build_msm(dtrajs, lag=1, k=5, reversible=True)
    q = np.asarray(model.committor_forward([0], [4]))
    half = tpt.committor_half_states(q, tol=0.25)
    assert 2 in half


def test_bridge_conditioning_beats_unconditional():
    _, P = _toy_dtrajs()
    Ph = bridge.doob_kernels(P, S_end=[4], horizon=20)
    cond = bridge.sample_bridge_paths(Ph, S_start=[0], n_paths=200, seed=0)
    unc = bridge.sample_unconditional_paths(
        P, S_start=[0], horizon=20, n_paths=200, seed=1
    )
    assert bridge.success_with_ci(cond, [4])["rate"] > 0.9
    assert bridge.success_with_ci(unc, [4])["rate"] < 0.5


def test_aib9_macrostate_classification():
    from scoobdoob.validation import aib9

    state_left = np.array(
        [
            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
            [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0, 0.5, 0.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 0.0, 0.5, 1.0, 1.0, 1.0, 1.0],
        ]
    )
    labels = aib9.classify_microstates(state_left)
    assert list(labels) == ["R", "L", "I_N", "I_C"]


def test_aib9_featurize_and_handedness_shapes():
    from scoobdoob.validation import aib9

    ang = [np.linspace(-3, 3, 18).reshape(1, 18).repeat(50, 0)]
    feat = aib9.featurize_sincos(ang)
    hand = aib9.handedness_indicator(ang)
    assert feat[0].shape == (50, 36)
    assert hand[0].shape == (50, 9)
    assert set(np.unique(hand[0])) <= {0.0, 1.0}


if __name__ == "__main__":
    test_connected_and_reversible()
    test_committor_solvers_agree()
    test_committor_half_localizes_in_middle()
    test_bridge_conditioning_beats_unconditional()
    test_aib9_macrostate_classification()
    test_aib9_featurize_and_handedness_shapes()
    print("all validation smoke tests passed")