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

EPOCHS = 25
BATCH = 128
GRID = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]


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


class MatchedReLU(nn.Module):
    def __init__(self, d, width=32, depth=2):
        super().__init__()
        self.layers = nn.ModuleList()
        q = d
        for _ in range(depth):
            self.layers.append(nn.Linear(q, width))
            q = width
        self.out = nn.Linear(q, 1)

    def forward(self, x):
        z = x
        for layer in self.layers:
            z = F.relu(layer(z))
        return self.out(z)


class ICNN(nn.Module):
    """Input-convex MLP with the same width/depth and direct input skips."""
    def __init__(self, d, width=32, depth=2):
        super().__init__()
        self.d, self.width, self.depth = d, width, depth
        self.A = nn.ParameterList()
        self.U = nn.ParameterList()
        self.b = nn.ParameterList()
        for k in range(depth):
            # z_0 is represented by a width-dimensional first affine embedding;
            # subsequent hidden-to-hidden matrices are nonnegative.
            self.A.append(nn.Parameter(torch.randn(width, width) * 0.12))
            self.U.append(nn.Parameter(torch.randn(width, d) * 0.12))
            self.b.append(nn.Parameter(torch.zeros(width)))
        self.aw = nn.Parameter(torch.zeros(width))
        self.u = nn.Parameter(torch.zeros(d))
        self.c = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        z = F.relu(x @ self.U[0].T + self.b[0])
        for k in range(1, self.depth):
            W = F.softplus(self.A[k]) + 1e-6
            z = F.relu(z @ W.T + x @ self.U[k].T + self.b[k])
        return (z @ F.softplus(self.aw) + x @ self.u + self.c).unsqueeze(-1)


def train_one(kind, cfg, seed, return_model=False):
    seed_all(seed)
    ds = get_dataset('tabular', seed=seed)
    d = int(ds['xtr'].shape[1])
    model = MatchedReLU(d) if kind == 'baseline' else ICNN(d)
    net, metric, hist = train_model(model, ds, epochs=EPOCHS,
                                    lr=cfg['lr'], batch=BATCH,
                                    weight_decay=0.0, log=lambda *_: None)
    if net is None:
        raise RuntimeError('training failed')
    if return_model:
        return float(metric), net, ds
    return float(metric)


def make_fn(kind, cfg):
    return lambda seed: train_one(kind, cfg, seed)


def mechanism_signature():
    # Re-test convexity on predictions from trained benchmark models, not toy weights.
    rows = []
    for seed in range(8):
        _, net, ds = train_one('idea', {'lr': 3e-3}, seed, return_model=True)
        device = next(net.parameters()).device
        x = ds['xte'][:96].to(device)
        rng = torch.Generator().manual_seed(9000 + seed)
        perm = torch.randperm(len(x), generator=rng)
        a, b = x, x[perm]
        t = 0.5
        with torch.no_grad():
            lhs = net(t*a + (1-t)*b).reshape(-1)
            rhs = (t*net(a) + (1-t)*net(b)).reshape(-1)
        gap = lhs - rhs
        rows.append((float(gap.max()), float((gap > 1e-5).float().mean())))
    max_gap = max(r[0] for r in rows)
    rate = float(np.mean([r[1] for r in rows]))
    return {
        'claim': 'trained ICNN predictions satisfy midpoint Jensen convexity',
        'predicted_max_violation': 0.0,
        'observed_max_violation': max_gap,
        'observed_violation_rate': rate,
        'confirmed': bool(max_gap <= 1e-5 and rate == 0.0),
        'n_models': 8
    }


def main():
    base = sweep_baseline(lambda cfg: make_fn('baseline', cfg), GRID)
    # Same three settings are run for the idea; best is selected on the same sweep seeds.
    idea_sweep = []
    for cfg in GRID:
        r = evaluate(make_fn('idea', cfg), seeds=(0,1,2,3))
        idea_sweep.append({'cfg': cfg, 'mean': r['mean']})
    best_cfg = min(idea_sweep, key=lambda x: x['mean'])['cfg']
    idea_full = evaluate(make_fn('idea', best_cfg))
    idea_res = {'best_cfg': best_cfg, 'sweep': idea_sweep,
                'full': idea_full, **idea_full}
    report = make_report('tabular', 'mlp_med', base, idea_full,
                         extra=mechanism_signature())
    report['idea'] = idea_res
    report['protocol_notes'] = {
        'structural_match': 'tabular is the built-in track for architecture/regularization interventions',
        'paired_seeds': list(range(8)),
        'epochs': EPOCHS, 'batch': BATCH,
        'baseline_and_idea_share_grid': True,
        'baseline_architecture': '2-layer width-32 ReLU MLP',
        'idea_architecture': '2-layer width-32 ICNN with softplus W>=0 and nonnegative output weights'
    }
    with open('bench_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    main()