import json
from pathlib import Path
import numpy as np
from scipy.stats import norm


def newey_west(samples, lag=20):
    """Columnwise mean, Bartlett HAC variance, SE, z and lower-tail p-value."""
    x = np.asarray(samples, dtype=float)
    if x.ndim != 2 or x.shape[0] < 2:
        raise ValueError("samples must have shape [time, components], T >= 2")
    t = x.shape[0]
    lag = max(0, min(int(lag), t - 1))
    mean = x.mean(axis=0)
    c = x - mean
    hac = np.mean(c * c, axis=0)
    for k in range(1, lag + 1):
        gamma = np.mean(c[k:] * c[:-k], axis=0)
        hac += 2.0 * (1.0 - k / (lag + 1.0)) * gamma
    return mean, np.maximum(hac, 1e-12)


def confidence_scores(samples, delta, lag=20):
    mean, hac = newey_west(samples, lag)
    se = np.sqrt(hac / samples.shape[0])
    z = (mean - delta) / se
    return {"mean": mean, "hac": hac, "se": se, "z": z, "p": norm.cdf(z)}


def select_components(samples, delta, keep, lag=20):
    s = confidence_scores(samples, delta, lag)
    # Large p means evidence is not against useful contribution.
    indices = np.argsort(s["p"])[-keep:]
    return np.sort(indices), s


def ar1(mu, rho, sigma, t, rng):
    e = rng.normal(0.0, sigma * np.sqrt(1.0 - rho * rho), size=t)
    y = np.empty(t)
    y[0] = mu + rng.normal(0.0, sigma)
    for i in range(1, t):
        y[i] = mu + rho * (y[i - 1] - mu) + e[i]
    return y


def mechanism_sweep(seed=7):
    """Numerical checks of the three quantitative predictions in the idea."""
    rng = np.random.default_rng(seed)
    # Prediction 1: SE approximately proportional to T^-1/2 for iid observations.
    ts = np.array([100, 200, 400, 800, 1600])
    se_obs = []
    for t in ts:
        vals = np.stack([ar1(0.3, 0.0, 1.0, t, rng) for _ in range(250)], axis=1)
        se_obs.append(float(np.mean(confidence_scores(vals, 0.0, lag=0)["se"])))
    slope = float(np.polyfit(np.log(ts), np.log(se_obs), 1)[0])

    # Prediction 2: AR(1) long-run variance inflation is (1+rho)/(1-rho).
    rhos = np.array([0.0, 0.3, 0.6, 0.8])
    infl_obs = []
    for rho in rhos:
        vals = np.stack([ar1(0.0, rho, 1.0, 3000, rng) for _ in range(35)], axis=1)
        infl_obs.append(float(np.mean(confidence_scores(vals, 0.0, lag=200)["hac"])))
    infl_obs = np.array(infl_obs)
    infl_obs /= infl_obs[0]
    infl_pred = (1 + rhos) / (1 - rhos)
    # normalize prediction consistently to the empirically estimated rho=0 baseline
    infl_pred /= infl_pred[0]

    # Prediction 3: when true mean falls below Delta, rejection probability rises.
    delta = 0.25
    means = np.array([0.45, 0.30, 0.20, 0.05])
    reject = []
    for mu in means:
        count = 0
        for _ in range(180):
            x = ar1(mu, 0.5, 1.0, 300, rng)[:, None]
            if confidence_scores(x, delta, lag=20)["p"][0] < 0.05:
                count += 1
        reject.append(count / 180.0)
    return {
        "prediction_1_se_slope": {"predicted": -0.5, "observed": slope, "se_by_T": se_obs},
        "prediction_2_hac_inflation": {"rho": rhos.tolist(), "predicted": infl_pred.tolist(), "observed": infl_obs.tolist()},
        "prediction_3_rejection_power": {"true_mean": means.tolist(), "delta": delta, "rejection_rate": reject},
    }


def pruning_comparison(seed=11, repeats=30):
    """Fixed-budget proxy: compare history inference with latest-minibatch ranking."""
    n, keep, t = 24, 10, 160
    truth = np.linspace(0.04, 0.48, n)
    stat_utils, latest_utils = [], []
    for r in range(repeats):
        rng = np.random.default_rng(seed + r)
        histories = np.stack([ar1(mu, 0.75, 0.55, t, rng) for mu in truth], axis=1)
        stat_idx, stat = select_components(histories, 0.20, keep, lag=20)
        latest_idx = np.argsort(histories[-1])[-keep:]
        stat_utils.append(float(truth[stat_idx].sum()))
        latest_utils.append(float(truth[latest_idx].sum()))
    # Include one run's allocation and p-values for inspectability.
    rng = np.random.default_rng(seed)
    histories = np.stack([ar1(mu, 0.75, 0.55, t, rng) for mu in truth], axis=1)
    stat_idx, stat = select_components(histories, 0.20, keep, lag=20)
    latest_idx = np.argsort(histories[-1])[-keep:]
    return {
        "keep": keep, "delta": 0.20, "repeats": repeats,
        "statistical_indices": stat_idx.tolist(),
        "latest_indices": np.sort(latest_idx).tolist(),
        "statistical_true_utility_mean": float(np.mean(stat_utils)),
        "latest_true_utility_mean": float(np.mean(latest_utils)),
        "statistical_win_rate": float(np.mean(np.array(stat_utils) > np.array(latest_utils))),
        "all_true_utility": float(np.sort(truth)[-keep:].sum()),
        "statistical_pvalues": stat["p"].tolist(),
    }


def main():
    result = {"mechanism": mechanism_sweep(), "pruning": pruning_comparison()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()