import json, math, random
from pathlib import Path
import numpy as np

SEED = 2206
np.random.seed(SEED)
random.seed(SEED)


def toy_sweep():
    # Exact quadratic value landscape. A step gamma moves from x=0 to x=gamma.
    # V(gamma)=a*gamma - b*gamma^2/2; KL(N(gamma,s2)||N(0,s2))=gamma^2/(2*s2).
    # Predictions: KL scales as gamma^2, Delta-V as gamma for small gamma,
    # eta=DeltaV/KL scales as gamma^-1, and eta is constant negative for nuisance.
    a, b, s2 = 1.7, 2.0, 0.8
    gammas = np.logspace(-3, math.log10(3.0), 48)
    kl = gammas**2 / (2.0 * s2)
    dv = a * gammas - 0.5 * b * gammas**2
    eta = dv / kl
    nuisance_dv = -0.5 * b * gammas**2
    nuisance_eta = nuisance_dv / kl

    def slope(x, y, n=10):
        return float(np.polyfit(np.log(x[:n]), np.log(np.abs(y[:n]) + 1e-30), 1)[0])

    # Exact zero crossing is gamma=2a/b. The grid estimate brackets it.
    changes = np.where(np.sign(dv[:-1]) != np.sign(dv[1:]))[0]
    bracket = None if len(changes) == 0 else [float(gammas[changes[0]]), float(gammas[changes[0] + 1])]
    return {
        'predictions': {
            'KL_loglog_slope': 2.0,
            'small_step_value_loglog_slope': 1.0,
            'small_step_eta_loglog_slope': -1.0,
            'aligned_value_zero_crossing': 2.0 * a / b,
            'nuisance_eta': -b * s2,
        },
        'observed': {
            'KL_loglog_slope': slope(gammas, kl),
            'small_step_value_loglog_slope': slope(gammas, dv),
            'small_step_eta_loglog_slope': slope(gammas, eta),
            'aligned_value_zero_crossing_bracket': bracket,
            'nuisance_eta_mean': float(np.mean(nuisance_eta)),
            'nuisance_eta_std': float(np.std(nuisance_eta)),
        },
    }


def neural_experiment():
    # A tiny linear predictor with useful x and nuisance z. The latter is
    # deliberately high dimensional and random, so fitting it acquires KL but
    # cannot improve held-out value. The monitor freezes when efficiency is low.
    import torch
    torch.manual_seed(SEED)
    try:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        torch.zeros(1, device=device)
    except Exception:
        device = torch.device('cpu')

    n_train, n_test, p = 512, 1024, 24
    rng = np.random.default_rng(SEED)
    xtr = rng.normal(size=(n_train, 1)).astype('float32')
    xte = rng.normal(size=(n_test, 1)).astype('float32')
    ztr = rng.normal(size=(n_train, p)).astype('float32')
    zte = rng.normal(size=(n_test, p)).astype('float32')
    ytr = (2.0 * xtr[:, 0] + 0.4 * rng.normal(size=n_train)).astype('float32')
    yte = (2.0 * xte[:, 0] + 0.4 * rng.normal(size=n_test)).astype('float32')
    Xtr = torch.tensor(np.concatenate([xtr, ztr], 1), device=device)
    Xte = torch.tensor(np.concatenate([xte, zte], 1), device=device)
    Ytr, Yte = torch.tensor(ytr, device=device), torch.tensor(yte, device=device)

    def run(monitored):
        w = torch.zeros(p + 1, device=device, requires_grad=True)
        opt = torch.optim.Adam([w], lr=0.08)
        prev = w.detach().clone()
        prev_val = None
        updates, infos, etas = 0, [], []
        for t in range(180):
            idx = torch.arange((t * 32) % n_train, (t * 32) % n_train + 32, device=device) % n_train
            loss = ((Xtr[idx] @ w - Ytr[idx]) ** 2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
            new = w.detach().clone()
            info = float(((new - prev) ** 2).sum().item() / (2 * 0.25))
            val = float(((Xte @ new - Yte) ** 2).mean().item())
            # Value is negative test MSE; a block is stopped after repeated
            # low/negative value-per-KL updates, while preserving equal steps.
            dv = 0.0 if prev_val is None else prev_val - val
            eta = dv / (info + 1e-12)
            infos.append(info); etas.append(eta)
            updates += 1
            if monitored and t > 12 and len(etas) >= 5 and np.mean(etas[-5:]) < 0.02:
                # trust-region response: undo the low-efficiency update
                with torch.no_grad(): w.copy_(prev)
            else:
                prev = new
                prev_val = val
        final_mse = float(((Xte @ w - Yte) ** 2).mean().item())
        return {'test_mse': final_mse, 'updates': updates,
                'total_kl': float(np.sum(infos)),
                'late_eta': float(np.mean(etas[-30:])),
                'device': str(device)}

    return {'baseline_adam': run(False), 'monitor_adam': run(True)}


def main():
    out = {'toy': toy_sweep(), 'neural': neural_experiment()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))


if __name__ == '__main__':
    main()