"""Run the Isaac Lab RL candy-pick policy on the real SO101 (state-based: joints + detected candy).

Same observation and action conventions as Isaac-SO101-CandyPick-Direct:
  obs  = arm joints [rad] + gripper opening [0..1], joint velocities x0.1, candy (x, y, sin2yaw,
         cos2yaw) from the phone at the start, towel (x, y), last action, episode time fraction
  act  = 5 arm-joint target deltas (x0.06 rad per 30 Hz step, clamped to the real-demo ranges)
         + gripper command (-1 closed .. +1 = LeRobot 45 open)

By default it only detects the candy and prints what it would do (DRY RUN, arm not moved).
Add --go to move the arm; you are asked to confirm first. Keep a hand near the power switch.

    cd ~/B1Hackathon/tasks/candy_sort/rl_deploy
    ~/miniconda3/envs/lerobot/bin/python deploy_rl.py --policy policy.pt            # dry run
    ~/miniconda3/envs/lerobot/bin/python deploy_rl.py --policy policy.pt --go       # moves the arm
Safety (real arm only): targets change by at most --max-step-deg per step (default 1 deg = 30 deg/s,
about a third of what the sim policy may command) and actions are low-pass filtered (--smooth).
Every run saves ../eval_videos/rl_<policy>_<time>.mp4 (phone + wrist) and a .csv of measured vs
commanded joints, unless --no-video.
"""

import argparse
import time

import numpy as np
import torch

from common import HERE, JOINTS, Phone, blue_blobs, candy_on_table, connect_robot, load_calib, read_joints

WRIST_DEV = "/dev/v4l/by-id/usb-Ausha_Webcam_Ausha_Webcam-video-index0"

# Must match so101_candy/candy_env_cfg.py
REAL_ARM_RANGE_DEG = np.array([(-44.9, 80.9), (-102.9, 32.6), (-12.2, 93.1), (0.2, 104.3), (-61.1, 99.7)])
ARM_SCALE_RAD = 0.06
GRIP_OPEN_LR, GRIP_CLOSED_LR = 45.0, 0.0
GRIP_LO, GRIP_HI = -0.174533, 1.74533  # sim gripper joint limits [rad] (LeRobot 0..100 maps onto them)
TOWEL_XY = (-0.20, -0.17)
TOWEL_HALF = (0.08, 0.12)
CANDY_X, CANDY_Y = (-0.06, 0.18), (-0.32, -0.12)  # spawn box of the sim task, with some margin
EPISODE_STEPS = 300  # 10 s at 30 Hz
HZ = 30.0

p = argparse.ArgumentParser()
p.add_argument("--policy", required=True, help="TorchScript policy exported from Isaac Lab (policy.pt)")
p.add_argument("--go", action="store_true", help="actually move the arm")
p.add_argument("--seconds", type=float, default=10.0)
p.add_argument("--max-step-deg", type=float, default=1.0,
               help="safety: max change of any joint target per 30 Hz step [deg] (sim policy allows ~3.4)")
p.add_argument("--smooth", type=float, default=0.5, help="safety: action low-pass, 0 = off, 0.5 = average with last")
p.add_argument("--no-video", action="store_true", help="don't save the phone/wrist video and the per-step log")
p.add_argument("--phone-url", default=None, help="phone stream URL (default: PHONE_URL env or http://127.0.0.1:4747/video)")
a = p.parse_args()

policy = torch.jit.load(a.policy, map_location="cpu").eval()
H = load_calib()

# 1) Detect the target candy (blue, not already on the towel, inside the trained workspace).
phone = Phone(a.phone_url) if a.phone_url else Phone()
img = phone.frame()
cands = []
for b in blue_blobs(img):
    xy, yaw = candy_on_table(H, b)
    on_towel = abs(xy[0] - TOWEL_XY[0]) < TOWEL_HALF[0] and abs(xy[1] - TOWEL_XY[1]) < TOWEL_HALF[1]
    in_ws = CANDY_X[0] <= xy[0] <= CANDY_X[1] and CANDY_Y[0] <= xy[1] <= CANDY_Y[1]
    print(f"blue blob at pixel {b['uv'].round(0)} -> table {xy.round(3)} m, yaw {np.degrees(yaw):.0f} deg"
          f"{' (on towel)' if on_towel else ''}{' (outside trained area)' if not in_ws else ''}")
    if not on_towel and in_ws:
        cands.append((xy, yaw))
if not cands:
    raise SystemExit("No blue candy inside the trained workspace. Move one closer to the arm and retry.")
xy, yaw = min(cands, key=lambda c: np.linalg.norm(c[0] - np.array([0.06, -0.22])))
candy_obs = np.array([xy[0], xy[1], np.sin(2 * yaw), np.cos(2 * yaw)])
print(f"target candy: {xy.round(3)} m, yaw {np.degrees(yaw):.0f} deg")

robot = connect_robot()
q = read_joints(robot)
print("current joints [deg]:", q.round(1))
if a.go:
    if input("The arm WILL move for %.0f s. Type 'go' to start: " % a.seconds).strip() != "go":
        robot.disconnect()
        raise SystemExit("aborted")
else:
    print("DRY RUN: computing actions without sending them (add --go to move the arm)")

# Documentation: phone (+ wrist if available) video and a per-step CSV of commanded vs measured joints.
import csv
import os
from datetime import datetime

import cv2

rec = None
if not a.no_video:
    os.makedirs(os.path.join(HERE, "..", "eval_videos"), exist_ok=True)
    stem = os.path.join(HERE, "..", "eval_videos", f"rl_{os.path.splitext(os.path.basename(a.policy))[0]}_{datetime.now():%Y%m%d_%H%M%S}"
                        + ("" if a.go else "_dryrun"))
    wrist = cv2.VideoCapture(WRIST_DEV)
    wrist = wrist if wrist.isOpened() else None
    vw = cv2.VideoWriter(stem + ".mp4", cv2.VideoWriter_fourcc(*"mp4v"), HZ, (640 + (320 if wrist else 0), 360))
    logf = open(stem + ".csv", "w", newline="")
    log = csv.writer(logf)
    log.writerow(["t"] + [f"meas_{n}" for n in JOINTS] + [f"target_{n}" for n in JOINTS] + [f"act_{i}" for i in range(6)])
    rec = (vw, wrist, log, logf)
    print(f"recording to {stem}.mp4 / .csv")

target = q.copy()  # LeRobot units
last_act = np.zeros(6)
prev_cmd = np.zeros(6)
prev_q, prev_t = q.copy(), time.perf_counter()
steps = int(a.seconds * HZ)
try:
    for k in range(steps):
        t0 = time.perf_counter()
        q = read_joints(robot)
        dt = max(t0 - prev_t, 1e-3)
        q_rad = np.deg2rad(q[:5])
        grip_open = q[5] / 100.0
        vel = np.r_[np.deg2rad(q[:5] - prev_q[:5]), (q[5] - prev_q[5]) / 100.0 * (GRIP_HI - GRIP_LO)] / dt
        prev_q, prev_t = q.copy(), t0
        obs = np.r_[q_rad, grip_open, 0.1 * vel, candy_obs, TOWEL_XY, last_act, min(k, EPISODE_STEPS) / EPISODE_STEPS]
        with torch.no_grad():
            act = policy(torch.tensor(obs, dtype=torch.float32)[None])[0].numpy()
        act = np.clip(act, -1.0, 1.0)
        last_act = act  # the policy sees its own (unsmoothed) action, as in sim
        cmd = a.smooth * prev_cmd + (1.0 - a.smooth) * act if a.smooth > 0 else act
        prev_cmd = cmd
        step = np.clip(np.rad2deg(ARM_SCALE_RAD * cmd[:5]), -a.max_step_deg, a.max_step_deg)
        target[:5] = np.clip(target[:5] + step, REAL_ARM_RANGE_DEG[:, 0], REAL_ARM_RANGE_DEG[:, 1])
        act = cmd
        target[5] = GRIP_CLOSED_LR + (GRIP_OPEN_LR - GRIP_CLOSED_LR) * (act[5] + 1.0) / 2.0
        if a.go:
            robot.send_action({f"{n}.pos": float(v) for n, v in zip(
                ["shoulder_pan", "shoulder_lift", "elbow_flex", "wrist_flex", "wrist_roll", "gripper"], target)})
        elif k % 15 == 0:
            print(f"t={k / HZ:4.1f}s target [deg]: {target.round(1)}")
        if rec:
            vw, wrist, log, _ = rec
            frame = cv2.cvtColor(phone.frame(max_age=1.0), cv2.COLOR_RGB2BGR)
            if wrist is not None:
                ok, w = wrist.read()
                w = cv2.resize(w, (320, 240)) if ok else np.zeros((240, 320, 3), np.uint8)
                frame = np.concatenate([frame, np.vstack([w, np.zeros((120, 320, 3), np.uint8)])], axis=1)
            cv2.putText(frame, f"t={k / HZ:4.1f}s {'GO' if a.go else 'DRY'}", (10, 25), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
            vw.write(frame)
            log.writerow([round(k / HZ, 3)] + q.round(2).tolist() + target.round(2).tolist() + np.round(act, 3).tolist())
        time.sleep(max(0.0, 1.0 / HZ - (time.perf_counter() - t0)))
except KeyboardInterrupt:
    print("\nstopped by Ctrl-C")
finally:
    robot.disconnect()
    phone.close()
    if rec:
        rec[0].release()
        if rec[1] is not None:
            rec[1].release()
        rec[3].close()
        print("saved video and per-step log next to", os.path.basename(stem))