#!/usr/bin/env python3
"""
Test OAT tokenizer with chunked encode/decode on raw parquet files.

This script:
1. Loads a trained OAT tokenizer from checkpoint
2. Reads raw parquet files directly
3. Encodes action chunks to tokens with different granularities (eval_keep_k)
4. Computes MSE for each granularity directly (no trajectory averaging)
5. Visualizes original vs reconstructed trajectories with chunk boundaries

Reference: eval_tokenizer_reconstruction.py

Usage:
    python tests/test_oat.py --model-path output/pat_OAT/checkpoints/checkpoint-000660
"""

import argparse
import sys
from pathlib import Path

# Add project root to path
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from typing import Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

from pat.models.pat_tokenizer import PATTokenizer, TokenizerType, ActionSpaceConfig
from pat.utils import eef_euler_to_6d, abs_6d_2_abs_euler, rotation_6d_to_matrix_torch


def compute_geodesic_distance(
    rot_6d_pred: np.ndarray, rot_6d_gt: np.ndarray
) -> np.ndarray:
    """
    Compute geodesic distance between two 6D rotation representations.
    
    This is the proper metric for rotation comparison, measuring the angle
    between two rotations in SO(3) space.
    
    Args:
        rot_6d_pred: Predicted 6D rotation (..., 6)
        rot_6d_gt: Ground truth 6D rotation (..., 6)
    
    Returns:
        Geodesic distance in radians (...,)
    """
    # Convert to torch tensors
    rot_6d_pred_torch = torch.from_numpy(rot_6d_pred).float()
    rot_6d_gt_torch = torch.from_numpy(rot_6d_gt).float()
    
    # Convert 6D to rotation matrices
    R_pred = rotation_6d_to_matrix_torch(rot_6d_pred_torch)
    R_gt = rotation_6d_to_matrix_torch(rot_6d_gt_torch)
    
    # Compute relative rotation: R_diff = R_pred^T @ R_gt
    R_diff = torch.matmul(R_pred.transpose(-2, -1), R_gt)
    
    # Extract trace and compute cos(theta)
    trace = R_diff[..., 0, 0] + R_diff[..., 1, 1] + R_diff[..., 2, 2]
    cos_theta = ((trace - 1) / 2).clamp(-1 + 1e-6, 1 - 1e-6)
    
    # Geodesic distance = arccos(cos(theta))
    geo_dist = torch.acos(cos_theta)
    
    return geo_dist.numpy()


def compute_rotation_metrics(
    actions_orig: np.ndarray,
    actions_recon: np.ndarray,
) -> Dict[str, float]:
    """
    Compute rotation-specific metrics between original and reconstructed actions.
    
    This properly evaluates rotation quality by:
    1. Extracting 6D rotation components
    2. Converting to rotation matrices
    3. Computing geodesic distance in SO(3)
    
    Args:
        actions_orig: Original actions (..., 20)
        actions_recon: Reconstructed actions (..., 20)
    
    Returns:
        Dictionary with rotation metrics:
        - geo_mean: Mean geodesic distance (radians)
        - geo_deg: Mean geodesic distance (degrees)
        - geo_max: Max geodesic distance (radians)
        - angle_mse: MSE of rotation angle (degrees^2)
    """
    # Extract 6D rotation components (left: 3:9, right: 13:19)
    left_6d_orig = actions_orig[..., 3:9]
    left_6d_recon = actions_recon[..., 3:9]
    right_6d_orig = actions_orig[..., 13:19]
    right_6d_recon = actions_recon[..., 13:19]
    
    # Compute geodesic distance for left and right arms
    geo_left = compute_geodesic_distance(left_6d_recon, left_6d_orig)
    geo_right = compute_geodesic_distance(right_6d_recon, right_6d_orig)
    
    # Combine both arms
    geo_all = np.concatenate([geo_left.flatten(), geo_right.flatten()])
    
    # Compute statistics in radians
    geo_mean = float(np.mean(geo_all))
    geo_max = float(np.max(geo_all))
    geo_std = float(np.std(geo_all))
    
    # Convert to degrees for better interpretability
    geo_mean_deg = np.degrees(geo_mean)
    geo_max_deg = np.degrees(geo_max)
    
    # Angle MSE (degrees^2) - measures squared error of rotation angle
    angle_mse_deg = float(np.mean(geo_all ** 2))
    angle_rmse_deg = np.sqrt(angle_mse_deg)
    
    # Per-arm statistics
    left_geo_mean = float(np.mean(geo_left))
    right_geo_mean = float(np.mean(geo_right))
    left_geo_mean_deg = np.degrees(left_geo_mean)
    right_geo_mean_deg = np.degrees(right_geo_mean)
    
    return {
        "geo_mean_rad": geo_mean,
        "geo_max_rad": geo_max,
        "geo_std_rad": geo_std,
        "geo_mean_deg": geo_mean_deg,
        "geo_max_deg": geo_max_deg,
        "angle_mse_deg2": angle_mse_deg,
        "angle_rmse_deg": angle_rmse_deg,
        "left_geo_mean_deg": left_geo_mean_deg,
        "right_geo_mean_deg": right_geo_mean_deg,
    }


def find_parquet_files(data_dir: str) -> List[str]:
    """Find all parquet files in directory."""
    import os

    parquet_files = []
    for root, dirs, files in os.walk(data_dir):
        for f in files:
            if f.endswith(".parquet"):
                parquet_files.append(os.path.join(root, f))
    return sorted(parquet_files)


def read_parquet_raw(parquet_path: str) -> pd.DataFrame:
    """Read parquet file directly."""
    df = pd.read_parquet(parquet_path)
    print(f"Loaded {parquet_path}: {len(df)} frames")
    return df


def extract_states(
    df: pd.DataFrame,
    state_key: str,
) -> np.ndarray:
    """Extract states from dataframe."""
    if state_key not in df.columns:
        raise ValueError(f"State key '{state_key}' not found in dataframe")

    values = df[state_key].values
    state_list = []
    for item in values:
        if hasattr(item, "tolist"):
            item = item.tolist()
        state_list.append(list(item))
    states_euler = np.array(state_list, dtype=np.float32)

    # Convert to 6D format
    states_6d = []
    for state_14d in states_euler:
        state_20d = eef_euler_to_6d(state_14d[0:7].tolist(), state_14d[7:14].tolist())
        states_6d.append(state_20d)
    states = np.stack(states_6d, axis=0)
    print(f"States shape: {states.shape} (6D format)")

    return states, states_euler


def convert_6d_to_euler(states_6d: np.ndarray) -> np.ndarray:
    """
    Convert state sequences from 6D rotation format back to euler (rxryrz).

    Reference: eval_tokenizer_reconstruction.py::_convert_6d_to_euler

    Args:
        states_6d: (T, 20) array with 6D rotation

    Returns:
        (T, 14) array with euler angles
    """
    states_euler = []
    for state_20d in states_6d:
        state_14d = abs_6d_2_abs_euler(state_20d)
        states_euler.append(state_14d)
    return np.stack(states_euler, axis=0)  # (T, 14)


def convert_states_to_6d(euler_states: np.ndarray) -> np.ndarray:
    """
    Convert state sequences from euler (rxryrz) to 6D rotation format.

    Reference: eval_tokenizer_reconstruction.py::_convert_states_to_6d

    Args:
        euler_states: (T, 14) array with euler angles

    Returns:
        (T, 20) array with 6D rotation representation
    """
    states_6d = []
    for state_14d in euler_states:
        state_20d = eef_euler_to_6d(state_14d[0:7].tolist(), state_14d[7:14].tolist())
        states_6d.append(state_20d)
    return np.stack(states_6d, axis=0)


def create_chunks(
    data: np.ndarray,
    horizon: int,
    overlap_rate: float = 0.0,
) -> Tuple[List[np.ndarray], List[int]]:
    """Create overlapping chunks from sequence."""
    overlap_length = int(round(horizon * overlap_rate))
    stride = horizon - overlap_length
    stride = max(1, stride)

    chunks = []
    chunk_starts = []
    for start in range(0, len(data) - horizon + 1, stride):
        chunk = data[start : start + horizon]
        chunks.append(chunk)
        chunk_starts.append(start)

    return chunks, chunk_starts


def compute_mse_per_dof(
    actions_orig: np.ndarray,
    actions_recon: np.ndarray,
) -> np.ndarray:
    """Compute MSE per degree of freedom."""
    mse_per_dof = np.mean((actions_orig - actions_recon) ** 2, axis=0)
    return mse_per_dof


def get_dof_names() -> List[str]:
    """Get DOF names for 14-dimensional Euler action space."""
    left_names = ["L_x", "L_y", "L_z", "L_rx", "L_ry", "L_rz", "L_grip"]
    right_names = ["R_x", "R_y", "R_z", "R_rx", "R_ry", "R_rz", "R_grip"]
    return left_names + right_names


def visualize_granularity_comparison_all_dofs(
    trajectory_orig_euler: np.ndarray,
    traj_by_granularity_euler: Dict[int, np.ndarray],
    chunk_starts: List[int],
    horizon: int,
    mse_by_granularity: Dict[int, float],
    output_dir: Path,
    prefix: str = "granularity_comparison",
    selected_granularities: Optional[List[int]] = None,
    selected_dofs: Optional[List[int]] = None,
):
    """Visualize granularity comparison for selected DOFs.

    Args:
        trajectory_orig_euler: Original trajectory in euler format (T, 14)
        traj_by_granularity_euler: Dict mapping k -> reconstructed trajectory
        chunk_starts: List of chunk start indices
        horizon: Horizon length
        mse_by_granularity: Dict mapping k -> MSE
        output_dir: Output directory
        prefix: Filename prefix
        selected_granularities: If None, use all; otherwise only show these k values
        selected_dofs: List of DOF indices to plot (0-13). If None, plot all DOFs (0-13)
    """
    output_dir.mkdir(parents=True, exist_ok=True)

    T = trajectory_orig_euler.shape[0]

    # Default: plot all DOFs (0-13)
    if selected_dofs is None:
        selected_dofs = list(range(14))  # All 14 DOFs

    all_granularities = sorted(traj_by_granularity_euler.keys())
    if selected_granularities is not None:
        granularities = [k for k in selected_granularities if k in traj_by_granularity_euler]
    else:
        granularities = all_granularities

    if len(granularities) == 0:
        print(f"  Warning: No valid granularities to visualize")
        return

    # ========== Bright, modern color palette ==========
    # Original trajectory: bright blue for strong visibility
    ORIGINAL_COLOR = "#0066CC"  # Bright blue
    
    # Bright, distinct colors for each granularity (lighter, more vibrant)
    granularity_colors = {
        2: "#FF6B6B",    # coral red - bright and warm
        4: "#4ECDC4",    # turquoise - fresh and light
        6: "#45B7D1",    # sky blue - clear and bright
        8: "#96CEB4",    # sage green - soft and light
        10: "#FFEAA7",   # soft yellow - light and cheerful
        12: "#DDA0DD",   # plum purple - light and distinct
        14: "#98D8C8",   # mint green - fresh
        16: "#F7DC6F",   # golden yellow - bright
    }
    
    # Linestyles for different granularities (more distinct)
    granularity_linestyles = {
        2: "-",           # solid
        4: "--",          # dashed
        6: "-.",         # dash-dot
        8: ":",          # dotted
        10: "-.",        # dash-dot
        12: "--",        # dashed
        14: ":",         # dotted
        16: "-",         # solid
    }
    
    # Default colors for any k not in the dict (bright palette)
    default_colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#98D8C8", "#F7DC6F"]
    default_linestyles = ["-", "--", "-.", ":", "-.", "--", ":", "-"]

    dof_names_full = [
        "L_x", "L_y", "L_z", "L_rx", "L_ry", "L_rz", "L_grip",
        "R_x", "R_y", "R_z", "R_rx", "R_ry", "R_rz", "R_grip"
    ]
    dof_display_names = {
        0: r"$L_{x}$", 1: r"$L_{y}$", 2: r"$L_{z}$", 3: r"$L_{rx}$", 4: r"$L_{ry}$", 5: r"$L_{rz}$", 6: r"$L_{grip}$",
        7: r"$R_{x}$", 8: r"$R_{y}$", 9: r"$R_{z}$", 10: r"$R_{rx}$", 11: r"$R_{ry}$", 12: r"$R_{rz}$", 13: r"$R_{grip}$"
    }

    # Calculate layout: 7 rows x 2 columns (left/right arm)
    num_dofs_per_arm = 7

    # Increased figure size for all DOFs - wider layout with 2 columns
    fig, axes = plt.subplots(num_dofs_per_arm, 2, figsize=(20, 28), facecolor='white')

    # Plot Left Arm DOFs (all 7: indices 0-6) - left column
    for row in range(num_dofs_per_arm):
        dof_idx = row  # 0-6 for left arm
        ax = axes[row, 0]
        ax.set_facecolor('#FAFAFA')  # Very light gray background for subplots
        
        # Plot original trajectory with strong visibility
        ax.plot(
            range(T),
            trajectory_orig_euler[:, dof_idx],
            ORIGINAL_COLOR,
            linewidth=1.5,
            alpha=1.0,
            label="Original",
            zorder=10,
        )

        # Plot reconstructed trajectories with enhanced visibility
        for i, k in enumerate(granularities):
            traj_k = traj_by_granularity_euler[k]
            mse_k = mse_by_granularity[k]
            color = granularity_colors.get(k, default_colors[i % len(default_colors)])
            linestyle = granularity_linestyles.get(k, default_linestyles[i % len(default_linestyles)])
            
            # Highlight selected granularities with thicker lines
            is_selected = selected_granularities is not None and k in selected_granularities
            linewidth = 1.8 if is_selected else 1.5
            
            ax.plot(
                range(T),
                traj_k[:, dof_idx],
                color=color,
                linewidth=linewidth,
                alpha=0.95,
                linestyle=linestyle,
                label=f"k={k}",
                zorder=8,
            )

        # Mark chunk boundaries with enhanced visibility
        for start in chunk_starts:
            if start < T:
                ax.plot(
                    start,
                    trajectory_orig_euler[start, dof_idx],
                    "go",
                    markersize=7,
                    alpha=0.8,
                    zorder=15,
                    markeredgewidth=1.5,
                    markeredgecolor='white',
                )

        ax.set_ylabel("Angle (rad)", fontsize=13, fontweight="bold")
        ax.set_xlabel("Timestep", fontsize=12)
        ax.set_title(f"Left Arm - {dof_display_names.get(dof_idx, f'DOF {dof_idx}')}",
                    fontsize=14, fontweight="bold")
        ax.grid(True, alpha=0.35, linestyle="--", linewidth=0.6, color='#CCCCCC')
        ax.tick_params(labelsize=11)
        ax.set_axisbelow(False)  # Grid lines behind data

    # Plot Right Arm DOFs (all 7: indices 7-13) - right column
    for row in range(num_dofs_per_arm):
        dof_idx = row + 7  # 7-13 for right arm
        ax = axes[row, 1]
        ax.set_facecolor('#FAFAFA')  # Very light gray background for subplots
        
        # Plot original trajectory with strong visibility
        ax.plot(
            range(T),
            trajectory_orig_euler[:, dof_idx],
            ORIGINAL_COLOR,
            linewidth=1.5,
            alpha=1.0,
            label="Original",
            zorder=10,
        )

        # Plot reconstructed trajectories with enhanced visibility
        for i, k in enumerate(granularities):
            traj_k = traj_by_granularity_euler[k]
            mse_k = mse_by_granularity[k]
            color = granularity_colors.get(k, default_colors[i % len(default_colors)])
            linestyle = granularity_linestyles.get(k, default_linestyles[i % len(default_linestyles)])
            
            # Highlight selected granularities with thicker lines
            is_selected = selected_granularities is not None and k in selected_granularities
            linewidth = 1.8 if is_selected else 1.5
            
            ax.plot(
                range(T),
                traj_k[:, dof_idx],
                color=color,
                linewidth=linewidth,
                alpha=0.95,
                linestyle=linestyle,
                label=f"k={k}",
                zorder=8,
            )

        # Mark chunk boundaries with enhanced visibility
        for start in chunk_starts:
            if start < T:
                ax.plot(
                    start,
                    trajectory_orig_euler[start, dof_idx],
                    "go",
                    markersize=7,
                    alpha=0.8,
                    zorder=15,
                    markeredgewidth=1.5,
                    markeredgecolor='white',
                )

        ax.set_ylabel("Angle (rad)", fontsize=13, fontweight="bold")
        ax.set_xlabel("Timestep", fontsize=12)
        ax.set_title(f"Right Arm - {dof_display_names.get(dof_idx, f'DOF {dof_idx}')}",
                    fontsize=14, fontweight="bold")
        ax.grid(True, alpha=0.35, linestyle="--", linewidth=0.6, color='#CCCCCC')
        ax.tick_params(labelsize=11)
        ax.set_axisbelow(False)  # Grid lines behind data

    # Build MSE summary
    mse_summary = ", ".join(
        [f"k={k}: {mse_by_granularity[k]:.6f}" for k in granularities]
    )

    # Title based on whether we're showing all or selected granularities
    if selected_granularities is None:
        title_suffix = "All Granularities"
    else:
        title_suffix = f"Selected Granularities {selected_granularities}"

    fig.suptitle(
        f"Rotation DOF Granularity Comparison: {title_suffix}\n"
        f"Non-rotation MSE: {mse_summary}",
        fontsize=16,
        fontweight="bold",
        y=0.995,
    )

    # Create legend at the bottom - separate from title, outside the plot area
    legend_elements = [
        plt.Line2D([0], [0], color=ORIGINAL_COLOR, linewidth=1.5, label="Original")
    ]
    for k in granularities:
        color = granularity_colors.get(k, default_colors[granularities.index(k) % len(default_colors)])
        linestyle = granularity_linestyles.get(k, default_linestyles[granularities.index(k) % len(default_linestyles)])
        is_selected = selected_granularities is not None and k in selected_granularities
        linewidth = 2.0 if is_selected else 1.5
        legend_elements.append(
            plt.Line2D(
                [0], [0],
                color=color,
                linewidth=linewidth,
                linestyle=linestyle,
                label=f"k={k} (MSE={mse_by_granularity[k]:.4f})",
            )
        )
    legend_elements.append(
        plt.Line2D(
            [0], [0],
            marker="o",
            color="w",
            markerfacecolor="g",
            markersize=10,
            label="Chunk start",
            linestyle="None",
            markeredgewidth=1.5,
            markeredgecolor='white',
        )
    )

    # Place legend BELOW the subplots, outside the plot area to avoid overlap
    fig.legend(
        handles=legend_elements,
        loc="lower center",
        ncol=min(len(granularities) + 2, 6),
        fontsize=12,
        bbox_to_anchor=(0.5, -0.02),
        framealpha=0.98,
        frameon=True,
        facecolor='white',
        edgecolor='#CCCCCC',
    )

    # Adjust layout to make room for legend at bottom (expanded bottom margin)
    plt.tight_layout(rect=[0.02, 0.03, 1, 0.96])
    suffix = "_selected" if selected_granularities else "_all"
    plt.savefig(output_dir / f"{prefix}_rotation_dofs{suffix}.png", dpi=200, bbox_inches="tight")
    plt.close(fig)

    # MSE vs Granularity for all DOFs - 2 rows (left/right arm)
    fig, axes = plt.subplots(2, 1, figsize=(14, 12), facecolor='white')

    # Bright color palette for all 7 DOFs per arm
    dof_colors = ["#FF6B6B", "#4ECDC4", "#45B7D1", "#96CEB4", "#FFEAA7", "#DDA0DD", "#98D8C8"]
    dof_markers = ["o", "s", "^", "D", "v", "<", "P"]

    for arm_idx, (arm_name, dof_offset) in enumerate([("Left Arm", 0), ("Right Arm", 7)]):
        ax = axes[arm_idx]
        ax.set_facecolor('#FAFAFA')  # Light gray background

        for dof_local_idx in range(7):  # All 7 DOFs per arm
            dof_global_idx = dof_offset + dof_local_idx
            mse_per_k = []
            for k in granularities:
                mse_dof = np.mean(
                    (
                        trajectory_orig_euler[:, dof_global_idx]
                        - traj_by_granularity_euler[k][:, dof_global_idx]
                    )
                    ** 2
                )
                mse_per_k.append(mse_dof)

            dof_color = dof_colors[dof_local_idx]
            marker = dof_markers[dof_local_idx % len(dof_markers)]
            ax.plot(
                granularities,
                mse_per_k,
                color=dof_color,
                marker=marker,
                linewidth=1.5,
                markersize=7,
                markeredgewidth=1.0,
                markeredgecolor='white',
                label=f"DOF {dof_local_idx}",
            )

        ax.set_xlabel("Number of Tokens (eval_keep_k)", fontsize=14, fontweight="bold")
        ax.set_ylabel("MSE (euler space)", fontsize=14, fontweight="bold")
        ax.set_title(f"{arm_name} - All DOFs", fontsize=15, fontweight="bold")
        ax.grid(True, alpha=0.35, linestyle="--", linewidth=0.7, color='#CCCCCC')
        ax.set_xticks(granularities)
        ax.tick_params(labelsize=12)
        # Move legend outside to the right
        ax.legend(loc="upper left", fontsize=10, bbox_to_anchor=(1.02, 1), framealpha=0.95)

    fig.suptitle(
        "Per-DOF MSE vs Granularity (All DOFs)",
        fontsize=16,
        fontweight="bold",
        y=0.995,
    )
    plt.tight_layout(rect=[0.02, 0.03, 0.98, 0.96])
    suffix = "_selected" if selected_granularities else "_all"
    plt.savefig(output_dir / f"{prefix}_mse_rotation{suffix}.png", dpi=200, bbox_inches="tight")
    plt.close(fig)

    print(f"  Saved granularity comparison visualizations to {output_dir}")


def visualize_full_trajectory_with_chunks(
    trajectory_orig_euler: np.ndarray,
    trajectory_recon_euler: np.ndarray,
    chunk_starts: List[int],
    horizon: int,
    mse_by_granularity: Dict[int, float],
    output_dir: Path,
    prefix: str = "trajectory",
):
    """Visualize full trajectory reconstruction with chunk boundaries."""
    output_dir.mkdir(parents=True, exist_ok=True)

    T = trajectory_orig_euler.shape[0]
    num_arms = 2
    dofs_per_arm = 7

    arm_names = ["Left Arm", "Right Arm"]
    dof_names = ["x", "y", "z", r"$r_x$", r"$r_y$", r"$r_z$", "Gripper"]

    mse_per_dof = compute_mse_per_dof(trajectory_orig_euler, trajectory_recon_euler)
    overall_mse = np.mean(mse_per_dof)

    fig, axes = plt.subplots(dofs_per_arm, num_arms, figsize=(20, 24), facecolor='white')

    for arm_idx in range(num_arms):
        for dof_idx in range(dofs_per_arm):
            ax = axes[dof_idx, arm_idx]
            dim_idx = arm_idx * 7 + dof_idx
            ax.set_facecolor('#FAFAFA')  # Light gray background

            ax.plot(
                range(T),
                trajectory_orig_euler[:, dim_idx],
                "#0066CC",  # Bright blue for original
                linewidth=1.5,
                alpha=1.0,
                label="Original",
            )
            ax.plot(
                range(T),
                trajectory_recon_euler[:, dim_idx],
                "#FF6B6B",  # Bright coral red for reconstructed
                linewidth=1.2,
                alpha=0.9,
                label="Reconstructed",
            )

            for start in chunk_starts:
                ax.axvline(
                    x=start, color="#999999", linestyle=":", linewidth=1.0, alpha=0.5
                )
                if start < T:
                    ax.plot(
                        start,
                        trajectory_orig_euler[start, dim_idx],
                        "go",
                        markersize=7,
                        alpha=0.8,
                        zorder=10,
                        markeredgewidth=1.5,
                        markeredgecolor='white',
                    )

            ax.set_ylabel(f"{dof_names[dof_idx]}", fontsize=13, fontweight="bold")
            ax.set_xlabel("Timestep", fontsize=12)
            ax.set_title(
                f"{arm_names[arm_idx]} - {dof_names[dof_idx]}\n"
                f"MSE: {mse_per_dof[dim_idx]:.6f}",
                fontsize=13,
                fontweight="bold",
            )
            ax.legend(loc="upper right", fontsize=10, framealpha=0.95)
            ax.grid(True, alpha=0.35, linestyle="--", linewidth=0.6, color='#CCCCCC')

    mse_summary = ", ".join([f"k={k}: {v:.6f}" for k, v in mse_by_granularity.items()])
    fig.suptitle(
        f"Full Trajectory Reconstruction: Original vs Reconstructed (k={list(mse_by_granularity.keys())[-1]})\n"
        f"Overall MSE: {overall_mse:.6f} | {mse_summary}",
        fontsize=16,
        fontweight="bold",
        y=0.995,
    )

    plt.tight_layout(rect=[0.02, 0.04, 1, 0.98])
    plt.savefig(
        output_dir / f"{prefix}_reconstruction.png", dpi=200, bbox_inches="tight"
    )
    plt.close(fig)

    print(f"  Saved trajectory reconstruction visualization to {output_dir}")


def main():
    parser = argparse.ArgumentParser(
        description="Test OAT tokenizer with chunked encode/decode"
    )
    parser.add_argument(
        "--model-path",
        type=str,
        default="output/pat_OAT/checkpoints/checkpoint-000660",
        help="Path to trained OAT checkpoint folder",
    )
    parser.add_argument(
        "--parquet-path",
        type=str,
        default="data/orbec02-towel8_28_0224",
        help="Path to parquet file or directory",
    )
    parser.add_argument(
        "--state-key",
        type=str,
        default="observation.state_eepose",
        help="Column name for states",
    )
    parser.add_argument(
        "--overlap-rate",
        type=float,
        default=0.0,  # Changed to 0.0 to match eval_tokenizer_reconstruction.py
        help="Overlap rate for chunks (0.0 to 1.0)",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default="output/oat_chunk_test",
        help="Output directory for visualizations",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda" if torch.cuda.is_available() else "cpu",
        help="Device for inference",
    )
    parser.add_argument(
        "--granularities",
        type=str,
        default="2,4,6,8,10,12,14,16",
        help="Comma-separated granularities to test (eval_keep_k values)",
    )
    args = parser.parse_args()

    granularities = [int(k.strip()) for k in args.granularities.split(",")]

    np.random.seed(args.seed)
    torch.manual_seed(args.seed)

    print("=" * 80)
    print("OAT Tokenizer: Chunked Encode/Decode Test")
    print("=" * 80)

    model_path = Path(args.model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"Model path does not exist: {model_path}")

    print(f"\nLoading tokenizer from {model_path}...")
    tokenizer = PATTokenizer.from_pretrained(str(model_path), device=args.device)
    tokenizer.to(args.device).eval()

    print(f"Tokenizer type: {tokenizer.config.tokenizer_type.value}")
    print(f"Horizon: {tokenizer.config.action_space.horizon}")
    print(f"Num registers: {tokenizer.config.oat_num_registers}")
    print(f"Device: {args.device}")
    print(f"Action normalizer fitted: {tokenizer.normalizer.action.is_fitted()}")

    parquet_path = Path(args.parquet_path)
    if parquet_path.is_dir():
        parquet_files = find_parquet_files(str(parquet_path))
        if len(parquet_files) == 0:
            raise ValueError(f"No parquet files found in {parquet_path}")
        print(f"\nFound {len(parquet_files)} parquet files in {parquet_path}")
    else:
        parquet_files = [str(parquet_path)]

    df = read_parquet_raw(parquet_files[0])
    states_6d, states_euler = extract_states(df, args.state_key)

    print(f"states_6d shape: {states_6d.shape}")
    print(f"states_euler shape: {states_euler.shape}")

    horizon = tokenizer.config.action_space.horizon
    chunks_6d, chunk_starts = create_chunks(states_6d, horizon, args.overlap_rate)
    chunks_euler, _ = create_chunks(states_euler, horizon, args.overlap_rate)

    print(
        f"\nCreated {len(chunks_6d)} chunks (horizon={horizon}, overlap={args.overlap_rate})"
    )
    print(f"Total trajectory length: {len(states_6d)} timesteps")

    # ========== Encode/Decode with different granularities ==========
    print("\nEncoding and decoding chunks with different granularities...")

    # Store results per granularity: list of (orig_6d, recon_6d, orig_euler, recon_euler) pairs
    # MSE is computed in 6D space, but euler is kept for visualization
    results_by_granularity = {
        k: {"orig_6d": [], "recon_6d": [], "orig_euler": [], "recon_euler": []}
        for k in granularities
    }

    with torch.no_grad():
        for i, (chunk_6d, chunk_euler) in enumerate(zip(chunks_6d, chunks_euler)):
            chunk_tensor = torch.from_numpy(chunk_6d).unsqueeze(0).to(args.device)

            # Encode - OAT now returns both latents and tokens
            latents, tokens = tokenizer.encode(
                chunk_tensor.cpu().numpy(), return_tokens=True
            )
            print(
                f"Chunk {i + 1}/{len(chunks_6d)}: latents shape: {latents.shape}, tokens shape: {tokens.shape}, tokens: {tokens}"
            )

            # Decode with different granularities from discrete tokens
            for k in granularities:
                if k <= horizon:
                    # Use is_tokens=True to indicate input is discrete tokens
                    actions_recon_k, _ = tokenizer.decode(
                        tokens, eval_keep_k=[k], is_tokens=True
                    )
                    # actions_recon_k, _ = tokenizer.decode(latents=latents, eval_keep_k=[k])

                    # Truncate to actual chunk length (chunk_euler may be shorter than horizon at sequence end)
                    actual_len = chunk_euler.shape[0]
                    actions_recon_k = actions_recon_k[0, :actual_len, :]  # (T, 20)

                    # Convert to euler for visualization
                    recon_euler = convert_6d_to_euler(actions_recon_k)

                    # Store both 6D (for MSE) and euler (for visualization)
                    results_by_granularity[k]["orig_6d"].append(chunk_6d)
                    results_by_granularity[k]["recon_6d"].append(actions_recon_k)
                    results_by_granularity[k]["orig_euler"].append(chunk_euler)
                    results_by_granularity[k]["recon_euler"].append(recon_euler)

            if (i + 1) % 10 == 0:
                print(f"  Processed {i + 1}/{len(chunks_6d)} chunks")

    # ========== Compute MSE for each granularity (non-rotation DOFs only) ==========
    # Note: Rotation quality is evaluated via geodesic distance metrics (see below)
    # Non-rotation DOFs: x, y, z, gripper (indices 0,1,2,6,7,8,9,13)
    print("\n" + "=" * 60)
    print("Reconstruction MSE by Granularity (non-rotation DOFs only):")
    print("=" * 60)

    mse_by_granularity = {}
    recon_trajectories_6d = {}
    recon_trajectories_euler = {}

    # Non-rotation DOF indices: position (x,y,z) + gripper for both arms
    non_rotation_dofs = [0, 1, 2, 6, 7, 8, 9, 13]

    for k in granularities:
        if k <= horizon and len(results_by_granularity[k]["orig_euler"]) > 0:
            # Concatenate all chunks to form full trajectory
            orig_traj_euler = np.concatenate(
                results_by_granularity[k]["orig_euler"], axis=0
            )
            recon_traj_euler = np.concatenate(
                results_by_granularity[k]["recon_euler"], axis=0
            )
            orig_traj_6d = np.concatenate(results_by_granularity[k]["orig_6d"], axis=0)
            recon_traj_6d = np.concatenate(
                results_by_granularity[k]["recon_6d"], axis=0
            )

            # Compute MSE for non-rotation DOFs only (rotation evaluated via geodesic distance)
            mse_non_rot = np.mean((orig_traj_euler[:, non_rotation_dofs] - recon_traj_euler[:, non_rotation_dofs]) ** 2)
            mse_by_granularity[k] = mse_non_rot
            recon_trajectories_6d[k] = recon_traj_6d
            recon_trajectories_euler[k] = recon_traj_euler

            print(f"  k={k:2d}: Non-rotation MSE = {mse_non_rot:.8f}")

    # Use original full trajectory for comparison
    orig_trajectory_6d = np.concatenate(chunks_6d, axis=0)
    orig_trajectory_euler = np.concatenate(chunks_euler, axis=0)

    # Per-DOF MSE for non-rotation DOFs only (rotation evaluated via geodesic distance)
    print(
        "\nPer-DOF MSE (non-rotation DOFs, Last granularity k={}):".format(max(granularities))
    )
    finest_k = max(granularities)
    if finest_k in recon_trajectories_euler:
        mse_per_dof_euler = compute_mse_per_dof(
            orig_trajectory_euler, recon_trajectories_euler[finest_k]
        )
        dof_names_euler = get_dof_names()
        # Only print non-rotation DOFs: x, y, z, gripper
        non_rotation_dofs = [0, 1, 2, 6, 7, 8, 9, 13]
        for dof_idx in non_rotation_dofs:
            arm = "Left" if dof_idx < 7 else "Right"
            print(f"  {arm} {dof_names_euler[dof_idx]}: {mse_per_dof_euler[dof_idx]:.8f}")

    # ========== Compute Rotation Metrics (Geodesic Distance) ==========
    print("\n" + "=" * 60)
    print("Rotation Metrics by Granularity (Geodesic Distance in SO(3)):")
    print("=" * 60)

    rotation_metrics_by_granularity = {}

    for k in granularities:
        if k <= horizon and len(results_by_granularity[k]["orig_6d"]) > 0:
            orig_traj_6d = np.concatenate(results_by_granularity[k]["orig_6d"], axis=0)
            recon_traj_6d = np.concatenate(
                results_by_granularity[k]["recon_6d"], axis=0
            )

            # Compute rotation metrics
            rot_metrics = compute_rotation_metrics(orig_traj_6d, recon_traj_6d)
            rotation_metrics_by_granularity[k] = rot_metrics

            print(f"\n  k={k:2d}:")
            print(f"    Mean Geodesic Distance = {rot_metrics['geo_mean_rad']:.6f} rad ({rot_metrics['geo_mean_deg']:.4f}°)")
            print(f"    Max  Geodesic Distance = {rot_metrics['geo_max_rad']:.6f} rad ({rot_metrics['geo_max_deg']:.4f}°)")
            print(f"    Std  Geodesic Distance = {rot_metrics['geo_std_rad']:.6f} rad")
            print(f"    Angle RMSE             = {rot_metrics['angle_rmse_deg']:.4f}°")
            print(f"    Left Arm Mean          = {rot_metrics['left_geo_mean_deg']:.4f}°")
            print(f"    Right Arm Mean         = {rot_metrics['right_geo_mean_deg']:.4f}°")

    # Rotation metrics for finest granularity
    print(
        "\n" + "=" * 60
    )
    print(f"Rotation Metrics Summary (Finest Granularity k={finest_k}):")
    print("=" * 60)
    
    if finest_k in rotation_metrics_by_granularity:
        metrics = rotation_metrics_by_granularity[finest_k]
        print(f"  Mean Rotation Error: {metrics['geo_mean_deg']:.4f}° ({metrics['geo_mean_rad']:.6f} rad)")
        print(f"  Max  Rotation Error: {metrics['geo_max_deg']:.4f}° ({metrics['geo_max_rad']:.6f} rad)")
        print(f"  RMSE Rotation Error: {metrics['angle_rmse_deg']:.4f}°")
        print(f"  Left Arm  Mean Error: {metrics['left_geo_mean_deg']:.4f}°")
        print(f"  Right Arm Mean Error: {metrics['right_geo_mean_deg']:.4f}°")

    # ========== Visualize ==========
    output_dir = Path(args.output_dir)
    print(f"\nGenerating visualizations to {output_dir}...")

    # Full trajectory reconstruction (using finest granularity)
    if finest_k in recon_trajectories_euler:
        visualize_full_trajectory_with_chunks(
            orig_trajectory_euler,
            recon_trajectories_euler[finest_k],
            chunk_starts,
            horizon,
            mse_by_granularity,
            output_dir,
            prefix="oat_trajectory",
        )

    # Granularity comparison - Set 1: All granularities
    if len(recon_trajectories_euler) > 0:
        visualize_granularity_comparison_all_dofs(
            orig_trajectory_euler,
            recon_trajectories_euler,
            chunk_starts,
            horizon,
            mse_by_granularity,
            output_dir,
            prefix="oat_granularity_comparison",
            selected_granularities=None,  # All granularities
        )
    
    # Granularity comparison - Set 2: Selected granularities (2, 6, 12, 16)
    selected_ks = [2, 6, 12, 16]
    selected_ks = [k for k in selected_ks if k in recon_trajectories_euler]
    if len(selected_ks) > 0:
        visualize_granularity_comparison_all_dofs(
            orig_trajectory_euler,
            recon_trajectories_euler,
            chunk_starts,
            horizon,
            mse_by_granularity,
            output_dir,
            prefix="oat_granularity_comparison",
            selected_granularities=selected_ks,  # Selected granularities only
        )
        print(f"  Also saved selected granularities visualization: {selected_ks}")

    # Save data
    np.save(output_dir / "trajectory_orig_euler.npy", orig_trajectory_euler)
    np.save(output_dir / "trajectory_orig_6d.npy", orig_trajectory_6d)
    for k in recon_trajectories_euler:
        np.save(
            output_dir / f"trajectory_recon_euler_k{k}.npy", recon_trajectories_euler[k]
        )
    for k in recon_trajectories_6d:
        np.save(output_dir / f"trajectory_recon_6d_k{k}.npy", recon_trajectories_6d[k])
    np.save(output_dir / "chunk_starts.npy", np.array(chunk_starts))

    import json

    # Helper to convert rotation metrics to serializable format
    def serialize_rotation_metrics(metrics_dict):
        result = {}
        for k, metrics in metrics_dict.items():
            result[str(k)] = {
                key: float(val) for key, val in metrics.items()
            }
        return result

    summary = {
        "total_frames": len(states_euler),
        "horizon": horizon,
        "num_chunks": len(chunks_6d),
        "overlap_rate": args.overlap_rate,
        "granularities_tested": granularities,
        "mse_by_granularity_6d": {
            str(k): float(v) for k, v in mse_by_granularity.items()
        },
        "rotation_metrics": serialize_rotation_metrics(rotation_metrics_by_granularity),
        "chunk_starts": chunk_starts,
    }

    with open(output_dir / "summary.json", "w") as f:
        json.dump(summary, f, indent=2, default=str)

    print(f"\nSaved all data and summary to {output_dir}")

    print("\n" + "=" * 80)
    print("Done!")
    print("=" * 80)


if __name__ == "__main__":
    main()
