import json
from pathlib import Path
import numpy as np

SEED = 2346


def projected_update(lam, vbar, alpha, tau, lam_max):
    return float(np.clip(lam + alpha * (vbar - tau), 0.0, lam_max))


def policy_target(lam, c=0.8, temperature=0.15):
    z = np.clip((c - lam) / temperature, -60.0, 60.0)
    return 1.0 / (1.0 + np.exp(-z))


def run(alpha, steps=400, beta=0.2, tau=0.2, lam_max=5.0,
        batch=64, noisy=False, fixed_penalty=None, seed=SEED):
    """Toy policy/dual loop. q is terminal violation probability.
    The policy relaxes toward the reward-minus-penalty optimum q*(lambda).
    """
    rng = np.random.default_rng(seed)
    q, lam = 0.8, 0.0
    qs, vs, ls, drifts = [], [], [], []
    for _ in range(steps):
        effective = lam if fixed_penalty is None else fixed_penalty
        q += beta * (policy_target(effective) - q)
        vbar = float(rng.binomial(batch, np.clip(q, 0.0, 1.0)) / batch) if noisy else q
        old = lam
        if fixed_penalty is None:
            lam = projected_update(lam, vbar, alpha, tau, lam_max)
        qs.append(q); vs.append(vbar); ls.append(lam); drifts.append(lam - old)
    return {"q": np.asarray(qs), "v": np.asarray(vs),
            "lam": np.asarray(ls), "drift": np.asarray(drifts)}


def math_check():
    r = np.random.default_rng(SEED)
    max_err = 0.0
    for _ in range(10000):
        lam, v, a, tau, cap = r.uniform(0, 10), r.uniform(0, 2), r.uniform(0, 3), r.uniform(0, 1), r.uniform(.1, 10)
        expected = min(cap, max(0.0, lam + a * (v - tau)))
        max_err = max(max_err, abs(projected_update(lam, v, a, tau, cap) - expected))
    x = run(alpha=.08, steps=1500, beta=.2, tau=.2, lam_max=5, noisy=False)
    interior = (x["lam"] > 1e-5) & (x["lam"] < 4.99999)
    drift_error = abs(np.mean(x["drift"][interior]) - .08 * np.mean((x["v"] - .2)[interior]))
    return {"max_projection_formula_error": max_err,
            "interior_drift_identity_error": float(drift_error),
            "final_lambda": float(x["lam"][-1]), "final_violation": float(x["v"][-1])}


def sweep():
    out = {}
    # Prediction 1: with fixed positive violation error, lambda reaches cap in
    # approximately (cap-lambda0)/(alpha*(v-tau)) steps.
    cap_rows = []
    for alpha in [.02, .05, .1, .2]:
        x = run(alpha, steps=300, beta=0.0, tau=0.1, lam_max=2.0, noisy=False)
        # beta=0 leaves q at .8, so v-tau=.7 exactly.
        hit = np.flatnonzero(x["lam"] >= 2.0 - 1e-10)
        observed = int(hit[0] + 1) if len(hit) else None
        predicted = int(np.ceil(2.0 / (alpha * .7)))
        cap_rows.append({"alpha": alpha, "predicted_hit_step": predicted,
                         "observed_hit_step": observed})
    out["cap_scaling"] = cap_rows

    # Prediction 2: away from boundaries, average violation error is zero in
    # a bounded drift-balanced regime.
    balance_rows = []
    for alpha in [.01, .05, .1, .2, .4]:
        x = run(alpha, steps=2000, beta=.2, tau=.2, lam_max=5, noisy=False)
        tail = slice(1000, None)
        balance_rows.append({"alpha": alpha,
                             "tail_mean_violation": float(np.mean(x["v"][tail])),
                             "tail_abs_mean_error": float(abs(np.mean(x["v"][tail]) - .2)),
                             "tail_lambda_mean": float(np.mean(x["lam"][tail])),
                             "tail_lambda_std": float(np.std(x["lam"][tail]))})
    out["drift_balance"] = balance_rows

    # Prediction 3: positive persistent drift saturates; negative drift projects
    # to zero. This tests both projection boundaries.
    boundary_rows = []
    for v in [0.0, .2, .8]:
        x = run(.1, steps=100, beta=0.0, tau=.2, lam_max=1.0, noisy=False)
        # beta=0 gives v=.8; use direct recurrence for arbitrary constant v.
        lam = 0.; hist=[]
        for _ in range(100):
            lam = projected_update(lam, v, .1, .2, 1.)
            hist.append(lam)
        boundary_rows.append({"constant_violation": v, "final_lambda": hist[-1],
                              "expected": 1.0 if v > .2 else (0.0 if v < .2 else 0.0)})
    out["projection_boundaries"] = boundary_rows

    # Secondary baseline: best fixed penalty chosen from a small grid, versus
    # adaptive dual controller at equal iterations.
    adaptive = run(.08, steps=400, beta=.2, tau=.2, lam_max=5, noisy=True, batch=64)
    fixed = []
    for p in np.linspace(0, 2, 21):
        z = run(0, steps=400, beta=.2, tau=.2, lam_max=5, fixed_penalty=float(p), noisy=True, batch=64, seed=SEED)
        fixed.append((abs(np.mean(z["v"][-100:]) - .2), p, np.mean(z["v"][-100:])))
    best = min(fixed)
    out["comparison"] = {
        "adaptive_tail_violation": float(np.mean(adaptive["v"][-100:])),
        "adaptive_tail_abs_error": float(abs(np.mean(adaptive["v"][-100:]) - .2)),
        "adaptive_tail_lambda": float(np.mean(adaptive["lam"][-100:])),
        "best_fixed_penalty": best[1], "best_fixed_tail_violation": best[2],
        "best_fixed_tail_abs_error": best[0]
    }
    return out


def main():
    result = {"seed": SEED, "math_check": math_check(), "sweeps": sweep()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()