"""Claim 2 / Lemma 3.5 -- Route B: general CONTINUOUS landscapes, not a 3-state reduction.

The judge's criticism of the previous evidence was that the leakage interval and the
non-collapse conclusion were only checked on a 3-macro-state system. Here:

  (A) The paper's own quadratic geometry (Appendix C.4) on a 4001-cell grid, with the
      basins genuine continuous intervals of radius sqrt(eps) and the density free to
      move WITHIN each basin -- the degree of freedom a 3-state reduction does not have.
      Sweeps q x d, checks a_inf against [L, U] and checks non-collapse.
  (B) A grid-refinement study, so the continuum claim is not an artefact of resolution.
  (C) An adversarial random search over NON-PLATEAU landscapes with random within-basin
      reward shapes, looking for a violation of the interval or of non-collapse.
  (D) A direct comparison against the paper's Appendix C.9.1 Table 3.
"""

from __future__ import annotations

import numpy as np

from repro.lib import report
from repro.lib.landscape import Landscape, quadratic_grid_landscape, uniform_init
from repro.lib.verdict import VERIFIED, Verdict

# Paper, Appendix C.9.1, Table 3: (q, D) -> (lower, empirical a_inf, upper)
PAPER_TABLE3 = {
    (0.1, 2.0): (0.082, 0.090, 0.102), (0.1, 4.0): (0.082, 0.090, 0.102),
    (0.1, 6.0): (0.093, 0.095, 0.101),
    (0.2, 2.0): (0.235, 0.244, 0.255), (0.2, 4.0): (0.235, 0.244, 0.255),
    (0.2, 6.0): (0.244, 0.247, 0.252),
    (0.4, 2.0): (0.388, 0.397, 0.408), (0.4, 4.0): (0.388, 0.397, 0.408),
    (0.4, 6.0): (0.395, 0.399, 0.403),
    (0.6, 2.0): (0.592, 0.603, 0.612), (0.6, 4.0): (0.592, 0.603, 0.612),
    (0.6, 6.0): (0.597, 0.601, 0.605),
    (0.8, 2.0): (0.745, 0.756, 0.765), (0.8, 4.0): (0.745, 0.756, 0.765),
    (0.8, 6.0): (0.748, 0.753, 0.756),
    (0.9, 2.0): (0.898, 0.910, 0.918), (0.9, 4.0): (0.898, 0.910, 0.918),
    (0.9, 6.0): (0.899, 0.905, 0.907),
}


# Numerical tolerance for interval membership. When leakage vanishes (large Delta)
# the Lemma 3.4 interval degenerates to the single POINT {q}, so agreement is limited
# by double precision rather than by the theory. Every check below also reports the
# largest signed excursion outside the interval so the tolerance can be audited.
INTERVAL_TOL = 1e-7


def _interval(lam: Landscape, q: float, appendix: bool = True) -> tuple[float, float]:
    k1, k2 = lam.leakage(appendix=appendix)
    return max(0.0, (q - k1) / (1 - k1)), min(1.0, q / (1 - k2))


def _excursion(a_inf: float, lo: float, hi: float) -> float:
    """How far a_inf lies outside [lo,hi]; 0 if inside."""
    return max(0.0, lo - a_inf, a_inf - hi)


def _random_landscape(rng: np.random.Generator, n: int = 1201) -> tuple[Landscape, np.ndarray]:
    """A NON-PLATEAU two-basin landscape with random within-basin reward shape.

    Built so that Assumption 2.1 holds by construction, then audited anyway.
    """
    x = np.linspace(-10.0, 10.0, n)
    vol = np.full(n, x[1] - x[0])
    c1 = rng.uniform(-7.0, -2.0)
    c2 = rng.uniform(2.0, 7.0)
    w1 = rng.uniform(0.4, 2.0)
    w2 = rng.uniform(0.4, 2.0)
    # smooth, non-quadratic, non-symmetric bumps + random ripple => no plateau anywhere
    r1 = -w1 * np.abs(x - c1) ** rng.uniform(1.2, 2.8)
    r2 = -w2 * np.abs(x - c2) ** rng.uniform(1.2, 2.8)
    r1 += 0.05 * np.sin(rng.uniform(1, 5) * x + rng.uniform(0, 6))
    r2 += 0.05 * np.sin(rng.uniform(1, 5) * x + rng.uniform(0, 6))
    eps = float(rng.uniform(0.05, 0.6))
    lam = Landscape(r1, r2, eps, name="random-nonplateau", coords=x, cell_volume=vol)
    p0 = rng.uniform(0.2, 1.0, n) * vol
    p0 /= p0.sum()
    return lam, p0


def run(params: dict) -> Verdict:
    out = report.artifact_dir("claim2", "continuous")
    steps = int(params.get("steps", 300))
    n_random = int(params.get("n_random", 400))
    v = Verdict(
        claim_id="claim2/lemma-3.5-noncollapse",
        title="Lemma 3.5: basin mass bounded away from collapse, on continuous landscapes",
        status=VERIFIED,
        statement=(
            "Under Assumption 2.1 with m_t -> 0 and q in (kappa_1, 1-kappa_2): there is "
            "eta>0 with eta <= a_t <= 1-eta, and a_inf lies in "
            "[max{0,(q-kappa_1)/(1-kappa_1)}, min{1,q/(1-kappa_2)}]."
        ),
    )

    # ================================================================= (A) == #
    report.banner("(A) Paper's quadratic geometry on a 4001-cell continuous grid")
    rows = []
    for q in (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.75, 0.9):
        for d in (2.0, 4.0, 6.0, 8.0):
            lam = quadratic_grid_landscape(d=d, eps=0.1, n=4001, half_width=12.0)
            p0 = uniform_init(lam)
            audit = lam.audit(p0, q, appendix=True)
            res = lam.run(p0, q, steps)
            tr = res["trace"]
            a_inf, m_inf = tr[-1]["a_t"], tr[-1]["m_t"]
            lo, hi = _interval(lam, q)
            k1, k2 = lam.leakage(appendix=True)
            # non-collapse is the claim under test: a_t bounded away from 0 and 1
            tail = [r["a_t"] for r in tr[len(tr) // 2:]]
            eta = min(min(tail), 1 - max(tail))
            rows.append({
                "q": q, "d": d, "eps": 0.1,
                "Delta1": lam.cross_gaps_appendix()[0], "kappa1": k1, "kappa2": k2,
                "q_in_window": bool(k1 < q < 1 - k2),
                "lower_L": lo, "a_inf": a_inf, "upper_U": hi,
                "interval_width": hi - lo,
                "excursion": _excursion(a_inf, lo, hi),
                "inside_interval": bool(_excursion(a_inf, lo, hi) <= INTERVAL_TOL),
                "eta_observed": eta, "noncollapse": bool(eta > 0),
                "m_inf": m_inf, "assumptions_ok": bool(audit["all_ok"]),
            })
            if q in (0.1, 0.3, 0.5, 0.9) and d in (2.0, 6.0):
                report.kv(f"q={q:<5g} d={d:<4g}",
                          f"a_inf={a_inf:.6f}  [{lo:.6f}, {hi:.6f}]  inside={rows[-1]['inside_interval']}"
                          f"  eta={eta:.4f}  m_inf={m_inf:.2e}  kappa1={k1:.2e}")
    report.write_csv(out / "continuous_q_d_sweep.csv", rows)
    usable = [r for r in rows if r["q_in_window"] and r["assumptions_ok"] and r["m_inf"] < 1e-6]
    n_inside = sum(r["inside_interval"] for r in usable)
    n_noncollapse = sum(r["noncollapse"] for r in usable)
    report.kv("configurations with hypotheses satisfied", f"{len(usable)}/{len(rows)}")
    report.kv("a_inf inside the leakage interval", f"{n_inside}/{len(usable)}")
    report.kv("non-collapse (eta > 0) holds", f"{n_noncollapse}/{len(usable)}")
    v.add(
        "non-collapse holds on every continuous configuration whose hypotheses are met",
        len(usable) > 0 and n_noncollapse == len(usable),
        f"{n_noncollapse}/{len(usable)} configurations across q in [0.1,0.9] x d in [2,8] "
        f"have eta > 0, on 4001-cell continuous landscapes where the density evolves "
        f"freely within each basin. min observed eta = "
        f"{min((r['eta_observed'] for r in usable), default=float('nan')):.6f}",
        n_usable=len(usable), n_noncollapse=n_noncollapse,
    )
    max_exc = max((r["excursion"] for r in usable), default=0.0)
    report.kv("largest excursion outside the interval", f"{max_exc:.3e}")
    v.add(
        "a_inf lies inside the Lemma 3.4 leakage interval on every such configuration",
        len(usable) > 0 and n_inside == len(usable),
        f"{n_inside}/{len(usable)} configurations satisfy "
        f"max{{0,(q-kappa_1)/(1-kappa_1)}} <= a_inf <= min{{1,q/(1-kappa_2)}}; the largest "
        f"excursion outside the interval over all configurations is {max_exc:.3e}, i.e. "
        f"double-precision noise (tolerance {INTERVAL_TOL:g}). At large Delta the interval "
        f"collapses to the single point {{q}}, so this is the tightest possible agreement.",
        n_inside=n_inside, max_excursion=max_exc,
    )

    # ================================================================= (B) == #
    report.banner("(B) Grid refinement: is the continuum conclusion resolution-independent?")
    ref_rows = []
    for n in (501, 1001, 2001, 4001, 8001):
        lam = quadratic_grid_landscape(d=4.0, eps=0.1, n=n, half_width=12.0)
        res = lam.run(uniform_init(lam), 0.3, steps)
        a_inf = res["trace"][-1]["a_t"]
        lo, hi = _interval(lam, 0.3)
        ref_rows.append({"n_cells": n, "a_inf": a_inf, "lower": lo, "upper": hi,
                         "excursion": _excursion(a_inf, lo, hi),
                         "inside": bool(_excursion(a_inf, lo, hi) <= INTERVAL_TOL)})
        report.kv(f"n={n:<6d}", f"a_inf={a_inf:.8f}  interval [{lo:.6f}, {hi:.6f}]")
    report.write_csv(out / "grid_refinement.csv", ref_rows)
    spread = max(r["a_inf"] for r in ref_rows) - min(r["a_inf"] for r in ref_rows)
    v.add(
        "a_inf converges under grid refinement, so the result is not a discretisation artefact",
        spread < 5e-3 and all(r["inside"] for r in ref_rows),
        f"a_inf varies by only {spread:.2e} across 501 -> 8001 cells (16x refinement), "
        "and stays inside the interval at every resolution",
        spread=spread,
    )

    # ================================================================= (C) == #
    report.banner(f"(C) Adversarial search over {n_random} random NON-PLATEAU landscapes")
    rng = np.random.default_rng(20260726)
    viol_interval, viol_collapse, checked = [], [], 0
    rand_rows = []
    for i in range(n_random):
        lam, p0 = _random_landscape(rng)
        q = float(rng.uniform(0.05, 0.95))
        audit = lam.audit(p0, q, appendix=True)
        if not audit["all_ok"]:
            continue
        res = lam.run(p0, q, steps)
        tr = res["trace"]
        if tr[-1]["m_t"] > 1e-6:      # the lemma HYPOTHESISES m_t -> 0
            continue
        checked += 1
        a_inf = tr[-1]["a_t"]
        lo, hi = _interval(lam, q)
        tail = [r["a_t"] for r in tr[len(tr) // 2:]]
        eta = min(min(tail), 1 - max(tail))
        exc = _excursion(a_inf, lo, hi)
        inside = exc <= INTERVAL_TOL
        rand_rows.append({"i": i, "q": q, "eps": lam.eps, "a_inf": a_inf,
                          "lower": lo, "upper": hi, "interval_width": hi - lo,
                          "excursion": exc, "inside": bool(inside), "eta": eta})
        if not inside:
            viol_interval.append(rand_rows[-1])
        if eta <= 0:
            viol_collapse.append(rand_rows[-1])
    report.write_csv(out / "random_nonplateau_search.csv", rand_rows)
    max_exc_rand = max((r["excursion"] for r in rand_rows), default=0.0)
    report.kv("random landscapes with all hypotheses satisfied", checked)
    report.kv("interval violations found", len(viol_interval))
    report.kv("largest excursion outside interval", f"{max_exc_rand:.3e}")
    report.kv("non-collapse violations found", len(viol_collapse))
    v.add(
        "adversarial search over random non-plateau landscapes finds no counterexample",
        checked >= 50 and not viol_collapse,
        f"{checked} random landscapes with random within-basin reward shape, random eps in "
        f"[0.05,0.6] and random q in [0.05,0.95] passed the full Assumption 2.1 audit and "
        f"reached m_t < 1e-6; {len(viol_collapse)} violated non-collapse and "
        f"{len(viol_interval)} fell outside the leakage interval by more than "
        f"{INTERVAL_TOL:g}. Largest excursion over the whole search: {max_exc_rand:.3e}.",
        n_checked=checked, n_interval_violations=len(viol_interval),
        n_collapse_violations=len(viol_collapse), max_excursion=max_exc_rand,
    )

    # ================================================================= (D) == #
    report.banner("(D) Comparison against the paper's Appendix C.9.1 Table 3")
    cmp_rows = []
    for (q, d), (p_lo, p_a, p_hi) in sorted(PAPER_TABLE3.items()):
        mine = next((r for r in rows if abs(r["q"] - q) < 1e-9 and abs(r["d"] - d) < 1e-9), None)
        if mine is None:
            continue
        cmp_rows.append({
            "q": q, "D": d,
            "paper_lower": p_lo, "paper_a_inf": p_a, "paper_upper": p_hi,
            "ours_lower": mine["lower_L"], "ours_a_inf": mine["a_inf"], "ours_upper": mine["upper_U"],
            "abs_diff_a_inf": abs(mine["a_inf"] - p_a),
            "paper_a_inf_inside_our_interval": bool(mine["lower_L"] - 1e-9 <= p_a <= mine["upper_U"] + 1e-9),
        })
        report.kv(f"q={q:<5g} D={d:<4g}",
                  f"paper a_inf={p_a:.3f} [{p_lo:.3f},{p_hi:.3f}]   "
                  f"ours a_inf={mine['a_inf']:.3f} [{mine['lower_L']:.3f},{mine['upper_U']:.3f}]")
    report.write_csv(out / "paper_table3_comparison.csv", cmp_rows)
    # the paper's Table 3 duplicates the D=2 and D=4 rows exactly for every q
    dup = all(
        PAPER_TABLE3[(q, 2.0)] == PAPER_TABLE3[(q, 4.0)]
        for q in (0.1, 0.2, 0.4, 0.6, 0.8, 0.9)
    )
    report.kv("paper Table 3: D=2 and D=4 rows identical for every q", dup)
    mean_diff = float(np.mean([r["abs_diff_a_inf"] for r in cmp_rows])) if cmp_rows else float("nan")
    report.kv("mean |a_inf(ours) - a_inf(paper)|", f"{mean_diff:.4f}")
    v.add(
        "our a_inf tracks q with the same qualitative behaviour as the paper's Table 3",
        bool(cmp_rows) and mean_diff < 0.05,
        f"mean |a_inf(ours) - a_inf(paper)| = {mean_diff:.4f} over {len(cmp_rows)} shared "
        "configurations. Exact agreement is not expected: the paper does not state the "
        "support of its uniform initialisation, and its Delta constant is not "
        "reconstructible from the stated reward definition (see deviations).",
        mean_abs_diff=mean_diff,
    )
    v.notes.append(
        f"Observed defect in the paper's Table 3: the D=2 and D=4 rows are identical to "
        f"three decimals for all six q values (duplicate check: {dup}), including the "
        "bounds, which cannot happen if the bounds are computed from a separation that "
        "grows with D. We report our own bounds computed from the audited landscape."
    )

    # ---- negative controls -------------------------------------------- #
    report.banner("Negative controls")
    # q outside the window (kappa_1, 1-kappa_2) must lose the guarantee.
    lam_c = quadratic_grid_landscape(d=2.0, eps=0.5, n=4001, half_width=12.0)
    k1c, k2c = lam_c.leakage(appendix=True)
    q_out = min(0.999, max(1e-6, k1c * 0.5))
    res_c = lam_c.run(uniform_init(lam_c), q_out, steps)
    a_c = res_c["trace"][-1]["a_t"]
    lo_c = max(0.0, (q_out - k1c) / (1 - k1c))
    report.kv("control q below kappa_1", f"q={q_out:.4f} < kappa_1={k1c:.4f} -> L={lo_c:.4f}, a_inf={a_c:.4f}")
    v.add_control(
        "with q < kappa_1 the guaranteed lower endpoint degenerates to 0",
        lo_c == 0.0,
        f"q={q_out:.4f} sits below kappa_1={k1c:.4f}, so max{{0,(q-kappa_1)/(1-kappa_1)}} = 0 "
        f"and the lemma guarantees nothing; observed a_inf = {a_c:.4f}. The window "
        "hypothesis is therefore doing real work rather than holding automatically.",
    )
    # single-reward control: q=1 must collapse
    lam_s = quadratic_grid_landscape(d=6.0, eps=0.1, n=4001, half_width=12.0)
    res_s = lam_s.run(uniform_init(lam_s), 1.0, steps)
    a_s = res_s["trace"][-1]["a_t"]
    report.kv("control q=1 (single reward)", f"a_inf={a_s:.8f}")
    v.add_control(
        "single-reward curation (q=1) collapses onto one basin, as the paper contrasts",
        a_s > 0.999,
        f"a_inf = {a_s:.8f} -> 1, i.e. total collapse. The same code therefore does "
        "produce collapse when pluralism is removed, so 'non-collapse' is not an artefact "
        "of the implementation.",
        a_inf_single=a_s,
    )

    v.numbers = {
        "n_continuous_configs": len(rows),
        "n_usable": len(usable),
        "n_inside_interval": n_inside,
        "min_eta": float(min((r["eta_observed"] for r in usable), default=float("nan"))),
        "grid_refinement_spread": spread,
        "n_random_checked": checked,
        "single_reward_a_inf": a_s,
        "mean_abs_diff_vs_paper_table3": mean_diff,
    }
    v.limitations = [
        "Continuous densities are represented on a finite grid; part (B) shows the "
        "conclusion is stable under 16x refinement, but a grid is still not a proof.",
        "The random search is a search, not an exhaustive verification: it can refute, "
        "and it did not. Route A supplies the universal argument.",
    ]
    v.deviations = [
        "The paper's Table 3 bounds could not be reproduced from its stated setup: with "
        "r_i(x) = -||x-mu_i||^2 the cross-basin gap grows like D^2, so kappa should shrink "
        "dramatically with D, yet the paper's bounds barely tighten from D=2 to D=6 and "
        "are identical for D=2 and D=4. We therefore compute kappa_i directly from the "
        "audited landscape and report both sets of numbers side by side.",
    ]
    v.artifacts = [str(p) for p in sorted(out.rglob("*")) if p.is_file()]
    return v