#!/usr/bin/env python
"""MuJoCo Tello simulator — real rigid-body physics + a real quadrotor controller.

Why MuJoCo and not Gazebo: the existing Tello Gazebo package (clydemcqueen/tello_ros) needs
ROS2 Foxy + Gazebo Classic, both EOL, and neither installs sanely on Apple Silicon. MuJoCo
ships native arm64 wheels (`pip install mujoco`) and needs no ROS. There is no official Tello
MJCF, so `assets/tello.xml` is written from the published spec (98×92.5×41 mm, 80 g, 3" props,
white shell + prop guards).

The drone is flown the way the real one is: you send high-level commands ("forward 30 cm") and
an onboard controller does the work. Here that controller is a standard cascade —

    position PD  ->  desired acceleration  ->  desired tilt + thrust
                 ->  attitude PD  ->  torques  ->  4 rotor forces

so the model actually banks to translate and settles like a real quadrotor, instead of
teleporting like the kinematic `SimDrone`.

Exposes the SAME interface as SimDrone (connect/takeoff/land/hover/move/flip), so `SafePilot`
and the whole BCI stack work unchanged. Rendering is OFFSCREEN (returns RGB arrays), which
avoids macOS's `mjpython` requirement for the interactive viewer and lets us drop the FPV
image straight into the Qt stimulus window.

  python sim_mujoco.py --demo            # scripted flight, saves a contact sheet
  python sim_mujoco.py --demo --live     # + an OpenCV window if cv2 is available
"""
from __future__ import annotations

import argparse
import threading
import time
from pathlib import Path

import numpy as np

ASSETS = Path(__file__).resolve().parents[2] / "assets" / "tello.xml"
G = 9.81


class MuJoCoDrone:
    """Physics-backed Tello with the SimDrone command surface."""

    def __init__(self, model_path=ASSETS, realtime=True, verbose=True):
        import mujoco
        self.mj = mujoco
        self.model = mujoco.MjModel.from_xml_path(str(model_path))
        self.data = mujoco.MjData(self.model)
        self.mass = float(self.model.body_subtreemass[
            mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "tello")])
        self.bid = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "tello")
        self.realtime, self.verbose = realtime, verbose
        self.flying = False
        self.target = np.array([0.0, 0.0, 0.05])      # world-frame setpoint (m)
        self.yaw_target = 0.0
        self._lock = threading.Lock()
        self._run = False
        self._thread = None
        self._renderers = {}
        self.log = []
        # control gains (tuned for an 80 g craft; gentle, BCI-friendly)
        self.kp_pos, self.kd_pos = np.array([3.0, 3.0, 12.0]), np.array([2.6, 2.6, 5.0])
        self.kp_att, self.kd_att = 0.06, 0.012
        self.tilt_max = np.deg2rad(22)

    # ------------------------------------------------------------------ state
    @property
    def pos(self):
        return self.data.xpos[self.bid].copy()

    @property
    def vel(self):
        return self.data.cvel[self.bid][3:].copy()

    def rpy(self):
        R = self.data.xmat[self.bid].reshape(3, 3)
        return (np.arctan2(R[2, 1], R[2, 2]),                       # roll
                -np.arcsin(np.clip(R[2, 0], -1, 1)),                # pitch
                np.arctan2(R[1, 0], R[0, 0]))                       # yaw

    # ------------------------------------------------------------- controller
    def _control(self):
        """Cascade: position PD -> tilt + thrust -> attitude PD -> rotor mix."""
        p, v = self.pos, self.vel
        roll, pitch, yaw = self.rpy()
        if not self.flying:
            self.data.ctrl[:] = 0.0
            return

        a_des = self.kp_pos * (self.target - p) - self.kd_pos * v      # world accel
        a_des[:2] = np.clip(a_des[:2], -6.0, 6.0)
        # thrust magnitude: vertical component + gravity, projected on body z
        T = self.mass * (G + a_des[2]) / max(np.cos(roll) * np.cos(pitch), 0.5)
        T = float(np.clip(T, 0.0, 3.0))
        # desired tilt (small-angle): +pitch tips thrust to +x, +roll tips it to -y
        cy, sy = np.cos(self.yaw_target), np.sin(self.yaw_target)
        ax, ay = a_des[0] * cy + a_des[1] * sy, -a_des[0] * sy + a_des[1] * cy
        pitch_des = np.clip(ax / G, -self.tilt_max, self.tilt_max)
        roll_des = np.clip(-ay / G, -self.tilt_max, self.tilt_max)

        w = self.data.cvel[self.bid][:3]                               # angular velocity
        tx = self.kp_att * (roll_des - roll) - self.kd_att * w[0]
        ty = self.kp_att * (pitch_des - pitch) - self.kd_att * w[1]
        yaw_err = (self.yaw_target - yaw + np.pi) % (2 * np.pi) - np.pi
        tz = 0.02 * yaw_err - 0.006 * w[2]

        L, k = 0.0455 * np.sqrt(2), 0.0075       # arm length (x-config), yaw reaction coeff
        f = np.array([                            # order matches actuators: FL, FR, BR, BL
            T / 4 + tx / (4 * L) - ty / (4 * L) + tz / (4 * k),
            T / 4 - tx / (4 * L) - ty / (4 * L) - tz / (4 * k),
            T / 4 - tx / (4 * L) + ty / (4 * L) + tz / (4 * k),
            T / 4 + tx / (4 * L) + ty / (4 * L) - tz / (4 * k)])
        self.data.ctrl[:] = np.clip(f, 0.0, 0.8)

    def step(self, n=1):
        with self._lock:
            for _ in range(n):
                self._control()
                self.mj.mj_step(self.model, self.data)

    # --------------------------------------------------------------- commands
    def _say(self, msg):
        self.log.append((time.time(), msg))
        if self.verbose:
            p = self.pos
            print(f"    [mujoco] {msg:<20} pos=({p[0]:+.2f},{p[1]:+.2f},{p[2]:.2f})m")

    def connect(self):
        self._say("connect"); return True

    def start(self, hz=500):
        """Run physics in a background thread (BCI loop stays responsive)."""
        if self._thread:
            return
        self._run = True

        def loop():
            dt = self.model.opt.timestep
            nxt = time.time()
            while self._run:
                self.step(1)
                if self.realtime:
                    nxt += dt
                    s = nxt - time.time()
                    if s > 0:
                        time.sleep(s)
                    elif s < -0.5:
                        nxt = time.time()
        self._thread = threading.Thread(target=loop, daemon=True); self._thread.start()

    def stop(self):
        self._run = False
        if self._thread:
            self._thread.join(timeout=1.0); self._thread = None

    def takeoff(self):
        self.flying = True
        self.target = np.array([self.pos[0], self.pos[1], 0.8])
        self._say("takeoff")

    def land(self):
        self.target = np.array([self.pos[0], self.pos[1], 0.02])
        self._say("land")
        t0 = time.time()
        while self.pos[2] > 0.06 and time.time() - t0 < 4.0:
            time.sleep(0.02)
        self.flying = False

    def hover(self):
        self.target = self.pos.copy()

    def move(self, direction, cm):
        if not self.flying:
            self._say(f"IGNORED {direction} (not flying)"); return
        d = cm / 100.0
        yaw = self.yaw_target
        fwd = np.array([np.cos(yaw), np.sin(yaw), 0.0])
        lft = np.array([-np.sin(yaw), np.cos(yaw), 0.0])
        delta = {"forward": fwd * d, "back": -fwd * d, "left": lft * d, "right": -lft * d,
                 "up": np.array([0, 0, d]), "down": np.array([0, 0, -d])}[direction]
        self.target = self.target + delta
        self.target[2] = max(0.15, self.target[2])
        self._say(f"{direction} {cm}cm")

    def rotate(self, direction, deg):
        if not self.flying:
            self._say(f"IGNORED {direction} (not flying)"); return
        d = np.deg2rad(deg)
        self.yaw_target += -d if direction == "rotate_cw" else d    # +yaw = CCW (right-hand)
        self._say(f"{direction} {deg}deg -> yaw {np.rad2deg(self.yaw_target):+.0f}")

    def flip(self, direction="f"):
        self._say(f"flip {direction} (visual only in sim)")

    def battery(self):
        return 100

    # ---------------------------------------------------------------- render
    def render(self, camera="chase", width=480, height=360):
        """Offscreen RGB frame — no mjpython needed, drops straight into a Qt widget."""
        key = (camera, width, height)
        if key not in self._renderers:
            self._renderers[key] = self.mj.Renderer(self.model, height=height, width=width)
        r = self._renderers[key]
        with self._lock:
            r.update_scene(self.data, camera=camera)
        return r.render()


# ----------------------------------------------------------------------- demo
def _demo(live=False, out="results/mujoco_flight.png"):
    from drone import SafePilot, SafetyConfig                  # noqa: E402
    d = MuJoCoDrone()
    print(f"model loaded: mass={d.mass*1000:.0f} g, hover thrust={d.mass*G:.3f} N total\n")
    d.connect(); d.start()
    pilot = SafePilot(d, SafetyConfig(step_cm=40, min_interval=0.0, hover_after=1e9))

    print("  BCI 起飞被拦截 ->", pilot.update("takeoff"))
    d.takeoff(); time.sleep(2.0)
    shots = [("takeoff", d.render("chase"))]
    for cmd in ["forward", "forward", "left", "up", "right", "back"]:
        pilot.update(cmd, 1.0)
        time.sleep(1.6)
        shots.append((cmd, d.render("chase")))
    fpv = d.render("fpv")
    d.land(); d.stop()
    p = d.pos
    print(f"\n  最终位置 x={p[0]:+.2f} y={p[1]:+.2f} z={p[2]:.2f} m")

    import matplotlib; matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    n = len(shots) + 1
    fig, ax = plt.subplots(2, (n + 1) // 2, figsize=(3.1 * ((n + 1) // 2), 5.2))
    fig.patch.set_facecolor("white")
    for a, (title, img) in zip(ax.ravel(), shots + [("FPV camera", fpv)]):
        a.imshow(img); a.set_title(title, fontsize=10); a.axis("off")
    for a in ax.ravel()[n:]:
        a.axis("off")
    fig.suptitle("MuJoCo Tello — physics sim (chase cam + onboard FPV)", fontweight="bold")
    fig.tight_layout()
    outp = Path(__file__).resolve().parents[2] / out
    outp.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(outp, dpi=120, bbox_inches="tight", facecolor="white")
    print("saved", outp)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--demo", action="store_true")
    ap.add_argument("--live", action="store_true", help="also show a live cv2 window")
    args = ap.parse_args()
    if args.demo:
        _demo(args.live)
    else:
        ap.print_help()


if __name__ == "__main__":
    main()