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

SEED = 2136
rng = np.random.default_rng(SEED)


def sample_deficits(n_trials, n_candidates, dim, kappa=2.0, curvature=1.0, batch=2000):
    """Q(u*)-max Q(U_i) for U uniform in [-1,1]^d and Q=-curvature*||u||^k."""
    out = np.empty(n_trials)
    for lo in range(0, n_trials, batch):
        m = min(batch, n_trials - lo)
        u = rng.uniform(-1.0, 1.0, size=(m, n_candidates, dim))
        r = np.linalg.norm(u, axis=-1)
        out[lo:lo+m] = curvature * np.min(r ** kappa, axis=1)
    return out


def fit_log_slope(ns, means):
    return float(np.polyfit(np.log(ns), np.log(means), 1)[0])


def main():
    # Prediction 1: E[Delta_N] scales as N^(-kappa/d).
    # Prediction 2: doubling N multiplies error by 2^(-kappa/d).
    # We check both d=1 and d=2, with kappa=2.
    n_trials = 12000
    kappa = 2.0
    scaling = {}
    for dim in (1, 2):
        ns = np.array([4, 8, 16, 32, 64, 128, 256])
        means = np.array([sample_deficits(n_trials, int(n), dim, kappa).mean() for n in ns])
        slope = fit_log_slope(ns, means)
        ratios = means[1:] / means[:-1]
        predicted_slope = -kappa / dim
        predicted_ratio = 2 ** predicted_slope
        scaling[str(dim)] = {
            "N": ns.tolist(), "mean_deficit": means.tolist(),
            "observed_loglog_slope": slope, "predicted_loglog_slope": predicted_slope,
            "observed_doubling_ratio_mean": float(ratios.mean()),
            "predicted_doubling_ratio": predicted_ratio,
            "doubling_ratios": ratios.tolist(),
        }

    # Prediction 3: scheduler N = ceil((C / ((1-beta) eps))^(d/kappa))
    # has N proportional to eps^(-d/kappa), and achieves the requested error scale.
    dim = 2
    beta = 0.9
    epsilons = np.array([0.04, 0.02, 0.01, 0.005])
    n_cal = 64
    cal = sample_deficits(30000, n_cal, dim, kappa)
    C_hat = float(cal.mean() * n_cal ** (kappa / dim))
    scheduled_ns = np.ceil((C_hat / ((1-beta) * epsilons)) ** (dim/kappa)).astype(int)
    scheduled_ns = np.maximum(scheduled_ns, 1)
    measured = np.array([sample_deficits(16000, int(n), dim, kappa).mean() for n in scheduled_ns])
    value_bound_scale = measured / (1-beta)
    eps_slope = fit_log_slope(epsilons, scheduled_ns)
    # For this geometry the asymptotic constant is known: area(pi r^2)/(area(box)) = pi r^2/4,
    # hence E[min ||U||^2] ~ 4/(pi N). This is also a check that C_hat is calibrated.
    C_exact = 4.0 / math.pi
    expected_at_schedule = C_exact * scheduled_ns ** (-dim / kappa) / (1-beta)
    scheduler = {
        "beta": beta, "C_hat": C_hat, "C_exact_asymptotic": C_exact,
        "C_hat_over_C_exact": C_hat / C_exact, "epsilons": epsilons.tolist(),
        "scheduled_N": scheduled_ns.tolist(), "measured_mean_delta": measured.tolist(),
        "measured_delta_over_1_minus_beta": value_bound_scale.tolist(),
        "expected_value_proxy_at_schedule": expected_at_schedule.tolist(),
        "mean_expected_proxy_over_epsilon": float(np.mean(expected_at_schedule / epsilons)),
        "observed_loglog_N_vs_epsilon_slope": eps_slope,
        "predicted_loglog_N_vs_epsilon_slope": -dim/kappa,
        "fraction_meeting_epsilon_bound": float(np.mean(value_bound_scale <= epsilons)),
    }

    # Mini comparison: fixed pools versus the adaptive pool over heterogeneous states.
    # Curvature changes the local C, so fixed N is deliberately wasteful on easy states.
    n_states = 3000
    curvatures = rng.uniform(0.35, 2.0, size=n_states)
    eps = 0.015
    fixed_ns = [8, 32, 118, 128]
    rows = []
    # Use common state curvature and independent candidate draws; report action error
    # and Bellman-value error proxy Delta/(1-beta), with mean Q evaluations.
    for n in fixed_ns:
        ds = np.array([sample_deficits(1, n, dim, kappa, float(a))[0] for a in curvatures])
        rows.append({"method": f"fixed_N={n}", "mean_evaluations": n,
                     "mean_delta": float(ds.mean()), "mean_value_error_proxy": float((ds/(1-beta)).mean()),
                     "fraction_value_proxy_le_eps": float(np.mean(ds/(1-beta) <= eps))})
    # Estimate C per state from a cheap N=8 pilot, then use the formula and cap at 128.
    pilot_n = 8
    pilot = np.array([sample_deficits(1, pilot_n, dim, kappa, float(a))[0] for a in curvatures])
    local_C = pilot * pilot_n ** (kappa/dim)
    adaptive_n = np.ceil((local_C / ((1-beta)*eps)) ** (dim/kappa)).astype(int)
    adaptive_n = np.clip(adaptive_n, pilot_n, 128)
    adaptive_ds = np.array([sample_deficits(1, int(n), dim, kappa, float(a))[0]
                            for a, n in zip(curvatures, adaptive_n)])
    rows.append({"method": "adaptive_scheduler", "mean_evaluations": float(adaptive_n.mean()),
                 "mean_delta": float(adaptive_ds.mean()), "mean_value_error_proxy": float((adaptive_ds/(1-beta)).mean()),
                 "fraction_value_proxy_le_eps": float(np.mean(adaptive_ds/(1-beta) <= eps)),
                 "pool_min": int(adaptive_n.min()), "pool_median": float(np.median(adaptive_n)),
                 "pool_max": int(adaptive_n.max())})

    result = {"seed": SEED, "theory": {"action": "uniform box, Q(u)=-||u||^2", "kappa": kappa},
              "scaling_checks": scaling, "scheduler_scaling_check": scheduler,
              "mini_comparison": rows}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()