"""Process video at 4fps and save temporally smoothed kinematics to file."""
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

import cv2
import numpy as np
import mujoco
from tqdm import tqdm

from ExoConfigs.so100_adhesive import SO100AdhesiveConfig 
from exo_utils import (estimate_robot_state, detect_and_set_link_poses, 
                       position_exoskeleton_meshes, render_from_camera_pose)
import matplotlib.pyplot as plt

# Configuration
video_path = "/Users/cameronsmith/Downloads/IMG_9546.MOV"
output_path = "scratch/kinematics_exp_test.npy"
target_fps = 4

# Setup robot config
SO100AdhesiveConfig.exo_alpha = 0.2
SO100AdhesiveConfig.aruco_alpha = 0.2
SO100AdhesiveConfig.exo_link_alpha = 0.2
robot_config = SO100AdhesiveConfig()

print("=" * 60)
print("Processing video for temporally smoothed kinematics")
print("=" * 60)

# Setup video capture
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv2.CAP_PROP_FPS)
frame_skip = int(fps / target_fps)
print(f"Video: {video_path}")
print(f"  {total_frames} frames @ {fps:.2f} fps")
print(f"  Processing at {target_fps} fps (every {frame_skip} frames)")

# Load MuJoCo model
model = mujoco.MjModel.from_xml_string(robot_config.xml)
data = mujoco.MjData(model)

# Process frames at target fps
smoothed_estimates = []
frames_to_process = list(range(0, total_frames, frame_skip))

for frame_idx in tqdm(frames_to_process, desc="Processing frames"):
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
    ret, frame = cap.read()
    if not ret:
        break
    
    rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_time = frame_idx / fps
    
    link_poses, camera_pose_world, cam_K, _, _, _ = detect_and_set_link_poses( rgb, model, data, robot_config)
    configuration = estimate_robot_state(model, data, robot_config, link_poses, ik_iterations=35)

    # Update robot state (warm start for next frame)
    data.qpos[:] = configuration.q
    data.ctrl[:] = configuration.q[:len(data.ctrl)]
    mujoco.mj_forward(model, data)
    position_exoskeleton_meshes(robot_config, model, data, link_poses)

    if frame_idx % 100 == 0 and 1:
        rendered = render_from_camera_pose(model, data, camera_pose_world, cam_K, *rgb.shape[:2])
        overlay = (rgb * 0.5 + rendered * 0.5).astype(np.uint8)
        plt.imshow(overlay)
        plt.show()
    
    smoothed_estimates.append({
        'frame_idx': frame_idx,
        'time': frame_time,
        'qpos': configuration.q.copy(),
        'camera_pose': camera_pose_world.copy(),
        'cam_K': cam_K.copy()
    })

cap.release()

print(f"\nSuccessfully processed {len(smoothed_estimates)}/{len(frames_to_process)} frames")

# Save to file
os.makedirs(os.path.dirname(output_path), exist_ok=True)
np.save(output_path, smoothed_estimates, allow_pickle=True)
print(f"\nSaved kinematics to: {output_path}")
print(f"File size: {os.path.getsize(output_path) / 1024 / 1024:.2f} MB")