import json, math, random
from pathlib import Path
import numpy as np
import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader

SEED = 1337
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
try:
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    DEVICE = torch.device('cpu')

# Smoothed R-function-like disk support: positive inside, negative outside.
def disk_score(x, cx=0.0, cy=0.0, radius=0.65, smooth=0.035):
    r = torch.sqrt((x[:, 0]-cx)**2 + (x[:, 1]-cy)**2 + 1e-8)
    return torch.sigmoid((radius-r)/smooth)

def make_data(kind, n, seed, train=True):
    g = torch.Generator().manual_seed(seed)
    x = (torch.rand(n, 2, generator=g) * 2.4 - 1.2)
    true_g = disk_score(x)
    wrong_g = disk_score(x, cx=0.48, cy=-0.30)
    if kind == 'correct':
        # Geometry is the dominant signal, with a small smooth nuisance term.
        y = true_g + 0.08*x[:, 0] - 0.04*x[:, 1]
        gp = true_g
    elif kind == 'wrong':
        y = torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2
        gp = wrong_g
    elif kind == 'none':
        y = torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2
        gp = true_g
    elif kind == 'half':
        # The geometry rule applies to the left half only; the right half is unrelated.
        obey = (x[:,0] < 0).float()
        y = obey*true_g + (1-obey)*(torch.sin(2.2*x[:,0]) + 0.35*x[:,1]**2)
        y = y + 0.05*x[:,0]
        gp = true_g
    else: raise ValueError(kind)
    if train:
        y = y + 0.035*torch.randn(n, generator=g)
    return x, y[:,None], gp[:,None]

class Free(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(2,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,x,g=None): return self.net(x)

class Geom(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(3,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,x,g): return self.net(torch.cat([x,g],1))

class Concat(Geom):
    pass

class Mixture(nn.Module):
    def __init__(self, input_gate=False):
        super().__init__(); self.free=Free(); self.geom=Geom()
        self.input_gate=input_gate
        self.beta=nn.Parameter(torch.tensor(-4.0))
        self.mixer=nn.Sequential(nn.Linear(2,16),nn.Tanh(),nn.Linear(16,1)) if input_gate else None
    def forward(self,x,g):
        hf=self.free(x); hg=self.geom(x,g)
        a=torch.sigmoid(self.mixer(x) if self.input_gate else self.beta)
        return a*hg+(1-a)*hf, hf, hg, a

def train_model(kind, mode, epochs=160, n=640, lam=0.0):
    x,y,g=make_data(kind,n,100+SEED)
    if mode=='free': model=Free().to(DEVICE)
    elif mode in ('geom','concat'): model=Geom().to(DEVICE)
    elif mode=='mixer': model=Mixture(input_gate=True).to(DEVICE)
    elif mode=='global': model=Mixture(input_gate=False).to(DEVICE)
    else: raise ValueError(mode)
    opt=torch.optim.Adam(model.parameters(),lr=0.012,weight_decay=1e-4)
    x,y,g=x.to(DEVICE),y.to(DEVICE),g.to(DEVICE)
    for _ in range(epochs):
        opt.zero_grad()
        if mode=='free': pred=model(x); loss=((pred-y)**2).mean()
        elif mode in ('geom','concat'): pred=model(x,g); loss=((pred-y)**2).mean()
        else:
            pred,hf,hg,a=model(x,g); loss=((pred-y)**2).mean()+lam*a.mean()
        loss.backward(); opt.step()
    return model

def evaluate(model, kind, mode):
    x,y,g=make_data(kind,1800,900+SEED,False)
    x,y,g=x.to(DEVICE),y.to(DEVICE),g.to(DEVICE)
    with torch.no_grad():
        if mode=='free': p=model(x); a=torch.zeros_like(p); hf=p; hg=p
        elif mode in ('geom','concat'): p=model(x,g); a=torch.ones_like(p); hf=p; hg=p
        else: p,hf,hg,a=model(x,g)
        rmse=torch.sqrt(((p-y)**2).mean()).item()
        branch_f=torch.sqrt(((hf-y)**2).mean()).item(); branch_g=torch.sqrt(((hg-y)**2).mean()).item()
        ent=(-(a.clamp(1e-7,1-1e-7)*torch.log(a.clamp(1e-7,1-1e-7))+(1-a).clamp(1e-7,1-1e-7)*torch.log((1-a).clamp(1e-7,1-1e-7))).mean()).item()
        return {'rmse':rmse,'free_branch_rmse':branch_f,'geom_branch_rmse':branch_g,'mean_alpha':a.mean().item(),'gate_entropy':ent}

def math_check():
    torch.manual_seed(7)
    h1=torch.randn(10000); h2=torch.randn(10000); a=torch.rand(10000)
    lo=torch.minimum(h1,h2); hi=torch.maximum(h1,h2)
    inside=bool(((a*h1+(1-a)*h2 >= lo-1e-6)&(a*h1+(1-a)*h2 <= hi+1e-6)).all())
    beta=torch.tensor(0.,requires_grad=True); alpha=torch.sigmoid(beta); (alpha).backward()
    deriv=float(beta.grad); beta2=torch.tensor(-4.)
    return {'convex_interpolation_inside_bounds':inside,'d_mean_alpha_d_beta_at_zero':deriv,'initial_alpha_beta_minus4':float(torch.sigmoid(beta2)),'suppression_gradient_sign':deriv>0}

def main():
    out={'device':str(DEVICE),'math_check':math_check(),'results':{}}
    for kind in ['correct','wrong','none','half']:
        out['results'][kind]={}
        for mode in ['free','geom','concat','global','mixer']:
            # concat is architecturally the same geometry-feature baseline; included explicitly.
            model=train_model(kind,mode)
            out['results'][kind][mode]=evaluate(model,kind,mode)
        # Repeat the global gate with mild suppression to test the stated regularizer behavior.
        m=train_model(kind,'global',lam=0.02)
        out['results'][kind]['global_lambda_0.02']=evaluate(m,kind,'global')
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()