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

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


def cosine(a, b):
    den = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / den) if den else 1.0


def make_population(n=12000):
    x = rng.uniform(-3.0, 3.0, size=(n, 2))
    z = 2.0 * x[:, 0] + 0.7 * x[:, 1] - 0.3
    fail_prob = 0.02 + 0.93 / (1.0 + np.exp(-z))
    y = rng.binomial(1, fail_prob)
    c = np.clip(fail_prob, 0.01, 0.99)
    features = np.column_stack([np.ones(n), x[:, 0], x[:, 1], x[:, 0] ** 2,
                                 x[:, 0] * x[:, 1], x[:, 1] ** 2])
    target = np.sin(x[:, 0]) + 0.35 * x[:, 1] + 0.15 * x[:, 0] ** 2
    return features, target, y, c


def fit_critic(x, y, epochs=800, lr=0.08):
    # Lightweight auxiliary predictor C_phi(s), fit only from state/failure pairs.
    X = np.column_stack([np.ones(len(x)), x])
    phi = np.zeros(X.shape[1])
    for _ in range(epochs):
        pred = 1.0 / (1.0 + np.exp(-np.clip(X @ phi, -30, 30)))
        phi -= lr * (X.T @ (pred - y)) / len(y)
    pred = 1.0 / (1.0 + np.exp(-np.clip(X @ phi, -30, 30)))
    bce = float(-np.mean(y * np.log(pred + 1e-8) + (1-y) * np.log(1-pred + 1e-8)))
    return np.clip(pred, 0.01, 0.99), bce


def proposal(c, alpha, eps=0.02):
    a = (eps + c) ** alpha
    z = float(a.mean())
    q = a / (len(a) * z)
    w = z / a
    ess_frac = float(1.0 / np.dot(q, w * w))
    return a, q, w, ess_frac


def exact_quantities(y, c, alpha, eps=0.02):
    a, q, w, ess_frac = proposal(c, alpha, eps)
    fail_p = float(y.mean())
    fail_q = float(np.dot(q, y))
    return a, q, w, fail_q / fail_p, ess_frac


def verify(alpha_grid, y, features, target, c, batches=400, batch_size=96):
    per_grad = -2.0 * target[:, None] * features
    uniform_grad = per_grad.mean(axis=0)
    rows = []
    for alpha in alpha_grid:
        _, q, w, enrich_exact, ess_exact = exact_quantities(y, c, alpha)
        failures, weighted_cos, weighted_err, unweighted_cos = [], [], [], []
        for _ in range(batches):
            ids = rng.choice(len(y), size=batch_size, replace=True, p=q)
            g = per_grad[ids]
            gw = (w[ids, None] * g).mean(axis=0)
            gu = g.mean(axis=0)
            failures.append(float(y[ids].mean()))
            weighted_cos.append(cosine(gw, uniform_grad))
            unweighted_cos.append(cosine(gu, uniform_grad))
            weighted_err.append(float(np.linalg.norm(gw - uniform_grad)))
        # Since q and w are explicitly known on this finite population, this is
        # an exact check of E_q[w g] = E_p[g], independent of Monte Carlo noise.
        exact_weighted = np.sum((q * w)[:, None] * per_grad, axis=0)
        exact_identity_error = float(np.linalg.norm(exact_weighted - uniform_grad))
        rows.append({"alpha": alpha, "predicted_enrichment": enrich_exact,
                     "observed_enrichment": float(np.mean(failures) / y.mean()),
                     "predicted_ess_fraction": ess_exact,
                     "weighted_gradient_cosine": float(np.mean(weighted_cos)),
                     "unweighted_gradient_cosine": float(np.mean(unweighted_cos)),
                     "weighted_gradient_l2_error": float(np.mean(weighted_err)),
                     "exact_identity_l2_error": exact_identity_error})
    return rows


def train_replay(features, target, c, alpha, weighted, steps=450, batch_size=96, lr=0.003):
    _, q, w, _ = proposal(c, alpha)
    theta = np.zeros(features.shape[1])
    for _ in range(steps):
        ids = rng.choice(len(target), size=batch_size, replace=True, p=q)
        xb, tb = features[ids], target[ids]
        err = xb @ theta - tb
        grad = 2.0 * (err[:, None] * xb)
        if weighted:
            grad *= w[ids, None]
        theta -= lr * grad.mean(axis=0)
    return float(np.mean((features @ theta - target) ** 2))


def main():
    random.seed(SEED)
    np.random.seed(SEED)
    features, target, y, oracle_c = make_population()
    x = features[:, 1:3]
    c, critic_bce = fit_critic(x, y)
    alpha_grid = [0.0, 0.5, 1.0, 2.0, 4.0]
    verification = verify(alpha_grid, y, features, target, c)
    uniform = train_replay(features, target, c, 0.0, False)
    training = []
    for alpha in [0.5, 1.0, 2.0]:
        training.append({"alpha": alpha, "uniform_mse": uniform,
                         "weighted_mse": train_replay(features, target, c, alpha, True),
                         "unweighted_mse": train_replay(features, target, c, alpha, False)})
    result = {"seed": SEED, "n_states": len(y), "base_failure_rate": float(y.mean()),
        "critic": {"bce": critic_bce, "mean_prediction": float(c.mean()),
                   "mean_label": float(y.mean())},
        "predictions": {"enrichment": "rho_fail=E_q[y]/E_p[y]",
                         "ess": "ESS/B approaches 1/E_q[w^2]",
                         "gradient_identity": "E_q[w g]=E_p[g]"},
        "verification": verification, "training": training,
        "notes": "Predictions are finite-population exact; observed values are Monte Carlo.",
        "importance_identity_check": {"max_exact_l2_error": float(max(r["exact_identity_l2_error"] for r in verification)),
                                     "criterion": "exact error should be numerical roundoff; finite-batch cosine degrades as ESS falls"}}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()