import sys, json, random
import numpy as np
import torch
from torch import nn
from scipy.stats import wasserstein_distance
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
from bench import get_dataset, train_model, sweep_baseline, evaluate, make_report

SEED = 461
LR_GRID = [1.5e-3, 3e-3, 6e-3]
EPOCHS = 12
BATCH = 128
K = 8

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

class PoolNet(nn.Module):
    """Matched end-to-end systems; only the set pooling map differs."""
    def __init__(self, kind, k=K):
        super().__init__()
        self.kind = kind
        self.k = k
        # A shared-size output head makes the comparison parameter-matched.
        self.head = nn.Sequential(nn.Linear(k, 32), nn.Tanh(), nn.Linear(32, 1))
        if kind == 'idea':
            # scalar probes with |slope|<=1; subtracting phi(0) enforces phi(theta)=0
            self.raw_slope = nn.Parameter(torch.randn(k) * 0.15)
            self.bias = nn.Parameter(torch.zeros(k))
        elif kind != 'baseline':
            raise ValueError(kind)

    def probe_values(self, x):
        slope = torch.tanh(self.raw_slope)
        return torch.tanh(x.unsqueeze(-1) * slope + self.bias) - torch.tanh(self.bias)

    def pooled(self, x):
        if self.kind == 'baseline':
            z = x.mean(dim=1, keepdim=True).expand(-1, self.k)
        else:
            z = self.probe_values(x).mean(dim=1)
        return z

    def forward(self, x):
        return self.head(self.pooled(x))

def train_one(kind, seed, lr, retain=False):
    seed_all(seed)
    d = get_dataset('sequence', seed, n_train=400, n_test=400)
    model = PoolNet(kind)
    net, metric, hist = train_model(model, d, epochs=EPOCHS, lr=lr,
                                    batch=BATCH, weight_decay=0.0)
    if net is None: raise RuntimeError('benchmark training failed')
    return float(metric), net, d

def fn(kind, lr):
    return lambda seed: train_one(kind, seed, lr)[0]

def mechanism_signature(net, d, n_pairs=100):
    net.eval()
    dev = next(net.parameters()).device
    x = d['xte'].to(dev)
    with torch.no_grad():
        vals = net.probe_values(x).mean(dim=1).detach().cpu().numpy()
    # For equal-weight 1D empirical measures, W1 is an independent task-free
    # geometric reference; D_hat is computed from the trained probes.
    rng = np.random.default_rng(9001)
    dh, w1 = [], []
    xx = d['xte'].cpu().numpy()
    for _ in range(n_pairs):
        a, b = rng.integers(0, len(xx), 2)
        # Re-evaluate individual probe pooled values, not model predictions.
        va, vb = vals[a], vals[b]
        dh.append(float(np.max(np.abs(va - vb))))
        w1.append(float(wasserstein_distance(xx[a], xx[b])))
    dh, w1 = np.asarray(dh), np.asarray(w1)
    corr = float(np.corrcoef(dh, w1)[0,1]) if np.std(dh)>1e-12 and np.std(w1)>1e-12 else 0.0
    ratio = float(np.mean(dh / np.maximum(w1, 1e-8)))
    violations = float(np.mean(dh > w1 + 1e-5))
    # Prediction: a valid 1-Lipschitz dual estimate should not exceed W1.
    return {'quantity': 'trained probe D_hat versus independent 1D Wasserstein-1',
            'predicted': 'D_hat <= W1 for every pair; positive geometric association',
            'observed_mean_D_hat': float(np.mean(dh)),
            'observed_mean_W1': float(np.mean(w1)),
            'mean_ratio_D_hat_over_W1': ratio,
            'pearson_correlation': corr,
            'violation_fraction': violations,
            'confirmed': bool(violations <= 0.05 and corr > 0.2)}

def main():
    # Baseline sweep and idea sweep use identical learning-rate union and seeds.
    grid = [{'lr': lr, 'epochs': EPOCHS, 'batch': BATCH} for lr in LR_GRID]
    base = sweep_baseline(lambda cfg: fn('baseline', cfg['lr']), grid)
    idea_sweep = []
    for cfg in grid:
        r = evaluate(fn('idea', cfg['lr']))
        idea_sweep.append({'cfg': cfg, 'mean': r['mean'], 'std': r['std'], 'per_seed': r['per_seed']})
    best_i = min(idea_sweep, key=lambda q: q['mean'])['cfg']
    idea_full = evaluate(fn('idea', best_i['lr']))
    # Re-train one paired model for the behavioral signature.
    _, sig_net, sig_d = train_one('idea', 0, best_i['lr'])
    sig = mechanism_signature(sig_net, sig_d)
    report = make_report('sequence', 'transformer_tiny_pooling_mvp', base, idea_full,
                         {'mechanism_signature': sig,
                          'idea_sweep': idea_sweep,
                          'selection_seeds': [0,1,2,3],
                          'epochs': EPOCHS, 'batch': BATCH,
                          'parameter_counts': {'baseline': sum(p.numel() for p in PoolNet('baseline').parameters()),
                                               'idea': sum(p.numel() for p in PoolNet('idea').parameters())},
                          'structural_match': 'sequence-level window forecasting with token-set pooling'})
    report['idea']['selected_cfg'] = best_i
    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()