import json
from pathlib import Path
import numpy as np
import torch

SEED = 7
np.random.seed(SEED)
torch.manual_seed(SEED)


def toy_check():
    a = 0.8
    ns = np.array([2, 4, 8, 16, 32, 64, 128, 256])
    cumulative = np.array([(1-a**int(n))/(1-a) for n in ns])
    mean_gap = cumulative / ns
    slope = float(np.polyfit(np.log(ns[2:]), np.log(mean_gap[2:]), 1)[0])
    return {'a': a, 'N': ns.tolist(), 'cumulative_gap': cumulative.tolist(),
            'mean_gap': mean_gap.tolist(), 'loglog_slope': slope,
            'theoretical_C': 1/(1-a)}


def make_data(T=240, a=.82, b=.55, noise=.025):
    rng = np.random.default_rng(SEED)
    u = rng.normal(size=T).astype(np.float32)
    x = np.zeros(T+1, dtype=np.float32); y = np.zeros(T, dtype=np.float32)
    x[0] = .7
    for k in range(T):
        y[k] = x[k] + noise*rng.normal(); x[k+1] = a*x[k] + b*u[k]
    return torch.tensor(u), torch.tensor(y)


def rollout(theta, x0, u):
    a, b = theta[0], theta[1]; x, out = x0, []
    for uk in u:
        out.append(x); x = a*x + b*uk
    return torch.stack(out)


def estimate_gap(theta, u, y, x0, horizon, inner_steps=20):
    z = torch.nn.Parameter(x0.detach().clone())
    inner = torch.optim.Adam([z], lr=.08)
    uu, yy = u[:horizon], y[:horizon]
    for _ in range(inner_steps):
        inner.zero_grad(); loss = ((rollout(theta.detach(), z, uu)-yy)**2).mean()
        loss.backward(); inner.step()
    with torch.no_grad():
        fixed, free = rollout(theta, x0, uu), rollout(theta, z, uu)
        return (fixed-free).abs().sum().item(), ((free-yy)**2).mean().item(), float(z)


def fit(method, u, y, steps=600, base_horizon=40, eps=.05, every=25):
    theta = torch.nn.Parameter(torch.tensor([.35, .25])); opt = torch.optim.Adam([theta], lr=.018)
    fixed_x0 = torch.tensor(0.0); chosen, chat = base_horizon, 0.0; hist = []
    for step in range(steps):
        if method == 'calibrated' and step % every == 0:
            c, _, _ = estimate_gap(theta, u, y, fixed_x0, base_horizon)
            chat = .8*chat + .2*c
            chosen = min(base_horizon, max(4, int(np.ceil(chat/eps))))
        n = base_horizon if method == 'full' else (8 if method == 'fixed_8' else chosen)
        start = (3*step) % (len(u)-n+1)
        loss = ((rollout(theta, fixed_x0, u[start:start+n])-y[start:start+n])**2).mean()
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([theta], 5.0); opt.step()
        if step % 50 == 0: hist.append([step, float(loss), int(n), float(chat)])
    with torch.no_grad(): mse = ((rollout(theta, fixed_x0, u)-y)**2).mean().item()
    gap, free_loss, free_x0 = estimate_gap(theta, u, y, fixed_x0, base_horizon)
    return {'mse': mse, 'theta': theta.detach().numpy().tolist(), 'mean_fixed_free_gap': gap/base_horizon,
            'free_loss': free_loss, 'free_x0': free_x0, 'selected_horizon': int(chosen), 'history': hist}


def empirical_sweep(theta, u, y):
    ns = [2, 4, 8, 16, 32, 40]
    gaps = [estimate_gap(theta, u, y, torch.tensor(0.0), n)[0]/n for n in ns]
    slope = float(np.polyfit(np.log(ns[1:]), np.log(gaps[1:]), 1)[0])
    return {'N': ns, 'mean_gap': gaps, 'loglog_slope': slope}


def main():
    u, y = make_data()
    full = fit('full', u, y)
    fixed = fit('fixed_8', u, y)
    calibrated = fit('calibrated', u, y, eps=.05)
    relaxed = fit('calibrated', u, y, eps=.20)
    out = {'toy_check': toy_check(), 'training': {'full_40': full, 'fixed_8': fixed,
           'calibrated_eps_.05': calibrated, 'calibrated_eps_.20': relaxed},
           'empirical_gap_sweep_full_model': empirical_sweep(torch.tensor(full['theta']), u, y)}
    Path('results.json').write_text(json.dumps(out, indent=2)); print(json.dumps(out, indent=2))

if __name__ == '__main__': main()