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

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

NTRAIN, NTEST = 400, 200
EPOCHS = 10
BATCH = 128
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 6e-3}]


def riesz_matrix(n, lam=1.0, eps=1e-5):
    D = np.zeros((n, n), dtype=np.float32)
    for i in range(n):
        D[i, i] = -1.0
        D[i, (i + 1) % n] = 1.0
    L = D.T @ D
    H = L + eps * np.eye(n, dtype=np.float32)
    R = lam * D @ np.linalg.inv(np.eye(n, dtype=np.float32) + lam * lam * H)
    return torch.tensor(R), torch.tensor(D)


class RieszTransformer(nn.Module):
    def __init__(self, win, out_dim, lam=1.0):
        super().__init__()
        d = 64
        self.inp = nn.Linear(1, d)
        self.pos = nn.Parameter(torch.zeros(1, win, d))
        nn.init.normal_(self.pos, std=.02)
        R, D = riesz_matrix(win, lam)
        self.register_buffer('R', R)
        self.register_buffer('D', D)
        self.branch = nn.Linear(d, d, bias=False)
        self.g = nn.Parameter(torch.tensor(0.1))
        layer = nn.TransformerEncoderLayer(d, nhead=2, dim_feedforward=128,
                                           batch_first=True, dropout=0.0)
        self.enc = nn.TransformerEncoder(layer, 2)
        self.head = nn.Linear(win * d, out_dim)

    def forward(self, x):
        h = self.inp(x.unsqueeze(-1)) + self.pos[:, :x.shape[1]]
        z = torch.einsum('ij,bjd->bid', self.R, h)
        z = z / (z.square().mean(dim=(1, 2), keepdim=True).sqrt() + 1e-5)
        h = h + self.g * self.branch(z)
        return self.head(self.enc(h).reshape(x.shape[0], -1))


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


def train_one(kind, cfg, seed, return_model=False):
    seed_all(seed)
    ds = get_dataset('sequence', seed, n_train=NTRAIN, n_test=NTEST)
    if kind == 'baseline':
        model = make_model('transformer_tiny', ds['input_shape'], ds['out_dim'])
    else:
        model = RieszTransformer(ds['input_shape'][0], ds['out_dim'])
    net, metric, hist = train_model(model, ds, epochs=EPOCHS,
                                    lr=cfg['lr'], batch=BATCH, log=lambda _: None)
    if net is None:
        raise RuntimeError('training failed')
    if return_model:
        return float(metric), net, ds
    return float(metric)


def signature():
    vals = []
    for seed in range(8):
        b, bm, ds = train_one('baseline', {'lr': 3e-3}, seed, True)
        r, rm, _ = train_one('idea', {'lr': 3e-3}, seed, True)
        devb = next(bm.parameters()).device
        devi = next(rm.parameters()).device
        xb_in = ds['xte'][:64].to(devb)
        ri_in = ds['xte'][:64].to(devi)
        noise_b = torch.randn_like(xb_in) * 0.10
        noise_i = noise_b.to(devi)
        with torch.no_grad():
            xb = bm(xb_in); xbn = bm(xb_in + noise_b)
            ri = rm(ri_in); rin = rm(ri_in + noise_i)
        vals.append({'baseline_output_noise_ratio': float((xbn-xb).norm()/(xb.norm()+1e-8)),
                     'idea_output_noise_ratio': float((rin-ri).norm()/(ri.norm()+1e-8)),
                     'baseline_metric': b, 'idea_metric': r})
    br = np.mean([v['baseline_output_noise_ratio'] for v in vals])
    ir = np.mean([v['idea_output_noise_ratio'] for v in vals])
    return {'mean_baseline_output_noise_ratio': float(br),
            'mean_idea_output_noise_ratio': float(ir),
            'predicted': 'idea should attenuate feature perturbation',
            'confirmed': bool(ir < br), 'per_seed': vals}


def main():
    seed_all(425)
    def baseline_fn(cfg):
        return lambda seed: train_one('baseline', cfg, seed)
    base = sweep_baseline(baseline_fn, GRID, seeds=(0, 1, 2, 3))
    # Full eight-seed idea run at best baseline lr and two nearby union-parity settings.
    idea_runs = []
    for cfg in GRID:
        idea_runs.append((cfg, evaluate(lambda seed, c=cfg: train_one('idea', c, seed))))
    best_cfg, idea_res = min(idea_runs, key=lambda z: z[1]['mean'])
    sig = signature()
    report = make_report('sequence', 'transformer_tiny', base, idea_res,
                         {'mechanism_signature': sig,
                          'track_match': 'sequence-level correlated multi-token forecast',
                          'idea_best_cfg': best_cfg,
                          'idea_grid': [{'cfg': c, 'full': r} for c, r in idea_runs],
                          'epochs': EPOCHS, 'n_train': NTRAIN, 'n_test': NTEST})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()