from pathlib import Path
import argparse
import csv
import hashlib
import importlib.metadata
import itertools
import json
import platform
import sys
import time

import numpy as np

OUT = Path(__file__).resolve().parent
ROOT = OUT.parents[1]
sys.path.insert(0, str(ROOT))
from scoobdoob import _conditioning

CONDITIONS = [
    ("doob", None),
    ("none", 0),
    ("p10", 10),
    ("p25", 25),
    ("p50", 50),
    ("p75", 75),
]
SETTINGS = dict(
    cluster_seeds=list(range(20)),
    k=100,
    tica_lag=20,
    msm_lag=5,
    horizon=16,
    alpha_per_transition=7.5,
    floor_percentiles=[10, 25, 50, 75],
    no_floor_control=True,
    start_law="uniform on the original source set; fixed-start mixture",
    s1_radius=0.25,
)


def dump(path, obj):
    path.write_text(json.dumps(obj, indent=2, allow_nan=False) + "\n")


def digest(path):
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def exact_metrics(kernels, mu, region, targets, corridor):
    """Forward full law and killed laws, including the state at step zero."""
    m, n = len(mu), len(kernels)
    mass = mu.copy()
    survival = mu.copy()
    survival[region] = 0
    avoiding = np.broadcast_to(mu, (m, m)).copy()
    np.fill_diagonal(avoiding, 0)
    s1_occupancy = float(mass[region].sum())
    self_steps = 0.0
    max_outside = float(mass[~corridor].sum())
    max_outside_arrivals = 0.0
    for K in kernels:
        self_steps += float(mass @ np.diag(K))
        mass = mass @ K
        np.testing.assert_allclose(mass.sum(), 1, atol=2e-12, rtol=0)
        max_outside = max(max_outside, float(mass[~corridor].sum()))
        max_outside_arrivals = max(max_outside_arrivals, float(mass[~corridor].sum()))
        s1_occupancy += float(mass[region].sum())
        survival = survival @ K
        survival[region] = 0
        avoiding = avoiding @ K
        np.fill_diagonal(avoiding, 0)
    hit_each = np.clip(1 - avoiding.sum(axis=1), 0, 1)
    return dict(
        s1_probability=float(np.clip(1 - survival.sum(), 0, 1)),
        expected_distinct_states=float(hit_each.sum()),
        expected_self_transitions=self_steps,
        self_transition_fraction=self_steps / n,
        expected_s1_state_observations=s1_occupancy,
        terminal_probability=float(mass[targets].sum()),
        max_outside_corridor_probability=max_outside,
        max_outside_arrival_mask_probability=max_outside_arrivals,
    )


def reference_hit_probability(P, w, mu, region, targets, horizon):
    """Independent ratio of weighted path sums, separately for each start."""
    full = np.zeros(len(P))
    full[targets] = 1
    avoided = full.copy()
    avoided[region] = 0
    w_avoid = w.copy()
    w_avoid[region] = 0
    for _ in range(horizon):
        full = P @ (w * full)
        avoided = P @ (w_avoid * avoided)
        scale = float(full.max())
        if scale <= 0:
            raise RuntimeError("No feasible weighted continuation")
        full /= scale
        avoided /= scale
    initial = mu > 0
    if np.any(full[initial] <= 0):
        raise RuntimeError("Infeasible prescribed start")
    avoided[region] = 0
    return float(1 - mu[initial] @ (avoided[initial] / full[initial]))


def validate_toy():
    """Brute-force the weighted reference law, not a second kernel simulation."""
    P = np.array([[0.5, 0.4, 0.1], [0.1, 0.6, 0.3], [0.2, 0.3, 0.5]])
    w = np.array([0.7, 0.3, 1.2])
    mu = np.array([0.3, 0.7, 0.0])
    region = np.array([False, True, False])
    targets = np.array([2])
    horizon = 4
    kernels, _ = _conditioning.build(P, w, targets, horizon, starts=[0, 1])
    got = exact_metrics(kernels, mu, region, targets, np.ones(3, dtype=bool))
    expected = np.zeros(4)
    for i in [0, 1]:
        norm = 0.0
        moments = np.zeros(4)
        for rest in itertools.product(range(3), repeat=horizon):
            if rest[-1] != 2:
                continue
            path = (i,) + rest
            prob = float(np.prod([P[a, b] * w[b] for a, b in zip(path[:-1], path[1:])]))
            norm += prob
            moments += prob * np.array(
                [
                    any(region[s] for s in path),
                    len(set(path)),
                    sum(a == b for a, b in zip(path[:-1], path[1:])),
                    sum(region[s] for s in path),
                ],
                dtype=float,
            )
        expected += mu[i] * moments / norm
    actual = [
        got[k]
        for k in [
            "s1_probability",
            "expected_distinct_states",
            "expected_self_transitions",
            "expected_s1_state_observations",
        ]
    ]
    np.testing.assert_allclose(actual, expected, atol=2e-12, rtol=0)
    other = reference_hit_probability(P, w, mu, region, targets, horizon)
    np.testing.assert_allclose(other, expected[0], atol=2e-12, rtol=0)

    masked_w = np.array([0.0, 0.3, 1.2])
    K, _ = _conditioning.build(P, masked_w, targets, 1, starts=[0])
    edge = exact_metrics(K, np.array([1.0, 0.0, 0.0]), region, targets, masked_w > 0)
    np.testing.assert_allclose(
        [
            edge["terminal_probability"],
            edge["expected_distinct_states"],
            edge["max_outside_corridor_probability"],
            edge["max_outside_arrival_mask_probability"],
        ],
        [1.0, 2.0, 1.0, 0.0],
        atol=2e-12,
        rtol=0,
    )
    return dict(
        exhaustive_reference_check=True,
        outside_mask_start_check=True,
        max_absolute_error=float(np.max(np.abs(np.array(actual) - expected))),
    )


def run(seed):
    start_time = time.time()
    toy = validate_toy()
    from scoobdoob.validation import msm_build as mb, tpt
    from scoobdoob.validation import mb_ground_truth as gt
    from scoobdoob.validation import mb_data as V

    files = sorted((ROOT / "data/mb_reactive").glob("traj_*.npy"))
    if len(files) != 12:
        raise RuntimeError(
            "Expected twelve trajectory files in data/mb_reactive; run python -m scoobdoob.validation.mb_data first"
        )
    trajs = [np.load(p).astype(np.float64) for p in files]
    model = mb.build_full_msm(
        trajs, trajs, lag=5, tica_lag=20, k=100, n_comp=2, seed=seed, use_tica=True
    )

    np.testing.assert_array_equal(model.active_symbols, np.arange(100))
    A, B, _ = mb.q2_sign_endpoints(model)
    starts, targets = V.orient_endpoints(model, A, B)
    starts, targets = np.array(starts), np.array(targets)
    xy = model.active_xy
    agreement, q, _ = tpt.committor_agreement(model, starts.tolist(), targets.tolist())
    flux = tpt.reactive_flux(model, starts.tolist(), targets.tolist())
    through = np.asarray(flux.gross_flux).sum(axis=1)
    corridor = through / through.sum() > 0.5 / model.n_active
    corridor[starts] = True
    corridor[targets] = True
    saddles = gt.saddle_points()
    S1 = saddles[np.argmin(saddles[:, 0])]
    landmarks = np.vstack([xy[starts].mean(0), xy[targets].mean(0), saddles])
    lo, hi = landmarks.min(0) - 0.35, landmarks.max(0) + 0.35
    corridor &= np.all((xy >= lo) & (xy <= hi), axis=1)

    counts = model.counts[model.active_symbols].astype(float)
    rho = np.maximum(counts, 1) / counts.mean()
    react = (q > 0.05) & (q < 0.95)
    region = np.linalg.norm(xy - S1, axis=1) < 0.25
    if not np.any(region) or not np.any(react):
        raise RuntimeError("Missing S1 or reference committor region")
    mu = np.zeros(model.n_active)
    mu[starts] = 1 / len(starts)
    saved_match = None
    if seed == 0:
        with np.load(
            ROOT / "data/mb_panels_data/panels_cache.npz", allow_pickle=True
        ) as z:
            saved_match = dict(
                max_state_coordinate_difference=float(np.max(np.abs(xy - z["xy"]))),
                counts_equal=bool(np.array_equal(counts, z["counts"])),
                corridor_equal=bool(np.array_equal(corridor, z["on"])),
            )
    np.savez_compressed(
        OUT / f"model_{seed:02d}.npz",
        P=model.P,
        pi=model.stationary,
        q=q,
        xy=xy,
        counts=counts,
        rho=rho,
        corridor=corridor,
        source=starts,
        target=targets,
        S1_region=region,
        mu=mu,
    )
    records = []
    for name, pct in CONDITIONS:
        floor = (
            None
            if pct is None
            else (0.0 if pct == 0 else float(np.percentile(rho[react], pct)))
        )
        w = (
            np.ones_like(rho)
            if name == "doob"
            else np.exp(-7.5 * np.maximum(rho, floor))
        )
        if np.any(w <= 0):
            raise FloatingPointError("A positive density weight underflowed")
        w[~corridor] = 0
        kernels, h = _conditioning.build(model.P, w, targets, 16, starts=starts)
        for K in kernels:
            if np.any((K > 0) & (model.P == 0)):
                raise RuntimeError("New edge introduced")
        record = exact_metrics(kernels, mu, region, targets, corridor)
        independent = reference_hit_probability(model.P, w, mu, region, targets, 16)
        error = abs(independent - record["s1_probability"])
        if error > 2e-10:
            raise AssertionError(f"Two path-law calculations disagree by {error}")
        record.update(
            condition=name,
            floor_percentile=pct,
            density_floor=floor,
            fraction_corridor_states_below_floor=(
                None if floor is None else float(np.mean(rho[corridor] < floor))
            ),
            independent_recursion_error=error,
        )
        records.append(record)
        print(
            f"seed={seed} {name} S1={record['s1_probability']:.8f} "
            f"distinct/path={record['expected_distinct_states']:.6f}",
            flush=True,
        )
    source_names = [
        "make_mb_panels.py",
        "mb_stats.py",
        "scoobdoob/validation/msm_build.py",
        "scoobdoob/_conditioning.py",
        "scoobdoob/validation/mb_data.py",
    ]
    sources = {str(p.relative_to(ROOT)): digest(p) for p in files}
    sources.update({p: digest(ROOT / p) for p in source_names})
    sources[str(Path(__file__).relative_to(ROOT))] = digest(__file__)
    versions = {
        p: importlib.metadata.version(p)
        for p in ["numpy", "scipy", "scikit-learn", "deeptime"]
    }
    dump(
        OUT / f"result_{seed:02d}.json",
        dict(
            seed=seed,
            settings=SETTINGS,
            records=records,
            n_frames=sum(len(t) for t in trajs),
            n_active=model.n_active,
            corridor_states=int(corridor.sum()),
            s1_states=int(region.sum()),
            source_states_outside_arrival_mask=starts[~corridor[starts]].tolist(),
            target_states_outside_arrival_mask=targets[~corridor[targets]].tolist(),
            reference_committor_agreement=float(agreement),
            seed0_saved_figure_check=saved_match,
            validation=toy,
            source_sha256=sources,
            versions=versions,
            python=platform.python_version(),
            elapsed_seconds=time.time() - start_time,
        ),
    )


def summarize():
    results = [
        json.loads((OUT / f"result_{s:02d}.json").read_text()) for s in range(20)
    ]
    if any(r["settings"] != SETTINGS for r in results):
        raise RuntimeError("Inconsistent settings")
    if len({json.dumps(r["source_sha256"], sort_keys=True) for r in results}) != 1:
        raise RuntimeError("Source/data hashes changed across jobs")
    records = [
        {**row, "cluster_seed": r["seed"]} for r in results for row in r["records"]
    ]
    with (OUT / "all_results.csv").open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=list(records[0]))
        writer.writeheader()
        writer.writerows(records)
    values = {
        name: np.array(
            [
                next(x for x in r["records"] if x["condition"] == name)[
                    "s1_probability"
                ]
                for r in results
            ]
        )
        * 100
        for name, _ in CONDITIONS
    }
    stats = {}
    for name, _ in CONDITIONS:
        y = values[name]
        delta = y - values["doob"]
        delta25 = y - values["p25"]
        by = [next(x for x in r["records"] if x["condition"] == name) for r in results]
        stats[name] = dict(
            s1_mean_pct=float(y.mean()),
            s1_sd_pct=float(y.std(ddof=1)),
            s1_min_pct=float(y.min()),
            s1_max_pct=float(y.max()),
            paired_gain_vs_doob_mean_pp=float(delta.mean()),
            paired_gain_vs_doob_min_pp=float(delta.min()),
            paired_gain_vs_doob_max_pp=float(delta.max()),
            paired_vs_p25_mean_pp=float(delta25.mean()),
            paired_vs_p25_min_pp=float(delta25.min()),
            paired_vs_p25_max_pp=float(delta25.max()),
            n_clusterings_better_than_doob=int(np.sum(delta > 1e-10)),
            distinct_per_path_mean=float(
                np.mean([x["expected_distinct_states"] for x in by])
            ),
            self_steps_mean=float(
                np.mean([x["expected_self_transitions"] for x in by])
            ),
        )
    archive = json.loads((ROOT / "data/mb_stats.json").read_text())
    comparison = {
        name: dict(
            archived_mean_pct=float(np.mean(archive[key])),
            mean_exact_minus_archived_pp=float(np.mean(values[name] - archive[key])),
            max_absolute_difference_pp=float(
                np.max(np.abs(values[name] - archive[key]))
            ),
        )
        for name, key in [("doob", "raw"), ("p25", "fk")]
    }
    check = dict(
        toy_exhaustive_passed=all(
            r["validation"]["exhaustive_reference_check"] for r in results
        ),
        max_independent_recursion_error=max(
            x["independent_recursion_error"] for x in records
        ),
        max_terminal_error=max(abs(x["terminal_probability"] - 1) for x in records),
        max_outside_corridor_probability=max(
            x["max_outside_corridor_probability"] for x in records
        ),
        max_outside_arrival_mask_probability=max(
            x["max_outside_arrival_mask_probability"] for x in records
        ),
        endpoint_mask_exceptions=[
            dict(
                seed=r["seed"],
                source=r["source_states_outside_arrival_mask"],
                target=r["target_states_outside_arrival_mask"],
            )
            for r in results
            if r["source_states_outside_arrival_mask"]
            or r["target_states_outside_arrival_mask"]
        ],
        seed0_saved_figure_check=results[0]["seed0_saved_figure_check"],
    )
    summary = dict(
        settings=SETTINGS,
        conditions=stats,
        checks=check,
        archive_comparison=comparison,
        scope="Exact probabilities on refitted MSMs from fixed MD data. SD is clustering sensitivity, not MD uncertainty.",
        completed_clusterings=len(results),
    )
    dump(OUT / "summary.json", summary)
    import matplotlib

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

    labels = ["No floor", "10th", "25th (original)", "50th", "75th"]
    names = ["none", "p10", "p25", "p50", "p75"]
    fig, axes = plt.subplots(1, 2, figsize=(10, 4.3))
    x = np.arange(5)
    for s in range(20):
        axes[0].plot(
            x, [values[n][s] for n in names], color="#93b6cc", alpha=0.42, lw=0.8
        )
        axes[1].plot(
            x,
            [values[n][s] - values["doob"][s] for n in names],
            color="#93b6cc",
            alpha=0.42,
            lw=0.8,
        )
    for ax, metric in zip(axes, ["s1_mean_pct", "paired_gain_vs_doob_mean_pp"]):
        ax.plot(
            x,
            [stats[n][metric] for n in names],
            "o-",
            color="#164e70",
            lw=2,
            label="Mean of 20 clusterings",
        )
        ax.set_xticks(x, labels, rotation=18)
        ax.set_xlabel("Density-floor percentile")
        ax.spines[["top", "right"]].set_visible(False)
        ax.grid(axis="y", alpha=0.2)
    axes[0].axhline(
        stats["doob"]["s1_mean_pct"], ls="--", color="#777777", label="Pure-Doob mean"
    )
    axes[1].axhline(0, ls="--", color="#777777")
    axes[0].set_ylabel("Exact probability of S1 visitation (%)")
    axes[1].set_ylabel("Paired gain over pure Doob (percentage points)")
    axes[0].legend(frameon=False, fontsize=8)
    fig.suptitle("MB floor sensitivity at fixed density strength and horizon")
    fig.tight_layout()
    (ROOT / "figures").mkdir(parents=True, exist_ok=True)
    for ext in ["pdf", "png", "svg"]:
        fig.savefig(
            ROOT / "figures" / f"floor_sensitivity.{ext}", dpi=180, bbox_inches="tight"
        )
    lines = [
        "MB FLOOR SENSITIVITY — completed",
        "",
        summary["scope"],
        "",
        "Condition       S1 mean +/- SD (%)      gain vs Doob (pp)   better /20",
    ]
    for name, _ in CONDITIONS:
        z = stats[name]
        lines.append(
            f"{name:12s} {z['s1_mean_pct']:7.3f} +/- {z['s1_sd_pct']:6.3f} "
            f"          {z['paired_gain_vs_doob_mean_pp']:+7.3f}             "
            f"{z['n_clusterings_better_than_doob']}"
        )
    lines += [
        "",
        "Checks:",
        json.dumps(check, indent=2),
        "",
        "The original 25th-percentile setting is retained. No optimum was selected.",
        "Use per-clustering paired differences in all_results.csv; original results were not overwritten.",
        "Exact evaluation has no rollout Monte Carlo error. Refitting may differ from archived finite samples.",
        "Aib9 and the p53 held-out density experiment used no density floor; this test concerns MB.",
    ]
    (OUT / "RESULTS.txt").write_text("\n".join(lines) + "\n")
    print("\n".join(lines), flush=True)


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    group = p.add_mutually_exclusive_group(required=True)
    group.add_argument("--seed", type=int, choices=range(20))
    group.add_argument("--summarize", action="store_true")
    group.add_argument("--validate", action="store_true")
    args = p.parse_args()
    if args.validate:
        print(json.dumps(validate_toy(), indent=2))
    elif args.summarize:
        summarize()
    else:
        run(args.seed)