# Copyright 2026 Toyota Research Institute.  All rights reserved.
"""
Visualization script with minimal dependencies to load a single LBM episode in unified format and visualize
the arm's end-effector poses + action trajectories projected onto RGB images,
using camera intrinsics and extrinsics and generate a video.

Usage:
First convert a spartan episode to unified format:
    uv run anydata/converters/lbm.py /data/cv_downloaded/spartan_tiny

Then run:
    uv run anydata/visualize/vis_lbm_unified.py \
        --episode_path data/cv_unified/spartan_tiny/BimanualPutRedBellPepperInBin/riverway/sim/teleop/2025-01-06T08-58-31-05-00/episode_131 \
        --output_dir ./validation_videos \
        --camera scene_right_0 \
        --max_frames 200
"""

import os
import json
import argparse

import cv2
import numpy as np
import torch
from tqdm import tqdm

from anydata.geometry.camera_utils import points_to_camera_frame, project_to_image
from anydata.utils.viz import draw_points_on_image
from anydata.utils.write import write_video


def _extract_tarballs_if_needed(sequence_path):
    """Extract rgb.tar.gz and lowdim.tar.gz into sequence_path if not already extracted.

    Skips extraction when the expected output directory already exists.
    Validates all member paths to reject absolute paths and '..' components
    to prevent path traversal attacks.
    """
    import tarfile

    # tarname -> expected extracted directory name
    targets = {
        'rgb.tar.gz':    'rgb',
        'lowdim.tar.gz': 'lowdim',
    }

    # Flat layout: camera dirs sit directly under sequence_path
    # Check for any camera dir as a proxy for "already extracted"
    cameras = ['scene_left_0', 'scene_right_0', 'wrist_left_plus', 'wrist_right_minus']
    flat_extracted = any(os.path.isdir(os.path.join(sequence_path, c)) for c in cameras)

    for tarname, extracted_dir in targets.items():
        tarpath = os.path.join(sequence_path, tarname)
        if not os.path.isfile(tarpath):
            continue

        # Skip if already extracted (nested layout)
        if os.path.isdir(os.path.join(sequence_path, extracted_dir)):
            continue

        # Skip if already extracted (flat layout)
        if flat_extracted:
            continue

        # Safe extraction: reject absolute paths and '..' components
        abs_dest = os.path.realpath(sequence_path)
        with tarfile.open(tarpath) as tf:
            for member in tf.getmembers():
                member_path = os.path.realpath(os.path.join(abs_dest, member.name))
                if not member_path.startswith(abs_dest + os.sep):
                    raise RuntimeError(
                        f"Refusing to extract '{member.name}': path traversal detected"
                    )
            print(f"Extracting {tarname}...")
            tf.extractall(sequence_path)


def visualize_unified_sequence(sequence_path, output_dir, camera='scene_right_0',
                                max_frames=None):
    """
    Load a single unified-format sequence and visualize end-effector poses on RGB images.

    Args:
        sequence_path: Path to the unified sequence directory (contains rgb/, lowdim/, metadata.json)
        output_dir: Directory to save the visualization video
        camera: Camera name to visualize (default: scene_right_0)
        max_frames: Maximum number of frames to process (None = all frames)
    """
    os.makedirs(output_dir, exist_ok=True)

    print(f"Loading sequence from: {sequence_path}")

    # Load metadata
    metadata_file = os.path.join(sequence_path, 'metadata.json')
    if not os.path.isfile(metadata_file):
        raise FileNotFoundError(f"Metadata file not found: {metadata_file}")

    with open(metadata_file) as f:
        metadata = json.load(f)

    available_cameras = metadata.get('cameras', [])
    if camera not in available_cameras:
        raise ValueError(f"Camera '{camera}' not found. Available: {available_cameras}")

    framerate = metadata.get('framerate', 10)
    print(f"Using camera: {camera}, framerate: {framerate} FPS")

    # Auto-extract tarballs if needed (rgb.tar.gz, lowdim.tar.gz)
    _extract_tarballs_if_needed(sequence_path)

    # Discover frames — support both flat layout (episode/<cam>/<frame>.jpg)
    # and nested layout (episode/rgb/<cam>/<frame>.jpg + episode/lowdim/<cam>/<frame>.npz)
    flat_dir = os.path.join(sequence_path, camera)
    if os.path.isdir(flat_dir):
        rgb_dir    = flat_dir
        lowdim_dir = flat_dir
    else:
        rgb_dir    = os.path.join(sequence_path, 'rgb',    camera)
        lowdim_dir = os.path.join(sequence_path, 'lowdim', camera)

    if not os.path.isdir(rgb_dir):
        raise FileNotFoundError(f"Camera directory not found: {flat_dir} (also tried rgb/ and lowdim/)")

    frame_ids = sorted([
        os.path.splitext(f)[0]
        for f in os.listdir(rgb_dir) if f.endswith('.jpg')
    ])

    N = len(frame_ids)
    if max_frames is not None:
        N = min(N, max_frames)
    frame_ids = frame_ids[:N]

    print(f"Found {N} frames")

    # First pass: load all frames
    # The unified format stores pre-chunked trajectories in each npz:
    #   action['raw'][key] has shape (chunk_size, D) where index 0 = current frame
    frames_rgb       = []
    ee_left_actual   = []   # (N, 3)   — actual EE at each frame
    ee_right_actual  = []
    left_xyz_traj    = []   # (N, chunk_size, 3) — desired future trajectory
    right_xyz_traj   = []

    intrinsics_cam = None
    extrinsics_cam = None

    EE_LEFT_ACT_KEY  = 'robot__actual__poses__left::panda__xyz'
    EE_RIGHT_ACT_KEY = 'robot__actual__poses__right::panda__xyz'
    EE_LEFT_DES_KEY  = 'robot__desired__poses__left::panda__xyz'
    EE_RIGHT_DES_KEY = 'robot__desired__poses__right::panda__xyz'

    for frame_id in tqdm(frame_ids, desc="Loading frames"):
        # RGB
        img_path = os.path.join(rgb_dir, f'{frame_id}.jpg')
        img = cv2.imread(img_path)
        if img is None:
            raise FileNotFoundError(f"Failed to read image at path: {img_path}")
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        frames_rgb.append(img)

        # Lowdim
        lowdim = np.load(
            os.path.join(lowdim_dir, f'{frame_id}.npz'), allow_pickle=True
        )

        if intrinsics_cam is None:
            intrinsics_cam = lowdim['intrinsics']   # (3, 3)
            extrinsics_cam = lowdim['extrinsics']   # (4, 4) — use frame 0 for static cameras

        action_dict = lowdim['action'].item()
        # Support both legacy schema (with nested "raw") and newer unified schema (flat dict).
        raw = action_dict.get('raw', action_dict)

        # Index 0 of the chunk = current frame's actual EE position
        ee_left_actual.append(raw[EE_LEFT_ACT_KEY][0])    # (3,)
        ee_right_actual.append(raw[EE_RIGHT_ACT_KEY][0])  # (3,)

        # Full chunk = desired future trajectory (chunk_size, 3)
        left_xyz_traj.append(raw[EE_LEFT_DES_KEY])
        right_xyz_traj.append(raw[EE_RIGHT_DES_KEY])

    # Stack to arrays — pad trajectories to uniform chunk size (last frames may be shorter)
    pose_xyz_left  = np.stack(ee_left_actual)   # (N, 3)
    pose_xyz_right = np.stack(ee_right_actual)  # (N, 3)

    chunk_size = max(t.shape[0] for t in left_xyz_traj)
    def pad_traj(trajs, chunk_size):
        padded = []
        for t in trajs:
            if t.shape[0] < chunk_size:
                pad = np.tile(t[-1:], (chunk_size - t.shape[0], 1))
                t = np.concatenate([t, pad], axis=0)
            padded.append(t)
        return np.stack(padded)

    left_xyz_act  = pad_traj(left_xyz_traj,  chunk_size)   # (N, chunk_size, 3)
    right_xyz_act = pad_traj(right_xyz_traj, chunk_size)   # (N, chunk_size, 3)
    print(f"Left EE pose shape:  {pose_xyz_left.shape}")
    print(f"Right EE pose shape: {pose_xyz_right.shape}")
    print(f"Trajectory chunk size: {chunk_size}")

    # Project all trajectory points to camera frame (pre-compute)
    left_xyz_act_cam  = points_to_camera_frame(
        left_xyz_act.reshape(-1, 3), extrinsics_cam
    ).reshape(N, chunk_size, 3)
    right_xyz_act_cam = points_to_camera_frame(
        right_xyz_act.reshape(-1, 3), extrinsics_cam
    ).reshape(N, chunk_size, 3)
    combined_xyz_act_cam = np.concatenate(
        [left_xyz_act_cam, right_xyz_act_cam], axis=2
    )  # (N, chunk_size, 6)

    # Project EE poses to camera frame
    pose_xyz_left_cam  = points_to_camera_frame(pose_xyz_left,  extrinsics_cam)  # (N, 3)
    pose_xyz_right_cam = points_to_camera_frame(pose_xyz_right, extrinsics_cam)  # (N, 3)
    ee_pose = np.hstack([pose_xyz_left_cam, pose_xyz_right_cam])                  # (N, 6)

    # Extend intrinsics to (3, 4) projection matrix
    intrinsics_proj = intrinsics_cam
    if intrinsics_proj.shape == (3, 3):
        intrinsics_proj = np.hstack([intrinsics_proj, np.zeros((3, 1))])

    # Per-frame drawing loop
    video_frames = []

    H, W = frames_rgb[0].shape[:2]
    text_y_ee_left  = max(H - 40, 20)
    text_y_ee_right = max(H - 20, 40)

    for t in tqdm(range(N), desc="Rendering frames"):
        img = frames_rgb[t].copy()
        if img.max() <= 1.0:
            img = (img * 255).astype(np.uint8)
        else:
            img = img.astype(np.uint8)

        # Draw current EE pose as red dots
        ee_pose_t_proj = project_to_image(
            ee_pose[t, None, :].reshape(-1, 3), intrinsics_proj
        )
        img = draw_points_on_image(img, ee_pose_t_proj, palette="Reds", dot_size=5)

        left_ee_px  = ee_pose_t_proj[0]
        right_ee_px = ee_pose_t_proj[1]
        cv2.putText(img, 'Left_EE',  (int(left_ee_px[0])  + 5, int(left_ee_px[1])  - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
        cv2.putText(img, 'Right_EE', (int(right_ee_px[0]) + 5, int(right_ee_px[1]) - 5),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)

        # Draw action trajectory as green dots
        ee_action_proj = project_to_image(
            combined_xyz_act_cam[t].reshape(-1, 3), intrinsics_proj
        )
        img = draw_points_on_image(img, ee_action_proj, palette="Greens", dot_size=3)

        # Overlay text
        cv2.putText(img, f'Frame {t}/{N}', (10, 30),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2)
        cv2.putText(img, 'Red: EE Pose', (10, 60),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
        cv2.putText(img, 'Green: EE Action Trajectory', (10, 85),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)

        left_pos  = pose_xyz_left[t]
        right_pos = pose_xyz_right[t]
        cv2.putText(img,
                    f'Left EE:  ({left_pos[0]:.2f}, {left_pos[1]:.2f}, {left_pos[2]:.2f})',
                    (10, text_y_ee_left),  cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        cv2.putText(img,
                    f'Right EE: ({right_pos[0]:.2f}, {right_pos[1]:.2f}, {right_pos[2]:.2f})',
                    (10, text_y_ee_right), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

        video_frames.append(img)

    # Save video
    if video_frames:
        video_tensor = torch.stack([torch.from_numpy(f) for f in video_frames])
        seq_name  = os.path.basename(sequence_path.rstrip('/'))
        video_path = os.path.join(output_dir, f'{seq_name}_{camera}_ee_visualization.mp4')
        print(f"Saving video: {video_path}")
        write_video(video_path, video_tensor, fps=framerate)
        print("Video saved successfully!")
    else:
        print("No frames to save!")


def main():
    parser = argparse.ArgumentParser(
        description='Visualize end-effector poses from a unified-format LBM sequence')
    parser.add_argument('--episode_path', type=str, required=True,
                        help='Path to a single unified episode directory (contains metadata.json)')
    parser.add_argument('--output_dir', type=str, default='./validation_videos',
                        help='Output directory for visualization video')
    parser.add_argument('--camera', type=str, default='scene_right_0',
                        help='Camera to visualize (default: scene_right_0)')
    parser.add_argument('--max_frames', type=int, default=None,
                        help='Maximum frames to process (default: all frames)')

    args = parser.parse_args()

    # Validate camera against names available in this episode's metadata
    import json as _json
    metadata_file = os.path.join(args.episode_path, 'metadata.json')
    with open(metadata_file) as f:
        available_cameras = _json.load(f).get('cameras', [])
    if args.camera not in available_cameras:
        parser.error(f"--camera '{args.camera}' not found. Available: {available_cameras}")

    print(f"Episode path:     {args.episode_path}")
    print(f"Output directory: {args.output_dir}")
    print(f"Camera:           {args.camera}")

    visualize_unified_sequence(
        args.episode_path,
        args.output_dir,
        camera=args.camera,
        max_frames=args.max_frames,
    )

    print("Visualization complete!")


if __name__ == '__main__':
    main()