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

SEED = 2230
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(8, torch.get_num_threads()))


def sinkhorn_np(G, tau=1.0, steps=10):
    P = np.exp(np.clip(G / tau, -60, 60))
    for _ in range(steps):
        P = P / P.sum(axis=1, keepdims=True)
        P = P / P.sum(axis=0, keepdims=True)
    return P


def sinkhorn_torch(G, tau=1.0, steps=8):
    P = torch.exp(torch.clamp(G / tau, -30, 30))
    for _ in range(steps):
        P = P / P.sum(dim=-1, keepdim=True)
        P = P / P.sum(dim=-2, keepdim=True)
    return P


def toy_verification():
    rng = np.random.default_rng(SEED)
    S, d = 6, 13
    # Prediction 1: alternating normalization improves both constraint residuals.
    G = rng.normal(size=(S, S)) * 2.5
    ks = [0, 1, 2, 3, 5, 10, 20]
    constraint_rows = []
    for k in ks:
        A = sinkhorn_np(G, tau=0.7, steps=k)
        constraint_rows.append({"K": k,
            "max_row_error": float(np.max(np.abs(A.sum(1)-1))),
            "max_col_error": float(np.max(np.abs(A.sum(0)-1))),
            "min_entry": float(A.min())})
    # Prediction 2: DS mixing is non-expansive in Frobenius norm.
    ratios = []
    for scale in [0.0, 0.5, 1.0, 2.0, 5.0, 10.0]:
        for _ in range(40):
            A = sinkhorn_np(rng.normal(size=(S,S))*scale, tau=1.0, steps=100)
            X = rng.normal(size=(S,d))
            ratios.append((scale, float(np.linalg.norm(A@X)/np.linalg.norm(X))))
    ratio_summary = {str(s): {"max": max(r for x,r in ratios if x==s),
                              "mean": float(np.mean([r for x,r in ratios if x==s]))}
                     for s in [0.0,0.5,1.0,2.0,5.0,10.0]}
    # Prediction 3: exact uniform input stream mass is preserved by row-stochastic A.
    mass_errors = []
    for _ in range(100):
        A = sinkhorn_np(rng.normal(size=(S,S))*4, steps=100)
        x = rng.normal(size=S)
        mass_errors.append(float(abs((A@x).sum()-x.sum())))
    # Spectral sweep: DS matrices should have spectral norm <= 1 (up to numerical tol).
    spectral = []
    for scale in [0, .5, 1, 2, 5, 10]:
        vals=[]
        for _ in range(40):
            A=sinkhorn_np(rng.normal(size=(S,S))*scale, steps=100)
            vals.append(np.linalg.svd(A,compute_uv=False)[0])
        spectral.append({"scale":scale,"max_sigma":float(max(vals)),"mean_sigma":float(np.mean(vals))})
    return {"sinkhorn_constraints": constraint_rows, "frobenius_ratio": ratio_summary,
            "arbitrary_total_mass_max_error": max(mass_errors), "spectral_norm": spectral}


class StreamNet(nn.Module):
    def __init__(self, constrained, S=4, D=8, hidden=32, tau=1.0):
        super().__init__(); self.S=S; self.D=D; self.constrained=constrained; self.tau=tau
        self.inp=nn.Linear(1,S*D); self.g=nn.Parameter(torch.zeros(S,S))
        self.mlp=nn.Sequential(nn.LayerNorm(D),nn.Linear(D,hidden),nn.Tanh(),nn.Linear(hidden,D))
        self.out=nn.Linear(S*D,1)
    def forward(self,x):
        X=self.inp(x).view(-1,self.S,self.D)
        if self.constrained: A=sinkhorn_torch(self.g,self.tau,10)
        else: A=self.g # deliberately unconstrained control, initialized near zero
        # normalize control's scale only to make optimization numerically comparable
        if not self.constrained: A=A/(A.abs().sum(dim=1,keepdim=True)+1e-6)
        Xm=torch.einsum('ij,bjd->bid',A,X)
        F=self.mlp(X)
        Y=Xm + 0.1*F
        return self.out(Y.reshape(x.shape[0],-1)), A, X, Xm


def train_compare():
    torch.manual_seed(SEED)
    n=512; x=torch.linspace(-1,1,n).unsqueeze(1); y=torch.sin(5*x)+0.15*x*x
    results={}
    for constrained in [False,True]:
        torch.manual_seed(SEED)
        model=StreamNet(constrained)
        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        losses=[]; grad_spikes=0; mix_ratios=[]; cvs=[]
        for step in range(250):
            ix=torch.randperm(n)[:64]; pred,A,X,Xm=model(x[ix]); loss=((pred-y[ix])**2).mean()
            opt.zero_grad(); loss.backward()
            gn=float(torch.nn.utils.clip_grad_norm_(model.parameters(),1e9))
            if gn>10: grad_spikes+=1
            opt.step(); losses.append(float(loss.detach()))
            with torch.no_grad():
                mix_ratios.append(float(torch.linalg.norm(Xm)/torch.linalg.norm(X)))
                rms=Xm.pow(2).mean(dim=(0,2)).sqrt(); cvs.append(float(rms.std()/(rms.mean()+1e-8)))
        results['idea' if constrained else 'baseline']={
            'final_loss':float(np.mean(losses[-25:])), 'best_loss':float(min(losses)),
            'grad_spikes_gt10':grad_spikes, 'mix_ratio_mean':float(np.mean(mix_ratios)),
            'stream_rms_cv_mean':float(np.mean(cvs)),
            'row_error':float((A.sum(1)-1).abs().max()), 'col_error':float((A.sum(0)-1).abs().max())}
    return results


def main():
    out={'seed':SEED,'verification':toy_verification(),'training':train_compare()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()