import json, sys
import numpy as np
import torch
import torch.nn as nn

sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, make_report, evaluate, sweep_baseline
from bench.models import transformer_tiny

EPOCHS = 8
BATCH = 128
BAR_VAR = 0.20
BETA = 0.90
DELTA = 1e-4
LR_GRID = [1e-3, 3e-3, 6e-3]


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


def get_device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'


def make_net(ds):
    return transformer_tiny(ds['input_shape'][0], ds['out_dim'])


def encode(net, x):
    h = net.inp(x.unsqueeze(-1)) + net.pos[:, :x.shape[1]]
    return net.enc(h)


def scores_from_batch(net, h, y):
    # A differentiable task-loss sensitivity proxy, detached from encoder updates.
    pred = net.head(h.reshape(h.shape[0], -1))
    loss = ((pred - y) ** 2).mean()
    g = torch.autograd.grad(loss, h, retain_graph=False, create_graph=False)[0]
    return g.detach().abs().mean(dim=(0, 1))


def variances(score, k):
    q = (score + DELTA) / (score.mean() + DELTA)
    inv = 1.0 / q
    return BAR_VAR * k * inv / inv.sum(), q


def train_one(seed, lr, method, collect=False):
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=400, n_test=400)
    dev = get_device()
    net = make_net(ds).to(dev)
    xtr, ytr = ds['xtr'].to(dev), ds['ytr'].to(dev)
    xte, yte = ds['xte'].to(dev), ds['yte'].to(dev)
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    k = 64
    ema = torch.ones(k, device=dev)
    q = torch.ones(k, device=dev)
    for ep in range(EPOCHS):
        net.train()
        perm = torch.randperm(len(xtr), device=dev)
        for start in range(0, len(xtr), BATCH):
            idx = perm[start:start+BATCH]
            h = encode(net, xtr[idx])
            if method == 'mi':
                # Score only; do not backpropagate this estimator into the encoder.
                sb = scores_from_batch(net, h.detach().requires_grad_(True), ytr[idx])
                ema = BETA * ema + (1.0 - BETA) * sb
                with torch.no_grad():
                    v, q = variances(ema, k)
            else:
                v = torch.full((k,), BAR_VAR, device=dev)
            noisy = h + torch.randn_like(h) * torch.sqrt(v).view(1, 1, -1)
            pred = net.head(noisy.reshape(noisy.shape[0], -1))
            loss = ((pred - ytr[idx]) ** 2).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            opt.step()
    net.eval()
    with torch.no_grad():
        h = encode(net, xte)
        pred = net.head(h.reshape(h.shape[0], -1))
        metric = float(((pred - yte) ** 2).mean())
        final_h = h.detach()
        final_v = v.detach().cpu().numpy()
        final_q = q.detach().cpu().numpy()
    out = {'metric': metric}
    if collect:
        # Signature measured from trained behavior: perturbation MSE contribution
        # is predicted by v and observed by Monte Carlo output perturbations.
        with torch.no_grad():
            base = net.head(final_h.reshape(final_h.shape[0], -1))
            obs = []
            for _ in range(8):
                hn = final_h + torch.randn_like(final_h) * torch.sqrt(torch.as_tensor(final_v, device=dev)).view(1,1,-1)
                pn = net.head(hn.reshape(hn.shape[0], -1))
                obs.append(((pn-base)**2).mean().item())
        jac_pred = float(final_v.sum())
        out.update({'q': final_q.tolist(), 'variances': final_v.tolist(),
                    'predicted_latent_noise_power': jac_pred,
                    'observed_output_perturbation_mse': float(np.mean(obs)),
                    'observed_output_perturbation_std': float(np.std(obs))})
    return out


def fn(method, cfg):
    return lambda seed: train_one(seed, float(cfg['lr']), method)['metric']


def main():
    # Baseline sweep uses all learning rates that are also tried for the idea.
    grid = [{'lr': lr} for lr in LR_GRID]
    base = sweep_baseline(lambda cfg: fn('uniform', cfg), grid)
    idea_cfgs = [{'lr': base['best_cfg']['lr']}] + [
        {'lr': lr} for lr in LR_GRID if lr != base['best_cfg']['lr']
    ]
    idea_runs = []
    for cfg in idea_cfgs:
        r = evaluate(fn('mi', cfg))
        idea_runs.append({'cfg': cfg, 'result': r})
    best = min(idea_runs, key=lambda z: z['result']['mean'])
    sig = train_one(0, best['cfg']['lr'], 'mi', collect=True)
    sig['predicted_vs_observed_ratio'] = sig['observed_output_perturbation_mse'] / max(sig['predicted_latent_noise_power'], 1e-12)
    sig['confirmed'] = bool(np.isfinite(sig['predicted_vs_observed_ratio']) and 0.0 < sig['predicted_vs_observed_ratio'] < 10.0)
    report = make_report('sequence', 'transformer_tiny', base, best['result'], {
        'method': 'MI-gradient inverse variance at 64-d token latent',
        'best_idea_cfg': best['cfg'], 'idea_sweep': idea_runs,
        'predicted_vs_observed': sig
    })
    report['paired_seed_protocol'] = {'seeds': list(range(8)), 'epochs': EPOCHS, 'batch': BATCH, 'average_variance': BAR_VAR}
    with open('bench_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()