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

SEEDS = tuple(range(8))
EPOCHS = 14
BATCH = 64
WEIGHT_DECAY = 1e-5


def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)


def project_input(x):
    p, v = x[..., :3], x[..., 4:]
    d = torch.cdist(p, p)
    k = min(8, p.shape[1]-1)
    idx = d.topk(k+1, largest=False).indices[..., 1:]
    src = p.unsqueeze(1).expand(-1, p.shape[1], -1, -1)
    nb = torch.gather(src, 2, idx.unsqueeze(-1).expand(-1, -1, -1, 3))
    c = nb - p.unsqueeze(2)
    cov = torch.einsum('bmki,bmkj->bmij', c, c) / float(k)
    _, e = torch.linalg.eigh(cov)
    n = e[..., 0]
    n = n / (n.norm(dim=-1, keepdim=True) + 1e-8)
    eye = torch.eye(3, dtype=x.dtype, device=x.device).view(1,1,3,3)
    P = eye - n.unsqueeze(-1) * n.unsqueeze(-2)
    vp = torch.einsum('bmij,bmj->bmi', P, v)
    r = p - p[:, :1]
    rq = torch.einsum('bij,bmj->bmi', P[:, 0], r)
    out = torch.cat([rq, x[..., 3:4], vp], dim=-1)
    return out.reshape(x.shape[0], -1), n, vp


class SharedNet(nn.Module):
    def __init__(self, intrinsic=False, width=64):
        super().__init__(); self.intrinsic = intrinsic
        dim = 20 * (7 if not intrinsic else 7)
        self.net = nn.Sequential(nn.Linear(dim, width), nn.ReLU(),
                                 nn.Linear(width, width), nn.ReLU(),
                                 nn.Linear(width, 1))

    def forward(self, x):
        if self.intrinsic:
            z, n, vp = project_input(x)
        else:
            z = x.reshape(x.shape[0], -1); n = None; vp = None
        return self.net(z), n, vp


def run_one(seed, intrinsic, lr):
    seed_all(seed)
    raw = track.get_dataset(seed, 400, 100)
    ds = {k: torch.as_tensor(raw[k], dtype=torch.float32) for k in ('xtr','ytr','xte','yte')}
    ds.update(task='regression', metric='mse')
    model = SharedNet(intrinsic=intrinsic)
    class Wrapper(nn.Module):
        def __init__(self, m): super().__init__(); self.m=m
        def forward(self, x): return self.m(x)[0]
    net, metric, hist = train_model(Wrapper(model), ds, epochs=EPOCHS, lr=lr,
                                    batch=BATCH, weight_decay=WEIGHT_DECAY, log=lambda *_: None)
    if net is None: return float('nan'), None
    return float(metric), model


def factory(intrinsic, lr):
    return lambda seed: run_one(seed, intrinsic, lr)[0]


def main():
    # Both methods are evaluated at every learning rate in the shared union.
    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
    base = sweep_baseline(lambda c: factory(False, c['lr']), grid, seeds=(0,1,2,3))
    # Ensure baseline is evaluated on all candidate settings as required.
    baseline_candidates = []
    for c in grid:
        baseline_candidates.append({'cfg': c, 'full': evaluate(factory(False, c['lr']), SEEDS)})
    best_cfg = min(baseline_candidates, key=lambda z: z['full']['mean'])['cfg']
    base['union_full'] = baseline_candidates
    idea_candidates = []
    for c in grid:
        idea_candidates.append({'cfg': c, 'full': evaluate(factory(True, c['lr']), SEEDS)})
    idea = min(idea_candidates, key=lambda z: z['full']['mean'])['full']
    best_idea_cfg = min(idea_candidates, key=lambda z: z['full']['mean'])['cfg']

    # Re-test the proposed mechanism on independently trained benchmark models.
    raw_leak, proj_leak, rot_change = [], [], []
    for s in SEEDS:
        metric, model = run_one(s, True, best_idea_cfg['lr'])
        if model is None: continue
        raw = track.get_dataset(s+9000, 400, 20)
        x = torch.as_tensor(raw['xte'], dtype=torch.float32)
        device = next(model.parameters()).device
        x = x.to(device)
        with torch.no_grad():
            _, n, vp = model(x)
            rawv = x[...,4:]
            raw_leak.append(float(torch.abs((n*rawv).sum(-1)).mean()))
            proj_leak.append(float(torch.abs((n*vp).sum(-1)).max()))
            R = torch.tensor([[0.,-1.,0.],[1.,0.,0.],[0.,0.,1.]], device=device)
            xr=x.clone(); xr[...,:3]=x[...,:3]@R.T; xr[...,4:]=x[...,4:]@R.T
            pred1=model(x)[0]; pred2=model(xr)[0]
            rot_change.append(float(torch.abs(pred1-pred2).mean()/(torch.abs(pred1).mean()+1e-8)))
    signature = {'raw_mean_normal_component': float(np.mean(raw_leak)),
                 'projected_max_normal_component': float(np.max(proj_leak)),
                 'rotation_relative_prediction_change': float(np.mean(rot_change)),
                 'predicted_projected_leakage': '<1e-5',
                 'confirmed': bool(np.max(proj_leak) < 1e-5)}
    report = make_report('tangent_surface_vector_regression', 'shared_mlp',
                         {'best_cfg': best_cfg, 'sweep': base['sweep'],
                          'union_full': baseline_candidates, 'full': base['full']},
                         idea, {'mechanism_signature': signature,
                                'custom_track': {'name': track.META['name'],
                                                 'file': 'tangent_surface_track.py',
                                                 'domain': track.META['domain']},
                                'idea_sweep': idea_candidates,
                                'best_idea_cfg': best_idea_cfg})
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()