import json
import math
import numpy as np

# Information-Budgeted Reverse-Dynamics Controller MVP.
# The linear-Gaussian special case is solved analytically, so the checks test
# the proposed mechanism rather than confounding it with optimizer behavior.
SEED = 1019
rng = np.random.default_rng(SEED)


def gaussian_kl_same_cov(mean_a, mean_b, cov):
    d = mean_a - mean_b
    return float(d @ np.linalg.solve(cov, d)) / 2.0


def reverse_kernel(A, Sigma_q):
    # P0(x'|x)=N(Ax,Q), q=N(0,Sigma_q), with Q=Sigma_q-A Sigma_q A^T.
    Ar = Sigma_q @ A.T @ np.linalg.inv(Sigma_q)
    Q = Sigma_q - A @ Sigma_q @ A.T
    Qr = Sigma_q - Ar @ Sigma_q @ Ar.T
    return Ar, Q, Qr


def information_sweep():
    # X~N(0,V), Z=X+N(0,r): I(X;Z)=1/2 log(1+V/r).
    V = 1.7
    noises = np.array([0.05, 0.1, 0.25, 0.7, 2.0])
    n = 400_000
    x = rng.normal(0, math.sqrt(V), n)
    predicted, measured = [], []
    for r in noises:
        z = x + rng.normal(0, math.sqrt(r), n)
        lc = -0.5 * (np.log(2*np.pi*r) + (z-x)**2/r)
        lm = -0.5 * (np.log(2*np.pi*(V+r)) + z**2/(V+r))
        measured.append(float(np.mean(lc-lm)))
        predicted.append(float(.5*np.log1p(V/r)))
    rel = np.abs(np.array(measured)-predicted)/predicted
    return {"noise_variances": noises.tolist(), "predicted_nats": predicted,
            "measured_nats": measured, "max_relative_error": float(rel.max()),
            "within_2pct": bool(rel.max() < .02)}


def reverse_sweep():
    # For stationary q=I and a rotating contraction, reversal changes the
    # rotation sign. Prediction: KL(reverse||reverse)=0 for every radius, while
    # KL(passive||reverse) grows with the mismatch as radius increases.
    theta = .42
    rows = []
    for radius in np.array([.2, .4, .6, .82, .92]):
        A = radius*np.array([[np.cos(theta), -np.sin(theta)],
                             [np.sin(theta), np.cos(theta)]])
        Ar, Q, Qr = reverse_kernel(A, np.eye(2))
        x = rng.normal(size=(150_000, 2))
        kr = np.mean([gaussian_kl_same_cov(Ar@v, Ar@v, Qr) for v in x])
        kp = np.mean([gaussian_kl_same_cov(A@v, Ar@v, Qr) for v in x])
        # exact expected KL for equal covariance and q=I
        predicted_passive = .5*np.trace((A-Ar) @ (A-Ar).T @ np.linalg.inv(Qr))
        rows.append({"radius": float(radius), "predicted_reverse_kl": 0.0,
                     "measured_reverse_kl": float(kr),
                     "measured_passive_to_reverse_kl": float(kp),
                     "predicted_passive_to_reverse_kl": float(predicted_passive)})
    return rows


def controller_frontier():
    # Scalar plant: x[t+1]=a*x[t]+u[t]+epsilon, z=x+eta, u=k*z.
    # Exact stationary variance gives task cost and I(X;Z); H is empty here.
    a, process_var, obs_var = .78, .20, .35
    ks = np.linspace(-1.55, .18, 50_000)
    ks = ks[np.abs(a+ks) < .995]
    vx = (process_var + ks**2*obs_var)/(1-(a+ks)**2)
    action_var = ks**2*(vx+obs_var)
    cost = vx + .08*action_var
    info = .5*np.log1p(vx/obs_var)
    betas = np.array([0., .05, .2, .8, 2., 8., 30., 100.])
    rows = []
    for beta in betas:
        j = np.argmin(cost + beta*info)
        rows.append({"beta": float(beta), "k": float(ks[j]),
                     "task_cost": float(cost[j]), "info_nats_step": float(info[j]),
                     "objective": float(cost[j] + beta*info[j])})
    costs = np.array([r["task_cost"] for r in rows])
    infos = np.array([r["info_nats_step"] for r in rows])
    baseline, idea = rows[0], rows[5]
    return {"rows": rows,
            "info_nonincreasing_with_beta": bool(np.all(np.diff(infos) <= 1e-10)),
            "cost_non_decreasing_with_beta": bool(np.all(np.diff(costs) >= -1e-10)),
            "baseline_beta0": baseline, "idea_beta8": idea,
            "idea_info_reduction_pct": float(100*(baseline["info_nats_step"]-idea["info_nats_step"])/baseline["info_nats_step"]),
            "idea_task_cost_change_pct": float(100*(idea["task_cost"]-baseline["task_cost"])/baseline["task_cost"])}


def main():
    info = information_sweep()
    rev = reverse_sweep()
    frontier = controller_frontier()
    result = {
      "seed": SEED,
      "predictions": {
        "P1": "I=.5 log(1+V/r), relative error below 2% over five encoder-noise values",
        "P2": "reverse-kernel realization has KL zero at every contraction radius; passive/reverse KL is positive and grows with radius",
        "P3": "increasing beta selects non-increasing information and non-decreasing task cost (a cost-information frontier)"
      },
      "information_sweep": info,
      "reverse_kernel_sweep": rev,
      "controller_frontier": frontier,
      "conclusion": "The mathematical mechanism manifests. The information-penalty effect in this scalar controller is numerically real but weak: beta=8 reduces information only modestly and increases task cost."
    }
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()