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

SEED = 1122
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def sinkhorn(cost, eps=0.12, iters=120):
    # Uniform marginals; log-domain updates are stable for this small toy problem.
    n, m = cost.shape
    logK = -cost / eps
    loga = torch.full((n,), -math.log(n), device=cost.device)
    logb = torch.full((m,), -math.log(m), device=cost.device)
    u = torch.zeros_like(loga); v = torch.zeros_like(logb)
    for _ in range(iters):
        u = loga - torch.logsumexp(logK + v[None, :], dim=1)
        v = logb - torch.logsumexp(logK + u[:, None], dim=0)
    return torch.exp(logK + u[:, None] + v[None, :])

def sample_xy(n):
    s = torch.rand(n, device=device) * 4 - 2
    nuisance = torch.randn(n, device=device)
    # Bimodal conditional law, with a mild heteroscedastic component.
    sign = torch.where(torch.rand(n, device=device) < .5, -1., 1.)
    y = sign * (1.0 + .30*s) + (.12 + .025*s.abs()) * torch.randn(n, device=device)
    return torch.stack([s, nuisance], 1), s[:, None], y[:, None]

class Flow(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(3,48), nn.Tanh(), nn.Linear(48,48), nn.Tanh(), nn.Linear(48,1))
    def forward(self,t,y,z): return self.net(torch.cat([t,y,z],1))

class MeanHead(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,z): return self.net(z)

def train_flow(lam, steps=260, n=64):
    model=Flow().to(device); opt=torch.optim.Adam(model.parameters(), lr=2e-3)
    for step in range(steps):
        _, z, y = sample_xy(n); y0=torch.randn_like(y)
        # R is normalized representation-locality cost as in the proposal.
        R=(z-z.T).pow(2); R=R/(R.mean().detach()+1e-6)
        C=(y0-y.T).pow(2)
        P=sinkhorn(C + lam*R, eps=.16, iters=55).detach()
        t=torch.rand(n,n,device=device)
        yt=(1-t)*y0 + t*y.T
        u=y.T-y0
        pred=model(t.reshape(-1,1),yt.reshape(-1,1),z.T.repeat(n,1).reshape(-1,1))
        loss=(P.reshape(-1,1)*(pred-u.reshape(-1,1)).pow(2)).sum()
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def train_mean(steps=260,n=64):
    model=MeanHead().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    for _ in range(steps):
        _,z,y=sample_xy(n); loss=(model(z)-y).pow(2).mean()
        opt.zero_grad();loss.backward();opt.step()
    return model

def true_samples(s, n=1600):
    s=torch.full((n,),float(s),device=device); sign=torch.where(torch.arange(n,device=device)%2==0,-1.,1.)
    return (sign*(1+.30*s)+(.12+.025*abs(s))*torch.randn(n,device=device)).cpu().numpy()

def w1(a,b):
    a=np.sort(a); b=np.sort(b); return float(np.mean(np.abs(a-b)))

def evaluate(flow, mean):
    grid=[-1.5,-.5,.5,1.5]; wf=[]; wm=[]; cover=[]
    for s in grid:
        z=torch.full((512,1),s,device=device); target=true_samples(s,512)
        with torch.no_grad():
            y=torch.randn(512,1,device=device)
            for k in range(30):
                t=torch.full_like(y,(k+.5)/30)
                y=y+flow(t,y,z)/30
            yp=y[:,0].cpu().numpy(); mp=mean(z)[:,0].cpu().numpy()
        wf.append(w1(yp,target)); wm.append(w1(mp,target))
        # fraction in either true mode neighborhood, a simple mode-coverage proxy
        cover.append(float(((yp < -.35) | (yp > .35)).mean()))
    return {'flow_w1':float(np.mean(wf)),'mse_w1':float(np.mean(wm)),
            'flow_mode_coverage':float(np.mean(cover)), 'per_s_w1':wf}

def main():
    # Prediction 1: lambda=0 has no representation-locality pressure.
    # Prediction 2: increasing lambda lowers paired representation distance.
    # Prediction 3: increasing epsilon washes out locality (higher paired distance).
    z=torch.linspace(-2,2,40,device=device)[:,None]; y0=torch.randn(40,1,device=device)
    rows=[]
    for eps in [.06,.16,.40]:
        for lam in [0.,.5,2.,8.]:
            R=(z-z.T).pow(2); R=R/(R.mean()+1e-6); C=(y0-y0.T).pow(2)
            P=sinkhorn(C+lam*R,eps=eps,iters=180)
            locality=float((P*R).sum().cpu()); marginal=float(max((P.sum(1)-1/40).abs().max(),(P.sum(0)-1/40).abs().max()).cpu())
            rows.append({'epsilon':eps,'lambda':lam,'paired_R':locality,'marginal_err':marginal})
    # Exact 1D quadratic OT check: sorted coupling is no worse than a random permutation.
    a=torch.tensor([-.8,.1,1.4,2.0]); b=torch.tensor([-1.1,.4,1.0,2.5])
    sorted_cost=float(((a-b)**2).mean()); random_cost=float(((a-b[torch.tensor([2,0,3,1])])**2).mean())
    # Interpolation derivative check.
    t=.37; y_t=(1-t)*a+t*b; y_t2=(1-(t+1e-4))*a+(t+1e-4)*b
    deriv_err=float(((y_t2-y_t)/1e-4-(b-a)).abs().max())
    math_check={'sorted_ot_cost':sorted_cost,'random_coupling_cost':random_cost,'derivative_max_error':deriv_err}
    mean=train_mean(); flow0=train_flow(0.0); flow8=train_flow(8.0)
    eval0=evaluate(flow0,mean); eval8=evaluate(flow8,mean)
    out={'device':device,'math_check':math_check,'locality_sweep':rows,
         'comparison_lambda0':eval0,'comparison_lambda8':eval8,
         'predictions':{
          'lambda_effect_at_eps_0.16': 'paired_R should decrease monotonically with lambda',
          'epsilon_effect_at_lambda_8': 'paired_R should increase with epsilon',
          'flow_distributional_effect': 'flow W1 should be lower than deterministic MSE W1 and mode coverage near 1'}}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()