"""
PAT Tokenizer - Quick Usage Examples

This file demonstrates how to use PAT Tokenizer with UAT and VQVAE backends.
"""

import numpy as np
from pat.models.pat_tokenizer import (
    create_uat_tokenizer,
    create_vqvae_tokenizer,
    ArmMode,
)


def example_uat_tokenizer():
    """
    Example: Using UAT (ActionCodec) Tokenizer

    UAT is Transformer-based and supports:
    - Multi-embodiment (different robots)
    - Variable duration actions
    - VQ (single quantizer) or RVQ (multiple quantizers)
    """
    print("=" * 60)
    print("UAT (ActionCodec) Tokenizer Example")
    print("=" * 60)

    # Step 1: Create tokenizer
    # - n_tokens: Number of discrete tokens to output
    # - n_quantizers: 1 for VQ, >=2 for RVQ
    # - horizon: Time horizon of input actions
    tokenizer = create_uat_tokenizer(
        n_tokens=16,
        n_quantizers=1,  # VQ mode
        horizon=50,
        arm_mode=ArmMode.BILATERAL,  # or ArmMode.RIGHT_ONLY for single-arm
    )
    tokenizer.eval()

    # Step 2: Prepare input actions
    # Shape: (batch_size, horizon, action_dim)
    # action_dim = 20 (left_arm: 10 + right_arm: 10)
    batch_size = 1
    horizon = 50
    action_dim = 20
    actions = np.random.randn(batch_size, horizon, action_dim).astype(np.float32)
    print(f"\nInput actions shape: {actions.shape}")

    # Step 3: Encode actions to tokens
    # embodiment_ids: Specifies which robot configuration to use (0 = default)
    tokens = tokenizer.encode(actions, embodiment_ids=0)
    print(f"Encoded tokens: {len(tokens)} sequences x {len(tokens[0])} tokens")
    print(
        f"Token range: [{min(min(t) for t in tokens)}, {max(max(t) for t in tokens)}]"
    )

    # Step 4: Decode tokens back to actions
    # durations: Duration of each sequence in seconds (horizon / freq)
    # For 20Hz and 50 steps: duration = 50/20 = 2.5 seconds
    actions_recon, mask = tokenizer.decode(
        tokens, embodiment_ids=0, durations=[2.5] * batch_size
    )
    print(f"Decoded actions shape: {actions_recon.shape}")
    print(f"Mask shape: {mask.shape}")

    # Step 5: Check reconstruction quality
    mse = np.mean((actions - actions_recon) ** 2)
    print(f"Reconstruction MSE: {mse:.6f}")

    print()


def example_vqvae_tokenizer():
    """
    Example: Using VQVAE Tokenizer

    VQVAE is CNN-based and features:
    - Efficient encoding/decoding
    - Fixed compression ratio (8x downsampling)
    - RVQ (Residual Vector Quantization)
    """
    print("=" * 60)
    print("VQVAE Tokenizer Example")
    print("=" * 60)

    # Step 1: Create tokenizer
    # - n_codebooks: Number of RVQ codebooks
    # - horizon: Must be divisible by 8 (due to 3x 2x downsampling)
    tokenizer = create_vqvae_tokenizer(
        n_codebooks=4,
        horizon=48,  # Must be divisible by 8
        arm_mode=ArmMode.BILATERAL,
    )
    tokenizer.eval()

    # Step 2: Prepare input actions
    batch_size = 1
    horizon = 48  # VQVAE requires horizon % 8 == 0
    action_dim = 20
    actions = np.random.randn(batch_size, horizon, action_dim).astype(np.float32)
    print(f"\nInput actions shape: {actions.shape}")

    # Step 3: Encode actions to tokens
    # Output shape: (batch, horizon/8 * n_codebooks)
    tokens = tokenizer.encode(actions)
    print(f"Encoded tokens shape: {tokens.shape}")
    print(f"Token range: [{tokens.min()}, {tokens.max()}]")

    # Step 4: Decode tokens back to actions
    actions_recon, mask = tokenizer.decode(tokens)
    print(f"Decoded actions shape: {actions_recon.shape}")
    print(f"Mask shape: {mask.shape}")

    # Step 5: Check reconstruction quality
    mse = np.mean((actions - actions_recon) ** 2)
    print(f"Reconstruction MSE: {mse:.6f}")

    print()


def example_single_arm_mode():
    """
    Example: Using Single-Arm Mode

    In single-arm mode, one arm is zeroed out automatically.
    """
    print("=" * 60)
    print("Single-Arm Mode Example (Right Arm Only)")
    print("=" * 60)

    # Create tokenizer with RIGHT_ONLY mode
    tokenizer = create_uat_tokenizer(
        n_tokens=16,
        n_quantizers=1,
        horizon=50,
        arm_mode=ArmMode.RIGHT_ONLY,  # Left arm will be zeroed
    )
    tokenizer.eval()

    # Prepare input (full 20-dim actions)
    actions = np.random.randn(1, 50, 20).astype(np.float32)

    # Encode and decode
    tokens = tokenizer.encode(actions, embodiment_ids=0)
    actions_recon, _ = tokenizer.decode(tokens, embodiment_ids=0, durations=[2.5])

    # Verify left arm is zeroed
    left_arm_zero = np.allclose(actions_recon[..., 0:10], 0, atol=1e-5)
    right_arm_active = not np.allclose(actions_recon[..., 10:20], 0)

    print(f"Left arm zeroed: {left_arm_zero}")
    print(f"Right arm active: {right_arm_active}")
    print()


def main():
    """Run all examples"""
    print("\n" + "=" * 60)
    print("PAT Tokenizer Usage Examples")
    print("=" * 60 + "\n")

    # Example 1: UAT Tokenizer
    example_uat_tokenizer()

    # Example 2: VQVAE Tokenizer
    example_vqvae_tokenizer()

    # Example 3: Single-Arm Mode
    example_single_arm_mode()

    print("=" * 60)
    print("All examples completed!")
    print("=" * 60)


if __name__ == "__main__":
    main()
