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

SEED = 2235
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(12, torch.get_num_threads()))
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type == 'cuda':
        torch.cuda.set_device(0)
        torch.empty(1, device=device)
except Exception:
    device = torch.device('cpu')

# Smooth exchangeable module: f(W)=sum_i v_i phi(u_i^T x), phi(0)=0.
# The quadratic term is phi'(0) sum_i v_i u_i^T x.
def phi(s):
    return s + 0.5*s*s + 0.1*s*s*s

def branch_output(W, x):
    # W is p x d, last row is v and first k rows are u.
    u, v = W[:-1, :], W[-1, :]
    return (v * phi(u.T @ x)).sum()

def structure_A(x):
    k = x.numel()
    A = torch.zeros((k+1, k+1), dtype=x.dtype, device=x.device)
    A[:-1, -1] = 0.5*x
    A[-1, :-1] = 0.5*x
    return A

def quadratic_output(W, x):
    M = W @ W.T
    return torch.sum(M * structure_A(x))

def fit_slope(xs, ys):
    return float(np.polyfit(np.log(xs), np.log(np.maximum(ys, 1e-30)), 1)[0])

def math_checks():
    k, d = 5, 13
    x = torch.randn(k)
    base = torch.randn(k+1, d)
    eps = np.logspace(-3, -0.5, 8)
    residuals=[]; magnitudes=[]
    for e in eps:
        W=e*base
        residuals.append(abs(float(branch_output(W,x)-quadratic_output(W,x))))
        magnitudes.append(float(torch.linalg.norm(W)))
    slope=fit_slope(np.array(magnitudes), np.array(residuals))

    # Exact exchangeability: arbitrary branch permutation leaves both outputs unchanged.
    perm=torch.randperm(d)
    perm_err=abs(float(branch_output(base[:,perm],x)-branch_output(base,x)))
    # Gram-only quadratic invariance: orthogonal mixing leaves WW^T exactly unchanged,
    # while the original module differs only by higher-order terms near zero.
    Q,_=torch.linalg.qr(torch.randn(d,d))
    gram_err=abs(float(quadratic_output(base@Q,x)-quadratic_output(base,x)))
    ortho_errs=[]
    for e in eps:
        ortho_errs.append(abs(float(branch_output(e*base@Q,x)-branch_output(e*base,x))))
    ortho_slope=fit_slope(np.array(magnitudes), np.array(ortho_errs))

    # Width sweep: independent branch cubic contributions predict RMS residual ~ d^1/2
    # for random signs, while a fixed aligned construction predicts linear d. Measure both.
    widths=np.array([4,8,16,32,64])
    rms=[]; aligned=[]
    # Use nested prefixes for random widths and one fixed column/input for the
    # coherent construction; otherwise resampling would add unrelated noise.
    fixed_x=torch.randn(k)
    fixed_col=torch.randn(k+1,1)
    for dd in widths:
        vals=[]
        for rep in range(30):
            X=torch.randn(k+1,dd); xx=torch.randn(k)
            vals.append(abs(float(branch_output(0.03*X,xx)-quadratic_output(0.03*X,xx))))
        rms.append(np.mean(vals))
        X=fixed_col.repeat(1,dd)
        aligned.append(abs(float(branch_output(0.03*X,fixed_x)-quadratic_output(0.03*X,fixed_x))))
    width_slope_rms=fit_slope(widths, np.array(rms))
    width_slope_aligned=fit_slope(widths, np.array(aligned))
    return {
      'device':str(device), 'epsilon_values':eps.tolist(),
      'remainder_values':residuals, 'remainder_loglog_slope':slope,
      'predicted_remainder_slope':3.0,
      'permutation_error':perm_err, 'quadratic_orthogonal_error':gram_err,
      'original_orthogonal_difference_slope':ortho_slope,
      'predicted_orthogonal_difference_slope':3.0,
      'widths':widths.tolist(), 'width_mean_abs_remainder':rms,
      'width_slope_random':width_slope_rms, 'predicted_width_slope_random':0.5,
      'width_slope_aligned':width_slope_aligned, 'predicted_width_slope_aligned':1.0
    }

class Branch(torch.nn.Module):
    def __init__(self,k,d,init=0.05):
        super().__init__(); self.W=torch.nn.Parameter(init*torch.randn(k+1,d))
    def forward(self,x):
        u,v=self.W[:-1],self.W[-1]
        return (phi(x @ u) * v).sum(1)
class Quadratic(torch.nn.Module):
    def __init__(self,k,d,init=0.05):
        super().__init__(); self.W=torch.nn.Parameter(init*torch.randn(k+1,d))
    def forward(self,x):
        # Equivalent low-cost factorized trace: sum_i w_i^T A(x) w_i.
        u,v=self.W[:-1],self.W[-1]
        return (x @ u * v).sum(1)

def train_model(cls, X, y, k, d, steps=500):
    torch.manual_seed(SEED+7)
    m=cls(k,d).to(device); opt=torch.optim.Adam(m.parameters(),lr=0.03)
    Xt,yt=X.to(device),y.to(device)
    losses=[]
    for t in range(steps):
        opt.zero_grad(); pred=m(Xt); loss=((pred-yt)**2).mean(); loss.backward(); opt.step()
        if t in (0,99,249,499): losses.append(float(loss.detach().cpu()))
    return losses

def mini_experiment():
    torch.manual_seed(SEED+99)
    k,d,N=6,16,128
    X=torch.randn(N,k)
    teacher=torch.randn(k+1,d)*0.18
    y=branch_output(teacher,X.T if False else X[0]) # overwritten vectorized below
    u,v=teacher[:-1],teacher[-1]
    y=(phi(X@u)*v).sum(1)
    b=train_model(Branch,X,y,k,d)
    q=train_model(Quadratic,X,y,k,d)
    return {'dataset':'synthetic cubic exchangeable branch teacher','k':k,'width':d,
            'steps':[0,99,249,499],'baseline_branch_mse':b,'idea_quadratic_mse':q,
            'parameter_count_each':(k+1)*d}

def main():
    checks=math_checks(); comparison=mini_experiment()
    # Mechanism passes only if all requested scaling predictions are numerically close.
    checks['prediction_pass'] = (abs(checks['remainder_loglog_slope']-3)<0.15 and
        abs(checks['original_orthogonal_difference_slope']-3)<0.2 and
        abs(checks['width_slope_random']-0.5)<0.25 and
        abs(checks['width_slope_aligned']-1)<0.2 and
        checks['quadratic_orthogonal_error']<1e-5 and checks['permutation_error']<1e-5)
    out={'math_checks':checks,'mini_experiment':comparison}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()