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

SEED=2380
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.get_device_name(0)
except Exception:
    device='cpu'

# Smooth operators from the proposal.
def softmin(a,tau,dim=-1):
    return -tau*torch.logsumexp(-a/tau,dim=dim)
def softmax_reduce(a,tau,dim=-1):
    return tau*torch.logsumexp(a/tau,dim=dim)

# First, verify the finite-set approximation claim numerically.
def math_check():
    rng=np.random.default_rng(SEED)
    rows=[]
    max_softmin_bound=0.; max_softmax_bound=0.
    for n in [2,4,8,16,32]:
        a=torch.tensor(rng.normal(size=(500,n)),dtype=torch.float32)
        amin=a.min(1).values; amax=a.max(1).values
        for tau in [0.2,0.1,0.05,0.02]:
            sm=softmin(a,tau); sx=softmax_reduce(a,tau)
            # softmin is below min by at most tau log n; softmax above max likewise.
            e1=(amin-sm).max().item(); e2=(sx-amax).max().item()
            b=tau*math.log(n)
            max_softmin_bound=max(max_softmin_bound,e1/b if b else 0)
            max_softmax_bound=max(max_softmax_bound,e2/b if b else 0)
            rows.append({'n':n,'tau':tau,'softmin_error':float((amin-sm).mean()),
                         'softmax_error':float((sx-amax).mean()),'bound':b,
                         'min_bound_ratio':e1/b,'max_bound_ratio':e2/b})
    # Direct monotone scaling check for one fixed candidate set.
    fixed=torch.tensor([[0.1,0.7,1.4,-0.2]],dtype=torch.float32)
    taus=torch.tensor([.4,.2,.1,.05,.025])
    errs=[float((fixed.min()-softmin(fixed,float(t))).abs()) for t in taus]
    return {'rows':rows,'max_bound_ratio_min':max_softmin_bound,
            'max_bound_ratio_max':max_softmax_bound,'fixed_tau':taus.tolist(),
            'fixed_softmin_abs_error':errs}

class Dynamics(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,x): return self.net(x)

def make_data(seed):
    g=torch.Generator().manual_seed(seed)
    x=-1+2*torch.rand(256,1,generator=g)
    # Smooth, mildly nonlinear concrete map; noisy observations make abstraction useful.
    y=.72*x+.18*torch.sin(3*x)+.025*torch.randn(x.shape,generator=g)
    return x,y

def train(seed,beta,steps=500):
    torch.manual_seed(seed)
    x,y=make_data(seed+100)
    # cells are learnable, but assignments and candidate successor lists are fixed for a fair comparison.
    K=16
    centers0=torch.linspace(-1.15,1.15,K).view(-1,1)
    centers=nn.Parameter(centers0+0.015*torch.randn_like(centers0))
    model=Dynamics()
    opt=torch.optim.Adam(list(model.parameters())+[centers],lr=3e-3)
    # Fixed local successor topology: three cells around the nominal target map.
    with torch.no_grad():
        nominal=.72*centers0+.18*torch.sin(3*centers0)
        nearest=torch.cdist(nominal,centers0).topk(3,largest=False).indices
        assign=torch.cdist(x,centers0).argmin(1)
    tau=.06
    for step in range(steps):
        pred=model(x)
        mse=((pred-y)**2).mean()
        ci=centers[assign]
        dcur=(x-ci).abs().squeeze(1)
        succ=centers[nearest[assign]].reshape(x.shape[0],-1) # B,3
        dnext=(pred-succ).abs()
        # Smooth over successors then samples, as in the proposal.
        sim=softmax_reduce(softmin(dnext,tau,1)-dcur,tau,0)
        violation=torch.relu(dcur-.10)
        eps_loss=softmax_reduce(violation,tau,0)
        loss=mse+beta*(sim+0.15*eps_loss)
        opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad(): centers.clamp_(-1.3,1.3)
    with torch.no_grad():
        pred=model(x); ci=centers[assign]; cur=(x-ci).abs().squeeze(1)
        succ=centers[nearest[assign]].reshape(x.shape[0],-1); dn=(pred-succ).abs()
        # Independent hard reverse-simulation residual and coverage at epsilon=.10.
        hard=(dn.min(1).values-cur).max().item()
        hard_mean=(dn.min(1).values-cur).mean().item()
        covered=((dn.min(1).values<=.10)&(cur<=.10)).float().mean().item()
        pred_mse=((pred-y)**2).mean().item()
        # Number of candidate successors needed for each sample at radius epsilon.
        counts=(dn<=.10).sum(1).float()
        avg_count=counts.mean().item(); zero=(counts==0).float().mean().item()
        sim_value=float(softmax_reduce(softmin(dn,tau,1)-cur,tau,0))
    return {'seed':seed,'beta':beta,'prediction_mse':pred_mse,'hard_max_residual':hard,
            'hard_mean_residual':hard_mean,'coverage_at_eps':covered,'avg_successors_at_eps':avg_count,
            'zero_coverage_fraction':zero,'smooth_sim':sim_value}

def main():
    check=math_check()
    results=[]
    for beta in [0.,0.15,0.5]:
        for seed in [11,22,33]:
            try: results.append(train(seed,beta))
            except RuntimeError as e:
                if 'CUDA' in str(e) or 'out of memory' in str(e).lower():
                    global device; device='cpu'; results.append(train(seed,beta))
                else: raise
    out={'device':device,'math_check':check,'results':results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    for b in [0.,.15,.5]:
        r=[z for z in results if z['beta']==b]
        vals=np.mean([[q['prediction_mse'],q['hard_max_residual'],q['coverage_at_eps'],q['avg_successors_at_eps'],q['smooth_sim']] for q in r],axis=0); print('beta',b,'mse %.5f hard %.5f coverage %.3f count %.3f sim %.5f'%tuple(vals))
    print('math ratios',check['max_bound_ratio_min'],check['max_bound_ratio_max'],'errors',check['fixed_softmin_abs_error'])
if __name__=='__main__': main()