import json, math, os
import numpy as np
import torch

SEED = 1267
np.random.seed(SEED)
torch.manual_seed(SEED)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    if device == 'cuda':
        torch.cuda.set_device(0)
        _ = torch.zeros(1, device='cuda')
except Exception:
    device = 'cpu'


def solve_local(R, B, a, g, lam, u0=None, bdiff=None):
    """Exact minimizer of 1/2 u'Ru + g'u + lam/2 (b+Bu-b0-Bu0)'a^-1(...)."""
    R, B, a, g = map(np.asarray, (R, B, a, g))
    n = R.shape[0]
    u0 = np.zeros(n) if u0 is None else np.asarray(u0)
    bdiff = np.zeros(B.shape[0]) if bdiff is None else np.asarray(bdiff)
    ai = np.linalg.pinv(a)
    H = R + lam * B.T @ ai @ B
    rhs = -g - lam * B.T @ ai @ (bdiff - B @ u0)
    return np.linalg.solve(H, rhs), H


def toy_verification():
    # Two drift coordinates with equal actuation but very different diffusion.
    R = np.eye(2)
    B = np.eye(2)
    a = np.diag([0.25, 4.0])
    g = np.array([-2.0, -2.0])
    u0 = np.zeros(2)
    lambdas = np.array([0., .1, .3, 1., 3., 10., 30., 100.])
    us = np.array([solve_local(R, B, a, g, x, u0)[0] for x in lambdas])
    weights = np.diag(np.linalg.pinv(a))
    drift_kl = .5 * np.sum((us @ B.T) ** 2 * weights, axis=1)
    # Prediction 1: KL decreases monotonically as lambda increases.
    monotone = bool(np.all(np.diff(drift_kl) <= 1e-10))
    # Prediction 2: each coordinate's deviation has asymptotic lambda*u -> -a*g.
    tail = lambdas >= 10
    asymp_expected = -np.diag(a) * g
    asymp_observed = np.mean(us[tail] * lambdas[tail, None], axis=0)
    rel_asymp_err = float(np.max(np.abs(asymp_observed - asymp_expected) /
                                (np.abs(asymp_expected) + 1e-12)))
    # Prediction 3: regularization curvature ratio equals inverse-noise ratio.
    curv0 = np.diag(R)
    curv_lam = np.diag(R + 1.0 * B.T @ np.linalg.pinv(a) @ B)
    observed_ratio = curv_lam[0] / curv_lam[1]
    predicted_ratio = (1 + 1/a[0, 0]) / (1 + 1/a[1, 1])
    return {
        'lambdas': lambdas.tolist(), 'solutions': us.tolist(), 'trajectory_kl': drift_kl.tolist(),
        'prediction_1_kl_monotone': monotone,
        'prediction_2_expected_lambda_u_tail': asymp_expected.tolist(),
        'prediction_2_observed_lambda_u_tail': asymp_observed.tolist(),
        'prediction_2_relative_error': rel_asymp_err,
        'prediction_3_predicted_curvature_ratio': float(predicted_ratio),
        'prediction_3_observed_curvature_ratio': float(observed_ratio),
        'prediction_3_relative_error': float(abs(observed_ratio-predicted_ratio)/predicted_ratio),
    }


def finite_difference_check():
    R = np.diag([1.3, .7]); B = np.array([[1., .2], [.1, 1.]])
    a = np.diag([.4, 2.5]); g = np.array([-.8, .3]); lam = 2.7
    u, H = solve_local(R, B, a, g, lam)
    def f(x):
        d = B @ x
        return .5*x@R@x + g@x + .5*lam*d@np.linalg.pinv(a)@d
    eps = 1e-5
    grad = np.array([(f(u+eps*np.eye(2)[i])-f(u-eps*np.eye(2)[i]))/(2*eps) for i in range(2)])
    hess = np.empty((2,2))
    for i in range(2):
        for j in range(2):
            ei, ej = np.eye(2)[i]*eps, np.eye(2)[j]*eps
            hess[i,j] = (f(u+ei+ej)-f(u+ei-ej)-f(u-ei+ej)+f(u-ei-ej))/(4*eps*eps)
    return {'stationarity_gradient_norm': float(np.linalg.norm(grad)),
            'hessian_error_norm': float(np.linalg.norm(hess-H)), 'analytic_u': u.tolist()}


def train_policy():
    # Same synthetic states and task for action-norm baseline versus noise-whitened drift KL.
    torch.manual_seed(SEED)
    n, d, act = 4096, 3, 2
    x = torch.randn(n, d, device=device)
    target = x @ torch.tensor([[1.0, -.4], [.2, .8], [-.7, .3]], device=device)
    # State-dependent but known diagonal diffusion; one direction is highly noisy.
    var = torch.tensor([.2, 4.0], device=device)
    B = torch.eye(act, device=device)
    lam = 2.0
    def fit(kind):
        torch.manual_seed(SEED + (0 if kind == 'action' else 1))
        W = torch.nn.Parameter(torch.zeros(d, act, device=device))
        opt = torch.optim.Adam([W], lr=.08)
        for _ in range(700):
            u = x @ W
            task = .5 * ((u-target)**2).mean()
            if kind == 'action':
                reg = .5 * (u**2).sum(1).mean()
            else:
                reg = .5 * ((u**2)/var).sum(1).mean()
            loss = task + lam*reg
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            u = x @ W
            task = .5*((u-target)**2).mean().item()
            action_reg = .5*(u**2).sum(1).mean().item()
            traj_kl = .5*((u**2)/var).sum(1).mean().item()
        return {'task_mse_half': task, 'action_penalty': action_reg,
                'trajectory_kl': traj_kl, 'W': W.detach().cpu().numpy().tolist()}
    return {'device': device, 'lambda': lam, 'action_regularization': fit('action'),
            'noise_whitened_drift_regularization': fit('drift')}


def main():
    out = {'seed': SEED, 'math_check': finite_difference_check(),
           'toy_predictions': toy_verification(), 'policy_experiment': train_policy()}
    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()