import json, math, os
import numpy as np
from scipy.linalg import solve_discrete_are

SEED = 7
rng_global = np.random.default_rng(SEED)

# Small unstable LTI plant. The controller knows this model; process noise is external.
A = np.array([[1.10, 0.10], [0.00, 0.95]])
B = np.array([[0.0], [0.10]])
Q = np.diag([1.0, 0.2])
R = np.array([[0.08]])
P = solve_discrete_are(A, B, Q, R)
K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)
Acl = A - B @ K
D = P - Acl.T @ P @ Acl
alpha_P = np.min(np.linalg.eigvalsh(D))

# A small learned residual surrogate: deliberately useful near the origin but imperfect.
def residual(x):
    return np.array([0.18 * np.tanh(1.5*x[0]) - 0.03*x[1]])

def lqr(x):
    return (-K @ x).reshape(1)

def rollout_model(x, u):
    return A @ x + (B @ np.asarray(u).reshape(1))[..., 0]

def mppi_action(x, M, lam, H=8, sigma=0.55, rng=None):
    # Nominal residual policy around LQR, with receding-horizon sampled corrections.
    if rng is None: rng = np.random.default_rng()
    unom = float(lqr(x)[0] + residual(x)[0])
    # Constant nominal sequence is sufficient for the tiny sanity test.
    eps = rng.normal(0.0, sigma, size=(M, H))
    us = unom + eps
    xs = np.repeat(x[None, :], M, axis=0)
    costs = np.zeros(M)
    for t in range(H):
        costs += np.einsum('bi,ij,bj->b', xs, Q, xs) + R[0,0] * us[:,t]**2
        xs = xs @ A.T + us[:,t,None] * B.T
    costs += np.einsum('bi,ij,bj->b', xs, P, xs)
    # numerically stable importance weights
    z = -(costs - costs.min()) / lam
    w = np.exp(np.clip(z, -80, 0))
    sw = w.sum() + 1e-12
    vals = us[:,0]
    uhat = float(np.dot(w, vals) / sw)
    var = float(np.dot(w, (vals-uhat)**2) / sw)
    ess = float(sw*sw / (np.dot(w,w) + 1e-12))
    se = math.sqrt(max(var, 0.0) / (ess + 1e-12))
    return uhat, se, ess, unom

def budgeted_action(x, lam, M0=16, Mmax=128, c=0.45, q=0.35, rng=None):
    if rng is None: rng = np.random.default_rng()
    # Conservative sample scaling from the supplied corollary (constant hidden factor omitted).
    scale = (np.linalg.norm(P)**2 * np.linalg.norm(B)**2 * np.linalg.norm(Acl)**2 /
             (alpha_P**2)) * math.log(2.0 / 0.05)
    # Start small, then spend computation if the estimated perturbation exceeds budget.
    M = M0
    attempts = 0
    while True:
        uhat, se, ess, unom = mppi_action(x, M, lam, rng=rng)
        ehat = se + abs(uhat-unom)*q
        budget = c * alpha_P * np.linalg.norm(x)
        accepted = np.linalg.norm(B[:,0]) * abs(ehat) <= budget
        attempts += 1
        if accepted or M >= Mmax:
            # At Mmax, safety takes precedence over refinement.
            if not accepted:
                return float(lqr(x)[0]), True, M, attempts, se, ess, scale
            return uhat, False, M, attempts, se, ess, scale
        M *= 2

def one_episode(kind, rng, T=100, lam=0.35, fixed_M=64):
    x = np.array([1.8, 0.0]) + rng.normal(0, .05, 2)
    total = 0.; unstable = False; fallbacks = 0; samples = 0; sevals=[]
    for _ in range(T):
        if kind == 'lqr':
            u = float(lqr(x)[0]); Mused = 0
        elif kind == 'residual':
            u = float(lqr(x)[0] + residual(x)[0]); Mused = 0
        elif kind == 'fixed':
            u, se, _, _ = mppi_action(x, fixed_M, lam, rng=rng); Mused=fixed_M; sevals.append(se)
        else:
            u, fb, Mused, _, se, _, _ = budgeted_action(x, lam, rng=rng)
            fallbacks += int(fb); sevals.append(se)
        u = float(np.clip(u, -5, 5))
        total += float(x @ Q @ x + R[0,0]*u*u)
        x = A @ x + B[:,0]*u + rng.normal(0, .035, 2)
        samples += Mused
        if np.linalg.norm(x) > 8 or not np.all(np.isfinite(x)):
            unstable=True; total += 5000.; break
    return total, unstable, fallbacks, samples, np.mean(sevals) if sevals else 0.

def scaling_check():
    # Repeated independent MPPI estimates at one state; report RMS error to a high-M reference.
    x=np.array([.2, -.05]); lam=.35
    ref=np.mean([mppi_action(x, 4096, lam, rng=np.random.default_rng(10000+i))[0] for i in range(20)])
    rows=[]
    for M in [16,32,64,128,256]:
        vals=np.array([mppi_action(x,M,lam,rng=np.random.default_rng(20000+100*M+j))[0]
                       for j in range(30)])
        rms=float(np.sqrt(np.mean((vals-ref)**2)))
        rows.append((M,rms))
    slope=float(np.polyfit(np.log([r[0] for r in rows]), np.log([r[1] for r in rows]), 1)[0])
    return {'reference_action':float(ref), 'rows':rows, 'loglog_slope':slope}

def main():
    print('geometry', json.dumps({'P_norm':float(np.linalg.norm(P)), 'K':K.tolist(),
      'Acl':Acl.tolist(), 'alpha_P':float(alpha_P),
      'rho_Acl':float(max(abs(np.linalg.eigvals(Acl))))}, sort_keys=True))
    scaling=scaling_check(); print('scaling',json.dumps(scaling))
    n=40; results={}
    for kind in ['lqr','residual','fixed','budgeted']:
        vals=[]
        for i in range(n): vals.append(one_episode(kind,np.random.default_rng(5000+i),lam=.35))
        arr=np.array(vals)
        results[kind]={'mean_cost':float(arr[:,0].mean()), 'std_cost':float(arr[:,0].std()),
          'escape_rate':float(arr[:,1].mean()), 'fallback_rate_per_step':float(arr[:,2].sum()/(n*100)),
          'mean_rollout_samples_per_step':float(arr[:,3].mean()/100), 'mean_action_se':float(arr[:,4].mean())}
    # Temperature sweep for fixed and budgeted, same states/noise seeds per temperature.
    temps={}
    for lam in [.15,.35,.7]:
        for kind in ['fixed','budgeted']:
            vals=[one_episode(kind,np.random.default_rng(9000+i),lam=lam) for i in range(25)]
            a=np.array(vals); temps[kind+str(lam)]={'cost':float(a[:,0].mean()),'escape':float(a[:,1].mean()),'samples':float(a[:,3].mean()/100)}
    out={'geometry':{'P':P.tolist(),'K':K.tolist(),'Acl':Acl.tolist(),'alpha_P':float(alpha_P)},
         'corollary_scale_without_hidden_constant':float((np.linalg.norm(P)**2*np.linalg.norm(B)**2*np.linalg.norm(Acl)**2/alpha_P**2)*math.log(2/.05)),
         'scaling_check':scaling,'episodes':results,'temperature_sweep':temps}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()