import sys, json, random
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

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

SEEDS = tuple(range(8))
EPOCHS = 12
BATCH = 128
# Baseline and idea share the complete learning-rate union.
GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0},
    {'lr': 3e-3, 'weight_decay': 0.0},
    {'lr': 6e-3, 'weight_decay': 0.0},
]
IDEA_GRID = [
    {'lr': 1e-3, 'weight_decay': 0.0, 'lambda_cycle': 0.01},
    {'lr': 3e-3, 'weight_decay': 0.0, 'lambda_cycle': 0.03},
    {'lr': 6e-3, 'weight_decay': 0.0, 'lambda_cycle': 0.10},
]


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(seed, cfg, intervention=False, return_signature=False):
    seed_all(seed)
    ds = get_dataset('dynamics', seed, n_train=400, n_test=200)
    net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    lam = float(cfg.get('lambda_cycle', 0.0)) if intervention else 0.0
    seq_len = int(np.prod(ds['input_shape']) // 3)
    last_h = None
    for device in (['cuda', 'cpu'] if torch.cuda.is_available() else ['cpu']):
        try:
            net = net.to(device)
            x, y = ds['xtr'].to(device), ds['ytr'].to(device)
            opt = torch.optim.Adam(net.parameters(), lr=float(cfg['lr']),
                                   weight_decay=float(cfg.get('weight_decay', 0.0)))
            mse = nn.MSELoss()
            for _ in range(EPOCHS):
                net.train()
                perm = torch.randperm(len(x), device=device)
                for start in range(0, len(x), BATCH):
                    idx = perm[start:start+BATCH]
                    xb, yb = x[idx], y[idx]
                    inp = xb.view(len(xb), seq_len, 3)
                    _, hseq = net.rnn(inp)
                    pred = net.head(hseq[-1])
                    loss = mse(pred, yb)
                    if lam:
                        # Re-run the GRU one prefix at a time to expose h_t.
                        # Penalize the exact period-two signature: adjacent
                        # states differ but two-step states are close.
                        hs = []
                        h = torch.zeros(1, len(xb), net.rnn.hidden_size,
                                         device=device, dtype=xb.dtype)
                        for t in range(seq_len):
                            _, h = net.rnn(inp[:, t:t+1], h)
                            hs.append(h[-1])
                        H = torch.stack(hs, dim=1)
                        if seq_len >= 3:
                            two = (H[:, 2:] - H[:, :-2]).pow(2).mean()
                            adjacent = (H[:, 1:] - H[:, :-1]).pow(2).mean()
                            cycle_pen = two / (adjacent.detach() + 1e-4)
                            loss = loss + lam * cycle_pen
                    opt.zero_grad(set_to_none=True)
                    loss.backward()
                    opt.step()
            net.eval()
            with torch.no_grad():
                pred = net(ds['xte'].to(device))
                metric = float(((pred - ds['yte'].to(device)) ** 2).mean())
                # Signature is measured from the trained model, not an analytic toy.
                inp = ds['xte'].to(device).view(len(ds['xte']), seq_len, 3)
                h = torch.zeros(1, len(inp), net.rnn.hidden_size, device=device)
                hs = []
                for t in range(seq_len):
                    _, h = net.rnn(inp[:, t:t+1], h); hs.append(h[-1])
                H = torch.stack(hs, 1)
                if seq_len >= 3:
                    two = float((H[:, 2:] - H[:, :-2]).pow(2).mean())
                    one = float((H[:, 1:] - H[:, :-1]).pow(2).mean())
                    ratio = two / (one + 1e-8)
                else:
                    two, one, ratio = 0.0, 0.0, 0.0
            if return_signature:
                return metric, {'observed_two_step_mse': two,
                                'observed_adjacent_mse': one,
                                'observed_two_to_one_ratio': ratio}
            return metric
        except RuntimeError:
            if device == 'cpu':
                raise
            net = make_model('rnn_small', ds['input_shape'], ds['out_dim'])
    raise RuntimeError('training failed')


def base_fn(cfg):
    return lambda seed: train_one(seed, cfg, False)


def idea_fn(cfg):
    return lambda seed: train_one(seed, cfg, True)


def main():
    baseline = sweep_baseline(base_fn, GRID, seeds=(0, 1, 2, 3))
    # Evaluate all three idea settings on the same full paired seeds; report best.
    idea_trials = []
    for cfg in IDEA_GRID:
        r = evaluate(idea_fn(cfg), seeds=SEEDS)
        idea_trials.append({'cfg': cfg, 'result': r})
    best = min(idea_trials, key=lambda z: z['result']['mean'])
    sig_rows = []
    for s in SEEDS:
        _, sig = train_one(s, best['cfg'], True, True)
        _, bsig = train_one(s, baseline['best_cfg'], False, True)
        sig_rows.append({'seed': s, 'baseline': bsig, 'idea': sig})
    b2 = np.mean([r['baseline']['observed_two_to_one_ratio'] for r in sig_rows])
    i2 = np.mean([r['idea']['observed_two_to_one_ratio'] for r in sig_rows])
    signature = {
        'prediction': 'anti-oscillation penalty should reduce trained recurrent two-step similarity relative to adjacent change',
        'baseline_mean_two_to_one_ratio': float(b2),
        'idea_mean_two_to_one_ratio': float(i2),
        'predicted_direction': 'idea lower than baseline',
        'confirmed': bool(i2 < b2),
        'per_seed': sig_rows,
    }
    rep = make_report('dynamics', 'rnn_small', baseline, best['result'], extra=signature)
    rep['idea_trials'] = idea_trials
    rep['structural_match'] = 'Dynamics track: recurrent GRU predicts actuated pendulum rollout; anti-oscillation targets recurrent hidden-state period-two behaviour.'
    rep['protocol'] = {'paired_seeds': list(SEEDS), 'epochs': EPOCHS, 'batch': BATCH,
                       'baseline_grid': GRID, 'idea_grid': IDEA_GRID}
    Path('bench_report.json').write_text(json.dumps(rep, indent=2))
    print(json.dumps(rep, indent=2))


if __name__ == '__main__':
    main()