#!/usr/bin/env python3
"""Finite-horizon discrete LQG cart exercise. Python 3.10+, standard library only."""

import argparse
import json
import math

DT, STEPS = 0.1, 80
Q = [[4.0, 0.0], [0.0, 1.0]]
W = [[0.0004, 0.0], [0.0, 0.0025]]
INITIAL_COVARIANCE = [[0.04, 0.0], [0.0, 0.04]]
PRESETS = ("baseline", "noisy-sensor", "wrong-mass")


def dot(a, b):
    return a[0] * b[0] + a[1] * b[1]


def add(a, b):
    return [a[0] + b[0], a[1] + b[1]]


def mv(a, x):
    return [dot(a[0], x), dot(a[1], x)]


def transpose(a):
    return [[a[0][0], a[1][0]], [a[0][1], a[1][1]]]


def mm(a, b):
    bt = transpose(b)
    return [[dot(a[0], bt[0]), dot(a[0], bt[1])],
            [dot(a[1], bt[0]), dot(a[1], bt[1])]]


def ma(a, b):
    return [[a[i][j] + b[i][j] for j in range(2)] for i in range(2)]


def symmetric(a):
    off = (a[0][1] + a[1][0]) / 2
    return [[a[0][0], off], [off, a[1][1]]]


def cart_model(mass=2.0, damping=1.0, dt=DT):
    """Exact ZOH dynamics. W below is per-step, not a continuous spectral density."""
    if not all(math.isfinite(x) and x > 0 for x in (mass, damping, dt)):
        raise ValueError("Positive finite mass, damping and sample time required.")
    one_minus_decay = -math.expm1(-damping * dt / mass)
    return {"A": [[1.0, mass / damping * one_minus_decay], [0.0, 1 - one_minus_decay]],
            "B": [dt / damping - mass / damping**2 * one_minus_decay,
                  one_minus_decay / damping]}


def riccati(model, steps=STEPS, q=Q, r=1.0, terminal=Q, dt=DT):
    """Stage weights dt*Q, dt*R; unscaled terminal Qf. Gains indexed 0..N-1."""
    matrices, gains = [None] * (steps + 1), [None] * steps
    matrices[steps] = [row[:] for row in terminal]
    a, b = model["A"], model["B"]
    at = transpose(a)
    for k in range(steps - 1, -1, -1):
        pb = mv(matrices[k + 1], b)
        denominator = dt * r + dot(b, pb)
        cross = mv(at, pb)
        gains[k] = [c / denominator for c in cross]
        propagated = mm(mm(at, matrices[k + 1]), a)
        matrices[k] = symmetric([
            [dt * q[i][j] + propagated[i][j] - cross[i] * cross[j] / denominator
             for j in range(2)] for i in range(2)])
    return {"model": model, "gains": gains, "riccati": matrices}


def kalman_update(prior, covariance, measurement, variance):
    """Only y enters: simulator state is deliberately absent from this API."""
    innovation_variance = covariance[0][0] + variance
    gain = [covariance[0][0] / innovation_variance, covariance[1][0] / innovation_variance]
    innovation = measurement - prior[0]
    estimate = [prior[i] + gain[i] * innovation for i in range(2)]
    residual = [[1 - gain[0], 0.0], [-gain[1], 1.0]]
    # Joseph form preserves the covariance's positive-semidefinite structure.
    posterior = symmetric(ma(mm(mm(residual, covariance), transpose(residual)),
                             [[variance * gain[i] * gain[j] for j in range(2)] for i in range(2)]))
    return {"estimate": estimate, "covariance": posterior, "gain": gain,
            "innovation": innovation, "innovationVariance": innovation_variance}


def kalman_predict(estimate, covariance, control, model, w=W):
    return {"estimate": add(mv(model["A"], estimate), [b * control for b in model["B"]]),
            "covariance": symmetric(ma(mm(mm(model["A"], covariance), transpose(model["A"])), w))}


def normal_random(seed):
    """32-bit seed avalanche, xorshift32, Box-Muller. Matches the browser's draws."""
    if not isinstance(seed, int) or not 1 <= seed <= 0xffffffff:
        raise ValueError("Seed must be an integer from 1 through 4294967295.")
    mask = 0xffffffff
    mixed = (seed + 0x9e3779b9) & mask
    mixed = ((mixed ^ (mixed >> 16)) * 0x21f0aaad) & mask
    mixed = ((mixed ^ (mixed >> 15)) * 0x735a2d97) & mask
    state = (mixed ^ (mixed >> 15)) & mask or 0x6d2b79f5
    spare = None

    def uniform():
        nonlocal state
        state ^= (state << 13) & mask
        state ^= state >> 17
        state ^= (state << 5) & mask
        state &= mask
        return (state + 0.5) / 4294967296

    def normal():
        nonlocal spare
        if spare is not None:
            value, spare = spare, None
            return value
        radius = math.sqrt(-2 * math.log(uniform()))
        angle = 2 * math.pi * uniform()
        spare = radius * math.sin(angle)
        return radius * math.cos(angle)

    return normal


def make_draws(seed):
    normal = normal_random(seed)
    initial = [1 + 0.2 * normal(), 0.2 * normal()]
    steps = [{"measurement": normal(), "process": [normal() * 0.02, normal() * 0.05]}
             for _ in range(STEPS)]
    return {"initial": initial, "steps": steps}


def run(controller, preset, draws, design):
    """y[k] -> posterior -> u[k] -> x[k+1], next prior. No terminal observation."""
    variance = 0.09 if preset == "noisy-sensor" else 0.01
    plant = cart_model(4.0 if preset == "wrong-mass" else 2.0)
    state, prior = draws["initial"][:], [1.0, 0.0]
    covariance = [row[:] for row in INITIAL_COVARIANCE]
    previous_measurement = None
    cost = effort = position_squared = 0.0
    rows = []
    for k in range(STEPS):
        measurement = state[0] + math.sqrt(variance) * draws["steps"][k]["measurement"]
        update = kalman_update(prior, covariance, measurement, variance)
        if controller == "oracle":
            estimate = state[:]
        elif controller == "raw":
            estimate = [measurement, 0.0 if previous_measurement is None else
                        (measurement - previous_measurement) / DT]
        else:
            estimate = update["estimate"]
        control = -dot(design["gains"][k], estimate)
        effort += DT * control * control
        position_squared += state[0]**2
        cost += DT * (dot(state, mv(Q, state)) + control * control)
        rows.append({"step": k, "time": k * DT, "state": state[:], "estimate": estimate,
                     "measurement": measurement, "control": control, "cumulativeCost": cost,
                     "covariance": update["covariance"] if controller == "lqg" else None})
        prediction = kalman_predict(update["estimate"], update["covariance"], control, design["model"])
        prior, covariance = prediction["estimate"], prediction["covariance"]
        state = add(add(mv(plant["A"], state), [b * control for b in plant["B"]]),
                    draws["steps"][k]["process"])
        previous_measurement = measurement
    cost += dot(state, mv(Q, state))
    terminal_estimate = state[:] if controller == "oracle" else prior if controller == "lqg" else rows[-1]["estimate"][:]
    rows.append({"step": STEPS, "time": STEPS * DT, "state": state[:], "estimate": terminal_estimate,
                 "measurement": None, "control": None, "cumulativeCost": cost,
                 "covariance": covariance if controller == "lqg" else None})
    return {"controller": controller, "rows": rows, "cost": cost,
            "positionRmse": math.sqrt(position_squared / STEPS), "effort": effort}


def experiment(preset="baseline", seed=7):
    if preset not in PRESETS:
        raise ValueError("Unknown preset.")
    draws, design = make_draws(seed), riccati(cart_model())
    return {"preset": preset, "seed": seed, "design": design,
            "runs": {controller: run(controller, preset, draws, design)
                     for controller in ("lqg", "raw", "oracle")}}


def summary(preset="baseline", count=64):
    costs = [experiment(preset, seed)["runs"] for seed in range(1, count + 1)]

    def paired(controller):
        differences = [runs[controller]["cost"] - runs["lqg"]["cost"] for runs in costs]
        mean = sum(differences) / count
        variance = sum((x - mean)**2 for x in differences) / (count - 1)
        return {"mean": mean, "standardError": math.sqrt(variance / count)}

    return {"count": count,
            "means": {controller: sum(runs[controller]["cost"] for runs in costs) / count
                      for controller in ("lqg", "raw", "oracle")},
            "rawMinusLqg": paired("raw"), "oracleMinusLqg": paired("oracle")}


def check():
    """Check hand calculation, independent Bellman identities and reported costs."""
    def close(actual, expected, tolerance=1e-10):
        if not math.isclose(actual, expected, rel_tol=tolerance, abs_tol=tolerance):
            raise AssertionError(f"Expected {expected}, got {actual}")

    update = kalman_update([1.0, 0.0], INITIAL_COVARIANCE, 1.2, 0.01)
    close(update["estimate"][0], 1.16)
    close(update["covariance"][0][0], 0.008)
    close(update["covariance"][1][1], 0.04)
    design = riccati(cart_model())
    close(design["gains"][0][0], 1.8947971217370019)
    close(design["gains"][0][1], 2.092005724620362)
    # Bellman completion of the square for arbitrary state and input at several k.
    for k in (0, 20, 79):
        x, u = [0.7, -0.4], 0.3
        model, next_s = design["model"], design["riccati"][k + 1]
        nx = add(mv(model["A"], x), [b * u for b in model["B"]])
        left = DT * (dot(x, mv(Q, x)) + u*u) + dot(nx, mv(next_s, nx))
        curvature = DT + dot(model["B"], mv(next_s, model["B"]))
        right = dot(x, mv(design["riccati"][k], x)) + curvature * (u + dot(design["gains"][k], x))**2
        close(left, right)
    for preset in PRESETS:
        for result in experiment(preset, 7)["runs"].values():
            rows = result["rows"]
            measured = sum(DT * (4*r["state"][0]**2 + r["state"][1]**2 + r["control"]**2)
                           for r in rows[:-1])
            measured += 4 * rows[-1]["state"][0]**2 + rows[-1]["state"][1]**2
            close(result["cost"], measured)
    print("PASS: hand Kalman update, finite-horizon gains, Bellman identities, and three-preset costs.")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--preset", choices=PRESETS, default="baseline")
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--check", action="store_true", help="run reproducibility checks before the results")
    parser.add_argument("--json", action="store_true", help="emit full selected-run traces for comparison")
    args = parser.parse_args()
    if not 1 <= args.seed <= 0xffffffff:
        parser.error("seed must be between 1 and 4294967295")
    if args.check:
        check()
    result = experiment(args.preset, args.seed)
    if args.json:
        print(json.dumps(result, separators=(",", ":")))
        return
    averages = summary(args.preset)
    print(f"Preset: {args.preset}; seed: {args.seed}; horizon: {STEPS * DT:.1f} s")
    print("Controller       selected cost    64-seed mean")
    for controller in ("lqg", "raw", "oracle"):
        print(f"{controller:16s} {result['runs'][controller]['cost']:13.6f} {averages['means'][controller]:15.6f}")
    difference = averages["rawMinusLqg"]
    print(f"Paired raw - LQG: {difference['mean']:.6f}; Monte Carlo SE: {difference['standardError']:.6f}")
    print("Means use seeds 1..64 and shared noise draws. SE is a sampling estimate, not a performance guarantee.")
    if args.preset == "wrong-mass":
        print("All controllers use the nominal 2 kg gain. The full-state baseline is not redesigned for 4 kg.")


if __name__ == "__main__":
    main()