"""Toy MVP for certified ambiguity gating.
Run: /home/maxwelhelp/main/bin/python3 ambiguity_experiment.py
"""
import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 611
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
try:
    torch.set_num_threads(4)
except Exception:
    pass


def audit(probs, tau=0.30, delta=0.05):
    probs = np.asarray(probs)
    admissible = probs >= tau
    amb = admissible.sum(axis=1) >= 2
    dhat = float(amb.mean()); n = len(amb)
    eps = math.sqrt(math.log(1/delta)/(2*n))
    floor_lcb = max(0., dhat-eps)/2
    return {"Dhat": dhat, "epsilon": eps, "floor_LCB": floor_lcb,
            "ambiguous": amb, "admissible": admissible}


def verify_minimax_floor():
    # Finite population: 40% have two admissible labels, others have one.
    rng = np.random.default_rng(SEED)
    n = 2000
    amb = rng.random(n) < .40
    # All deterministic predictors are represented by their predictions;
    # adversary chooses the opposite admissible target pointwise.
    # For a randomized q=P(pred=1), worst target risk is max(q,1-q).
    q = rng.random(n)
    admissible_risk = np.where(amb, np.maximum(q, 1-q),
                              np.where(q >= .5, 0., 1.))
    # On unique points q is set to the unique admissible label (0 here).
    q_opt = np.where(amb, .5, 0.)
    opt_risk = np.where(amb, .5, 0.).mean()
    empirical_D = amb.mean()
    return {"D": float(empirical_D), "minimax_risk_fair": float(opt_risk),
            "D_over_2": float(empirical_D/2),
            "random_predictor_worst_risk": float(admissible_risk.mean())}


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2, 32), nn.Tanh(), nn.Linear(32, 2))
    def forward(self, x): return self.net(x)


def make_data(n, rng):
    x = rng.uniform(-1, 1, size=(n, 2)).astype(np.float32)
    # Ground-truth boundary is nonlinear but learnable by the small MLP.
    score = x[:, 0] + .55*np.sin(3*x[:, 1])
    y = (score > 0).astype(np.int64)
    # Supervisor is confident away from boundary and ambiguous near it.
    amb = np.abs(score) < .38
    p1 = np.where(amb, .50 + .03*np.sign(score),
                  np.where(score > 0, .92, .08))
    # Independent generated hard label, as with stochastic LLM decoding.
    yl = (rng.random(n) < p1).astype(np.int64)
    probs = np.column_stack([1-p1, p1]).astype(np.float32)
    return x, y, yl, probs, amb


def train(x, yl, probs, keep, soft=False, epochs=80):
    torch.manual_seed(SEED + (3 if soft else 0) + int(keep.sum()))
    model = MLP(); opt = torch.optim.Adam(model.parameters(), lr=.012)
    xt = torch.from_numpy(x); yt = torch.from_numpy(yl)
    pt = torch.from_numpy(probs)
    idx = np.flatnonzero(keep)
    # Full-batch makes equal-step comparison deterministic; equal retained data.
    for _ in range(epochs):
        opt.zero_grad(); logits = model(xt[idx])
        if soft:
            loss = -(pt[idx] * torch.log_softmax(logits, 1)).sum(1).mean()
        else:
            loss = nn.functional.cross_entropy(logits, yt[idx])
        loss.backward(); opt.step()
    return model


def metrics(model, x, y):
    with torch.no_grad():
        z = model(torch.from_numpy(x)); p = torch.softmax(z, 1)[:,1].numpy()
    pred = (p >= .5).astype(np.int64)
    acc = float((pred == y).mean())
    # Expected calibration error, 10 fixed bins.
    ece = 0.
    conf = np.maximum(p, 1-p); corr = (pred == y).astype(float)
    for lo, hi in zip(np.linspace(0,1,11)[:-1], np.linspace(0,1,11)[1:]):
        m = (conf >= lo) & (conf < hi if hi < 1 else conf <= hi)
        if m.any(): ece += m.mean()*abs(corr[m].mean()-conf[m].mean())
    return acc, float(ece)


def main():
    rng = np.random.default_rng(SEED)
    xtr, ytr, yl, probs, latent_amb = make_data(1800, rng)
    xte, yte, _, _, _ = make_data(5000, rng)
    au = audit(probs, tau=.30, delta=.05)
    amb = au["ambiguous"]
    # Gate keeps certified-unambiguous examples. Confidence control retains
    # exactly the same number, using largest max-label probability.
    gate_keep = ~amb
    confidence = probs.max(1)
    conf_keep = np.zeros(len(confidence), dtype=bool)
    conf_keep[np.argsort(confidence)[-gate_keep.sum():]] = True
    hard_all = np.ones(len(yl), dtype=bool)
    models = {
      "hard_all": train(xtr, yl, probs, hard_all, False),
      "ambiguity_gate": train(xtr, yl, probs, gate_keep, False),
      "confidence_matched": train(xtr, yl, probs, conf_keep, False),
      "soft_all": train(xtr, yl, probs, hard_all, True),
    }
    results = {k: {"accuracy": metrics(v,xte,yte)[0], "ECE": metrics(v,xte,yte)[1]}
               for k,v in models.items()}
    out = {"seed": SEED, "tau": .30, "n_train": len(xtr),
           "audit": {k:(float(v) if np.isscalar(v) else None) for k,v in au.items() if k not in ("ambiguous","admissible")},
           "retained_gate": int(gate_keep.sum()), "retention": float(gate_keep.mean()),
           "minimax_verification": verify_minimax_floor(), "results": results}
    print(json.dumps(out, indent=2, sort_keys=True))
    Path("results.json").write_text(json.dumps(out, indent=2, sort_keys=True))

if __name__ == '__main__': main()