import os, sys, json, random
import numpy as np
import torch
import torch.nn as nn

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

SEED0 = 2145
EPOCHS = 18
NTR, NTE = 400, 200
N_CAND, ACTIVE = 8, 3

class CandidateMLP(nn.Module):
    """Shared architecture: incumbent MLP plus a bank of gated feature blocks."""
    def __init__(self, input_dim, out_dim=1, active=(0,1,2)):
        super().__init__()
        self.base = nn.Sequential(nn.Linear(input_dim, 32), nn.Tanh(), nn.Linear(32, 16), nn.Tanh())
        self.base_head = nn.Linear(16, out_dim)
        self.cands = nn.ModuleList([nn.Sequential(nn.Linear(input_dim, 8), nn.Tanh(), nn.Linear(8, out_dim)) for _ in range(N_CAND)])
        self.active = tuple(int(x) for x in active)
    def forward(self, x):
        y = self.base_head(self.base(x))
        for j in self.active:
            y = y + self.cands[j](x)
        return y
    def incumbent(self, x):
        return self.base_head(self.base(x))
    def candidate_response(self, x, j):
        return self.cands[j](x)

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 projected_scores(net, x):
    """Q_j=(I-P_B) C_j for scalar responses, Y=I; return gains and singular values."""
    with torch.no_grad():
        b = net.incumbent(x).detach().cpu().numpy().reshape(-1, 1)
        cs = [net.candidate_response(x, j).detach().cpu().numpy().reshape(-1, 1) for j in range(N_CAND)]
    # Stable Y-weighted projector via least squares, no explicit inverse assumption.
    scores = []
    qs = []
    for c in cs:
        coef, *_ = np.linalg.lstsq(b, c, rcond=1e-8)
        q = c - b @ coef
        sv = np.linalg.svd(q, compute_uv=False)
        scores.append(float(np.sum(q*q))); qs.append(q)
    return scores, qs

def train_one(track, seed, lr, mode, return_net=False):
    seed_all(seed)
    d = get_dataset(track, seed, n_train=NTR, n_test=NTE)
    dim = int(np.prod(d['input_shape']))
    # Keep exact same initial bank for paired baseline/idea; only activation differs.
    probe = CandidateMLP(dim, d['out_dim'], active=())
    scores, qs = projected_scores(probe, d['xtr'][:min(128, NTR)])
    if mode == 'idea':
        ranked = sorted(range(N_CAND), key=lambda j: scores[j], reverse=True)
        # Riesz filter: reject nearly null and badly scaled responses; fallback preserves budget.
        accepted = [j for j in ranked if np.linalg.norm(qs[j]) >= 1e-4 and np.linalg.norm(qs[j]) <= 30][:ACTIVE]
        if len(accepted) < ACTIVE: accepted = ranked[:ACTIVE]
    else:
        accepted = list(range(ACTIVE))
    net = CandidateMLP(dim, d['out_dim'], active=accepted)
    # train_model is the canonical path for the trained system
    net, metric, hist = train_model(net, d, epochs=EPOCHS, lr=lr, batch=128, log=lambda *_: None)
    if net is None: return float('nan') if not return_net else (float('nan'), None, accepted, scores)
    if return_net: return float(metric), net, accepted, scores
    return float(metric)

def make_baseline(track, cfg):
    return lambda seed: train_one(track, int(seed), float(cfg['lr']), 'baseline')

def make_idea(track, cfg):
    return lambda seed: train_one(track, int(seed), float(cfg['lr']), 'idea')

def signature(track, cfg, seeds=(0,1,2,3,4,5,6,7)):
    pred, obs = [], []
    for s in seeds:
        metric, net, selected, scores = train_one(track, s, cfg['lr'], 'idea', True)
        d = get_dataset(track, s, n_train=NTR, n_test=NTE)
        x = d['xte'].to(next(net.parameters()).device)
        with torch.no_grad():
            full = net(x).cpu().numpy().reshape(-1)
            inc = net.incumbent(x).cpu().numpy().reshape(-1)
        # predicted residual energy on calibration; observed response energy on held-out trained model
        pred.append(float(sum(scores[j] for j in selected)))
        obs.append(float(np.mean((full-inc)**2)))
    corr = float(np.corrcoef(pred, obs)[0,1]) if np.std(pred)>0 and np.std(obs)>0 else 0.0
    return {'predicted_projected_gain_mean': float(np.mean(pred)), 'observed_heldout_selected_response_mse_mean': float(np.mean(obs)), 'predicted_per_seed': pred, 'observed_per_seed': obs, 'correlation': corr, 'confirmed': bool(corr > 0.3)}

def main():
    track, model = 'tabular', 'candidate_mlp_shared'
    # Union of all step sizes is shared by both sides; three settings satisfy idea nearby sweep.
    grid = [{'lr': 1e-3}, {'lr': 3e-3}, {'lr': 1e-2}]
    base = sweep_baseline(lambda cfg: make_baseline(track, cfg), grid)
    idea_trials = []
    for cfg in grid:
        r = evaluate(make_idea(track, cfg))
        idea_trials.append({'cfg': cfg, 'result': r})
    best = min(idea_trials, key=lambda z: z['result']['mean'])
    report = make_report(track, model, base, best['result'], extra=signature(track, best['cfg']))
    report['idea_sweep'] = idea_trials
    report['selection'] = {'candidates': N_CAND, 'active_blocks': ACTIVE, 'epochs': EPOCHS, 'n_train': NTR, 'n_test': NTE, 'track_rationale': 'tabular regression is the built-in optimizer/regularizer/architecture-compatible small MLP track; no PDE/control structure is required by this adapter-bank idea'}
    with open('bench_report.json','w') as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()