import json
import numpy as np


def score_distribution(a, u, n, kappa=0.05, epsilon=0.05):
    score = (np.abs(a) + kappa) * u / np.sqrt(np.maximum(n, 1.0))
    q = score / score.sum()
    return (1.0 - epsilon) * q + epsilon / len(q)


def allocate_controller(a, u, budget, batch=16, kappa=0.05, epsilon=0.05, seed=0):
    rng = np.random.default_rng(seed)
    c = len(u)
    n = np.ones(c, dtype=int)
    while n.sum() < budget:
        q = score_distribution(a, u, n, kappa, epsilon)
        take = min(batch, budget - n.sum())
        chosen = rng.choice(c, size=take, p=q)
        n += np.bincount(chosen, minlength=c)
    return n


def allocate_fixed(weights, budget):
    weights = np.asarray(weights, dtype=float)
    raw = budget * weights / weights.sum()
    n = np.floor(raw).astype(int)
    n = np.maximum(n, 1)
    while n.sum() < budget:
        n[np.argmax(raw - n)] += 1
    while n.sum() > budget:
        eligible = np.where(n > 1)[0]
        j = eligible[np.argmax(n[eligible] - raw[eligible])]
        n[j] -= 1
    return n


def variance_prediction(u, n):
    return float(np.sum(np.asarray(u) ** 2 / np.asarray(n)))


def empirical_variance(u, n, reps=4000, seed=11):
    rng = np.random.default_rng(seed)
    errors = np.zeros(reps)
    for c in range(len(u)):
        errors += rng.normal(0.0, u[c], size=(reps, int(n[c]))).mean(axis=1)
    return float(np.var(errors, ddof=1))


def main():
    rng = np.random.default_rng(947)
    C = 32
    budget = 4096
    # Rare families have small probability but remain important signed terms.
    P = np.exp(np.linspace(0.0, -5.0, C)); P /= P.sum()
    s = np.where(np.arange(C) % 3 == 0, -1.0, 1.0)
    E = 0.8 + 1.2 * rng.random(C)
    a = s * P * E
    u = 0.15 * (0.4 + 2.5 * np.sqrt(P / P.min()))
    u *= (0.8 + 0.4 * rng.random(C))

    # Prediction 1: independent-noise variance should equal sum u^2/n.
    n_test = allocate_fixed(np.ones(C), budget)
    pred_var = variance_prediction(u, n_test)
    obs_var = empirical_variance(u, n_test)
    rel_err = abs(obs_var - pred_var) / pred_var

    # Prediction 2: optimal allocation n proportional to u; compare a sweep.
    budgets = [512, 1024, 2048, 4096]
    opt_rows = []
    for b in budgets:
        nu = allocate_fixed(u, b)
        uniform = allocate_fixed(np.ones(C), b)
        opt_rows.append({
            "budget": b,
            "uniform_predicted_variance": variance_prediction(u, uniform),
            "u_proportional_predicted_variance": variance_prediction(u, nu),
            "ratio_uniform_over_u_proportional": variance_prediction(u, uniform) / variance_prediction(u, nu),
        })

    # Prediction 3: controller score should preferentially increase count of
    # high |a|u families, while epsilon enforces a nonzero floor.
    eps_rows = []
    target = (np.abs(a) + 0.05) * u
    for eps in [0.0, 0.01, 0.05, 0.2, 0.5]:
        n_ctrl = allocate_controller(a, u, budget, epsilon=eps, seed=947)
        corr = float(np.corrcoef(n_ctrl, target)[0, 1])
        min_share = float(np.min(n_ctrl / n_ctrl.sum()))
        eps_rows.append({"epsilon": eps, "count_score_correlation": corr, "minimum_family_share": min_share, "predicted_variance": variance_prediction(u, n_ctrl)})

    # Mode-collapse test: deliberately hide a high-uncertainty family from the
    # controller. With epsilon=0 it receives only its initial sample; an
    # exploration mixture should restore samples and reduce true variance.
    hidden = int(np.argmax(u))
    u_bad = u.copy()
    u_bad[hidden] = u.min() * 0.01
    collapse_rows = []
    for eps in [0.0, 0.01, 0.05, 0.2]:
        n_bad = allocate_controller(a, u_bad, budget, epsilon=eps, seed=947)
        collapse_rows.append({
            "epsilon": eps,
            "hidden_family": hidden,
            "hidden_family_count": int(n_bad[hidden]),
            "true_predicted_variance": variance_prediction(u, n_bad),
        })

    # Secondary equal-budget comparison: uniform, magnitude-only, and controller.
    n_uniform = allocate_fixed(np.ones(C), budget)
    n_mag = allocate_fixed(np.abs(a) + 0.05, budget)
    n_ctrl = allocate_controller(a, u, budget, epsilon=0.05, seed=947)
    comparison = {
        "uniform": {"predicted_variance": variance_prediction(u, n_uniform), "empirical_variance": empirical_variance(u, n_uniform, seed=21)},
        "magnitude_only": {"predicted_variance": variance_prediction(u, n_mag), "empirical_variance": empirical_variance(u, n_mag, seed=22)},
        "uncertainty_guided_controller": {"predicted_variance": variance_prediction(u, n_ctrl), "empirical_variance": empirical_variance(u, n_ctrl, seed=23)},
    }
    result = {
        "setup": {"families": C, "budget_including_initial_counts": budget, "P_min": float(P.min()), "P_max": float(P.max())},
        "prediction_1_variance_identity": {"predicted": pred_var, "observed": obs_var, "relative_error": rel_err},
        "prediction_2_optimal_allocation_sweep": opt_rows,
        "prediction_3_exploration_sweep": eps_rows,
        "prediction_3_mode_collapse_sweep": collapse_rows,
        "equal_budget_comparison": comparison,
    }
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()