import sys, json, random
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import evaluate, sweep_baseline, make_report
from gql_track import get_dataset

SEEDS = tuple(range(8))
GRID = [{'lr': 1e-3, 'epochs': 12}, {'lr': 3e-3, 'epochs': 12}, {'lr': 6e-3, 'epochs': 12}]
EPS = 1e-4

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(8, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 4))
    def forward(self, x):
        return self.net(x)

def gql_limit(u0, cand, eps=EPS, iters=60):
    delta = cand - u0
    def margin(t):
        v = u0 + t.unsqueeze(-1) * delta
        return v[:, 3] - torch.sqrt(torch.sum(v[:, :3] * v[:, :3], dim=1) + 1e-12)
    one = torch.ones(u0.shape[0], device=u0.device)
    bad = (cand[:, 0] < eps) | (margin(one) < eps)
    lo = torch.zeros_like(one); hi = one
    for _ in range(iters):
        mid = (lo + hi) * .5
        ok = margin(mid) >= eps
        lo = torch.where(ok, mid, lo); hi = torch.where(ok, hi, mid)
    theta = torch.where(bad, lo, one)
    db = cand[:, 0] < eps
    td = torch.clamp((u0[:, 0] - eps) / torch.clamp(u0[:, 0] - cand[:, 0], min=1e-12), 0., 1.)
    theta = torch.where(db, torch.minimum(theta, td), theta)
    return u0 + theta.unsqueeze(-1) * delta, theta

def train(seed, cfg, limited):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    d = get_dataset(seed, 400, 100)
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model = MLP().to(device)
        xtr, ytr = torch.as_tensor(d['xtr'], device=device), torch.as_tensor(d['ytr'], device=device)
        xte, yte = torch.as_tensor(d['xte'], device=device), torch.as_tensor(d['yte'], device=device)
        opt = torch.optim.Adam(model.parameters(), lr=cfg['lr'])
        model.train()
        for _ in range(cfg['epochs']):
            p = model(xtr)
            if limited:
                p, _ = gql_limit(xtr[:, :4], p)
            loss = ((p-ytr)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        model.eval()
        with torch.no_grad():
            raw = model(xte)
            pred, theta = gql_limit(xte[:, :4], raw) if limited else (raw, torch.ones(raw.shape[0], device=device))
            mse = ((pred-yte)**2).mean().item()
            margin = pred[:,3]-torch.sqrt(torch.sum(pred[:,:3]**2,dim=1)+1e-12)
            invalid = ((pred[:,0] < EPS) | (margin < EPS)).float().mean().item()
            limited_fraction = (theta < .999999).float().mean().item()
        return mse, {'invalid_rate':invalid,'limited_fraction':limited_fraction,'mean_theta':theta.mean().item()}
    except Exception:
        device=torch.device('cpu'); torch.manual_seed(seed)
        model=MLP(); xtr=torch.as_tensor(d['xtr']); ytr=torch.as_tensor(d['ytr'])
        opt=torch.optim.Adam(model.parameters(),lr=cfg['lr'])
        for _ in range(cfg['epochs']):
            p=model(xtr); p,_=gql_limit(xtr[:,:4],p) if limited else (p,None)
            loss=((p-ytr)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            raw=model(torch.as_tensor(d['xte'])); pred,th=gql_limit(torch.as_tensor(d['xte'][:,:4]),raw) if limited else (raw,torch.ones(100))
            margin=pred[:,3]-torch.sqrt(torch.sum(pred[:,:3]**2,1)+1e-12)
            return ((pred-torch.as_tensor(d['yte']))**2).mean().item(), {'invalid_rate':float(((pred[:,0]<EPS)|(margin<EPS)).float().mean()),'limited_fraction':float((th<.999999).float().mean()),'mean_theta':float(th.mean())}

def fn(cfg, limited):
    return lambda seed: train(seed,cfg,limited)[0]

def mechanism_signature():
    rows=[]
    for s in SEEDS:
        cfg=GRID[1]; d=get_dataset(s,400,100); random.seed(s); np.random.seed(s); torch.manual_seed(s)
        model=MLP(); opt=torch.optim.Adam(model.parameters(),lr=cfg['lr']); x=torch.as_tensor(d['xtr']); y=torch.as_tensor(d['ytr'])
        for _ in range(cfg['epochs']):
            p,_=gql_limit(x[:,:4],model(x)); loss=((p-y)**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            raw=model(torch.as_tensor(d['xte'])); safe,th=gql_limit(torch.as_tensor(d['xte'][:,:4]),raw)
            margin=safe[:,3]-torch.sqrt(torch.sum(safe[:,:3]**2,1)+1e-12)
            rows.append({'seed':s,'invalid_raw':float(((raw[:,0]<EPS)|((raw[:,3]-torch.sqrt(torch.sum(raw[:,:3]**2,1)+1e-12))<EPS)).float().mean()),'invalid_limited':float(((safe[:,0]<EPS)|(margin<EPS)).float().mean()),'mean_theta':float(th.mean())})
    pred=0.0; observed=float(np.mean([r['invalid_limited'] for r in rows])); raw=float(np.mean([r['invalid_raw'] for r in rows]))
    return {'claim':'limiting a valid baseline candidate along the complete residual enforces admissibility','predicted_invalid_rate_after_limit':pred,'observed_invalid_rate_after_limit':observed,'observed_raw_invalid_rate':raw,'trained_model_measurements':rows,'confirmed':bool(observed <= 1e-5 and raw > observed)}

def main():
    base=sweep_baseline(lambda c: fn(c,False), GRID, seeds=(0,1,2,3))
    idea_sweep=[{'cfg':c,'mean':evaluate(fn(c,True),seeds=(0,1,2,3))['mean']} for c in GRID]
    best=min(idea_sweep,key=lambda z:z['mean'])['cfg']
    idea=evaluate(fn(best,True),seeds=SEEDS)
    rep=make_report('custom_relativistic_conservative_regression','mlp_tiny',base,idea,{'mechanism_signature':mechanism_signature(),'custom_track':{'name':'relativistic_conservative_regression','file':'gql_track.py','domain':'conservative_state_admissibility'},'sweep_parity':{'union_grid':GRID,'idea_sweep':idea_sweep}})
    with open('bench_report.json','w') as f: json.dump(rep,f,indent=2)
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()