import itertools, json, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED = 2028
np.random.seed(SEED)
random.seed(SEED)
torch.manual_seed(SEED)


def subsets(n, B):
    return [s for k in range(B + 1) for s in itertools.combinations(range(n), k)]


def exact_regret(c, pred, d, subset):
    x = np.zeros(len(c), dtype=float)
    x[list(subset)] = 1.0
    ct, cp = c + d * x, pred + d * x
    yt, yp = int(np.argmin(ct)), int(np.argmin(cp))
    return float(ct[yp] - ct[yt]), int(yp != yt)


def mechanism_check():
    # True costs are c=[2,1], prediction is [3,.5]. Both choose edge 1 nominally.
    # Interdicting edge 1 by delay z: true switch is z=1; predicted switch z=2.5.
    c = np.array([2.0, 1.0])
    pred = np.array([3.0, 0.5])
    zs = np.linspace(0, 3, 61)
    observed = []
    for z in zs:
        r, flip = exact_regret(c, pred, np.array([0., z]), (1,))
        observed.append((float(z), r, flip))
    positive = [z for z, r, _ in observed if r > 1e-9]
    # In the open interval (1, 2.5), regret is z-1, then vanishes.
    onset = min(positive)
    end = max(z for z, r, _ in observed if r > 1e-9)
    linear = [(z, r) for z, r, _ in observed if 1.05 <= z <= 2.45]
    slope = np.polyfit([z for z, _ in linear], [r for _, r in linear], 1)[0]
    max_abs_linear_error = max(abs(r - (z - 1.0)) for z, r in linear)

    # Scaling prediction: multiplying all delays and costs by s scales regret by s,
    # while the switching thresholds scale by s.
    scale_rows = []
    for s in [0.5, 1., 2., 3.]:
        # use a fixed relative interdiction z=1.5*s, inside the disagreement band
        r, _ = exact_regret(s*c, s*pred, np.array([0., 1.5*s]), (1,))
        scale_rows.append((s, r, 0.5*s))
    scale_error = max(abs(r - expected) for s, r, expected in scale_rows)

    # Zero-delay prediction: no adversarial effect at d=0.
    zero_regret = max(exact_regret(c, pred, np.zeros(2), x)[0] for x in subsets(2, 1))
    return {
        "predictions": {
            "onset_boundary": "regret starts when delay exceeds true margin 1.0",
            "linear_scaling": "regret = delay - 1 in the disagreement interval (slope 1)",
            "scale_law": "scaling costs and delays by s scales regret by s and thresholds by s",
            "zero_perturbation": "d=0 has zero extra regret for a nominally equivalent predictor"
        },
        "observed": {
            "onset_grid": onset,
            "last_positive_grid": end,
            "fitted_regret_vs_delay_slope": float(slope),
            "max_linear_absolute_error": float(max_abs_linear_error),
            "scale_rows_s_regret_expected": scale_rows,
            "max_scale_absolute_error": float(scale_error),
            "zero_delay_worst_regret": float(zero_regret)
        },
        "theory": {"onset": 1.0, "end": 2.5, "slope": 1.0}
    }


class Predictor(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(1, 12), nn.Tanh(), nn.Linear(12, 2))
    def forward(self, w):
        return F.softplus(self.net(w))


def train(kind, steps=900, B=1):
    torch.manual_seed(SEED + {"mse": 11, "nominal": 12, "adversarial": 13}[kind])
    model = Predictor()
    opt = torch.optim.Adam(model.parameters(), lr=0.025)
    # Training distribution includes varying margins and delays are known scenarios.
    w = torch.linspace(0.35, 2.0, 96).reshape(-1, 1)
    c = torch.cat([2*w, w], dim=1)
    d = torch.cat([torch.zeros_like(w), 1.5*w], dim=1)
    scen = subsets(2, B)
    for step in range(steps):
        pred = model(w)
        if kind == "mse":
            loss = ((pred-c)**2).mean()
        else:
            losses = []
            # Differentiable decision loss: CE on soft shortest-path probabilities.
            for sub in scen:
                x = torch.zeros_like(c)
                if sub:
                    x[:, list(sub)] = 1.
                ct, cp = c+d*x, pred+d*x
                true_choice = torch.argmin(ct, dim=1)
                losses.append(F.cross_entropy(-cp / 0.12, true_choice, reduction="mean"))
            nominal = losses[0]
            if kind == "nominal":
                loss = nominal + 0.002*((pred-c)**2).mean()
            else:
                loss = torch.stack(losses).max() + 0.002*((pred-c)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return model


def evaluate(model, B=1):
    w = np.linspace(.35, 2.0, 120)
    c = np.stack([2*w, w], axis=1)
    # Delays cover weak to strong perturbations, and are proportional to w.
    d = np.stack([np.zeros_like(w), 1.5*w], axis=1)
    with torch.no_grad():
        pred = model(torch.tensor(w[:, None], dtype=torch.float32)).numpy()
    nominal = []; worst = []; flips = []; mse = []
    for i in range(len(w)):
        nominal.append(exact_regret(c[i], pred[i], d[i], ())[0])
        vals = [exact_regret(c[i], pred[i], d[i], s) for s in subsets(2, B)]
        worst.append(max(v[0] for v in vals)); flips.append(max(v[1] for v in vals)); mse.append(np.mean((pred[i]-c[i])**2))
    return {"nominal_regret_mean": float(np.mean(nominal)), "worst_regret_mean": float(np.mean(worst)), "worst_regret_max": float(np.max(worst)), "adversarial_flip_rate": float(np.mean(flips)), "cost_mse": float(np.mean(mse))}


def main():
    checks = mechanism_check()
    results = {}
    for kind in ["mse", "nominal", "adversarial"]:
        model = train(kind)
        results[kind] = evaluate(model)
    out = {"seed": SEED, "checks": checks, "results": results, "note": "Exact hard paths are used for evaluation; training uses differentiable cross-entropy over the exact two-route oracle."}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()