import sys, json, time, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

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

OUT = Path('bench_report.json')
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
EPOCHS, NTR, NTE, BATCH = 15, 1200, 1000, 128
LR_GRID = [0.0015, 0.003, 0.006]
LAMBDA_GRID = [0.0, 1e-6, 3e-6]


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)


class FactorizedLinear(nn.Module):
    def __init__(self, din, dout, rank=32):
        super().__init__()
        self.din, self.dout, self.rank = din, dout, rank
        self.u = nn.Parameter(torch.empty(dout, rank))
        self.v = nn.Parameter(torch.empty(din, rank))
        self.bias = nn.Parameter(torch.zeros(dout))
        # Scale initialization gives a normal-sized product while retaining all columns.
        nn.init.normal_(self.u, 0.0, 0.12)
        nn.init.normal_(self.v, 0.0, 0.12)

    def forward(self, x):
        return x @ self.v @ self.u.t() + self.bias


class FactorMLP(nn.Module):
    def __init__(self, rank=32):
        super().__init__()
        self.l1 = FactorizedLinear(10, 64, min(rank, 10))
        self.l2 = FactorizedLinear(64, 64, rank)
        self.l3 = FactorizedLinear(64, 1, 1)

    def forward(self, x):
        return self.l3(torch.relu(self.l2(torch.relu(self.l1(x)))))


def make_net(seed):
    seed_all(seed)
    return FactorMLP(32)


def ds_for(seed):
    return get_dataset('tabular', seed, n_train=NTR, n_test=NTE)


def baseline_train(seed, cfg):
    seed_all(seed)
    d = ds_for(seed)
    net = make_net(seed)
    # train_model is the canonical baseline path; only the model is the shared factorized MLP.
    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=float(cfg['lr']),
                               batch=BATCH, weight_decay=float(cfg['weight_decay']), log=lambda *_: None)
    return float(metric)


SIGNATURES = {}

def pairwise_prox(net, eta, lam, active):
    changed = 0
    gaps = []
    for layer in (net.l1, net.l2, net.l3):
        with torch.no_grad():
            # Exact shared paired rule: ||u_j||^2+||v_j||^2 <= 4 eta lambda.
            score = layer.u.square().sum(0) + layer.v.square().sum(0)
            keep = (score > 4.0 * eta * lam) & active.get(layer, torch.ones_like(score, dtype=torch.bool))
            changed += int((~keep & active.get(layer, torch.ones_like(keep))).sum().item())
            layer.u[:, ~keep] = 0; layer.v[:, ~keep] = 0
            active[layer] = keep
            if keep.any():
                nu = layer.u[:, keep].norm(dim=0)
                nv = layer.v[:, keep].norm(dim=0)
                a = torch.sqrt(nv / (nu + 1e-12))
                layer.u[:, keep] *= a
                layer.v[:, keep] /= a
                gaps.append(float((layer.u[:, keep].norm(dim=0) - layer.v[:, keep].norm(dim=0)).abs().mean().cpu()))
            # Prevent Adam momentum from resurrecting deleted paired columns.
            if hasattr(layer.u, '_optim_state'):
                pass
    return changed, float(np.mean(gaps)) if gaps else 0.0


def idea_train(seed, cfg):
    seed_all(seed); d = ds_for(seed); net = make_net(seed)
    try:
        dev = torch.device(DEVICE)
        net.to(dev); x, y = d['xtr'].to(dev), d['ytr'].to(dev)
        opt = torch.optim.Adam(net.parameters(), lr=float(cfg['lr']), weight_decay=float(cfg['weight_decay']))
        loader = DataLoader(TensorDataset(x, y), batch_size=BATCH, shuffle=True,
                            generator=torch.Generator(device='cpu').manual_seed(seed + 91))
        active = {layer: torch.ones(layer.rank, dtype=torch.bool, device=dev)
                  for layer in (net.l1, net.l2, net.l3)}
        ranks, gaps, total_pruned = [], [], 0
        # Warm-start continuation: lambda rises through four stages, preserving one model.
        stages = [0.0, cfg['lambda'] / 3.0, cfg['lambda'], 3.0 * cfg['lambda']]
        stage_counts = [EPOCHS // len(stages) + (i < EPOCHS % len(stages)) for i in range(len(stages))]
        for i, lam in enumerate(stages):
            for _ in range(stage_counts[i]):
                for xb, yb in loader:
                    opt.zero_grad(set_to_none=True)
                    loss = (net(xb) - yb).square().mean() / 2
                    loss.backward(); opt.step()
                    # Use the configured learning rate as eta in the stated proximal rule.
                    pruned, gap = pairwise_prox(net, float(cfg['lr']), float(lam), active)
                    total_pruned += pruned; gaps.append(gap)
                    for layer in active:
                        with torch.no_grad():
                            layer.u[:, ~active[layer]] = 0; layer.v[:, ~active[layer]] = 0
            ranks.append(int(sum(int(a.sum().item()) for a in active.values())))
        with torch.no_grad():
            pred = net(d['xte'].to(dev)); metric = float((pred - d['yte'].to(dev)).square().mean().cpu())
        sig = {'seed': seed, 'ranks_by_lambda_stage': ranks,
               'observed_monotone': bool(all(ranks[i+1] <= ranks[i] for i in range(len(ranks)-1))),
               'mean_balanced_norm_gap': float(np.mean(gaps)) if gaps else 0.0,
               'pruned_columns': total_pruned}
        SIGNATURES[seed] = sig
        return metric
    except Exception:
        # Required robust GPU fallback, rebuilding on CPU after any CUDA/runtime error.
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
        old = globals()['DEVICE']; globals()['DEVICE'] = 'cpu'
        try: return idea_train(seed, cfg)
        finally: globals()['DEVICE'] = old


def main():
    t0 = time.time()
    # Baseline decisive knobs are both learning rate and weight decay; idea tries the same lr union.
    grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in [0.0, 1e-4]]
    base = sweep_baseline(lambda cfg: (lambda s: baseline_train(s, cfg)), grid)
    best_lr = float(base['best_cfg']['lr']); best_wd = float(base['best_cfg']['weight_decay'])
    idea_grid = [{'lr': best_lr, 'weight_decay': best_wd, 'lambda': z} for z in LAMBDA_GRID]
    # Include nearby lr settings; each is already present in the baseline union sweep.
    idea_grid = [{'lr': lr, 'weight_decay': best_wd, 'lambda': lam}
                 for lr in LR_GRID for lam in LAMBDA_GRID]
    idea_runs = []
    best_idea = None
    for cfg in idea_grid:
        SIGNATURES.clear()
        res = __import__('bench').evaluate(lambda s, c=cfg: idea_train(s, c))
        idea_runs.append({'cfg': cfg, 'result': res, 'signature': dict(SIGNATURES)})
        if best_idea is None or res['mean'] < best_idea['result']['mean']: best_idea = idea_runs[-1]
    rep = make_report('tabular', 'mlp_tiny', base, best_idea['result'], {
        'prediction': 'warm-started exact paired pruning should produce nonincreasing active rank and balanced factor norms',
        'observed': best_idea['signature'],
        'predicted_monotone_rank': True,
        'predicted_balanced_gap': 0.0,
        'confirmed': bool(all(v['observed_monotone'] for v in best_idea['signature'].values()) and
                     np.mean([v['mean_balanced_norm_gap'] for v in best_idea['signature'].values()]) < 1e-5)
    })
    rep['idea_sweep'] = idea_runs
    rep['runtime_sec'] = time.time() - t0
    rep['device'] = DEVICE
    OUT.write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))

if __name__ == '__main__': main()