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

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

class SharedSymbolic(nn.Module):
    # Safe differentiable implementation of the proposed feature-bank bottleneck.
    def __init__(self,d=2,m=3,H=8,tau=1.0):
        super().__init__(); self.d=d; self.m=m; self.H=H; self.tau=tau
        self.w=nn.Parameter(torch.randn(H,6)*.15); self.b=nn.Parameter(torch.zeros(H))
        self.glogit=nn.Parameter(torch.zeros(H,6)); self.oplogit=nn.Parameter(torch.randn(H,7)*.1)
        self.a=nn.Parameter(torch.randn(m,H)*.15); self.v=nn.Parameter(torch.ones(m,H)*1.5)
        self.c=nn.Parameter(torch.zeros(m)); self.register_buffer('mu',torch.tensor([1.,1.]))
        self.register_buffer('sd',torch.tensor([1.,1.]))
    def set_norm(self,X):
        self.mu.copy_(X.mean(0)); self.sd.copy_(X.std(0).clamp_min(1e-5))
    def features(self,X):
        x=(X-self.mu)/self.sd
        # positive-domain safe primitives, as in the implementation plan
        raw=x; rec=1/x.clamp(-20,20).where(x.abs()>0.03, torch.sign(x)*.03)
        log=torch.log(x.abs()+1e-2)
        return torch.cat([raw,rec,log],1)
    def forward(self,X):
        ph=self.features(X); g=torch.sigmoid(self.glogit)
        u=torch.einsum('bd,hd->bh',ph,g*self.w)+self.b
        vals=torch.stack([torch.exp(u.clamp(-8,8)),u*u,u,1/u.clamp(-8,8),
                          torch.log1p(torch.relu(u)+1e-6),torch.sin(u),torch.sqrt(torch.relu(u)+1e-6)],-1)
        p=torch.softmax(self.oplogit/self.tau,-1); z=(vals*p[None,:,:]).sum(-1)
        q=torch.sigmoid(self.v)
        return self.c+z@(q*self.a).T, z, p, g

class MLP(nn.Module):
    def __init__(self,d=2,m=3):
        super().__init__(); self.net=nn.Sequential(nn.Linear(d,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,m))
    def forward(self,x): return self.net(x)

def data(n, noise, rng, lo=.15, hi=5.):
    X=rng.uniform(lo,hi,(n,2)).astype('float32'); K=np.array([.7,.35]); alpha=np.array([1.2,.8,1.5])
    z=1/(1+X@K); Y=np.stack([alpha[j]*X[:,j%2]*z for j in range(3)],1)
    if noise: Y += rng.normal(0,noise,Y.shape).astype('float32')
    return torch.tensor(X),torch.tensor(Y.astype('float32'))

def train(model,X,Y,epochs=500,lr=2e-3, symbolic=False):
    model.to(device); X=X.to(device); Y=Y.to(device)
    if symbolic: model.set_norm(X)
    opt=torch.optim.Adam(model.parameters(),lr=lr)
    for e in range(epochs):
        opt.zero_grad(); pred=model(X)
        if symbolic:
            pred,z,p,g=pred; loss=((pred-Y)**2).mean()+2e-4*(torch.sigmoid(model.v)*model.a).abs().sum()+1e-5*g.sum()+1e-5*(p*(p+1e-8).log()).sum()
            model.tau=max(.05,1-.95*e/epochs)
        else: loss=((pred-Y)**2).mean()
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),10); opt.step()
    return model

def mse(model,X,Y,symbolic=False):
    with torch.no_grad():
        out=model(X.to(device)); out=out[0] if symbolic else out
        return float(((out-Y.to(device))**2).mean().cpu())

def mechanism_checks():
    # Prediction 1: exact two-class softmax concentration and entropy versus temperature.
    gap=2.0; taus=[1.,.5,.25,.1,.05]; obs=[]; pred=[]
    for t in taus:
        q=np.exp(np.array([0.,gap])/t-np.max(np.array([0.,gap])/t)); q=q/q.sum(); h=float(-(q*np.log(np.maximum(q,1e-300))).sum()); obs.append(h); pred.append(float(math.log1p(math.exp(-gap/t))+(gap/t)/(1+math.exp(gap/t))))
    # Prediction 2: rank-one shared latent mechanism gives zero residual for every output.
    rng=np.random.default_rng(SEED); z=rng.uniform(.2,1.,1000); A=np.array([1.2,-.7,2.1]); Y=z[:,None]*A[None,:]
    fit=np.linalg.lstsq(z[:,None],Y,rcond=None)[0]; rank_res=float(np.mean((Y-z[:,None]@fit)**2))
    # Prediction 3: known additive noise has expected MSE sigma^2 (exact mechanism fit).
    noise_rows=[]
    for s in [0.,.01,.03,.1]:
        vals=[]
        for rep in range(30):
            eps=rng.normal(0,s,(1000,3)); vals.append(float(np.mean(eps**2)))
        noise_rows.append((s,float(np.mean(vals)),s*s))
    return {'softmax_temperature':{'gap':gap,'taus':taus,'observed_entropy':obs,'predicted_entropy':pred},
            'shared_latent_rank_residual':rank_res,'noise_scaling':noise_rows}

def main():
    checks=mechanism_checks(); rng=np.random.default_rng(SEED)
    Xtr,Ytr=data(1000,.01,rng,.15,5.); Xte,Yte=data(1000,0,rng,5.,10.)
    # symbolic model and standard shared-trunk MLP, same data and fixed seed
    torch.manual_seed(SEED); sym=train(SharedSymbolic(H=8),Xtr,Ytr,500,2e-3,True)
    torch.manual_seed(SEED); mlp=train(MLP(),Xtr,Ytr,500,2e-3,False)
    with torch.no_grad(): out,z,p,g=sym(Xte.to(device)); active=int((torch.sigmoid(sym.v)*sym.a).abs().mean(0).gt(.03).sum().cpu()); ops=p.argmax(1).cpu().tolist()
    result={'device':device,'checks':checks,'fit':{'symbolic_in_range_mse':mse(sym,Xtr,Ytr,True),'symbolic_extrapolation_mse':mse(sym,Xte,Yte,True),'mlp_in_range_mse':mse(mlp,Xtr,Ytr),'mlp_extrapolation_mse':mse(mlp,Xte,Yte),'symbolic_active_units':active,'operator_choices':ops}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()