"""Model, controller design and simulation for the Voltomni rotary inverted pendulum kit.

The stepper drives the arm with a commanded acceleration u (rad/s²), so the arm is a double
integrator and only the pendulum has dynamics. Angles: theta = arm, phi = pendulum measured from
UPRIGHT (phi = pi hanging down), both in radians.

    Jp * phi'' = m*lc*g*sin(phi) + Jp*sin(phi)*cos(phi)*theta'^2 - m*lc*r*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 furuta.py   (needs numpy only; matplotlib optional)
"""
from dataclasses import dataclass
import math

import numpy as np


@dataclass
class Rig:
    r: float = 0.120        # arm length: motor axis to pendulum pivot (m)
    rod_m: float = 0.010    # printed pendulum rod mass (kg); weigh yours
    rod_l: float = 0.150    # pivot to tip (m)
    tip_m: float = 0.002    # 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 = 60.0     # arm acceleration limit (rad/s²)
    v_max: float = 20.0     # arm speed limit (rad/s)

    @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: x = [theta, phi, theta', phi'], input u = theta''."""
    a = p.m * p.lc * p.g / p.Jp
    c = p.m * p.lc * p.r / 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([5.0, 50.0, 0.5, 1.0]), R=np.array([[0.02]])):
    """Discrete LQR by iterating the Riccati equation (no scipy needed). Returns K for u = -K x."""
    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(20000):
        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 pendulum is near upright."""

    # Swing-up defaults chosen in simulation to work across pivot friction 1e-5..1.5e-4 N·m·s and
    # pendulum mass ±20 % with gains from the nominal model.
    def __init__(self, p: Rig, K, k_energy=1.0, k_centre=80.0, k_damp=5.0, swing_max=60.0, catch=0.35, theta_limit=2.5):
        self.p, self.K, self.catch, self.theta_limit = p, K, catch, theta_limit
        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, theta, phi, dtheta, dphi):
        p = self.p
        if abs(theta) > self.theta_limit:  # cable/arm 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([theta, phi, dtheta, dphi]))
        else:
            # Energy relative to upright (zero there): E = phi'^2/2 - w0^2 (1 - cos phi)
            E = 0.5 * dphi**2 - self.w0sq * (1 - math.cos(phi))
            # Pump energy in phase with the swing (bounded), and keep the arm centred and slow so its
            # speed doesn't disturb the swing.
            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 * theta - self.k_damp * dtheta
        return max(-p.u_max, min(p.u_max, u))


def simulate(p: Rig, ctrl, T=10.0, phi0=math.pi - 0.05):
    theta, dtheta, phi, dphi = 0.0, 0.0, phi0, 0.0
    log = []
    n_sub = 10
    for k in range(int(T / p.dt)):
        u = ctrl(theta, wrap(phi), dtheta, dphi)
        if u is None:
            return log, "arm 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.Jp * math.sin(phi) * math.cos(phi) * dtheta**2 - p.m * p.lc * p.r * math.cos(phi) * u - p.b * dphi) / p.Jp
            dtheta = max(-p.v_max, min(p.v_max, dtheta + u * h))
            theta += dtheta * h
            dphi += ddphi * h
            phi += dphi * h
        log.append((k * p.dt, theta, wrap(phi), dtheta, dphi, u))
    return log, "ok"


if __name__ == "__main__":
    rig = Rig()
    K = dlqr(rig)
    print(f"Rig: m={rig.m*1000:.0f} g, lc={rig.lc*1000:.0f} mm, Jp={rig.Jp:.2e} kg·m², natural frequency {math.sqrt(rig.m*rig.lc*rig.g/rig.Jp)/(2*math.pi):.2f} Hz")
    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))
    up = [t for (t, th, ph, *_ ) in log if abs(ph) < 0.05]
    last = log[-500:]
    held = all(abs(ph) < 0.1 for (_, _, ph, *_ ) in last)
    print(f"Simulation: {status}; first upright at {up[0]:.2f} s" if up else f"Simulation: {status}; never reached upright")
    print(f"Balanced for the last second: {held}; arm range {min(l[1] for l in log):.2f}..{max(l[1] for l in log):.2f} rad")
    try:
        import matplotlib.pyplot as plt

        t = [l[0] for l in log]
        plt.plot(t, [l[1] for l in log], label="arm θ (rad)")
        plt.plot(t, [l[2] for l in log], label="pendulum φ (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