import json, math, os
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

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


def allocation(q, bar_sigma2):
    q = np.asarray(q, dtype=float)
    v = bar_sigma2 * len(q) * (1.0 / q) / np.sum(1.0 / q)
    return v


def toy_sweep():
    # Gaussian latent z~N(0,I), scalar task t=a^T z+eta.  In this toy,
    # s_k=|a_k| is the exact task sensitivity proxy.
    K = 8
    bar = 0.25
    contrasts = [1, 2, 4, 8, 16, 32]
    rows = []
    for c in contrasts:
        # Two high-relevance and six low-relevance coordinates.
        a = np.array([c, c] + [1.0] * (K - 2))
        q = (np.abs(a) + 1e-8) / (np.mean(np.abs(a)) + 1e-8)
        v = allocation(q, bar)
        # Exact predicted task-noise variance, and Monte Carlo observed value.
        pred = float(np.sum(a * a * v))
        uniform = float(np.sum(a * a * bar))
        z = rng.normal(size=(400000, K))
        eps = rng.normal(size=(400000, K)) * np.sqrt(v)
        observed = float(np.var((z + eps) @ a - z @ a))
        rows.append({
            "contrast": c,
            "predicted_total_variance": float(K * bar),
            "observed_total_variance": float(np.mean(np.sum(eps * eps, axis=1))),
            "predicted_task_noise": pred,
            "observed_task_noise": observed,
            "uniform_task_noise": uniform,
            "adaptive_over_uniform": pred / uniform,
            "high_relevance_noise_fraction": float(np.sum(v[:2]) / np.sum(v)),
        })
    # Separate direct checks of the formula over random relevance vectors.
    budget_errors = []
    ordering_ok = []
    for _ in range(100):
        s = np.exp(rng.normal(size=K))
        q = (s + 1e-6) / (np.mean(s) + 1e-6)
        v = allocation(q, bar)
        budget_errors.append(abs(np.mean(v) - bar))
        ordering_ok.append(np.all(np.argsort(s) == np.argsort(-v)))
    return {
        "K": K, "bar_sigma2": bar, "rows": rows,
        "budget_max_abs_error": float(max(budget_errors)),
        "inverse_ordering_fraction": float(np.mean(ordering_ok)),
    }


def make_data(n, d, weights, noise, seed):
    r = np.random.default_rng(seed)
    x = r.normal(size=(n, d))
    score = x @ weights + r.normal(scale=noise, size=n)
    y = (score > 0).astype(np.int64)
    return x, y


def gradient_relevance(x, y):
    # A task probe supplies the operational MI sensitivity estimate:
    # average absolute gradient of per-example BCE wrt each latent coordinate.
    probe = LogisticRegression(C=10.0, max_iter=300, random_state=SEED)
    probe.fit(x, y)
    p = probe.predict_proba(x)[:, 1]
    w = probe.coef_[0]
    grad = np.abs((p - y)[:, None] * w[None, :])
    return grad.mean(axis=0), np.abs(w)


def evaluate(xtr, ytr, xte, yte, method, bar_sigma2, repeats=8):
    s, magnitude = gradient_relevance(xtr, ytr)
    d = xtr.shape[1]
    if method == "mi_gradient":
        q = (s + 1e-8) / (np.mean(s) + 1e-8)
    elif method == "magnitude":
        q = (magnitude + 1e-8) / (np.mean(magnitude) + 1e-8)
    elif method == "random":
        q = np.random.default_rng(SEED + 17).permutation((magnitude + 1e-8) / (np.mean(magnitude) + 1e-8))
    else:
        q = np.ones(d)
    v = allocation(q, bar_sigma2)
    acc = []
    for j in range(repeats):
        r = np.random.default_rng(SEED + 1000 + j)
        noisy_tr = xtr + r.normal(size=xtr.shape) * np.sqrt(v)
        noisy_te = xte + r.normal(size=xte.shape) * np.sqrt(v)
        clf = LogisticRegression(C=10.0, max_iter=300, random_state=SEED)
        clf.fit(noisy_tr, ytr)
        acc.append(accuracy_score(yte, clf.predict(noisy_te)))
    return {
        "accuracy_mean": float(np.mean(acc)), "accuracy_std": float(np.std(acc)),
        "mean_variance": float(np.mean(v)), "variances": v.tolist(),
        "scores": s.tolist(), "probe_magnitude": magnitude.tolist(),
    }


def mini_experiment():
    d = 8
    true_w = np.array([3.0, 2.5, 1.8, 1.2, 0.7, 0.45, 0.25, 0.1])
    xtr, ytr = make_data(12000, d, true_w, 0.7, SEED)
    xte, yte = make_data(5000, d, true_w, 0.7, SEED + 1)
    out = {}
    for method in ["uniform", "random", "magnitude", "mi_gradient"]:
        out[method] = evaluate(xtr, ytr, xte, yte, method, bar_sigma2=0.8)
    return out


def main():
    result = {"seed": SEED, "toy": toy_sweep(), "mini": mini_experiment()}
    with open("results.json", "w") as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()