"""Model, controller design and simulation for the Voltomni cart-pole kit.

The stepper drives the cart through a GT2 belt with a commanded acceleration u (m/s²), so the cart is
a double integrator and only the pole has dynamics. x = cart position (m, 0 = centre of travel),
phi = pole angle measured from UPRIGHT (phi = pi hanging down), radians.

    Jp * phi'' = m*lc*g*sin(phi) - m*lc*cos(phi)*u - b*phi'

Run it to print LQR gains for the firmware (paste into the serial console with the `g` command)
and to simulate a swing-up and balance:  python3 cartpole.py   (needs numpy only; matplotlib optional)
"""
from dataclasses import dataclass
import math

import numpy as np


@dataclass
class Rig:
    rod_m: float = 0.015    # printed pole mass (kg); weigh yours
    rod_l: float = 0.250    # pivot to tip (m)
    tip_m: float = 0.003    # tip disc (kg)
    b: float = 0.00002      # pivot friction (N·m·s); fit from a free-swing decay (see guide)
    g: float = 9.81
    dt: float = 0.002       # control period (500 Hz, as in the firmware)
    u_max: float = 15.0     # cart acceleration limit (m/s²)
    v_max: float = 0.8      # cart speed limit (m/s): 20 motor rev/s
    x_limit: float = 0.15   # soft travel limit each side of centre (m); the switches sit further out

    @property
    def m(self):
        return self.rod_m + self.tip_m

    @property
    def lc(self):  # centre of mass from the pivot
        return (self.rod_m * self.rod_l / 2 + self.tip_m * self.rod_l) / self.m

    @property
    def Jp(self):  # inertia about the pivot
        return self.rod_m * self.rod_l**2 / 3 + self.tip_m * self.rod_l**2


def linear_model(p: Rig):
    """Continuous model about upright: s = [x, phi, x', phi'], input u = x''."""
    a = p.m * p.lc * p.g / p.Jp
    c = p.m * p.lc / p.Jp
    A = np.array([[0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0], [0, a, 0, -p.b / p.Jp]])
    B = np.array([[0], [0], [1], [-c]])
    return A, B


def dlqr(p: Rig, Q=np.diag([200.0, 50.0, 5.0, 1.0]), R=np.array([[0.02]])):
    """Discrete LQR by iterating the Riccati equation (no scipy needed). Returns K for u = -K s."""
    A, B = linear_model(p)
    Ad = np.eye(4) + A * p.dt + A @ A * p.dt**2 / 2
    Bd = B * p.dt + A @ B * p.dt**2 / 2
    P = Q.copy()
    for _ in range(50000):
        K = np.linalg.solve(R + Bd.T @ P @ Bd, Bd.T @ P @ Ad)
        Pn = Q + Ad.T @ P @ (Ad - Bd @ K)
        if np.max(np.abs(Pn - P)) < 1e-9:
            break
        P = Pn
    return K.ravel()


def wrap(a):
    return (a + math.pi) % (2 * math.pi) - math.pi


class Controller:
    """The firmware's logic: energy swing-up, then LQR once the pole is near upright."""

    def __init__(self, p: Rig, K, k_energy=0.2, k_centre=80.0, k_damp=2.0, swing_max=6.0, catch=0.3):
        self.p, self.K, self.catch = p, K, catch
        self.k_energy, self.k_centre, self.k_damp, self.swing_max = k_energy, k_centre, k_damp, swing_max
        self.w0sq = p.m * p.lc * p.g / p.Jp
        self.balancing = False

    def __call__(self, x, phi, dx, dphi):
        p = self.p
        if abs(x) > p.x_limit:  # travel limit: stop
            return None
        if abs(phi) < self.catch:
            self.balancing = True
        elif abs(phi) > 0.6:
            self.balancing = False
        if self.balancing:
            u = -float(self.K @ np.array([x, phi, dx, dphi]))
        else:
            # Energy relative to upright (zero there): E = phi'^2/2 - w0^2 (1 - cos phi). Pump in phase
            # with the swing (bounded) while a spring and damper keep the cart near the centre.
            E = 0.5 * dphi**2 - self.w0sq * (1 - math.cos(phi))
            pump = self.k_energy * E * math.copysign(1.0, dphi * math.cos(phi))
            u = max(-self.swing_max, min(self.swing_max, pump)) - self.k_centre * x - self.k_damp * dx
        return max(-p.u_max, min(p.u_max, u))


def simulate(p: Rig, ctrl, T=12.0, phi0=math.pi - 0.05):
    x, dx, phi, dphi = 0.0, 0.0, phi0, 0.0
    log = []
    n_sub = 10
    for k in range(int(T / p.dt)):
        u = ctrl(x, wrap(phi), dx, dphi)
        if u is None:
            return log, "travel limit"
        for _ in range(n_sub):  # integrate the plant finer than the controller
            h = p.dt / n_sub
            ddphi = (p.m * p.lc * p.g * math.sin(phi) - p.m * p.lc * math.cos(phi) * u - p.b * dphi) / p.Jp
            dx = max(-p.v_max, min(p.v_max, dx + u * h))
            x += dx * h
            dphi += ddphi * h
            phi += dphi * h
        log.append((k * p.dt, x, wrap(phi), dx, dphi, u))
    return log, "ok"


def summary(log, status):
    up = [t for (t, _, ph, *_) in log if abs(ph) < 0.05]
    held = status == "ok" and all(abs(ph) < 0.1 for (_, _, ph, *_) in log[-500:])
    return up[0] if up else None, held


if __name__ == "__main__":
    rig = Rig()
    K = dlqr(rig)
    print(f"Rig: m={rig.m*1000:.0f} g, lc={rig.lc*1000:.0f} mm, natural frequency {math.sqrt(rig.m*rig.lc*rig.g/rig.Jp)/(2*math.pi):.2f} Hz, w0^2={rig.m*rig.lc*rig.g/rig.Jp:.2f}")
    print("LQR gains (paste into the serial console):")
    print("g " + " ".join(f"{k:.3f}" for k in K))
    log, status = simulate(rig, Controller(rig, K))
    first, held = summary(log, status)
    print(f"Simulation: {status}; " + (f"first upright at {first:.2f} s" if first is not None else "never reached upright"))
    print(f"Balanced for the last second: {held}; cart range {min(l[1] for l in log)*1000:.0f}..{max(l[1] for l in log)*1000:.0f} mm")
    try:
        import matplotlib.pyplot as plt

        t = [l[0] for l in log]
        plt.plot(t, [l[1] * 10 for l in log], label="cart x (dm)")
        plt.plot(t, [l[2] for l in log], label="pole φ (rad, 0 = up)")
        plt.legend(); plt.xlabel("time (s)"); plt.title("Swing-up and balance (simulation)")
        plt.savefig("simulation.png", dpi=120)
        print("Plot saved to simulation.png")
    except ImportError:
        pass