import json, math, random
import numpy as np
import torch

SEED = 535
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')

# Scalar box-constrained output: y in [-1, 1]. The feature map is frozen;
# only the affine output head is trained, so the nominal objective is convex.
def make_data(n, noise, rng):
    x = rng.uniform(-1.0, 1.0, size=(n, 1)).astype(np.float32)
    y_clean = (0.82*x[:, 0] + 0.03*np.sin(3*x[:, 0])).astype(np.float32)
    y_obs = y_clean + rng.normal(0, noise, size=n).astype(np.float32)
    return x, y_obs, y_clean

def fit_nominal(x, y, steps=1000, lr=0.04):
    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
    Y = torch.tensor(y[:, None], device=device)
    w = torch.zeros(2, 1, device=device, requires_grad=True)
    opt = torch.optim.Adam([w], lr=lr)
    for _ in range(steps):
        opt.zero_grad(); pred = X @ w
        loss = ((pred-Y)**2).mean() + 1e-4*(w[0]**2)
        loss.backward(); opt.step()
    return w.detach()

def fit_scenario(x, y, residuals, mu=3.0, scen=48, steps=1200, lr=0.025):
    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
    Y = torch.tensor(y[:, None], device=device)
    R = torch.tensor(residuals[:, None], device=device)
    w = torch.zeros(2, 1, device=device, requires_grad=True)
    opt = torch.optim.Adam([w], lr=lr)
    g = torch.Generator(device=device); g.manual_seed(SEED+7)
    for _ in range(steps):
        opt.zero_grad(); pred = X @ w
        # Bootstrap residual scenarios independently for each minibatch item.
        idx = torch.randint(R.shape[0], (scen, X.shape[0]), generator=g, device=device)
        s = R[idx].squeeze(-1).T  # [batch, scenarios]
        ys = pred + s
        h = torch.relu(-1.0-ys) + torch.relu(ys-1.0)
        loss = ((pred-Y)**2).mean() + mu*h.mean() + 1e-4*(w[0]**2)
        loss.backward(); opt.step()
    return w.detach()

def predict(w, x):
    X = torch.tensor(np.c_[x, np.ones(len(x), dtype=np.float32)].astype(np.float32), device=device)
    return (X @ w).detach().cpu().numpy()[:, 0]

def violation_metrics(pred, clean, noise, rng):
    # Independent held-out disturbance distribution, with repeated draws to
    # estimate rare violations without conflating them with training labels.
    draws = clean[:, None] + rng.normal(0, noise, size=(len(clean), 80))
    # Predictor is nominal; each actual output is prediction error scenario.
    errors = draws - clean[:, None]
    actual = pred[:, None] + errors
    viol = np.maximum(0, -1-actual) + np.maximum(0, actual-1)
    return {
        'nominal_mse_to_clean': float(np.mean((pred-clean)**2)),
        'p95_abs_error': float(np.quantile(np.abs(pred[:, None]-draws), .95)),
        'trajectory_sample_violation_rate': float(np.mean(viol > 0)),
        'mean_slack': float(np.mean(viol)),
        'p99_slack': float(np.quantile(viol, .99)),
    }

def exact_slack_check():
    # For each scenario, min_{h>=0} mu*h subject to y+s <= 1+h and
    # -1-h <= y+s is attained at h=max(0, |y+s|-1).
    vals = np.array([-1.4, -1.0, -0.2, 0.7, 1.0, 1.25])
    mu = 2.7
    eliminated = mu*np.maximum(0, np.abs(vals)-1)
    # brute force grid minimization of the defining one-dimensional problem
    brute = []
    for v in vals:
        hs = np.linspace(0, 1.5, 30001)
        feasible = (v <= 1+hs) & (v >= -1-hs)
        brute.append(np.min(np.where(feasible, mu*hs, np.inf)))
    brute = np.array(brute)
    return {'max_abs_error': float(np.max(np.abs(brute-eliminated))),
            'slacks': eliminated.tolist(), 'passed': bool(np.max(np.abs(brute-eliminated)) < 1e-4)}

def main():
    rng = np.random.default_rng(SEED)
    xtr, ytr, cleantr = make_data(700, .10, rng)
    xte, yte, cleante = make_data(5000, .10, rng)
    # Residual buffer from observed prediction errors of a nominal fitted head.
    w0 = fit_nominal(xtr, ytr)
    residual_buffer = ytr - predict(w0, xtr)
    wb = fit_nominal(xtr, ytr)
    ws = fit_scenario(xtr, ytr, residual_buffer)
    result = {
      'device': str(device), 'exact_slack_check': exact_slack_check(),
      'baseline': violation_metrics(predict(wb,xte), cleante, .10, np.random.default_rng(SEED+2)),
      'residual_scenario': violation_metrics(predict(ws,xte), cleante, .10, np.random.default_rng(SEED+2)),
      'weights': {'baseline': wb.cpu().numpy().ravel().tolist(), 'scenario': ws.cpu().numpy().ravel().tolist()},
      'buffer': {'n': int(len(residual_buffer)), 'mean': float(residual_buffer.mean()), 'std': float(residual_buffer.std())}
    }
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()