#!/usr/bin/env python3
"""lonely_runner_sim.py — Lonely Runner Betti simulation.

Simulates k runners on S¹, computes scar (uncovered) region M_t,
tracks β₀(M_t) — connected components of uncovered set.
β₀(M_t) > 0 ⇔ lonely runner exists at time t.

Imports scar_density, threshold_scar_support, compute_betti_numbers
from betti_tracker for compatibility; uses direct 1D circular β₀
calculator for correctness on S¹ topology.

Usage:
  python3 lonely_runner_sim.py --test
  python3 lonely_runner_sim.py --k 3 --speeds 0 1 2 --time 5 --steps 1000
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

try:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    HAVE_MPL = True
except ImportError:
    HAVE_MPL = False

sys.path.insert(0, str(Path(__file__).resolve().parent))
from betti_tracker import threshold_scar_support, compute_betti_numbers

# ── simulation core ──────────────────────────────────────────────────


def runner_positions(
    speeds: list[float], times: np.ndarray
) -> np.ndarray:
    """Runner positions on S¹ at each time step.

    Returns (T, k) array in [0, 1).
    """
    speeds_arr = np.array(speeds, dtype=np.float64)
    pos = np.outer(times, speeds_arr) % 1.0
    return pos


def coverage_density(
    positions: np.ndarray, delta: float, N: int = 512
) -> np.ndarray:
    """Coverage density Φ(t,θ) = Σᵢ 𝟙(|θ - vᵢt| < δ mod 1).

    Returns (T, N) integer count array.
    """
    T, k = positions.shape
    theta = np.linspace(0, 1, N, endpoint=False)
    phi = np.zeros((T, N), dtype=np.int32)

    for t in range(T):
        diff = np.abs(theta[:, np.newaxis] - positions[t, np.newaxis, :])
        dist = np.minimum(diff, 1.0 - diff)
        phi[t, :] = (dist < delta).sum(axis=1)

    return phi


def scar_field(phi: np.ndarray) -> np.ndarray:
    """Scar (loneliness) field μ(t,θ) = 1 - min(Φ(t,θ), 1).

    Returns (T, N) float64 array where 1 = uncovered, 0 = covered.
    """
    return (1.0 - np.minimum(phi, 1)).astype(np.float64)


# ── 1D circular Betti-0 ─────────────────────────────────────────────


def betti_1d_circular(mask: np.ndarray) -> int:
    """β₀ of a 1D binary mask on S¹ (circular topology).

    Counts connected components of the foreground (1 = scar),
    accounting for wrap-around adjacency (first and last elements
    are neighbours on the circle).
    """
    if not mask.any():
        return 0
    if mask.all():
        return 1
    signed = 2 * mask.astype(np.int8) - 1
    padded = np.concatenate([[signed[-1]], signed])
    trans_up = int(((padded[1:] - padded[:-1]) > 0).sum())
    return trans_up


# ── simulate ────────────────────────────────────────────────────────


SIM_CASES: dict[str, dict] = {
    "k2_simple": {
        "k": 2,
        "speeds": [0, 1],
        "description": "k=2 with speeds [0, 1]",
    },
    "k3_012": {
        "k": 3,
        "speeds": [0, 1, 2],
        "description": "k=3 with speeds [0, 1, 2]",
    },
    "k4_0134": {
        "k": 4,
        "speeds": [0, 1, 3, 4],
        "description": "k=4 with speeds [0, 1, 3, 4] — known lonely times",
    },
    "k3_rational": {
        "k": 3,
        "speeds": [0, 0.5, 1.0],
        "description": "k=3 with rational multiples [0, 0.5, 1.0]",
    },
    "k4_random": {
        "k": 4,
        "speeds": [0, 1.0, 2.71, 3.14],
        "description": "k=4 with random-ish speeds",
    },
}


def simulate(
    speeds: list[float],
    total_time: float,
    time_steps: int,
    N: int = 512,
) -> dict:
    """Run a Lonely Runner simulation.

    Returns a results dict with positions, scar data, and Betti series.
    """
    k = len(speeds)
    delta = 1.0 / (k + 1.0)
    times = np.linspace(0, total_time, time_steps, endpoint=False)

    pos = runner_positions(speeds, times)
    phi = coverage_density(pos, delta, N=N)
    mu = scar_field(phi)
    masks = (mu > 0).astype(np.uint8)

    betti_series: list[int] = []
    scar_fractions: list[float] = []
    for t in range(time_steps):
        b0 = betti_1d_circular(masks[t])
        betti_series.append(b0)
        scar_fractions.append(float(masks[t].mean()))

    theta = np.linspace(0, 1, N, endpoint=False)

    return {
        "k": k,
        "speeds": speeds,
        "delta": delta,
        "total_time": total_time,
        "time_steps": time_steps,
        "N": N,
        "times": times.tolist(),
        "theta": theta.tolist(),
        "positions": pos.tolist(),
        "phi": phi.tolist(),
        "mu": mu.tolist(),
        "masks": masks.tolist(),
        "betti_series": betti_series,
        "scar_fractions": scar_fractions,
        "min_betti_0": min(betti_series),
        "max_betti_0": max(betti_series),
        "lonely_times_found": max(betti_series) > 0,
        "always_covered": max(betti_series) == 0,
        "fraction_with_scar": sum(1 for b in betti_series if b > 0) / time_steps,
    }


def build_report(results: dict) -> dict:
    """Build a JSON-serialisable report with concise summary."""
    b0 = results["betti_series"]
    summary_parts = []
    if results["lonely_times_found"]:
        summary_parts.append(
            f"β₀(M_t) > 0 at some t (max={results['max_betti_0']})"
        )
    else:
        summary_parts.append("β₀(M_t) = 0 ∀t — complete coverage")
    summary_parts.append(
        f"scar present {results['fraction_with_scar']*100:.1f}% of time"
    )
    summary_parts.append(
        f"min β₀ = {results['min_betti_0']}, max β₀ = {results['max_betti_0']}"
    )

    return {
        "schema": "lonely_runner_sim_v1",
        "claim_boundary": "numerical-simulation;no-lean-proof",
        "k": results["k"],
        "speeds": results["speeds"],
        "delta": results["delta"],
        "total_time": results["total_time"],
        "time_steps": results["time_steps"],
        "N": results["N"],
        "summary": "; ".join(summary_parts),
        "lonely_times_found": results["lonely_times_found"],
        "always_covered": results["always_covered"],
        "min_betti_0": results["min_betti_0"],
        "max_betti_0": results["max_betti_0"],
        "fraction_with_scar": results["fraction_with_scar"],
        "scar_fraction_mean": float(np.mean(results["scar_fractions"])),
        "betti_series": b0,
        "scar_fractions": results["scar_fractions"],
    }


# ── plotting ────────────────────────────────────────────────────────


def plot_results(
    results: dict, outdir: str, prefix: str = "lonely_runner"
) -> str | None:
    """Generate three-panel diagnostic plot.

    Returns path to saved PNG, or None if matplotlib unavailable.
    """
    if not HAVE_MPL:
        return None

    out_path = Path(outdir)
    out_path.mkdir(parents=True, exist_ok=True)
    ts = np.array(results["times"])
    theta = np.array(results["theta"])
    pos = np.array(results["positions"])
    phi = np.array(results["phi"])
    mu = np.array(results["mu"])
    masks = np.array(results["masks"])
    b0 = results["betti_series"]

    fig, axes = plt.subplots(4, 1, figsize=(12, 11), sharex=False)

    # Panel 1: runner trajectories + uncovered regions
    ax = axes[0]
    for i in range(results["k"]):
        ax.plot(ts, pos[:, i], linewidth=0.8, label=f"R{i} (v={results['speeds'][i]})")
    scar_t, scar_theta = np.where(masks > 0)
    if len(scar_t) > 0:
        ax.scatter(
            ts[scar_t], theta[scar_theta],
            s=0.5, c="red", alpha=0.3, label="uncovered",
        )
    ax.set_ylabel("position on S¹")
    ax.set_title(
        f"k={results['k']} speeds={results['speeds']}  δ={results['delta']:.4f}"
    )
    ax.legend(fontsize=7, ncol=min(results["k"], 4))
    ax.set_ylim(0, 1)

    # Panel 2: coverage heatmap
    ax = axes[1]
    extent = [0, results["total_time"], 0, 1]
    im = ax.imshow(
        phi.T, aspect="auto", origin="lower", extent=extent,
        cmap="YlOrRd", interpolation="nearest",
    )
    ax.set_ylabel("angle θ")
    ax.set_title("Coverage density Φ(t,θ)")
    plt.colorbar(im, ax=ax)

    # Panel 3: scar field heatmap
    ax = axes[2]
    im2 = ax.imshow(
        mu.T, aspect="auto", origin="lower", extent=extent,
        cmap="Blues", interpolation="nearest", vmin=0, vmax=1,
    )
    ax.set_ylabel("angle θ")
    ax.set_title("Scar (loneliness) field μ(t,θ)")
    plt.colorbar(im2, ax=ax)

    # Panel 4: β₀ over time
    ax = axes[3]
    ax.plot(ts, b0, "o-", color="C0", markersize=2, linewidth=0.8)
    ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)
    ax.set_ylabel("β₀")
    ax.set_xlabel("time")
    ax.set_title(f"β₀(M_t) — {'lonely times found' if results['lonely_times_found'] else 'always covered'}")
    ax.set_ylim(-0.5, max(b0) + 1.5)

    fig.suptitle(
        f"Lonely Runner Simulation  (k={results['k']}, δ={results['delta']:.4f})",
        fontsize=13,
    )
    plt.tight_layout(rect=[0, 0, 1, 0.97])
    png_path = out_path / f"{prefix}_k{results['k']}.png"
    fig.savefig(png_path, dpi=150)
    plt.close(fig)
    return str(png_path)


# ── case runner ─────────────────────────────────────────────────────


def run_case(
    name: str,
    config: dict,
    total_time: float = 5.0,
    time_steps: int = 1000,
    N: int = 512,
    outdir: str = ".",
) -> dict:
    """Run one simulation case and write output files."""
    out_path = Path(outdir) / name
    out_path.mkdir(parents=True, exist_ok=True)

    print(f"\n  Case: {name}  ({config['description']})")
    print(f"  k={config['k']}  speeds={config['speeds']}  δ={1.0/(config['k']+1):.4f}")

    results = simulate(
        config["speeds"],
        total_time=total_time,
        time_steps=time_steps,
        N=N,
    )
    report = build_report(results)

    png = plot_results(results, str(out_path), prefix=f"lonely_runner_{name}")

    json_path = out_path / f"{name}_report.json"
    with open(json_path, "w") as f:
        json.dump(report, f, indent=2)

    print(f"    min β₀ = {report['min_betti_0']}, max β₀ = {report['max_betti_0']}")
    print(f"    scar fraction: {report['scar_fraction_mean']:.4f}")
    print(f"    lonely times: {report['lonely_times_found']}")
    if png:
        print(f"    plot: {png}")
    print(f"    JSON: {json_path}")

    return report


# ── tests ────────────────────────────────────────────────────────────


def test_betti_1d():
    """Unit tests for the 1D circular Betti-0 calculator."""
    # all zeros → β₀ = 0
    assert betti_1d_circular(np.zeros(10, dtype=np.uint8)) == 0
    # all ones → β₀ = 1 (whole circle)
    assert betti_1d_circular(np.ones(10, dtype=np.uint8)) == 1
    # single block, no wrap
    assert betti_1d_circular(np.array([1, 1, 0, 0, 0], dtype=np.uint8)) == 1
    # two blocks, no wrap
    assert betti_1d_circular(np.array([1, 1, 0, 0, 1], dtype=np.uint8)) == 1
    # two blocks, wraps around
    assert betti_1d_circular(np.array([1, 0, 1, 0, 1], dtype=np.uint8)) == 2
    # alternating single
    assert betti_1d_circular(np.array([1, 0, 1, 0, 1, 0], dtype=np.uint8)) == 3
    # single isolated 1
    assert betti_1d_circular(np.array([0, 0, 1, 0, 0], dtype=np.uint8)) == 1
    print("  ✓ betti_1d_circular — all unit tests passed")


def test_coverage_sanity():
    """Sanity checks on coverage density calculation."""
    # k=3, at t=0 each runner covers 2δ = 0.5 of circle
    # with runners at 0, overlap means total covered < 3*0.5
    speeds = [0, 1, 2]
    delta = 1.0 / (3 + 1)  # 0.25
    times = np.array([0.0])
    pos = runner_positions(speeds, times)
    phi = coverage_density(pos, delta, N=100)
    frac_covered = (phi[0] > 0).mean()
    # At t=0 all runners at 0 → exactly δ on each side → 0.5
    assert abs(frac_covered - 0.5) < 0.06, f"Expected ~0.5, got {frac_covered}"
    # At t=0, coverage count at θ=0 should be 3 (all runners)
    assert phi[0, 0] == 3, f"Expected phi=3 at origin, got {phi[0, 0]}"
    print(f"  ✓ coverage_sanity — k=3 at t=0 covers {frac_covered:.3f} (expected ~0.5)")


def test_small_k_cases():
    """Check well-known small-k cases produce lonely times."""
    cases = [
        ("k=2 [0,1]", [0, 1]),
        ("k=3 [0,1,2]", [0, 1, 2]),
        ("k=4 [0,1,3,4]", [0, 1, 3, 4]),
    ]
    for label, speeds in cases:
        k = len(speeds)
        res = simulate(speeds, total_time=5.0, time_steps=1000, N=512)
        assert res["lonely_times_found"], f"{label}: expected lonely times, got always covered"
        print(f"  ✓ {label}: β₀ max = {res['max_betti_0']}, scar {res['fraction_with_scar']*100:.1f}%")


def _test():
    """Full test suite."""
    print("=" * 60)
    print("lonely_runner_sim — test suite")
    print("=" * 60)

    # Unit tests
    print("\n--- Unit tests ---")
    test_betti_1d()
    test_coverage_sanity()

    # Small-k verification
    print("\n--- Small-k lonely times ---")
    test_small_k_cases()

    # Rational-multiple case (more coverage overlap)
    print("\n--- Rational speeds (more overlap expected) ---")
    r = run_case(
        "k3_rational",
        SIM_CASES["k3_rational"],
        total_time=5.0, time_steps=1000, N=512,
        outdir="/tmp/lonely_runner_test",
    )
    print(f"    (rational case: scar {r['fraction_with_scar']*100:.1f}% of time)")

    # Random speeds
    print("\n--- Random-ish speeds ---")
    r = run_case(
        "k4_random",
        SIM_CASES["k4_random"],
        total_time=5.0, time_steps=1000, N=512,
        outdir="/tmp/lonely_runner_test",
    )

    print("\n" + "=" * 60)
    print("ALL TESTS PASSED")
    print("=" * 60)


# ── CLI ─────────────────────────────────────────────────────────────


def main():
    ap = argparse.ArgumentParser(
        description="Lonely Runner Betti simulation"
    )
    ap.add_argument("--k", type=int, default=3, help="Number of runners")
    ap.add_argument("--speeds", type=float, nargs="+", default=None, help="Runner speeds")
    ap.add_argument("--time", type=float, default=5.0, help="Total simulation time")
    ap.add_argument("--steps", type=int, default=1000, help="Number of time steps")
    ap.add_argument("--N", type=int, default=512, help="Spatial discretisation points")
    ap.add_argument("--outdir", default=".", help="Output directory")
    ap.add_argument("--name", default="lonely_runner", help="Run name")
    ap.add_argument("--test", action="store_true", help="Run test suite")
    args = ap.parse_args()

    if args.test:
        _test()
        return

    speeds = args.speeds
    if speeds is None:
        speeds = list(range(args.k))
    elif len(speeds) != args.k:
        print(f"Error: --speeds must have {args.k} values for k={args.k}")
        sys.exit(1)

    config = {
        "k": args.k,
        "speeds": speeds,
        "description": f"k={args.k} speeds={speeds}",
    }
    run_case(
        args.name, config,
        total_time=args.time,
        time_steps=args.steps,
        N=args.N,
        outdir=args.outdir,
    )


if __name__ == "__main__":
    main()