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

SEED=2711
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.cuda.empty_cache()
except Exception:
    device='cpu'

def true_x(u):
    return 0.2 + 0.8*u + 0.08*torch.sin(4*math.pi*u)

def cost(x,u):
    return (x-0.72)**2 + 0.035*(u-0.15)**2

def grid_opt(xfun, n=20001):
    u=torch.linspace(0,1,n,device=device)
    with torch.no_grad():
        j=cost(xfun(u),u); k=int(torch.argmin(j))
    return float(u[k]), float(j[k])

def stability_sweep():
    H=3.7; rows=[]
    for rho in np.linspace(.05,.7,14):
        e=1.0
        for _ in range(80): e=(1-rho*H)*e
        rows.append({'rho':float(rho),'rhoH':float(rho*H),'error80':float(abs(e)),
                     'stable_observed':bool(abs(e)<1e-4),
                     'stable_predicted':bool(abs(1-rho*H)<1)})
    return {'H':H,'predicted_rho_boundary_2_over_H':2/H,'rows':rows}

def minima_count(eps,k=6*math.pi,n=20001):
    u=np.linspace(0,1,n); us=.63
    j=(u-us)**2+eps*np.cos(k*u)
    idx=np.where((j[1:-1]<j[:-2])&(j[1:-1]<j[2:]))[0]+1
    return int(len(idx)),u[idx].tolist(),float(u[np.argmin(j)])

def minima_sweep():
    k=6*math.pi; predicted=2/(k*k)
    epses=np.array([0,.001,.003,.005,.01,.02,.04,.08])
    return {'k':k,'predicted_curvature_threshold':predicted,
            'rows':[{'eps':float(e),'count':minima_count(e,k)[0],
                     'global_u':minima_count(e,k)[2]} for e in epses]}

class Surrogate(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,u): return self.net(u)

def decision_unroll(model, starts, steps=35, lr=.08):
    u=starts.clone().requires_grad_(True)
    for _ in range(steps):
        j=cost(model(u[:,None]).squeeze(1),u)
        g=torch.autograd.grad(j.sum(),u,create_graph=True)[0]
        u=torch.clamp(u-lr*g,0,1)
    j=cost(model(u[:,None]).squeeze(1),u)
    return u,j

def train(decision=False, seed=2711, epochs=1800):
    torch.manual_seed(seed)
    model=Surrogate().to(device)
    # Sparse observations intentionally leave decision-relevant interpolation freedom.
    ud=torch.tensor([[0.0],[.22],[.48],[.76],[1.0]],dtype=torch.float32,device=device)
    yd=true_x(ud)
    ut,jt=grid_opt(true_x)
    ustar=torch.tensor(ut,device=device); jstar=torch.tensor(jt,device=device)
    opt=torch.optim.Adam(model.parameters(),lr=.008)
    starts=torch.tensor([.05,.25,.5,.75,.95],device=device)
    for ep in range(epochs):
        pred=model(ud).squeeze(1); loss=((pred-yd)**2).mean()
        if decision:
            us,js=decision_unroll(model,starts,steps=18,lr=.10)
            # soft selection: lowest-start solution is used as differentiable proxy
            q=torch.argmin(js.detach()); ld=(us[q]-ustar)**2 + 2.0*(js[q]-jstar)**2
            loss=loss+3.0*ld
        opt.zero_grad(); loss.backward(); opt.step()
    return model,ut,jt

def assess(model, ut, jt):
    us=np.linspace(0,1,101); u=torch.tensor(us,dtype=torch.float32,device=device)
    with torch.no_grad(): j=cost(model(u[:,None]).squeeze(1),u).cpu().numpy()
    ids=np.where((j[1:-1]<j[:-2])&(j[1:-1]<j[2:]))[0]+1
    # multistart projected gradient, no graph needed
    sols=[]
    for s in np.linspace(.01,.99,100):
        z=torch.tensor([s],dtype=torch.float32,device=device,requires_grad=True)
        for _ in range(120):
            jj=cost(model(z),z); g=torch.autograd.grad(jj,z)[0]
            with torch.no_grad(): z.clamp_(0,1); z.sub_(.08*g).clamp_(0,1)
            z.requires_grad_(True)
        sols.append(float(z.detach()))
    distinct=[]
    for z in sorted(sols):
        if not distinct or abs(z-distinct[-1])>.025: distinct.append(z)
    k=int(np.argmin(j));
    return {'grid_opt_u':float(us[k]),'grid_opt_j':float(j[k]),'displacement':abs(float(us[k])-ut),
            'grid_local_minima':int(len(ids)),'multistart_distinct':len(distinct),
            'multistart_solutions':distinct,'data_mse':float(((model(torch.tensor([[0.],[.22],[.48],[.76],[1.]],device=device))-true_x(torch.tensor([[0.],[.22],[.48],[.76],[1.]],device=device)))**2).mean().detach().cpu())}

def main():
    ut,jt=grid_opt(true_x)
    base,_a,_b=train(False); idea,_a,_b=train(True)
    out={'device':device,'true_optimum':{'u':ut,'j':jt},'math_stability':stability_sweep(),'math_extra_minima':minima_sweep(),
         'baseline':assess(base,ut,jt),'decision_aware':assess(idea,ut,jt)}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()