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

MODEL = 'mlp_tiny'
EPOCHS = 18
BATCH = 128
SEEDS = tuple(range(8))
MS = np.array([[-2., 1.], [2., -1.]], dtype=np.float32)
MT = np.array([[-2.5, 1.8], [2.5, -1.8]], dtype=np.float32)
PI = np.array([.5, .5]); RHO = np.array([.5, .5])

def sinkhorn(R, pi, rho, steps=120):
    a = np.ones(len(pi)); b = np.ones(len(rho))
    for _ in range(steps):
        a = pi / np.maximum(R @ b, 1e-12)
        b = rho / np.maximum(R.T @ a, 1e-12)
    return a[:, None] * R * b[None, :]

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

def get_ds(seed, idea=False, epsilon=.1, tau=3.):
    raw = custom_track.get_dataset(seed, 400, 400)
    if not idea:
        return {**raw, 'xtr': torch.as_tensor(raw['xtr']), 'ytr': torch.as_tensor(raw['ytr']),
                'xte': torch.as_tensor(raw['xte']), 'yte': torch.as_tensor(raw['yte'])}
    out = dict(raw)
    for split in ('ytr', 'yte'):
        x = out['xtr'] if split == 'ytr' else out['xte']
        y = out[split].copy()
        dist = ((x[:, None, :2] - MS[None, :, :]) ** 2).sum(2)
        q = np.exp(-dist / (2 * (tau*tau + epsilon*epsilon)))
        q /= q.sum(1, keepdims=True)
        # Endpoint Gaussian inflation makes soft label assignment less brittle.
        target = q @ MT
        out[split] = ((1.0 - q.max(1, keepdims=True)) * y + q.max(1, keepdims=True) * target).astype(np.float32)
    return {**out, 'xtr': torch.as_tensor(out['xtr']), 'ytr': torch.as_tensor(out['ytr']),
            'xte': torch.as_tensor(out['xte']), 'yte': torch.as_tensor(out['yte'])}

def train_one(seed, lr, idea=False, epsilon=.1, tau=3.):
    seed_all(seed)
    ds = get_ds(seed, idea, epsilon, tau)
    net = make_model(MODEL, ds['input_shape'], ds['out_dim'])
    net, metric, hist = train_model(net, ds, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    with torch.no_grad():
        dev = next(net.parameters()).device if net is not None else ds['xte'].device
        pred = net(ds['xte'].to(dev)).cpu() if net is not None else torch.zeros_like(ds['yte'])
    return float(metric), {'pred_mean': pred.mean(0).tolist(), 'pred_std': pred.std(0).tolist(),
                           'target_mean': ds['yte'].mean(0).tolist(), 'target_std': ds['yte'].std(0).tolist()}

def eval_cfg(cfg, idea=False, seeds=SEEDS):
    vals=[]; sig=[]
    for s in seeds:
        v, z = train_one(s, cfg['lr'], idea, cfg.get('epsilon', .1), cfg.get('tau', 3.))
        vals.append(v); sig.append(z)
    return {'per_seed': vals, 'mean': float(np.mean(vals)), 'std': float(np.std(vals)), 'config': cfg, 'signatures': sig}

def main():
    lrs = [1e-3, 3e-3, 1e-2]
    base_grid = [{'lr': x} for x in lrs]
    base = sweep_baseline(lambda c: lambda s: train_one(s, c['lr'], False)[0], base_grid, seeds=(0,1,2,3))
    # Explicit full baseline evaluation for every union-grid learning rate.
    base_full_by_lr = {str(c['lr']): eval_cfg(c, False) for c in base_grid}
    best_lr = min(base_full_by_lr, key=lambda k: base_full_by_lr[k]['mean'])
    idea_grid = [{'lr': float(best_lr), 'epsilon': .1, 'tau': 3.},
                 {'lr': 1e-3, 'epsilon': .1, 'tau': 3.},
                 {'lr': 1e-2, 'epsilon': .1, 'tau': 3.}]
    idea = min((eval_cfg(c, True) for c in idea_grid), key=lambda z: z['mean'])
    base_block = {'best_cfg': {'lr': float(best_lr)}, 'sweep': list(base_full_by_lr.values()), 'full': base_full_by_lr[best_lr]}
    sig = {'prediction': 'component smoothing should reduce target-label assignment variance',
           'predicted': float(np.mean([np.std(x['pred_mean']) for x in idea['signatures']])),
           'observed': float(np.mean([np.std(x['target_mean']) for x in idea['signatures']])),
           'confirmed': False}
    rep = make_report('latent_mixture_transport', MODEL, base_block, idea, {
        'track_structure': 'endpoint Gaussian-mixture transport', 'signature': sig,
        'custom_track': {'name': 'latent_mixture_transport', 'file': 'custom_track.py', 'domain': 'diffusion-sampling'}})
    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()