from __future__ import annotations
import json, sys
from pathlib import Path
import numpy as np
import torch
from torch import nn
sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
import bench

SEEDS = tuple(range(8))
DEPTHS = (1, 2, 3, 4)
LRS = (0.0015, 0.003, 0.006)
EPOCHS = 12
BATCH = 128
WIDTH = 64

class VariableMLP(nn.Module):
    def __init__(self, input_dim, depth, width=WIDTH):
        super().__init__()
        layers, d = [], input_dim
        for _ in range(depth):
            layers += [nn.Linear(d, width), nn.ReLU()]
            d = width
        layers.append(nn.Linear(d, 1))
        self.net = nn.Sequential(*layers)
    def forward(self, x):
        return self.net(x.reshape(x.shape[0], -1))

def make_model(ds, depth):
    return VariableMLP(int(np.prod(tuple(ds['xtr'].shape[1:]))), int(depth))

def train_metric(ds, depth, lr, return_model=False):
    model = make_model(ds, depth)
    net, metric, _ = bench.train_model(model, ds, epochs=EPOCHS, lr=float(lr),
                                       batch=BATCH, weight_decay=0.0, log=lambda *_: None)
    if metric is None:
        raise RuntimeError('bench training failed')
    return (float(metric), net) if return_model else float(metric)

def eval_cfg(cfg, seeds=SEEDS):
    vals = []
    for seed in seeds:
        torch.manual_seed(10000 + seed)
        np.random.seed(20000 + seed)
        ds = bench.get_dataset('tabular', seed, n_train=400, n_test=400)
        vals.append(train_metric(ds, cfg['depth'], cfg['lr']))
    return {'per_seed': vals, 'mean': float(np.mean(vals)),
            'std': float(np.std(vals, ddof=1))}

def baseline_train(cfg):
    return lambda seed: train_metric(bench.get_dataset('tabular', seed, 400, 400), cfg['depth'], cfg['lr'])

def calibration(seed=9001):
    ds = bench.get_dataset('tabular', seed, n_train=400, n_test=400)
    # Calibration discrepancy is validation MSE to a fixed high-precision teacher.
    torch.manual_seed(30000 + seed)
    teacher, _ = train_metric(ds, 4, 0.003, return_model=True)
    del teacher
    rows = []
    for d in DEPTHS:
        torch.manual_seed(31000 + d)
        err = train_metric(ds, d, 0.003)
        rows.append((d, err))
    # Fit E(d)=floor+C_syn/d; residual margin gives conservative [L,U].
    X = np.array([[1.0, 1.0/d] for d, _ in rows])
    y = np.array([e for _, e in rows])
    coef = np.linalg.lstsq(X, y, rcond=None)[0]
    resid = y - X @ coef
    margin = max(0.002, 2.0 * float(np.std(resid)))
    return {'floor_L': float(coef[0]-margin), 'floor_U': float(coef[0]+margin),
            'C_syn': float(max(0.0, coef[1])), 'rows': rows,
            'residual_std': float(np.std(resid))}

def main():
    # Search-space parity: every idea lr/depth pair is included in baseline grid.
    grid = [{'depth': d, 'lr': lr} for d in DEPTHS for lr in LRS]
    base = bench.sweep_baseline(baseline_train, grid, seeds=SEEDS[:4])
    base['full'] = eval_cfg(base['best_cfg'], SEEDS)
    cal = calibration()
    eps = float(base['full']['mean'])
    candidates = []
    for d in DEPTHS:
        radius = cal['C_syn'] / d + cal['floor_U'] - cal['floor_L']
        upper = cal['floor_U'] + radius
        lower = cal['floor_L'] - radius
        candidates.append((upper, d, lower))
    feasible = [x for x in candidates if x[0] <= eps]
    chosen_depth = min((x[1] for x in feasible), default=4)
    # Same-size idea sweep at the baseline best lr and two nearby union-grid lrs.
    idea_grid = [{'depth': chosen_depth, 'lr': lr} for lr in LRS]
    idea_trials = [{'cfg': c, 'result': eval_cfg(c, SEEDS[:4])} for c in idea_grid]
    best = min(idea_trials, key=lambda z: z['result']['mean'])['cfg']
    idea = eval_cfg(best, SEEDS)
    comparison = bench.compare_results(base['full'], idea)
    # NN-scale mechanism signature: observed fitted depth law from trained models.
    observed = []
    for d in DEPTHS:
        r = eval_cfg({'depth': d, 'lr': best['lr']}, SEEDS[:2])
        observed.append((d, r['mean']))
    slope = float(np.polyfit(np.log([d for d,_ in observed]), np.log(np.maximum([e for _,e in observed],1e-12)), 1)[0])
    signature = {'predicted_radius_depth_exponent': -1.0,
                 'observed_test_error_depth_slope': slope,
                 'predicted_vs_observed_tolerance': 0.35,
                 'confirmed': bool(abs(slope + 1.0) <= 0.35),
                 'observed_depth_errors': observed,
                 'note': 'Measured on independently trained benchmark models.'}
    report = bench.make_report('tabular', 'variable_mlp', base, idea,
        {'certificate_calibration': cal, 'candidate_certificates': candidates,
         'chosen_depth': chosen_depth, 'idea_grid': idea_trials,
         'mechanism_signature': signature})
    report['mechanism_signature'] = signature
    report['bench_report'] = {'track': 'tabular', 'model': 'variable_mlp',
                              'calibration': cal, 'chosen_depth': chosen_depth}
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()