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

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


def forward(x, w, b, a, v, c):
    return torch.relu(x @ w.T + b) @ a + x @ v + c


def canonical_merge(w, b, a, tol=1e-7):
    """Exact/tolerance merge, preserving orientation (u,beta), never sign-flips."""
    w = np.asarray(w, dtype=np.float64); b=np.asarray(b, dtype=np.float64); a=np.asarray(a, dtype=np.float64)
    groups=[]; zero=[]
    for j in range(len(a)):
        r=np.linalg.norm(w[j])
        if r <= tol:
            zero.append(j); continue
        u=w[j]/r; beta=b[j]/r; alpha=a[j]*r
        found=None
        for g in groups:
            if np.max(np.abs(np.r_[u,beta]-np.r_[g['u'],g['beta']])) <= tol:
                found=g; break
        if found is None:
            groups.append({'u':u.copy(),'beta':float(beta),'alpha':float(alpha)})
        else: found['alpha'] += float(alpha)
    if not groups:
        return np.zeros((0,w.shape[1])), np.zeros(0), np.zeros(0), zero
    W=np.array([g['u'] for g in groups]); B=np.array([g['beta'] for g in groups]); A=np.array([g['alpha'] for g in groups])
    return W,B,A,zero


def jacobian_norm(model_tuple, x):
    x=x.detach().clone().requires_grad_(True)
    y=forward(x,*model_tuple)
    grads=torch.autograd.grad(y.sum(),x)[0]
    return float(torch.linalg.vector_norm(grads,dim=1).mean())


def main():
    d=2; n=600
    x=torch.randn(n,d,device=device)
    y=(torch.sin(1.7*x[:,0])+0.35*x[:,1]**2-0.2*x[:,0]*x[:,1]+0.08*torch.randn(n,device=device))
    xv=torch.randn(400,d,device=device)
    yv=(torch.sin(1.7*xv[:,0])+0.35*xv[:,1]**2-0.2*xv[:,0]*xv[:,1])
    # Train a small ordinary network; this is the compact representation to be duplicated.
    m=24
    w=nn.Parameter(0.8*torch.randn(m,d,device=device)); b=nn.Parameter(torch.zeros(m,device=device))
    a=nn.Parameter(0.3*torch.randn(m,device=device)); v=nn.Parameter(torch.zeros(d,device=device)); c=nn.Parameter(torch.zeros((),device=device))
    opt=torch.optim.Adam([w,b,a,v,c],lr=0.025)
    t0=time.time()
    for step in range(900):
        opt.zero_grad(); pred=forward(x,w,b,a,v,c); loss=((pred-y)**2).mean(); loss.backward(); opt.step()
    train_sec=time.time()-t0
    base=tuple(q.detach().clone() for q in (w,b,a,v,c))
    # Make a width-48 nominal model by splitting every coefficient into two identical neurons.
    wb,bb,ab,vb,cb=base
    wr=torch.repeat_interleave(wb,2,dim=0); br=torch.repeat_interleave(bb,2); ar=torch.repeat_interleave(ab/2,2)
    dense=(wr,br,ar,vb,cb)
    # Merge in CPU canonical coordinates, then evaluate the resulting model.
    W,B,A,zero=canonical_merge(wr.cpu().numpy(),br.cpu().numpy(),ar.cpu().numpy(),tol=1e-9)
    merged=tuple([torch.tensor(z,dtype=torch.float32,device=device) for z in (W,B,A)])+(vb,cb)
    with torch.no_grad():
        probe=torch.randn(5000,d,device=device)
        yd=forward(probe,*dense); ym=forward(probe,*merged); ycompact=forward(probe,*base)
        maxerr=float((yd-ym).abs().max()); rmserr=float(torch.sqrt(((yd-ym)**2).mean()))
        val_dense=float(((forward(xv,*dense)-yv)**2).mean()); val_merge=float(((forward(xv,*merged)-yv)**2).mean())
        val_comp=float(((forward(xv,*base)-yv)**2).mean())
    # Independent algebra check over random neurons, including non-unit scales.
    ww=np.random.randn(17,d); bb=np.random.randn(17); aa=np.random.randn(17)
    W2,B2,A2,_=canonical_merge(ww,bb,aa,tol=1e-12)
    z=torch.randn(1000,d); lhs=torch.relu(z@torch.tensor(ww,dtype=torch.float32).T+torch.tensor(bb,dtype=torch.float32))@torch.tensor(aa,dtype=torch.float32)
    rhs=torch.relu(z@torch.tensor(W2,dtype=torch.float32).T+torch.tensor(B2,dtype=torch.float32))@torch.tensor(A2,dtype=torch.float32)
    algebra_err=float((lhs-rhs).abs().max())
    jac_dense=jacobian_norm(dense,probe[:512]); jac_merge=jacobian_norm(merged,probe[:512])
    raw_params=2*m + m + d + 1 # width 24 compact
    redundant_params=2*(2*m)+(2*m)+d+1
    merged_params=2*m+m+d+1
    result={
      'device':str(device), 'seed':SEED, 'train_seconds':train_sec,
      'math_algebra_max_error':algebra_err,
      'constructed_width':2*m, 'merged_width':len(A), 'width_reduction_factor':(2*m)/len(A),
      'probe_max_abs_function_error':maxerr, 'probe_rmse_function_error':rmserr,
      'val_mse_redundant':val_dense, 'val_mse_merged':val_merge, 'val_mse_compact':val_comp,
      'jacobian_mean_norm_redundant':jac_dense, 'jacobian_mean_norm_merged':jac_merge,
      'params_redundant':redundant_params, 'params_merged':merged_params,
      'note':'The redundant model is an exact coefficient-split reparameterization of a trained compact model; this isolates the claimed merger phenomenon rather than claiming spontaneous duplicate discovery.'
    }
    os.makedirs('results',exist_ok=True)
    with open('results/metrics.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

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