import json, math, time
import numpy as np
import torch
import torch.nn as nn

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

torch.set_default_dtype(torch.float64)


def simplex_map(z, eps):
    return torch.softmax(-z / eps, dim=-1)


def simplex_jacobian(w, eps):
    # derivative of softmax(-z/eps), evaluated at the corresponding w
    return -(torch.diag(w) - w[:, None] * w[None, :]) / eps


def verify_math():
    # Prediction 1: at z=0, tangent singular values are exactly 1/(n eps).
    n = 7
    z = torch.zeros(n)
    rows = []
    for eps in [0.05, 0.1, 0.2, 0.4]:
        w = simplex_map(z, eps)
        J = simplex_jacobian(w, eps).numpy()
        sv = np.linalg.svd(J, compute_uv=False)
        measured = float(sv[0]); predicted = 1.0 / (n * eps)
        rows.append({"eps": eps, "predicted_tangent_gain": predicted,
                     "measured_tangent_gain": measured,
                     "relative_error": abs(measured-predicted)/predicted})
    # Prediction 2: probability ratios obey log(w_i/w_j)=- (z_i-z_j)/eps.
    ratio_rows = []
    z = torch.tensor([0.0, 0.37, -0.21, 0.8])
    for eps in [0.05, 0.1, 0.25, 0.5]:
        w = simplex_map(z, eps)
        observed = float(torch.log(w[0]/w[1]))
        predicted = float(-(z[0]-z[1])/eps)
        ratio_rows.append({"eps": eps, "predicted_log_ratio": predicted,
                           "measured_log_ratio": observed,
                           "absolute_error": abs(observed-predicted)})
    # Prediction 3: scaling equivariance T_eps(c z)=T_{eps/c}(z).
    z = torch.tensor([0.2, -0.4, 0.9, -0.1, 0.5])
    equiv = []
    for eps, c in [(0.1, 2.), (0.2, 3.), (0.4, .5), (.07, 5.)]:
        lhs = simplex_map(c*z, eps)
        rhs = simplex_map(z, eps/c)
        equiv.append({"eps": eps, "c": c,
                      "max_abs_error": float((lhs-rhs).abs().max())})
    return {"jacobian_scaling": rows, "log_ratio_scaling": ratio_rows,
            "epsilon_score_equivariance": equiv}


class CapacityBarrier(nn.Module):
    """Entropy Legendre regularizer on the simplex plus log barriers Aw<b.
    Equality sum(w)=1 is handled by a KKT Newton solve; iterations are unrolled.
    """
    def __init__(self, A, b, eps=0.12, steps=14):
        super().__init__()
        self.register_buffer("A", A)
        self.register_buffer("b", b)
        self.eps, self.steps = eps, steps

    def forward(self, z):
        B, n = z.shape; m = self.A.shape[0]
        w = torch.full((B,n), 1.0/n, dtype=z.dtype, device=z.device)
        # fixed damped Newton with a feasibility-preserving step chosen per row
        for _ in range(self.steps):
            slack = self.b[None,:] - w @ self.A.T
            # strict interior is guaranteed by construction of the experiment
            g = z + self.eps * (torch.log(w)+1.0) + self.eps * (1.0/slack) @ self.A
            H = self.eps * (torch.diag_embed(1.0/w) +
                torch.einsum('mi,bm,mj->bij', self.A, 1.0/(slack**2), self.A))
            C = torch.ones((B,1,n), dtype=z.dtype, device=z.device)
            K = torch.cat([torch.cat([H,C.transpose(1,2)], dim=2),
                           torch.cat([C, torch.zeros(B,1,1,dtype=z.dtype,device=z.device)], dim=2)], dim=1)
            rhs = torch.cat([g, (w.sum(1)-1.0)[:,None]], dim=1)
            sol = torch.linalg.solve(K, rhs[...,None])[...,0]
            delta = sol[:,:n]
            # Backtracking only determines a scalar; detach it to retain a stable graph.
            alpha = torch.ones(B, dtype=z.dtype, device=z.device)
            neg = delta > 0
            alpha = torch.minimum(alpha, torch.where(neg, -0.9*w/(delta-1e-30), alpha[:,None]).min(1).values)
            dslack = delta @ self.A.T
            alpha = torch.minimum(alpha, torch.where(dslack < 0, -0.9*slack/(dslack-1e-30), alpha[:,None]).min(1).values)
            alpha = alpha.clamp(max=1.0, min=1e-4).detach()
            w = w - alpha[:,None] * delta
        return w


def make_data(N=768, n=8, m=3):
    rng = np.random.default_rng(SEED+4)
    A = rng.uniform(.05, 1.0, size=(m,n)); A /= A.sum(1,keepdims=True)
    b = np.full(m, .70) # uniform slack .30, while random policies often violate
    X = rng.normal(size=(N, 6))
    # Feasible targets are mixtures of uniform and random simplex points.
    raw = rng.dirichlet(np.ones(n), size=N)
    target = .72/n + .28*raw
    # use a fixed strictly-feasible target independent of rare random violations
    target = target / target.sum(1, keepdims=True)
    return map(torch.tensor,(X,target,A,b))


def train_compare():
    X, target, A, b = make_data(); X=X.to(device); target=target.to(device)
    A=A.to(device); b=b.to(device)
    split=600; n=target.shape[1]
    def model(): return nn.Sequential(nn.Linear(6,24), nn.Tanh(), nn.Linear(24,n)).to(device)
    def run(kind):
        torch.manual_seed(SEED+ (0 if kind=='penalty' else 1))
        net=model(); opt=torch.optim.Adam(net.parameters(), lr=.015)
        layer=CapacityBarrier(A,b).to(device)
        t0=time.perf_counter()
        for step in range(180):
            ix=torch.arange((step*32)%split, (step*32)%split+32, device=device) % split
            z=net(X[ix]); w=torch.softmax(z, -1) if kind=='penalty' else layer(z)
            viol=torch.relu(w@A.T-b)
            loss=((w-target[ix])**2).mean() + (50*viol.square().mean() if kind=='penalty' else 0)
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            z=net(X[split:]); w=torch.softmax(z,-1) if kind=='penalty' else layer(z)
            v=torch.relu(w@A.T-b)
            mse=((w-target[split:])**2).mean().item()
            maxv=v.max().item(); meanv=v.mean().item(); minsl=(b-w@A.T).min().item()
        return {"test_mse":mse, "max_capacity_violation":maxv,
                "mean_capacity_violation":meanv, "minimum_slack":minsl,
                "seconds":time.perf_counter()-t0}
    return {"penalty":run('penalty'), "legendre_barrier":run('legendre')}

if __name__ == '__main__':
    checks=verify_math()
    comparison=train_compare()
    out={"device":device, "checks":checks, "comparison":comparison}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))