import sys, json, math
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 make_model, train_model, evaluate, sweep_baseline, make_report

META = {'name': 'rare_event_markov', 'domain': 'dynamics',
        'description': 'Discrete Markov state-space event prediction with exact terminal reachability conditioning.'}
NSTATE, TARGET, MAX_H, Q = 10, 9, 12, 0.30
SEEDS = tuple(range(8))
LR_GRID = [1e-3, 3e-3, 1e-2]
EPOCHS = 20


def exact_message(state, horizon):
    need = TARGET - int(state)
    if need < 0:
        return 1.0
    if need > int(horizon):
        return 0.0
    return float(math.comb(int(horizon), need) * Q ** need * (1-Q) ** (int(horizon)-need))


def get_dataset(seed, n_train=400, n_test=400):
    rng = np.random.RandomState(seed)
    def make(n):
        state = rng.randint(0, 5, size=n)
        horizon = rng.randint(5, MAX_H + 1, size=n)
        y = np.zeros(n, dtype=np.int64)
        for i in range(n):
            z = int(state[i])
            for _ in range(int(horizon[i])):
                if z < TARGET and rng.rand() < Q:
                    z += 1
            y[i] = int(z == TARGET)
        x = np.stack([state / TARGET, horizon / MAX_H], axis=1).astype(np.float32)
        msg = np.array([exact_message(s, h) for s, h in zip(state, horizon)], dtype=np.float32)
        return x, y, msg
    xtr, ytr, mtr = make(n_train)
    xte, yte, mte = make(n_test)
    return {'xtr': torch.from_numpy(xtr), 'ytr': torch.from_numpy(ytr),
            'xte': torch.from_numpy(xte), 'yte': torch.from_numpy(yte),
            'message_tr': torch.from_numpy(mtr), 'message_te': torch.from_numpy(mte),
            'task': 'classification', 'metric': 'err', 'input_shape': (2,), 'out_dim': 2}


def baseline_train(cfg, seed):
    torch.manual_seed(seed); np.random.seed(seed)
    d = get_dataset(seed)
    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    _, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'],
                               weight_decay=cfg.get('weight_decay', 0.0), log=lambda *_: None)
    return float(metric)


def idea_train(cfg, seed, return_signature=False):
    torch.manual_seed(seed); np.random.seed(seed)
    d = get_dataset(seed)
    net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
    # This is the intervention: train the transition-event predictor toward
    # the exact backward feasibility message in addition to the event label.
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    try:
        net = net.to(device)
        x, y, msg = d['xtr'].to(device), d['ytr'].to(device), d['message_tr'].to(device)
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
        ce = nn.CrossEntropyLoss()
        for _ in range(EPOCHS):
            perm = torch.randperm(len(x), device=device)
            for i in range(0, len(x), 128):
                ix = perm[i:i+128]
                out = net(x[ix])
                prob = torch.softmax(out, 1)[:, 1]
                loss = ce(out, y[ix]) + cfg['alpha'] * ((prob - msg[ix]) ** 2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        net.eval()
        with torch.no_grad():
            out = net(d['xte'].to(device)); prob = torch.softmax(out, 1)[:, 1]
            metric = float((out.argmax(1) != d['yte'].to(device)).float().mean())
            rmse = float(torch.sqrt(((prob - d['message_te'].to(device)) ** 2).mean()))
            corr = float(torch.corrcoef(torch.stack([prob, d['message_te'].to(device)]))[0, 1])
    except RuntimeError:
        # CPU fallback mirrors the benchmark's robust device policy.
        net = make_model('mlp_tiny', d['input_shape'], d['out_dim'])
        net = net.to('cpu'); x, y, msg = d['xtr'], d['ytr'], d['message_tr']
        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']); ce = nn.CrossEntropyLoss()
        for _ in range(EPOCHS):
            for i in range(0, len(x), 128):
                out = net(x[i:i+128]); prob = torch.softmax(out, 1)[:, 1]
                loss = ce(out, y[i:i+128]) + cfg['alpha'] * ((prob-msg[i:i+128])**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            out = net(d['xte']); prob = torch.softmax(out, 1)
            metric = float((out.argmax(1) != d['yte']).float().mean())
            rmse = float(torch.sqrt(((prob-d['message_te'])**2).mean()))
            corr = float(torch.corrcoef(torch.stack([prob, d['message_te']]))[0,1])
    if return_signature:
        return float(metric), {'message_rmse': rmse, 'message_corr': corr}
    return float(metric)


def main():
    # Baseline sweep includes every lr evaluated by the idea; alpha is the
    # method knob and is swept symmetrically on the idea side (baseline alpha=0).
    grid = [{'lr': lr, 'weight_decay': wd} for lr in LR_GRID for wd in [0.0, 1e-4]]
    base = sweep_baseline(lambda cfg: (lambda seed: baseline_train(cfg, seed)), grid)
    idea_grid = [{'lr': lr, 'alpha': a} for lr in LR_GRID for a in [0.1, 0.3, 1.0]]
    candidates = []
    for cfg in idea_grid:
        r = evaluate(lambda seed, c=cfg: idea_train(c, seed), SEEDS[:4])
        candidates.append({'cfg': cfg, 'mean': r['mean']})
    best = min(candidates, key=lambda z: z['mean'])['cfg']
    idea = evaluate(lambda seed: idea_train(best, seed), SEEDS)
    sigs = [idea_train(best, s, True)[1] for s in SEEDS]
    signature = {'predicted': 'exact backward message should improve NN feasibility-probability fidelity',
                 'observed_message_rmse_mean': float(np.mean([x['message_rmse'] for x in sigs])),
                 'observed_message_corr_mean': float(np.mean([x['message_corr'] for x in sigs])),
                 'confirmed': bool(np.mean([x['message_corr'] for x in sigs]) > 0.9),
                 'idea_sweep': candidates}
    report = make_report('rare_event_markov', 'mlp_tiny', base, idea, signature)
    Path('bench_report.json').write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__':
    main()