#!/usr/bin/env python3
"""Click places on your map in RViz and the JEPA + PPO policy drives the robot there, one goal after another.

Needs (see README "Goals on the map"): the robot bringup, map localization (AMCL) with your saved map, and RViz with
rviz/map_goals.rviz. Then, on the laptop:

    python3 map_goals.py                    # waits for goals; Ctrl+C to quit

In RViz:  "2D Pose Estimate" = tell AMCL where the robot is (once, at the start)
          "Publish Point" (single click) or "2D Goal Pose" = send the robot to that spot
A new click while driving replaces the current goal. The robot only moves if AMCL is confident about its position
(position spread < 0.25 m, heading spread < 15 deg). Every run uses the same safety stops as run_policy.py.
The goal is kept in map coordinates and re-converted to odometry coordinates every step, so AMCL corrections for
wheel-odometry drift are applied while driving.
"""
import argparse
import math
import os
import signal
import sys
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from run_policy import PolicyRunner, build_parser  # noqa: E402

try:
    import rclpy
    from geometry_msgs.msg import PointStamped, PoseStamped, PoseWithCovarianceStamped
    from rclpy.executors import SingleThreadedExecutor
    from rclpy.node import Node
    from rclpy.time import Time
    from tf2_ros import Buffer, TransformListener
except ImportError:
    sys.exit('ROS 2 is not set up in this terminal. Run first:  cd to this folder and run:  source setup_terminal.sh')


def yaw_of(q):
    return math.atan2(2 * (q.w * q.z + q.x * q.y), 1 - 2 * (q.y * q.y + q.z * q.z))


class GoalManager(Node):
    def __init__(self, a):
        super().__init__('map_goal_manager')
        self.a = a
        self.tf = Buffer()
        self.tfl = TransformListener(self.tf, self)
        self.create_subscription(PoseStamped, a.goal_topic, self.on_goal_pose, 10)
        self.create_subscription(PointStamped, a.point_topic, self.on_point, 10)
        self.create_subscription(PoseWithCovarianceStamped, 'amcl_pose', self.on_amcl, 10)
        self.pending, self.goal_map, self.loc = None, None, None
        self.last_goal = (None, -1e9)

    def _new(self, frame, x, y):
        (g, t) = self.last_goal                                   # the same click delivered twice -> ignore the repeat
        if g is not None and g[0] == frame and math.hypot(g[1] - x, g[2] - y) < 0.05 and time.time() - t < 2.0:
            return
        self.last_goal = ((frame, x, y), time.time())
        self.pending = (frame, x, y)

    def on_goal_pose(self, m):
        self._new(m.header.frame_id or 'map', m.pose.position.x, m.pose.position.y)

    def on_point(self, m):
        self._new(m.header.frame_id or 'map', m.point.x, m.point.y)

    def on_amcl(self, m):
        c = m.pose.covariance
        self.loc = (time.time(), math.sqrt(max(c[0], c[7])), math.sqrt(c[35]))

    def localized(self):
        """-> (ok, reason)."""
        if self.a.no_loc_check:
            return True, 'check disabled'
        if self.loc is None:
            return False, 'no /amcl_pose yet: is localization running? Set the robot position with "2D Pose Estimate".'
        _, sxy, syaw = self.loc
        if sxy > 0.25 or syaw > math.radians(15):
            return False, (f'AMCL is unsure where the robot is (position spread {sxy:.2f} m, heading spread '
                           f'{math.degrees(syaw):.0f} deg). Use "2D Pose Estimate", or drive a little with teleop, then click again.')
        return True, f'position spread {sxy:.2f} m, heading spread {math.degrees(syaw):.0f} deg'

    def to_odom(self, frame, x, y):
        """Point in `frame` -> odom frame using the latest TF (map -> odom comes from AMCL)."""
        if frame == self.a.odom_frame:
            return x, y
        try:
            t = self.tf.lookup_transform(self.a.odom_frame, frame, Time())
        except Exception:
            return None
        tr, yaw = t.transform.translation, yaw_of(t.transform.rotation)
        return (tr.x + math.cos(yaw) * x - math.sin(yaw) * y, tr.y + math.sin(yaw) * x + math.cos(yaw) * y)


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('--speed-scale', type=float, default=0.5)
    ap.add_argument('--goal-topic', default='goal_pose')
    ap.add_argument('--point-topic', default='clicked_point')
    ap.add_argument('--odom-frame', default='odom')
    ap.add_argument('--no-loc-check', action='store_true', help=argparse.SUPPRESS)       # testing only
    ap.add_argument('--extra', default='', help=argparse.SUPPRESS)                       # extra run_policy options
    a = ap.parse_args()
    from v6_policy import DeployedPolicy
    policy = DeployedPolicy.load()
    rclpy.init()
    mgr = GoalManager(a)
    ex = SingleThreadedExecutor()
    ex.add_node(mgr)
    stop = {'flag': False}
    signal.signal(signal.SIGINT, lambda *_: stop.update(flag=True))
    signal.signal(signal.SIGTERM, lambda *_: stop.update(flag=True))
    runner, n = None, 0
    print('Ready. In RViz: "2D Pose Estimate" to localize, then "Publish Point" or "2D Goal Pose" to send the robot. Ctrl+C to quit.',
          flush=True)
    try:
        while rclpy.ok() and not stop['flag']:
            ex.spin_once(timeout_sec=0.02)
            if mgr.pending is not None:
                frame, gx, gy = mgr.pending
                mgr.pending = None
                ok, why = mgr.localized()
                g = mgr.to_odom(frame, gx, gy)
                if not ok or g is None:
                    print(f'\n[goal ({gx:.2f}, {gy:.2f}) in {frame}] NOT STARTED: '
                          + (why if not ok else f'no transform {frame} -> {a.odom_frame} yet (is localization running?)'), flush=True)
                    continue
                if runner is not None and runner.state != 'done':
                    runner.finish('NEW GOAL', 'replaced by a new click')
                    ex.remove_node(runner)
                    runner.destroy_node()
                n += 1
                print(f'\n[goal {n}] map ({gx:.2f}, {gy:.2f}) -> odom ({g[0]:.2f}, {g[1]:.2f}); localization OK ({why})', flush=True)
                args = build_parser().parse_args(['--goal-odom', str(g[0]), str(g[1]), '--speed-scale', str(a.speed_scale),
                                                  '--countdown', '1.0'] + a.extra.split())
                args.max_time = 60.0 / a.speed_scale
                target = (frame, gx, gy)
                runner = PolicyRunner(args, policy, name=f'map_goal_runner_{n}', log_prefix=f'mapgoal_{n}',
                                      goal_provider=lambda t=target: mgr.to_odom(*t))
                ex.add_node(runner)
            if runner is not None and runner.state == 'done' and runner.result != 'NEW GOAL':
                for _ in range(5):
                    runner.send(0.0, 0.0)
                print(f'[goal {n}] {runner.result}: {runner.detail}  | driven {runner.path:.2f} m, closest obstacle '
                      f'{runner.min_seen:.2f} m. Waiting for the next click.', flush=True)
                ex.remove_node(runner)
                runner.destroy_node()
                runner = None
    finally:
        if runner is not None:
            runner.finish('STOPPED BY YOU', 'Ctrl+C pressed')
            for _ in range(10):
                runner.send(0.0, 0.0)
                time.sleep(0.03)
        mgr.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main()