import sys, json, math, random
from pathlib import Path
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_model, train_model, evaluate, sweep_baseline, make_report

SEED0 = 2230
EPOCHS = 12
NTR, NTE = 800, 300
BATCH = 128


def sinkhorn(logits, tau=1.0, steps=10):
    # Stable log-domain alternating normalization; output is nonnegative DS.
    z = logits / tau
    for _ in range(steps):
        z = z - torch.logsumexp(z, dim=-1, keepdim=True)
        z = z - torch.logsumexp(z, dim=-2, keepdim=True)
    return torch.exp(z)


class StreamTransformer(nn.Module):
    """S streams of d features, with shared transformer-style feature processing."""
    def __init__(self, input_dim, out_dim, S=4, ds=16, depth=2, tau=1.0, sink_steps=10, alpha=.1, constrained=True):
        super().__init__()
        self.S, self.ds, self.tau, self.sink_steps, self.constrained = S, ds, tau, sink_steps, constrained
        self.inp = nn.Linear(1, S * ds)
        self.pos = nn.Parameter(torch.zeros(1, 32, S, ds))
        self.g = nn.Parameter(torch.eye(S).unsqueeze(0).repeat(depth, 1, 1) * 2.0)
        self.blocks = nn.ModuleList()
        for _ in range(depth):
            self.blocks.append(nn.ModuleDict({
                'norm': nn.LayerNorm(ds),
                'ff': nn.Sequential(nn.Linear(ds, 64), nn.GELU(), nn.Linear(64, ds))
            }))
        self.alpha = nn.Parameter(torch.tensor(float(alpha)))
        self.out = nn.Linear(32 * S * ds, out_dim)

    def forward(self, x):
        X = self.inp(x.unsqueeze(-1)).view(x.shape[0], x.shape[1], self.S, self.ds)
        X = X + self.pos[:, :x.shape[1]]
        for i, block in enumerate(self.blocks):
            A = sinkhorn(self.g[i], self.tau, self.sink_steps) if self.constrained else self.g[i]
            Xm = torch.einsum('ij,btjd->btid', A, X)
            # Shared streamwise residual transformation, preserving stream routing.
            F = block['ff'](block['norm'](Xm))
            X = Xm + self.alpha * F
        return self.out(X.reshape(x.shape[0], -1))


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


def train_one(kind, seed, lr, tau=1.0, alpha=.1, return_model=False):
    set_seed(seed)
    d = get_dataset('sequence', seed, n_train=NTR, n_test=NTE)
    if kind == 'baseline':
        model = StreamTransformer(d['input_shape'][0], d['out_dim'], tau=tau, alpha=alpha, constrained=False)
    else:
        model = StreamTransformer(d['input_shape'][0], d['out_dim'], tau=tau, alpha=alpha)
    net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr, batch=BATCH, log=lambda *_: None)
    if net is None: return (float('inf'), None, d) if return_model else float('inf')
    if return_model: return float(metric), net, d
    return float(metric)


def signature(seed=0, lr=3e-3):
    metric, net, d = train_one('idea', seed, lr, return_model=True)
    device = next(net.parameters()).device
    with torch.no_grad():
        xt = d['xte'][:64].to(device)
        X = net.inp(xt.unsqueeze(-1)).view(xt.shape[0], 32, net.S, net.ds)
        ratios, re, ce = [], [], []
        for g in net.g:
            A = sinkhorn(g, net.tau, net.sink_steps)
            Xm = torch.einsum('ij,btjd->btid', A, X)
            ratios.append(float(torch.linalg.norm(Xm) / torch.linalg.norm(X)))
            re.append(float((A.sum(-1)-1).abs().max())); ce.append(float((A.sum(-2)-1).abs().max()))
    # Prediction tested on trained model: DS mixing is non-expansive (allow tiny FP tolerance).
    return {'trained_model': True, 'predicted': {'max_frobenius_ratio': '<=1', 'row_col_error': '0'},
            'observed': {'max_frobenius_ratio': max(ratios), 'mean_frobenius_ratio': float(np.mean(ratios)),
                         'max_row_error': max(re), 'max_col_error': max(ce)},
            'confirmed': bool(max(ratios) <= 1.0001 and max(re) < 1e-4 and max(ce) < 1e-4),
            'metric_on_signature_model': float(metric)}


def main():
    # Union parity: every idea learning rate is also evaluated by baseline.
    grid = [{'lr': 7e-4}, {'lr': 1e-3}, {'lr': 1.5e-3}, {'lr': 2e-3}, {'lr': 3e-3}]
    base = sweep_baseline(lambda c: lambda s: train_one('baseline', s, c['lr']), grid)
    best_lr = base['best_cfg']['lr']
    idea_grid = [1e-3, 1.5e-3, 2e-3]
    idea_runs = []
    for lr in idea_grid:
        r = evaluate(lambda s, lr=lr: train_one('idea', s, lr))
        idea_runs.append({'lr': lr, 'result': r})
    best = min(idea_runs, key=lambda z: z['result']['mean'])
    rep = make_report('sequence', 'transformer_tiny', base, best['result'], {
        'idea_lr_sweep': idea_runs, 'selected_lr': best['lr'], **signature(0, best['lr'])})
    rep['custom_track'] = None
    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()