#!/usr/bin/env python3

"""SE(3) geometric control demo: circular trajectory tracking."""



from __future__ import annotations



import argparse

import sys

from pathlib import Path



import numpy as np



ROOT = Path(__file__).resolve().parents[1]

if str(ROOT) not in sys.path:

  sys.path.insert(0, str(ROOT / "src"))



from uavcatch.config import load_control_config, load_environment_config, load_uav_config

from uavcatch.control import TrackingController, circle_trajectory

from uavcatch.control.factory import build_cascaded_controller

from uavcatch.dynamics.state import QuadState

from uavcatch.sim.engine import SimEngine

from uavcatch.sim.mujoco_viz import MuJoCoViewer





def parse_args() -> argparse.Namespace:

  parser = argparse.ArgumentParser(

    description="Geometric SE(3) circular trajectory tracking demo",

  )

  parser.add_argument("--environment", type=Path, default=ROOT / "environment.toml")

  parser.add_argument("--uav", type=Path, default=ROOT / "uav.toml")

  parser.add_argument("--control", type=Path, default=ROOT / "control.toml")

  parser.add_argument("--radius", type=float, default=3.0, help="Circle radius (m)")

  parser.add_argument("--height", type=float, default=2.0, help="Flight height (m)")

  parser.add_argument("--period", type=float, default=12.0, help="One lap period (s)")

  parser.add_argument("--duration", type=float, default=24.0, help="Simulation duration (s)")

  parser.add_argument("--no-viewer", action="store_true")

  parser.add_argument(

    "--speed",

    type=float,

    default=1.0,

    help="Playback speed (1.0 = real time, <1 slow motion)",

  )

  return parser.parse_args()





def main() -> int:

  args = parse_args()

  env_cfg = load_environment_config(args.environment)

  uav_cfg = load_uav_config(args.uav)

  control_cfg = load_control_config(args.control)

  env_cfg.simulation.duration = args.duration



  if control_cfg.controller.mode != "geometric":

    print(

      f"Warning: control.toml mode is '{control_cfg.controller.mode}'; "

      "circle tracking works best with mode = \"geometric\"."

    )



  engine = SimEngine(env_cfg=env_cfg, uav_cfg=uav_cfg)

  engine.reset()



  center = np.array([0.0, 0.0, 0.0], dtype=np.float64)

  pos_fn, yaw_fn, accel_fn = circle_trajectory(

    center=center,

    radius=args.radius,

    height=args.height,

    period=args.period,

  )

  init_target = pos_fn(0.0)

  controller = build_cascaded_controller(

    control_cfg,

    engine.dynamics,

    uav_cfg,

    target_position=init_target.position,

    target_velocity=init_target.velocity,

    target_yaw=float(yaw_fn(0.0)),

  )

  tracker = TrackingController(

    controller=controller,

    target_fn=pos_fn,

    yaw_fn=yaw_fn,

    accel_fn=accel_fn,

  )



  duration = args.duration

  dt = env_cfg.simulation.dt

  use_viewer = env_cfg.viewer.enable and not args.no_viewer



  if use_viewer:

    with MuJoCoViewer(env_cfg, uav_cfg=uav_cfg) as viewer:

      def on_step(state: QuadState, t: float, omega: np.ndarray) -> None:

        viewer.render_realtime(state, omega, dt, speed=args.speed)

        if not viewer.is_running():

          raise KeyboardInterrupt



      try:

        engine.run(tracker, duration=duration, on_step=on_step)

      except KeyboardInterrupt:

        pass

  else:

    engine.run(tracker, duration=duration)



  final = engine.state

  final_target = pos_fn(duration)

  pos_err = float(np.linalg.norm(final.position - final_target.position))

  print(f"Simulation finished at t={engine.time:.2f}s")

  print(f"Final position:  {final.position}")

  print(f"Target position: {final_target.position}")

  print(f"Position error:  {pos_err:.4f} m")

  print(f"Final velocity:  {final.velocity}")



  if pos_err > 1.0:

    print("Warning: circular tracking error exceeds 1.0 m")

    return 1

  return 0





if __name__ == "__main__":

  try:

    raise SystemExit(main())

  except KeyboardInterrupt:

    raise SystemExit(130) from None