#!/usr/bin/env python3
"""Pre-flight check for the JEPA + PPO policy on a TurtleBot3 Burger with an LDS-02. NEVER moves the robot.

    python3 check_robot.py

1. Lidar: scans per second (want 4.5-7), points per scan, start angle / direction, share of invalid readings (0 / inf:
   no return or closer than 0.16 m). Odometry: messages per second (want >= 10).
2. Live direction test: nearest obstacle AHEAD / LEFT / BEHIND / RIGHT and the JEPA collision risk ahead / left.
   Hold a box ~30 cm in front of the robot: AHEAD must drop and risk ahead rise. Then on its left: LEFT must drop.
   If left and right are swapped, or ahead shows behind, check LDS_MODEL on the robot (or use --angle-offset 3.14159).
3. Times one full policy step (encoder + heads + policy) on this laptop (want < 50 ms).
Press Ctrl+C to quit.
"""
import argparse
import math
import os
import sys
import time

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
from lds02_adapter import RAY_ANGLES, LidarAdapter  # noqa: E402

try:
    import rclpy
    from nav_msgs.msg import Odometry
    from rclpy.qos import qos_profile_sensor_data
    from sensor_msgs.msg import LaserScan
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')

SECTORS = dict(AHEAD=0, LEFT=90, BEHIND=180, RIGHT=270)


def sector_min(clean, centre_deg, half=25):
    d = np.abs((np.degrees(RAY_ANGLES) - centre_deg + 180) % 360 - 180)
    return float(clean[d <= half].min())


def ok(flag):
    return 'OK  ' if flag else 'CHECK'


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--angle-offset', type=float, default=0.0)
    ap.add_argument('--seconds', type=float, default=60.0)
    a = ap.parse_args()
    rclpy.init()
    node = rclpy.create_node('tb3_check_v6')
    ad = LidarAdapter(a.angle_offset)
    st = dict(scan_t=[], odom_t=[], n=[], inv=[], info=None)

    def on_scan(m):
        r = np.asarray(m.ranges, np.float64)
        st['scan_t'].append(time.time())
        st['n'].append(len(r))
        st['inv'].append(float(np.mean(~np.isfinite(r) | (r <= 0) | (r < m.range_min))))
        st['info'] = (m.angle_min, m.angle_increment, m.range_min, m.range_max)
        ad.push(r, m.angle_min, m.angle_increment, m.header.stamp.sec + m.header.stamp.nanosec * 1e-9)

    node.create_subscription(LaserScan, 'scan', on_scan, qos_profile_sensor_data)
    node.create_subscription(Odometry, 'odom', lambda m: st['odom_t'].append(time.time()), 10)
    print('Listening for 5 s ...', flush=True)
    t0 = time.time()
    while time.time() - t0 < 5:
        rclpy.spin_once(node, timeout_sec=0.05)
    if len(st['scan_t']) < 3 or len(st['odom_t']) < 3:
        sys.exit('No lidar or odometry data. Is the robot bringup running, with the same ROS_DOMAIN_ID on both machines?')
    hz = (len(st['scan_t']) - 1) / (st['scan_t'][-1] - st['scan_t'][0])
    ohz = (len(st['odom_t']) - 1) / (st['odom_t'][-1] - st['odom_t'][0])
    amin, ainc, rmin, rmax = st['info']
    print(f'\n[{ok(4.5 <= hz <= 7.5)}] lidar: {hz:.1f} scans/s (training: 5)')
    print(f'[{ok(True)}] lidar: {int(np.median(st["n"]))} points per scan, {math.degrees(abs(ainc)):.2f} deg apart, '
          f'start {math.degrees(amin):.0f} deg, {"counter-clockwise" if ainc > 0 else "CLOCKWISE (handled)"}, '
          f'range {rmin:.2f}-{rmax:.1f} m')
    inv = float(np.mean(st['inv']))
    print(f'[{ok(inv < 0.25)}] lidar: {100 * inv:.0f} % invalid readings (no return / too close). Training saw 2-15 %; '
          'much more means open space beyond range or dark / shiny surfaces nearby.')
    print(f'[{ok(ohz >= 10)}] odometry: {ohz:.0f} messages/s')
    from v6_policy import DeployedPolicy
    pol = DeployedPolicy.load()
    t1 = time.perf_counter()
    for _ in range(5):
        pol._key = None
        pol.act(ad, 1.0, 0.0)
    ms = (time.perf_counter() - t1) / 5 * 1e3
    print(f'[{ok(ms < 50)}] one policy step on this laptop: {ms:.0f} ms (must be well under the 100 ms control period)')
    print(f'      model: v7 seed {pol.cfg["deployed"]["seed"]}, collision shield threshold {pol.agent.tau}\n')
    print('Live readout (metres; risk = JEPA probability of an obstacle within 0.25 m in the next 1 s). Ctrl+C to quit.')
    t0, last = time.time(), 0
    try:
        while time.time() - t0 < a.seconds:
            rclpy.spin_once(node, timeout_sec=0.05)
            if time.time() - last > 0.5 and ad.ready(1):
                last = time.time()
                c = ad.last_clean
                risk, _ = pol.heads(ad)
                print('  ' + '  '.join(f'{k} {sector_min(c, v):4.2f}' for k, v in SECTORS.items())
                      + f'  | risk ahead {risk[0]:.2f}  left {risk[2]:.2f}  | invalid rays held {ad.held_rays}', flush=True)
    except KeyboardInterrupt:
        pass
    node.destroy_node()
    rclpy.shutdown()


if __name__ == '__main__':
    main()