import os, 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')
import bench

SEEDS = tuple(range(8))
# Union is used on both sides: baseline and idea see every tested setting.
GRID = [
    {'lr': 0.001, 'weight_decay': 0.0},
    {'lr': 0.003, 'weight_decay': 0.0},
    {'lr': 0.006, 'weight_decay': 0.0},
]
EPOCHS = 18
BATCH = 128

class FullGRU(nn.Module):
    def __init__(self, hidden=45):
        super().__init__()
        self.rnn = nn.GRU(3, hidden, batch_first=True)
        self.head = nn.Linear(hidden, 1)
        self._no_cudnn = False
    def forward(self, x):
        q = x.view(x.shape[0], -1, 3)
        try:
            _, h = self.rnn(q)
        except RuntimeError:
            self._no_cudnn = True
        if self._no_cudnn:
            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
            try: _, h = self.rnn(q)
            finally: torch.backends.cudnn.enabled = old
        return self.head(h[-1])
    @torch.no_grad()
    def phase_predictions(self, x):
        q = x.view(x.shape[0], 8, 3)
        try:
            _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
        except RuntimeError:
            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
            try:
                _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
            finally: torch.backends.cudnn.enabled = old
        return self.head(hm[-1]), self.head(hf[-1])

class ChainedGRU(nn.Module):
    """Two event phases; phase-2 starts from phase-1's differentiable terminal h."""
    def __init__(self, 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, 1)
        self._no_cudnn = False
    def _run(self, module, x, h=None):
        try: return module(x, h)
        except RuntimeError:
            self._no_cudnn = True
            old = torch.backends.cudnn.enabled; torch.backends.cudnn.enabled = False
            try: return module(x, h)
            finally: torch.backends.cudnn.enabled = old
    def forward(self, x):
        q = x.view(x.shape[0], 8, 3)
        _, h1 = self._run(self.phase1, q[:, :4])
        _, h2 = self._run(self.phase2, q[:, 4:], h1)
        return self.head(h2[-1])
    @torch.no_grad()
    def phase_predictions(self, x):
        q = x.view(x.shape[0], 8, 3)
        _, h1 = self._run(self.phase1, q[:, :4])
        # This is an observable prediction at the event, before phase 2.
        mid = self.head(h1[-1])
        _, h2 = self._run(self.phase2, q[:, 4:], h1)
        return mid, self.head(h2[-1])

class InspectFull(FullGRU):
    @torch.no_grad()
    def phase_predictions(self, x):
        q = x.view(x.shape[0], 8, 3)
        _, hm = self.rnn(q[:, :4]); _, hf = self.rnn(q)
        return self.head(hm[-1]), self.head(hf[-1])

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)

def train_one(kind, cfg, seed, inspect=False):
    seed_all(seed)
    ds = bench.get_dataset('dynamics', seed, n_train=400, n_test=400)
    model = (ChainedGRU() if kind == 'idea' else FullGRU())
    # train_model is the benchmark's required robust CUDA/CPU path.
    net, metric, hist = bench.train_model(model, ds, epochs=EPOCHS,
        lr=cfg['lr'], batch=BATCH, weight_decay=cfg['weight_decay'], log=lambda *_: None)
    if net is None: return float('nan'), None
    if not inspect: return float(metric), None
    dev = next(net.parameters()).device
    x = torch.as_tensor(ds['xte'], dtype=torch.float32, device=dev)
    y = torch.as_tensor(ds['yte'], dtype=torch.float32, device=dev).reshape(-1,1)
    with torch.no_grad():
        mid, final = net.phase_predictions(x)
    # midpoint observed theta is input's fourth theta; final observed y is target.
    obs_mid = x.view(-1,8,3)[:,3,0:1]
    return float(metric), {'mid_pred_mae': float((mid-obs_mid).abs().mean().cpu()),
                           'mid_obs_mae_baseline_reference': float(obs_mid.abs().mean().cpu()),
                           'final_pred_mae': float((final-y).abs().mean().cpu()),
                           'final_obs_abs': float(y.abs().mean().cpu())}

def make_fn(kind, cfg):
    return lambda s: train_one(kind, cfg, s)[0]

def main():
    # Baseline sweep uses the same union of hyperparameters as the idea sweep.
    base = bench.sweep_baseline(lambda c: make_fn('base', c), GRID, seeds=(0,1,2,3))
    base['grid_union'] = GRID
    idea_trials = []
    for cfg in GRID:
        r = bench.evaluate(make_fn('idea', cfg), SEEDS)
        idea_trials.append({'cfg': cfg, **r})
    best = min(idea_trials, key=lambda z: z['mean'])
    idea_res = {'best_cfg': best['cfg'], 'sweep': idea_trials,
                'per_seed': best['per_seed'], 'mean': best['mean'],
                'std': best['std'], 'n': best['n']}
    # Re-train paired best systems to measure behaviour, not an identity.
    sig = {'baseline': [], 'idea': [], 'paired_seed': []}
    for s in SEEDS:
        bm, bs = train_one('base', best['cfg'], s, True)
        im, ins = train_one('idea', best['cfg'], s, True)
        sig['baseline'].append(bs); sig['idea'].append(ins); sig['paired_seed'].append(s)
    sig['mean_final_mae_baseline'] = float(np.mean([v['final_pred_mae'] for v in sig['baseline']]))
    sig['mean_final_mae_idea'] = float(np.mean([v['final_pred_mae'] for v in sig['idea']]))
    sig['mean_mid_mae_baseline'] = float(np.mean([v['mid_pred_mae'] for v in sig['baseline']]))
    sig['mean_mid_mae_idea'] = float(np.mean([v['mid_pred_mae'] for v in sig['idea']]))
    # Prediction being tested: chaining should help the post-boundary dynamics.
    sig['prediction'] = 'phase chaining improves final/post-event trajectory prediction'
    sig['confirmed'] = bool(sig['mean_final_mae_idea'] < sig['mean_final_mae_baseline'])
    report = bench.make_report('dynamics', 'rnn_small', base, idea_res,
        {'mechanism_signature': sig,
         'protocol_notes': '400 train/400 test, 8 paired seeds, 18 epochs; midpoint is the known phase boundary.'})
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    try: main()
    except Exception:
        if torch.cuda.is_available():
            torch.cuda.empty_cache(); os.environ['CUDA_VISIBLE_DEVICES']=''
            main()
        else: raise