import json, math, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 1599
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda":
        torch.zeros(1, device="cuda").sum().item()
except Exception:
    DEVICE = "cpu"


def step(w, dt, s, admission):
    return torch.relu(w - dt) + admission * s


def deterministic_check():
    # Baseline admits a small first job and rejects a large second job.
    # Controlled does the opposite: it is initially lower, then exceeds baseline.
    dt, s1, s2 = 0.1, 2.0, 10.0
    wb = step(torch.tensor(0.), 0., s1, torch.tensor(1.))
    wc = step(torch.tensor(0.), 0., s1, torch.tensor(0.))
    wb = step(wb, dt, s2, torch.tensor(0.))
    wc = step(wc, dt, s2, torch.tensor(1.))
    return {"baseline_after_second": float(wb), "controlled_after_second": float(wc),
            "violation": float(torch.relu(wc-wb)), "order_loss": float(torch.relu(wc-wb)**2)}


def math_sweeps():
    # With first-job residual b and controlled residual zero, the next-job
    # violation is max(S2-b,0), hence onset at S2=b and quadratic loss above it.
    b, s2 = 1.9, np.linspace(0, 5, 101)
    observed = np.maximum(s2-b, 0.)
    predicted = np.maximum(s2-b, 0.)
    onset = float(s2[np.argmax(observed > 1e-9)])
    # Fit slope of loss vs (S2-b)^2 only in the active region.
    active = observed > 0
    x = (s2[active]-b)**2; y = observed[active]**2
    slope = float(np.dot(x,y)/np.dot(x,x))
    # Scaling sweep directly changes the excess gap.
    scales = np.array([0.5, 1., 1.5, 2.])
    gaps = np.maximum(scales*4-b, 0.)
    ratios = gaps**2 / (gaps[1]**2)
    return {"threshold_prediction": b, "threshold_observed_grid": onset,
            "quadratic_prediction_slope": 1.0, "quadratic_observed_slope": slope,
            "scale_factors": scales.tolist(), "loss_ratio_observed": ratios.tolist(),
            "loss_ratio_predicted": ((gaps/gaps[1])**2).tolist()}


def traces(n, device):
    # Paired arrivals and service costs. Baseline is a state-dependent reference
    # gate; the learned gate sees the same workload, size, and position.
    dt = torch.rand(n, 6, device=device) * .35 + .05
    s = torch.exp(torch.randn(n, 6, device=device)*.55 + .35).clamp(.2, 8.)
    return dt, s

class Gate(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(4, 24), nn.Tanh(), nn.Linear(24, 1))
    def forward(self, x):
        return torch.sigmoid(self.net(x)).squeeze(-1)

def baseline_decision(w, s):
    # Reference production policy: avoid admitting a job if current workload
    # is high or if the job is large. This creates the feedback counterexample.
    return ((w < 1.5) & (s < 3.0)).float()

def rollout(model, dt, s, lam):
    n, k = s.shape
    wb = torch.zeros(n, device=dt.device); wc = torch.zeros(n, device=dt.device)
    violations=[]; probs=[]; accepted=[]; baseacc=[]
    for j in range(k):
        xb = torch.stack([wc, s[:,j], dt[:,j], torch.full_like(wc, j/(k-1))], 1)
        p = model(xb)
        ib = baseline_decision(wb, s[:,j])
        # Soft admission is the differentiable training surrogate.
        wb = step(wb, dt[:,j], s[:,j], ib)
        wc = step(wc, dt[:,j], s[:,j], p)
        v = torch.relu(wc-wb)
        violations.append(v); probs.append(p); accepted.append(p); baseacc.append(ib)
    v = torch.stack(violations,1); p = torch.stack(probs,1)
    # Throughput reward, mild workload regularization, plus proposed order term.
    objective = -p.mean() + .01*wc.mean() + lam*(v**2).mean()
    return objective, v.detach(), p.detach(), torch.stack(baseacc,1).detach(), wc.detach(), wb.detach()

def train(lam, device):
    model=Gate().to(device); opt=torch.optim.Adam(model.parameters(),lr=.025)
    for epoch in range(350):
        dt,s=traces(96,device)
        loss,*_=rollout(model,dt,s,lam)
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        dt,s=traces(2048,device)
        _,v,p,ba,wc,wb=rollout(model,dt,s,lam)
        return {"violation_frequency":float((v>1e-5).any(1).float().mean()),
                "mean_violation":float(v.mean()), "p99_violation":float(torch.quantile(v.flatten(),.99)),
                "mean_admission_probability":float(p.mean()), "mean_final_workload":float(wc.mean()),
                "baseline_mean_final_workload":float(wb.mean()),
                "baseline_admission_rate":float(ba.mean())}

def main():
    checks=math_sweeps(); exact=deterministic_check()
    results={}
    for lam in (0.0, 0.5, 2.0):
        results[str(lam)]=train(lam,DEVICE)
    out={"device":DEVICE,"exact_counterexample":exact,"math_sweeps":checks,"training":results}
    Path("results.json").write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=="__main__": main()