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

SEEDS = tuple(range(8))
# Union is shared: every idea learning rate is evaluated for baseline too.
GRID = [{'lr': 0.0015, 'weight_decay': 0.0},
        {'lr': 0.003,  'weight_decay': 0.0},
        {'lr': 0.006,  'weight_decay': 0.0}]
EPOCHS = 20

class ExactChainedRNN(nn.Module):
    """Two known phases; phase-2 recurrent state is differentiably initialized by phase 1."""
    def __init__(self, out_dim=1, hidden=32):
        super().__init__()
        self.phase1 = nn.GRU(3, hidden, batch_first=True)
        self.phase2 = nn.GRU(3, hidden, batch_first=True)
        self.head = nn.Linear(hidden, out_dim)

    def forward(self, x):
        x = x.to(next(self.parameters()).device)
        q = x.view(x.shape[0], -1, 3)
        _, h1 = self.phase1(q[:, :4])
        _, h2 = self.phase2(q[:, 4:], h1)
        return self.head(h2[-1])

    @torch.no_grad()
    def diagnostics(self, x):
        # CPU diagnostics avoid consuming the shared GPU/cuDNN workspace.
        self.cpu(); x = x.cpu()
        q = x.view(x.shape[0], -1, 3)
        _, h1 = self.phase1(q[:, :4])
        _, h2 = self.phase2(q[:, 4:], h1)
        pred = self.head(h2[-1]).squeeze(-1)
        jump = (q[:, 4, 2] - q[:, 3, 2]).abs()
        return pred, jump, h1, h2

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 train_one(kind, seed, cfg, keep=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=4000, n_test=1000)
    model = make_model('rnn_small', ds['input_shape'], ds['out_dim']) if kind == 'base' else ExactChainedRNN(ds['out_dim'])
    net, metric, hist = train_model(model, ds, epochs=EPOCHS, lr=cfg['lr'], batch=128,
                                    weight_decay=cfg['weight_decay'], log=lambda *_: None)
    if net is None: return float('nan'), None
    return (float(metric), (net, ds)) if keep else (float(metric), None)

def fn(kind, cfg):
    return lambda seed: train_one(kind, seed, cfg)[0]

def behavioral_signature(cfg):
    idea_rows, base_rows = [], []
    idea_param = base_param = None
    for s in SEEDS:
        _, ipack = train_one('idea', s, cfg, keep=True)
        _, bpack = train_one('base', s, cfg, keep=True)
        inet, ds = ipack; bnet, _ = bpack
        ip, jump, _, _ = inet.diagnostics(ds['xte'])
        bnet.cpu(); bp = bnet(ds['xte']).squeeze(-1).detach().cpu()
        obs = ds['yte'].squeeze(-1).cpu()
        high = jump >= torch.quantile(jump, 0.75)
        low = ~high
        idea_rows.append({'seed':s, 'high_jump_mse':float(((ip[high]-obs[high])**2).mean()), 'low_jump_mse':float(((ip[low]-obs[low])**2).mean())})
        base_rows.append({'seed':s, 'high_jump_mse':float(((bp[high]-obs[high])**2).mean()), 'low_jump_mse':float(((bp[low]-obs[low])**2).mean())})
        idea_param, base_param = count_params(inet), count_params(bnet)
    ih = float(np.mean([r['high_jump_mse'] for r in idea_rows])); bh = float(np.mean([r['high_jump_mse'] for r in base_rows]))
    return {'idea_by_seed':idea_rows, 'baseline_by_seed':base_rows,
            'idea_high_jump_mse':ih, 'baseline_high_jump_mse':bh,
            'idea_low_jump_mse':float(np.mean([r['low_jump_mse'] for r in idea_rows])),
            'baseline_low_jump_mse':float(np.mean([r['low_jump_mse'] for r in base_rows])),
            'high_jump_delta_idea_minus_baseline':ih-bh,
            'idea_params':int(idea_param), 'baseline_params':int(base_param),
            'confirmed': bool(ih < bh),
            'prediction':'phase chaining should reduce error on abrupt control changes'}

def main():
    base = sweep_baseline(lambda cfg: fn('base', cfg), GRID, seeds=(0,1,2,3))
    base['full'] = evaluate(fn('base', base['best_cfg']), SEEDS)
    idea_runs = [{'cfg':cfg, 'result':evaluate(fn('idea', cfg), SEEDS)} for cfg in GRID]
    chosen = min(idea_runs, key=lambda z:z['result']['mean'])
    idea = chosen['result']
    report = make_report('dynamics', 'rnn_small', base, idea,
        {'phase_chaining': behavioral_signature(chosen['cfg']),
         'idea_cfg':chosen['cfg'], 'idea_grid':idea_runs,
         'structural_match':'controlled pendulum rollout with an observed midpoint control event; primary metric is test MSE'})
    report['comparison']['selected_idea_cfg'] = chosen['cfg']
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()