"""Model, controller design and simulation for the Voltomni self-balancing robot kit.

The steppers drive the wheels with a commanded acceleration a (m/s² at the axle), so the wheels are a
double integrator and only the body tilts. x = distance travelled (m), theta = body tilt from
upright (rad, positive leaning forward), measured by the IMU.

    J * theta'' = m*l*g*sin(theta) - m*l*cos(theta)*a

J is the body's inertia about the axle and m*l its mass times the height of its centre of mass above
the axle. Run it to print LQR gains for the firmware (paste them into the robot's web page) and to
simulate recovering from a push and driving:  python3 balancer.py   (numpy; matplotlib optional)
"""
from dataclasses import dataclass, field
import math

import numpy as np


@dataclass
class Robot:
    # (mass kg, height of its centre above the axle m) for each part; motors sit on the axle.
    parts: list = field(default_factory=lambda: [
        (0.56, 0.0),    # two NEMA 17 motors
        (0.18, 0.130),  # 3S 2200 mAh battery on the top deck
        (0.25, 0.070),  # printed decks, brackets and standoffs
        (0.06, 0.070),  # electronics on the middle deck
    ])
    wheel_r: float = 0.045
    g: float = 9.81
    dt: float = 0.002       # control period (500 Hz, as in the firmware)
    a_max: float = 6.0      # wheel acceleration limit (m/s²)
    v_max: float = 1.0      # wheel speed limit (m/s)
    gyro_noise: float = 0.02  # rad/s, white noise on the gyro
    tilt_bias: float = 0.0  # rad, error in the IMU's upright calibration

    @property
    def ml(self):
        return sum(m * h for m, h in self.parts)

    @property
    def J(self):  # about the axle; the motors' own term is small
        return sum(m * h * h for m, h in self.parts) + 0.56 * 0.0423**2 / 6


def linear_model(r: Robot):
    """Continuous model about upright: s = [x, theta, x', theta'], input a = x''."""
    A = np.array([[0, 0, 1, 0], [0, 0, 0, 1], [0, 0, 0, 0], [0, r.ml * r.g / r.J, 0, 0]])
    B = np.array([[0], [0], [1], [-r.ml / r.J]])
    return A, B


def dlqr(r: Robot, Q=np.diag([2.0, 100.0, 4.0, 1.0]), R=np.array([[0.05]])):
    """Discrete LQR by iterating the Riccati equation (no scipy needed). Returns K for a = -K s."""
    A, B = linear_model(r)
    Ad = np.eye(4) + A * r.dt + A @ A * r.dt**2 / 2
    Bd = B * r.dt + A @ B * r.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 simulate(r: Robot, K, T=6.0, theta0=0.15, v_ref=lambda t: 0.0, seed=1):
    """Balance from a tilt of theta0, following the speed command v_ref(t). The controller sees the
    tilt through a complementary filter fed by a noisy gyro and the accelerometer, as on the robot."""
    rng = np.random.default_rng(seed)
    x = dx = dth = 0.0
    th = theta0
    est = theta0 + r.tilt_bias
    x_ref = 0.0
    log = []
    for k in range(int(T / r.dt)):
        t = k * r.dt
        vr = v_ref(t)
        x_ref += vr * r.dt
        gyro = dth + rng.normal(0, r.gyro_noise)
        acc_angle = th + r.tilt_bias + rng.normal(0, 0.02)
        est = 0.98 * (est + gyro * r.dt) + 0.02 * acc_angle
        a = -float(K @ np.array([x - x_ref, est, dx - vr, gyro]))
        a = max(-r.a_max, min(r.a_max, a))
        for _ in range(10):
            h = r.dt / 10
            ddth = (r.ml * r.g * math.sin(th) - r.ml * math.cos(th) * a) / r.J
            dx = max(-r.v_max, min(r.v_max, dx + a * h))
            x += dx * h
            dth += ddth * h
            th += dth * h
        log.append((t, x, th, dx, dth, a))
        if abs(th) > 0.8:
            return log, "fell"
    return log, "ok"


if __name__ == "__main__":
    r = Robot()
    K = dlqr(r)
    print(f"Robot: m*l={r.ml:.4f} kg·m, J={r.J:.5f} kg·m², fall time constant {math.sqrt(r.J/(r.ml*r.g))*1000:.0f} ms")
    print("LQR gains (enter them on the robot's web page):")
    print("K = " + " ".join(f"{k:.3f}" for k in K))
    log, status = simulate(r, K)
    print(f"Push recovery from 0.15 rad: {status}; settled tilt {abs(log[-1][2]):.3f} rad, moved {log[-1][1]*100:.0f} cm")
    log2, status2 = simulate(r, K, T=8, theta0=0.0, v_ref=lambda t: 0.4 if 1 < t < 5 else 0.0)
    print(f"Drive at 0.4 m/s for 4 s: {status2}; max tilt {max(abs(l[2]) for l in log2):.3f} rad, travelled {log2[-1][1]:.2f} m")
    try:
        import matplotlib.pyplot as plt

        t = [l[0] for l in log2]
        plt.plot(t, [l[1] for l in log2], label="distance (m)")
        plt.plot(t, [l[2] for l in log2], label="tilt (rad)")
        plt.legend(); plt.xlabel("time (s)"); plt.title("Driving at 0.4 m/s (simulation)")
        plt.savefig("simulation.png", dpi=120)
        print("Plot saved to simulation.png")
    except ImportError:
        pass