import json, math, os, random, time
import numpy as np
import torch
from torch import nn

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

# A deliberately nonconvex/nonmonotone synthetic energy on positive singular values.
def target(x):
    # Positive baseline plus localized oscillatory wells; the ICNN is trained as a lower fit.
    a, b = x[:, 0], x[:, 1]
    return 0.35*(a-1.05)**2 + 0.28*(b-1.25)**2 + 0.16*torch.sin(5.0*a)*torch.sin(4.0*b) + 0.18

class MLP(nn.Module):
    def __init__(self, width=24):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,1))
    def forward(self,x): return self.net(x).squeeze(-1)

class ICNN(nn.Module):
    """z_l=softplus(A_l z+B_l x+b), with A,B >= 0; positive output skips preserve monotonicity."""
    def __init__(self, width=20):
        super().__init__()
        self.w0 = nn.Parameter(torch.full((width,2), -4.0) + torch.randn(width,2)*0.05)
        self.b0 = nn.Parameter(torch.zeros(width))
        self.Araw = nn.ParameterList([nn.Parameter(torch.full((width,width), -4.0) + torch.randn(width,width)*0.05)])
        self.Braw = nn.ParameterList([nn.Parameter(torch.full((width,2), -4.0) + torch.randn(width,2)*0.05)])
        self.b1 = nn.Parameter(torch.zeros(width))
        self.craw = nn.Parameter(torch.full((width,), -4.0) + torch.randn(width)*0.05)
        self.draw = nn.Parameter(torch.full((2,), -4.0) + torch.randn(2)*0.05)
        self.outb = nn.Parameter(torch.tensor(0.15))
    @staticmethod
    def pos(v): return torch.nn.functional.softplus(v) + 1e-4
    def forward(self,x):
        # First layer has nonnegative direct input weights.
        z = torch.nn.functional.softplus(x @ self.pos(self.w0).T + self.b0)
        A, B = self.pos(self.Araw[0]), self.pos(self.Braw[0])
        z = torch.nn.functional.softplus(z @ A.T + x @ B.T + self.b1)
        return z @ self.pos(self.craw) + x @ self.pos(self.draw) + self.outb

def energy_from_F(model, F, eps=1e-6):
    s = torch.linalg.svdvals(F).clamp_min(eps)
    return model(s), s

def violation_rates(model, n=4000):
    x = 0.2 + 1.8*torch.rand(n,2,device=device)
    y = 0.2 + 1.8*torch.rand(n,2,device=device)
    t = torch.rand(n,1,device=device)
    z = t*x + (1-t)*y
    with torch.no_grad():
        convex_v = (model(z) > t.squeeze()*model(x)+(1-t.squeeze())*model(y)+1e-5).float().mean().item()
        mono_v = []
        for j in range(2):
            d = torch.zeros_like(x); d[:,j] = 0.05
            mono_v.append((model(x+d) < model(x)-1e-5).float().mean().item())
        q = target(x)
        lower_v = (model(x) > q+1e-4).float().mean().item()
    return convex_v, float(np.mean(mono_v)), lower_v

def train(model, lower=False, steps=1800):
    model.to(device); opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    for step in range(steps):
        x = 0.2 + 1.8*torch.rand(128,2,device=device)
        q = target(x); p = model(x)
        # Exact requested lower-envelope penalty, with a modest value term.
        if lower:
            loss = ((p-q)**2).mean() + 4.0*torch.relu(p-q).square().mean()
        else:
            loss = ((p-q)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return float(loss.detach().cpu())

def main():
    # Stage-1 math sanity: random ICNN parameters should satisfy the architectural claims.
    torch.manual_seed(SEED)
    icnn = ICNN().to(device)
    sanity = violation_rates(icnn, 12000)
    # Fit baseline and proposed model under identical sampled training protocol.
    torch.manual_seed(SEED+1); base = MLP().to(device)
    t0=time.perf_counter(); train(base, False, steps=1800); base_time=time.perf_counter()-t0
    torch.manual_seed(SEED+2); idea = ICNN().to(device)
    t0=time.perf_counter(); train(idea, True, steps=1800); idea_time=time.perf_counter()-t0
    base_v = violation_rates(base); idea_v = violation_rates(idea)
    grid1=torch.linspace(.2,2.0,100,device=device); grid2=torch.linspace(.2,2.0,100,device=device)
    xx,yy=torch.meshgrid(grid1,grid2,indexing='ij'); grid=torch.stack([xx.ravel(),yy.ravel()],1)
    with torch.no_grad():
        qb=target(grid); pb=base(grid); pi=idea(grid)
        metrics={
          'baseline_rmse': float(torch.sqrt(((pb-qb)**2).mean()).cpu()),
          'idea_rmse': float(torch.sqrt(((pi-qb)**2).mean()).cpu()),
          'baseline_grid_lower_violation': float((pb>qb+1e-4).float().mean().cpu()),
          'idea_grid_lower_violation': float((pi>qb+1e-4).float().mean().cpu()),
          'baseline_grid_mean_signed_error': float((pb-qb).mean().cpu()),
          'idea_grid_mean_signed_error': float((pi-qb).mean().cpu()),
          'baseline_seconds':base_time, 'idea_seconds':idea_time,
          'baseline_parameters':sum(p.numel() for p in base.parameters()),
          'idea_parameters':sum(p.numel() for p in idea.parameters()),
          'sanity_jensen_violation':sanity[0], 'sanity_monotonicity_violation':sanity[1],
          'sanity_lower_not_applicable':sanity[2],
          'baseline_random_jensen_violation':base_v[0], 'idea_random_jensen_violation':idea_v[0],
          'baseline_random_monotonicity_violation':base_v[1], 'idea_random_monotonicity_violation':idea_v[1],
          'baseline_random_lower_violation':base_v[2], 'idea_random_lower_violation':idea_v[2],
        }
    # Verify SVD wrapper shape and permutation invariance for a matrix energy.
    F=torch.randn(16,2,2,device=device); e,s=energy_from_F(idea,F)
    metrics['svd_energy_shape_ok']=bool(e.shape==(16,) and s.shape==(16,2))
    metrics['device']=str(device)
    os.makedirs('results',exist_ok=True)
    with open('results/metrics.json','w') as f: json.dump(metrics,f,indent=2)
    print(json.dumps(metrics,indent=2))

if __name__=='__main__':
    try: main()
    except Exception as exc:
        if str(device)=='cuda':
            print('CUDA failed, rerun on CPU:',repr(exc)); device=torch.device('cpu'); main()
        else: raise