"""Resolution-gated dual masking MVP with numerical check and synthetic sweep."""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np


def _codes(z):
    z = np.asarray(z, dtype=np.int64)
    if z.ndim == 1:
        return z
    # Binary candidate variables: encode a mask row as a mixed-radix integer.
    return z.dot(1 << np.arange(z.shape[1], dtype=np.int64))


def conditional_entropy(y, z=None):
    y = np.asarray(y, dtype=np.int64)
    if z is None:
        z = np.zeros(len(y), dtype=np.int64)
    keys = _codes(z)
    # H(Y|Z) from one joint bincount; this is substantially faster than loops.
    _, zi = np.unique(keys, return_inverse=True)
    ny = int(y.max()) + 1
    joint = np.bincount(zi * ny + y, minlength=int(zi.max() + 1) * ny).reshape(-1, ny)
    totals = joint.sum(axis=1)
    nz = joint > 0
    probs = np.divide(joint, totals[:, None], out=np.zeros_like(joint, dtype=float), where=totals[:, None] > 0)
    ent = -np.sum(np.where(nz, probs * np.log(np.maximum(probs, 1e-300)), 0.0), axis=1)
    return float(np.sum(totals / len(y) * ent))


def entropy_scores(y, X, mask=None):
    y, X = np.asarray(y), np.asarray(X, dtype=np.int64)
    mask = [] if mask is None else list(mask)
    h0 = conditional_entropy(y, np.zeros(len(y), dtype=np.int64) if not mask else X[:, mask])
    return np.array([h0 - conditional_entropy(y, X[:, mask + [j]])
                     for j in range(X.shape[1]) if j not in mask])


def conditional_mean_predict(train_x, train_y, test_x):
    train_x = np.asarray(train_x, dtype=np.int64); test_x = np.asarray(test_x, dtype=np.int64)
    if train_x.ndim == 1: train_x = train_x[:, None]
    if test_x.ndim == 1: test_x = test_x[:, None]
    # Mixed-radix encoding is safe here because all features are binary.
    kt = _codes(train_x); kv = _codes(test_x)
    sums = np.bincount(kt, weights=train_y)
    counts = np.bincount(kt)
    glob = float(np.mean(train_y))
    return np.divide(sums[kv], counts[kv], out=np.full(len(kv), glob), where=counts[kv] > 0)


def risk_scores(ytr, Xtr, yva, Xva, mask=None):
    mask = [] if mask is None else list(mask)
    bxtr = np.zeros((len(ytr), 1), dtype=np.int64) if not mask else Xtr[:, mask]
    bxva = np.zeros((len(yva), 1), dtype=np.int64) if not mask else Xva[:, mask]
    base = np.mean((yva - conditional_mean_predict(bxtr, ytr, bxva)) ** 2)
    out = []
    for j in range(Xtr.shape[1]):
        if j not in mask:
            pred = conditional_mean_predict(Xtr[:, mask + [j]], ytr, Xva[:, mask + [j]])
            out.append(base - np.mean((yva - pred) ** 2))
    return np.array(out), base


def bootstrap_gamma(y, X, B=30, seed=0):
    rng = np.random.default_rng(seed)
    vals = np.array([entropy_scores(y[ix], X[ix])
                     for ix in (rng.integers(0, len(y), len(y)) for _ in range(B))])
    observed = entropy_scores(y, X)
    j = int(np.argmax(observed))
    return float(observed[j] / (np.std(vals[:, j], ddof=1) + 1e-12)), observed


def gated_select(ytr, Xtr, yva, Xva, kappa_c=1.0, gamma_c=2.0, B=30, seed=0):
    kappa = conditional_entropy(ytr) / np.log(2.0)
    gamma, sh = bootstrap_gamma(ytr, Xtr, B=B, seed=seed)
    sr, _ = risk_scores(ytr, Xtr, yva, Xva)
    mh, mr = int(np.argmax(sh)), int(np.argmax(sr))
    chosen = mh if (kappa < kappa_c and gamma > gamma_c) else mr
    return dict(kappa=kappa, gamma=gamma, sh=sh, sr=sr, mh=mh, mr=mr, chosen=chosen)


def generate(n, p, d, seed):
    rng = np.random.default_rng(seed)
    lag = rng.integers(0, 2, n)
    X = np.column_stack([lag, rng.integers(0, 2, (n, d - 1))])
    y = lag ^ (rng.random(n) < p).astype(int)
    return X, y


def run():
    checks = []
    for p in (0.0, .1, .25, .4, .5):
        X, y = generate(100000, p, 3, 91)
        empirical = conditional_entropy(y) - conditional_entropy(y, X[:, 0])
        q = np.array([p, 1-p]); hb = -np.sum(q[q > 0] * np.log(q[q > 0]))
        checks.append((p, empirical, np.log(2)-hb))
    rows = []
    for n in (256, 1024, 4096):
      for p in (0.0, .1, .25, .4, .49):
       for d in (3, 9, 17):
        vals = {k: [] for k in ("entropy_recovery","risk_recovery","gate_recovery","risk_h","risk_r","risk_gate","kappa","gamma")}
        for rep in range(6):
            X,y = generate(n + n//2, p, d, 10000 + n + 100*d + rep)
            Xtr, Xva, ytr, yva = X[:n], X[n:], y[:n], y[n:]
            g = gated_select(ytr, Xtr, yva, Xva, B=30, seed=rep)
            pred = []
            for col in (g['mh'], g['mr'], g['chosen']):
                pred.append(conditional_mean_predict(Xtr[:, [col]], ytr, Xva[:, [col]]))
            vals["entropy_recovery"].append(g['mh'] == 0); vals["risk_recovery"].append(g['mr'] == 0); vals["gate_recovery"].append(g['chosen'] == 0)
            vals["risk_h"].append(np.mean((yva-pred[0])**2)); vals["risk_r"].append(np.mean((yva-pred[1])**2)); vals["risk_gate"].append(np.mean((yva-pred[2])**2)); vals["kappa"].append(g['kappa']); vals["gamma"].append(g['gamma'])
        rows.append(dict(n=n,p=p,d=d,**{k:float(np.mean(v)) for k,v in vals.items()}))
    out = {"math_checks":checks, "results":rows}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps({"math_checks":checks, "selected_summary": rows[:3]}, indent=2))

if __name__ == "__main__": run()