import json
import math
from pathlib import Path
import numpy as np

SEED = 1268


def rho_curve(m, benefit=0.28, length=2.0, stale=0.0, onset=3, p0=0.04):
    val = p0 + benefit * (1.0 - math.exp(-m / length)) - stale * max(0, m - onset)
    return float(np.clip(val, 1e-5, 0.95))


def escape(rho, M):
    return 1.0 - (1.0 - rho) ** M


def math_verification():
    rng = np.random.default_rng(SEED)
    # Prediction 1: exact M-fold escape formula.
    rho = 0.173
    formula = []
    for M in [1, 2, 4, 8, 16, 32]:
        n = 120000
        observed = rng.binomial(n, escape(rho, M)) / n
        predicted = escape(rho, M)
        formula.append({"M": M, "predicted": predicted, "observed": float(observed),
                        "abs_error": float(abs(observed - predicted))})

    # Prediction 2: monotone learning implies increasing memory improves escape.
    mono_r = [escape(rho_curve(m, benefit=.30, length=2.5, stale=0), 8)
              for m in range(1, 17)]
    monotone_fraction = float(np.mean(np.diff(mono_r) > 0))

    # Prediction 3: staleness creates a finite optimum; more drift shifts it left.
    opt_rows = []
    for stale in [0.0, .005, .015, .03, .05, .08]:
        rs = [escape(rho_curve(m, benefit=.30, length=2.5, stale=stale), 8)
              for m in range(1, 17)]
        # Independent high-sample sweep gives an observed optimum.
        obs = []
        for m in range(1, 17):
            p_m = rho_curve(m, benefit=.30, length=2.5, stale=stale)
            rho_obs = rng.binomial(200000, p_m) / 200000
            obs.append(escape(rho_obs, 8))
        opt_rows.append({"stale": stale, "predicted_opt_m": int(np.argmax(rs) + 1),
                         "observed_opt_m": int(np.argmax(obs) + 1),
                         "best_r": float(max(rs)),
                         "observed_best_r": float(max(obs))})
    # Prediction 4: in the low-rho regime, r is approximately M*rho.
    scaling = []
    for M in [1, 2, 4, 8, 16]:
        r = escape(.002, M)
        scaling.append({"M": M, "r": r, "r_over_Mrho": r / (M * .002)})
    return {"formula_check": formula, "monotone_fraction": monotone_fraction,
            "stale_optima": opt_rows, "low_rho_scaling": scaling}


def run_controller(rng, rounds=80, M=8, tau=.025, probe_every=4, mmax=8,
                    benefit=.30, length=2.5, stale=.045, noise_n=48):
    m = 1
    total_escape = 0
    trajectory = []
    for t in range(rounds):
        # Production event uses exactly M proposals, matching fixed_run.
        p = rho_curve(m, benefit, length, stale)
        total_escape += int(np.any(rng.random(M) < p))
        # Independent validation samples estimate rho for adaptation.
        rho_hat = rng.binomial(noise_n, p) / noise_n
        r_hat = escape(rho_hat, M)
        if t % probe_every == 0:
            estimates = {m: r_hat}
            for neighbor in (m - 1, m + 1):
                if 1 <= neighbor <= mmax:
                    q = rho_curve(neighbor, benefit, length, stale)
                    qhat = rng.binomial(noise_n, q) / noise_n
                    estimates[neighbor] = escape(qhat, M)
            if m + 1 in estimates and estimates[m + 1] - estimates[m] > tau:
                m += 1
            elif m - 1 in estimates and estimates[m] - estimates[m - 1] > tau:
                m -= 1
        trajectory.append(m)
    return {"escapes": total_escape, "mean_m": float(np.mean(trajectory)),
            "final_m": trajectory[-1], "trajectory": trajectory}

def fixed_run(rng, m, rounds=80, M=8, benefit=.30, length=2.5, stale=.045,
              noise_n=48):
    escapes = 0
    p = rho_curve(m, benefit, length, stale)
    for _ in range(rounds):
        rho_hat = rng.binomial(noise_n, p) / noise_n
        # Same estimated-proposal process and M-fold escape metric.
        escapes += int(rng.random() < escape(rho_hat, M))
    return escapes


def comparison():
    settings = {"rounds": 80, "M": 8, "stale": .045, "noise_n": 48}
    vals = {}
    for m in [1, 2, 4, 8, 16]:
        vals[str(m)] = fixed_run(np.random.default_rng(SEED + m), m, **settings)
    adaptive = run_controller(np.random.default_rng(SEED + 100), mmax=8, tau=.025,
                              probe_every=4, **settings)
    return {"fixed_escapes": vals, "adaptive": {k: v for k, v in adaptive.items()
                                                   if k != "trajectory"},
            "adaptive_trajectory": adaptive["trajectory"]}


def main():
    report = {"seed": SEED, "math": math_verification(), "comparison": comparison()}
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()