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

TRACK = 'two_view_gauge_localization_v2'

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)

def rot(a):
    c, s = np.cos(a), np.sin(a)
    return np.array([[c,-s],[s,c]], dtype=np.float32)

def two_view_init_numpy(x):
    q1,q2,ev1,ev2,et=[x[:,2*i:2*i+2] for i in range(5)]
    dq=q2-q1; dl=ev2-ev1
    psi=np.arctan2(dq[:,1],dq[:,0])-np.arctan2(dl[:,1],dl[:,0])
    out=[]
    for i,a in enumerate(psi):
        R=rot(float(a)); xx=q1[i]-R@ev1[i]; out.append(xx+R@et[i])
    return np.asarray(out,np.float32)[:, :1]

class SharedMLP(nn.Module):
    def __init__(self, bias=None):
        super().__init__()
        self.backbone=nn.Sequential(nn.Linear(10,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh())
        self.head=nn.Linear(64,1)
        if bias is not None:
            with torch.no_grad():
                self.head.weight.mul_(0.05); self.head.bias.copy_(torch.as_tensor(bias))
    def forward(self,x): return self.head(self.backbone(x))

def run_one(kind, seed, lr, epochs):
    seed_all(seed)
    ds=get_dataset(TRACK, seed=seed, n_train=400, n_test=400)
    bias=None
    if kind=='idea': bias=two_view_init_numpy(ds['xtr'].numpy()).mean(0)
    net, metric, hist=train_model(SharedMLP(bias), ds, epochs=epochs, lr=lr, batch=128, log=lambda *_:None)
    return float(metric), net, ds

def make_train_fn(kind, cfg):
    return lambda seed: run_one(kind, int(seed), cfg['lr'], cfg['epochs'])[0]

def main():
    assert TRACK in all_track_names(), all_track_names()
    seeds=tuple(range(8)); epochs=18
    grid=[{'lr':v,'epochs':epochs} for v in (0.001,0.003,0.01)]
    base=sweep_baseline(lambda cfg: make_train_fn('baseline',cfg), grid, seeds=seeds[:4])
    idea_cfg_results=[]
    for cfg in grid:
        r=evaluate(make_train_fn('idea',cfg), seeds=seeds)
        idea_cfg_results.append({'cfg':cfg,'mean':r['mean'],'std':r['std'],'per_seed':r['per_seed']})
    ib=min(idea_cfg_results,key=lambda z:z['mean'])
    idea={'config':ib['cfg'],'mean':ib['mean'],'std':ib['std'],'per_seed':ib['per_seed'],'n':8}
    rep=make_report(TRACK,'mlp_tiny',base,idea)
    # Trained-model signature: compare standard test MSE of trained models against
    # observed targets, plus the two-view estimate's prediction on the same test data.
    sig=[]
    for s in seeds:
        bm, bn, ds=run_one('baseline',s,base['best_cfg']['lr'],epochs)
        im, inn, _=run_one('idea',s,ib['cfg']['lr'],epochs)
        analytic=two_view_init_numpy(ds['xte'].numpy())
        with torch.no_grad():
            bd=next(bn.parameters()).device; idv=next(inn.parameters()).device
            with torch.no_grad():
                bp=bn(ds['xte'].to(bd)).detach().cpu().numpy(); ip=inn(ds['xte'].to(idv)).detach().cpu().numpy()
        y=ds['yte'].numpy()
        sig.append((float(np.mean((analytic-y)**2)),float(np.mean((bp-y)**2)),float(np.mean((ip-y)**2))))
    vals=np.asarray(sig)
    signature={'prediction':'two-view analytic gauge initialization reduces target prediction error','analytic_test_mse_mean':float(vals[:,0].mean()),'trained_baseline_test_mse_mean':float(vals[:,1].mean()),'trained_idea_test_mse_mean':float(vals[:,2].mean()),'confirmed':bool(vals[:,2].mean() < vals[:,1].mean())}
    rep['mechanism_signature']={'values':signature,'confirmed':signature['confirmed']}
    rep['custom_track']={'name':TRACK,'file':'two_view_registered_track_v2.py','domain':'geometry/localization'}
    Path('bench_report.json').write_text(json.dumps(rep,indent=2))
    print(json.dumps(rep,indent=2))

if __name__=='__main__': main()