import json, random, sys
import numpy as np
import torch

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

TRACK = 'poisson_dirichlet'
MODEL = 'mlp_tiny'
SEEDS = tuple(range(8))
EPOCHS = 20
BATCH = 128
# Shared union: every idea LR is also tested by baseline.
GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0},
    {'lr': 3e-3, 'weight_decay': 0.0},
    {'lr': 1e-2, 'weight_decay': 0.0},
]


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 basis_matrix():
    n_grid, n_modes = 32, 12
    x = np.arange(1, n_grid + 1, dtype=np.float32) / (n_grid + 1)
    return (np.sqrt(2.0) * np.sin(np.pi * np.outer(
        x, np.arange(1, n_modes + 1)))).astype(np.float32)

def get_ds(seed):
    d = get_dataset(TRACK, seed=int(seed), n_train=400, n_test=400)
    # bench flattens custom-track y arrays; restore one field per sample.
    nout = int(d['out_dim'])
    for k in ('xtr', 'xte'):
        d[k] = torch.as_tensor(d[k], dtype=torch.float32)
    for k, n in (('ytr', len(d['xtr'])), ('yte', len(d['xte']))):
        a = np.asarray(d[k]).reshape(n, nout)
        d[k] = torch.as_tensor(a, dtype=torch.float32)
    return d


class NullspaceModel(torch.nn.Module):
    def __init__(self, core, basis):
        super().__init__()
        self.core = core
        self.register_buffer('V', torch.as_tensor(basis, dtype=torch.float32))

    def forward(self, x):
        # core predicts z; decoder always lies in ker(C).
        interior = self.core(x) @ self.V.T
        z = torch.zeros((interior.shape[0], 1), device=interior.device, dtype=interior.dtype)
        return torch.cat((z, interior, z), dim=1)


def fit(seed, cfg, constrained, return_model=False):
    seed_all(seed)
    d = get_ds(seed)
    # Baseline predicts the full field, idea predicts the 12 nullspace coordinates.
    out_dim = int(basis_matrix().shape[1]) if constrained else int(d['out_dim'])
    core = make_model(MODEL, tuple(d['xtr'].shape[1:]), out_dim)
    net = NullspaceModel(core, basis_matrix()) if constrained else core
    net, metric, history = train_model(
        net, d, epochs=EPOCHS, lr=float(cfg['lr']), batch=BATCH,
        weight_decay=float(cfg['weight_decay']), log=lambda *_: None)
    if net is None or metric is None:
        return (float('inf'), None, d)
    return (float(metric), net if return_model else None, d)


def make_fn(constrained):
    def fn(cfg):
        return lambda seed: fit(seed, cfg, constrained)[0]
    return fn


def run():
    # Baseline sweep uses four seeds as specified, then re-evaluates best on eight.
    base = sweep_baseline(make_fn(False), GRID, seeds=(0, 1, 2, 3))
    # Explicitly run the idea at all three shared settings; report best using the
    # same four-seed selection budget, then evaluate its selected config on 8.
    idea_trials = []
    best_cfg, best_mean = None, float('inf')
    for cfg in GRID:
        r = evaluate(make_fn(True)(cfg), seeds=(0, 1, 2, 3))
        idea_trials.append({'cfg': cfg, 'mean': r['mean']})
        if r['mean'] < best_mean:
            best_cfg, best_mean = cfg, r['mean']
    idea_full = evaluate(make_fn(True)(best_cfg), seeds=SEEDS)
    base['idea_union_trials'] = idea_trials

    # Signature is measured from trained systems: boundary residual and decoded
    # nullspace violation on actual held-out inputs. C selects endpoints.
    probe_seed = 0
    bmetric, bnet, bd = fit(probe_seed, base['best_cfg'], False, True)
    imetric, inet, idata = fit(probe_seed, best_cfg, True, True)
    C = torch.zeros((2, 34), dtype=torch.float32)
    C[0, 0] = 1.0; C[1, -1] = 1.0
    with torch.no_grad():
        db = next(bnet.parameters()).device
        di = next(inet.parameters()).device
        ub = bnet(idata['xte'].to(db)).detach().cpu()
        ui = inet(idata['xte'].to(di)).detach().cpu()
    rb = torch.linalg.norm(ub @ C.T, dim=1)
    ri = torch.linalg.norm(ui @ C.T, dim=1)
    signature = {
        'prediction': 'nullspace decoder should make C u exactly zero while ambient output violates endpoints',
        'predicted_boundary_residual_idea': 0.0,
        'observed_boundary_residual_idea_max': float(ri.max()),
        'observed_boundary_residual_baseline_max': float(rb.max()),
        'observed_boundary_residual_baseline_mean': float(rb.mean()),
        'observed_boundary_residual_idea_mean': float(ri.mean()),
        'ratio_baseline_to_idea_max': float((rb.max() + 1e-12) / (ri.max() + 1e-12)),
        'confirmed': bool(float(ri.max()) < 1e-5 and float(rb.max()) > 1e-3),
        'trained_models': True,
    }
    report = make_report(TRACK, MODEL, base, idea_full, {
        'selection': 'best of shared 3-point LR/weight-decay grid on seeds 0-3',
        'idea_best_cfg': best_cfg,
        **signature,
    })
    report['custom_track'] = {
        'name': 'poisson_dirichlet',
        'file': '/home/maxwelhelp/all/math2nn/bench/custom_tracks/poisson_dirichlet.py',
        'domain': 'pde',
        'reason': 'built-in custom PDE track has fixed-grid field outputs and homogeneous Dirichlet constraints',
    }
    report['protocol'] = {'epochs': EPOCHS, 'batch': BATCH, 'seeds': list(SEEDS), 'grid': GRID}
    with open('bench_report.json', 'w') as f:
        json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))


if __name__ == '__main__':
    run()