import sys, json, math, 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 make_model, train_model, evaluate, sweep_baseline, make_report
from bloch_track import get_dataset

K = 2.0 * math.pi
SEEDS = tuple(range(8))
SWEEP_SEEDS = tuple(range(4))

class GaugeModel(nn.Module):
    """Same MLP as baseline, with canonical q and exact output gauge."""
    def __init__(self, base):
        super().__init__(); self.base = base
    def forward(self, z):
        c, x, q = z[:, 0], z[:, 1], z[:, 2]
        m = torch.round(q / K)
        q0 = q - m * K
        raw = self.base(torch.stack([c, x, q0], dim=1))
        phase = torch.exp(-1j * m * K * x)
        out = torch.complex(raw[:, 0], raw[:, 1]) * phase
        return torch.stack([out.real, out.imag], dim=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 make_system(kind, ds):
    base = make_model('mlp_tiny', ds['input_shape'], 2)
    return GaugeModel(base) if kind == 'idea' else base

def run_one(kind, cfg, seed):
    seed_all(seed)
    ds = get_dataset(seed, 400, 400)
    ds['xtr'] = torch.as_tensor(ds['xtr'], dtype=torch.float32); ds['ytr'] = torch.as_tensor(ds['ytr'], dtype=torch.float32)
    ds['xte'] = torch.as_tensor(ds['xte'], dtype=torch.float32); ds['yte'] = torch.as_tensor(ds['yte'], dtype=torch.float32)
    ds['input_shape'] = tuple(ds['xtr'].shape[1:]); ds['out_dim'] = 2
    net = make_system(kind, ds)
    trained, metric, hist = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None)
    return float(metric) if metric is not None else float('nan')

def model_for_signature(kind, cfg, seed):
    seed_all(seed); ds = get_dataset(seed, 400, 400)
    ds['xtr'] = torch.as_tensor(ds['xtr'], dtype=torch.float32); ds['ytr'] = torch.as_tensor(ds['ytr'], dtype=torch.float32)
    ds['xte'] = torch.as_tensor(ds['xte'], dtype=torch.float32); ds['yte'] = torch.as_tensor(ds['yte'], dtype=torch.float32)
    ds['input_shape'] = tuple(ds['xtr'].shape[1:]); ds['out_dim'] = 2
    net = make_system(kind, ds)
    net, _, _ = train_model(net, ds, epochs=cfg['epochs'], lr=cfg['lr'], batch=128, weight_decay=cfg['weight_decay'], log=lambda *_: None)
    return net, ds

def signature(cfg):
    rows = {'baseline': [], 'idea': []}
    # Query the trained systems at q0 and q0+K, then compare observed shift with exact phase prediction.
    rng = np.random.default_rng(991)
    c = torch.tensor(rng.uniform(-.9, .9, 96), dtype=torch.float32)
    x = torch.tensor(rng.uniform(0, 1, 96), dtype=torch.float32)
    q0 = torch.tensor(rng.uniform(-.8*math.pi, .8*math.pi, 96), dtype=torch.float32)
    for kind in ('baseline', 'idea'):
        net, _ = model_for_signature(kind, cfg, 0)
        net.eval()
        dev = next(net.parameters()).device
        with torch.no_grad():
            probe0 = torch.stack([c, x, q0], 1).to(dev)
            probe1 = torch.stack([c, x, q0 + K], 1).to(dev)
            p0 = net(probe0)
            p1 = net(probe1)
            cdev, xdev, qdev = c.to(dev), x.to(dev), q0.to(dev)
        z0 = torch.complex(p0[:,0], p0[:,1]); z1 = torch.complex(p1[:,0], p1[:,1])
        expected = z0 * torch.exp(-1j*K*xdev)
        err = torch.sqrt(torch.mean(torch.abs(z1-expected)**2)).item()
        scale = torch.sqrt(torch.mean(torch.abs(z0)**2)).item()
        rows[kind] = {'shift_m': 1, 'observed_gauge_rmse': err, 'relative_rmse': err/(scale+1e-8)}
    # Quantitative stage-1 prediction: exact covariance should be near zero for the wrapped model,
    # while the raw model should have a nonzero alias discrepancy.
    confirmed = rows['idea']['relative_rmse'] < 0.08 and rows['idea']['relative_rmse'] < rows['baseline']['relative_rmse']
    return {'prediction': 'canonicalized model has much smaller trained-model reciprocal-shift gauge discrepancy', 'rows': rows, 'confirmed': bool(confirmed)}

def main():
    grid = [
        {'lr': 0.001, 'epochs': 25, 'weight_decay': 0.0},
        {'lr': 0.003, 'epochs': 25, 'weight_decay': 0.0},
        {'lr': 0.006, 'epochs': 25, 'weight_decay': 0.0},
    ]
    base = sweep_baseline(lambda cfg: lambda s: run_one('baseline', cfg, s), grid, seeds=SWEEP_SEEDS)
    idea_sweep = []
    for cfg in grid:
        r = evaluate(lambda s, cfg=cfg: run_one('idea', cfg, s), seeds=SWEEP_SEEDS)
        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
    best_idea_cfg = min(idea_sweep, key=lambda x: x['mean'])['cfg']
    idea_full = evaluate(lambda s: run_one('idea', best_idea_cfg, s), seeds=SEEDS)
    rep = make_report('bloch_gauge_pde', 'mlp_tiny',
                      {'best_cfg': base['best_cfg'], 'sweep': base['sweep'], 'full': base['full']},
                      dict(idea_full, best_cfg=best_idea_cfg, sweep=idea_sweep),
                      extra=signature(best_idea_cfg))
    rep['custom_track'] = {'name': 'bloch_gauge_pde', 'file': 'bloch_track.py', 'domain': 'pde'}
    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))
if __name__ == '__main__': main()