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

SEED = 1029
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    if device.type == "cuda":
        torch.cuda.manual_seed_all(SEED)
        # A failed CUDA allocation should not invalidate the experiment.
        torch.zeros(1, device=device)
except Exception:
    device = torch.device("cpu")

class Coupling(nn.Module):
    """Additive coupling C_f,g and its analytically computed inverse."""
    def __init__(self, d, scale=0.1):
        super().__init__()
        h = d // 2
        self.f = nn.Linear(h, h, bias=False)
        self.g = nn.Linear(h, h, bias=False)
        nn.init.normal_(self.f.weight, std=scale / math.sqrt(h))
        nn.init.normal_(self.g.weight, std=scale / math.sqrt(h))

    def forward(self, x):
        a, b = x.chunk(2, dim=-1)
        ap = a + self.f(b)
        bp = b + self.g(ap)
        return torch.cat((ap, bp), dim=-1)

    def inverse(self, y):
        ya, yb = y.chunk(2, dim=-1)
        b = yb - self.g(ya)
        a = ya - self.f(b)
        return torch.cat((a, b), dim=-1)

class InverseCopy(nn.Module):
    def __init__(self, source):
        super().__init__()
        d = source.f.in_features * 2
        self.block = Coupling(d, 0.01)
        self.block.load_state_dict(source.state_dict())
    def forward(self, x):
        return self.block.inverse(x)

class Pair(nn.Module):
    def __init__(self, source):
        super().__init__()
        self.forward_block = Coupling(source.f.in_features * 2, 0.01)
        self.forward_block.load_state_dict(source.state_dict())
        self.inverse_block = InverseCopy(source)
    def forward(self, x):
        return self.inverse_block(self.forward_block(x))

class Net(nn.Module):
    def __init__(self, d, source=None):
        super().__init__()
        self.block = Coupling(d, 0.15) if source is None else source
        self.out = nn.Linear(d, 1)
    def forward(self, x):
        return self.out(self.block(x))
    def append_pair(self):
        old = self.block
        self.block = nn.Sequential(old, Pair(old).to(next(old.parameters()).device))


def coupling_error(scale, dtype=torch.float32, n=4096):
    torch.manual_seed(7)
    d = 8
    c = Coupling(d, 0.01).to(device=device, dtype=dtype)
    with torch.no_grad():
        c.f.weight.mul_(scale / 0.01); c.g.weight.mul_(scale / 0.01)
        x = torch.randn(n, d, device=device, dtype=dtype)
        z = c.inverse(c(x))
        err = (z-x).abs().max().item()
        rel = (z-x).norm().item() / max(x.norm().item(), 1e-30)
    return err, rel

def perturbation_sweep():
    # At an identity pair, a small untied parameter step produces O(eta) output change.
    torch.manual_seed(11)
    d=8; x=torch.randn(512,d,device=device)
    source=Coupling(d,0.08).to(device)
    pair=Pair(source).to(device)
    with torch.no_grad(): base=pair(x).clone()
    loss=(pair(x)**2).mean(); loss.backward()
    grads=[p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p) for p in pair.parameters()]
    grad_norm=math.sqrt(sum((g**2).sum().item() for g in grads))
    vals=[]
    for eta in [1e-5, 3e-5, 1e-4, 3e-4, 1e-3]:
        q=Pair(source).to(device); q.load_state_dict(pair.state_dict())
        with torch.no_grad():
            for p,g in zip(q.parameters(), grads):
                p.sub_(eta*g)
            delta=(q(x)-base).norm().item()/x.norm().item()
        vals.append((eta,delta))
    slope=np.polyfit(np.log([a for a,b in vals if b>0]),np.log([b for a,b in vals if b>0]),1)[0]
    return {"gradient_norm":grad_norm,"sweep":vals,"loglog_slope":float(slope)}

def train(model, x, y, steps, lr=2e-3):
    opt=torch.optim.Adam(model.parameters(),lr=lr)
    losses=[]
    for _ in range(steps):
        opt.zero_grad(set_to_none=True)
        loss=((model(x)-y)**2).mean()
        loss.backward(); opt.step(); losses.append(loss.item())
    return losses

def mini_experiment():
    torch.manual_seed(SEED)
    d=8; n=2048
    x=torch.randn(n,d,device=device)
    # A target requiring a nontrivial but small transformation.
    target=Coupling(d,0.22).to(device)
    with torch.no_grad(): y=(target(target(x))).sum(dim=1,keepdim=True)/2.0
    torch.manual_seed(SEED)
    progressive=Net(d).to(device)
    t0=time.time(); pre=train(progressive,x,y,250); before=((progressive(x)-y)**2).mean().item()
    progressive.append_pair()
    with torch.no_grad(): after=((progressive(x)-y)**2).mean().item()
    insertion_rel=(after-before)/max(abs(before),1e-30)
    post=train(progressive,x,y,300)
    # Standard shallow continuation, and a randomly initialized two-block model.
    torch.manual_seed(SEED); shallow=Net(d).to(device); train(shallow,x,y,250); shallow_post=train(shallow,x,y,300)
    torch.manual_seed(SEED+1); random_deep=Net(d).to(device); random_deep.block=nn.Sequential(random_deep.block,Coupling(d,0.15).to(device)); random_post=train(random_deep,x,y,300)
    return {"device":str(device),"before_insertion_loss":before,"after_insertion_loss":after,"insertion_relative_change":insertion_rel,
            "paired_loss_after_300":post[-1],"shallow_loss_after_300":shallow_post[-1],"random_deep_loss_after_300":random_post[-1],
            "paired_best_300":min(post),"wall_seconds":time.time()-t0}

def main():
    scales=[0.01,0.03,0.1,0.3,1.0]
    cancel=[]
    for s in scales:
        e,r=coupling_error(s); cancel.append({"scale":s,"max_abs_error":e,"relative_error":r,"predicted":"identity; fp error increases with operator products"})
    pert=perturbation_sweep()
    mini=mini_experiment()
    result={"seed":SEED,"math_predictions":[
        "C^{-1}(C(x))=x, so insertion loss change is zero up to floating point error.",
        "Composition roundoff should increase with coupling scale (products of f and g operators).",
        "After untie, a parameter step of size eta changes the pair output proportionally to eta (log-log slope 1)."],
        "cancellation_sweep":cancel,"untied_step_sweep":pert,"mini_experiment":mini}
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__ == "__main__": main()