"""
PAT (Policy Action Tokenizer) - Unified Action Tokenizer Interface

Supports multiple tokenizer backends:
- UAT (ActionCodec): Transformer-based, supports multi-embodiment
- OAT (Object-Action Transformer): Register-based transformer encoder
- VQVAE: CNN-based, efficient encoding
- MultiVQVAE: Multi-action space separated encoding

Action Space Definition (20 dimensions):
- Left arm (10 dims): xyz(3) + 6d_rotation(6) + gripper(1) [0:10]
- Right arm (10 dims): xyz(3) + 6d_rotation(6) + gripper(1) [10:20]
- Combined grippers (2 dims): left_gripper(1) + right_gripper(1) [14:16]
- Reserved (4 dims): [16:20]

Single-arm mode: Left arm is zeroed, only right arm is used.
"""

import logging
from abc import ABC, abstractmethod
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Literal, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn

from pat.normalization import ActionNormalizer, TensorLinearNormalizer

logger = logging.getLogger(__name__)


class TokenizerType(str, Enum):
    """Supported Tokenizer Types"""

    UAT = "uat"
    OAT = "oat"
    VQVAE = "vqvae"
    MULTIVQVAE = "multivqvae"


class ArmMode(str, Enum):
    """Robot Arm Mode"""

    BILATERAL = "bilateral"  # Dual-arm
    RIGHT_ONLY = "right_only"  # Right arm only
    LEFT_ONLY = "left_only"  # Left arm only


@dataclass
class ActionSpaceConfig:
    """
    Action Space Configuration

    20-dimensional action space layout:
    ┌────────────────────────────────────────────────────────────┐
    │ Index  │  Dimension         │  Description                │
    ├────────────────────────────────────────────────────────────┤
    │ 0-2    │ left_xyz          │ Left arm end-effector pos    │
    │ 3-8    │ left_6d_rot       │ Left arm 6D rotation         │
    │ 9      │ left_gripper      │ Left gripper (0=close/1=open)│
    ├────────────────────────────────────────────────────────────┤
    │ 10-12  │ right_xyz         │ Right arm end-effector pos   │
    │ 13-18  │ right_6d_rot      │ Right arm 6D rotation        │
    │ 19     │ right_gripper     │ Right gripper (0=close/1=open)│
    ├────────────────────────────────────────────────────────────┤
    │ 14-15  │ grippers          │ Combined grippers [L, R]     │
    │ 16-19  │ reserved          │ Reserved/padding             │
    └────────────────────────────────────────────────────────────┘

    Note: Gripper appears in two locations:
    - Individually within arm (index 9, 19)
    - Combined in gripper channel (index 14, 15)
    """

    # Action dimension configuration
    action_dim: int = 20
    horizon: int = 24  # Default time horizon

    # Arm mode
    arm_mode: ArmMode = ArmMode.BILATERAL

    # Action space slice definitions
    left_arm_slice: slice = field(default=slice(0, 10), init=False)
    right_arm_slice: slice = field(default=slice(10, 20), init=False)

    # Detailed action dimension slices
    left_xyz_slice: slice = field(default=slice(0, 3), init=False)
    left_rot_slice: slice = field(default=slice(3, 9), init=False)
    left_gripper_idx: int = 9

    right_xyz_slice: slice = field(default=slice(10, 13), init=False)
    right_rot_slice: slice = field(default=slice(13, 19), init=False)
    right_gripper_idx: int = 19

    gripper_slice: slice = field(default=slice(14, 16), init=False)
    reserved_slice: slice = field(default=slice(16, 20), init=False)

    # MultiVQVAE specific configuration
    pos_dim: int = 6  # xyz(3) * 2 arms
    rot_dim: int = 12  # 6d_rot(6) * 2 arms
    grip_dim: int = 2  # gripper(1) * 2 arms

    def __post_init__(self):
        if isinstance(self.arm_mode, str):
            self.arm_mode = ArmMode(self.arm_mode)
        if self.arm_mode == ArmMode.RIGHT_ONLY:
            # Single-arm mode: zero out left arm
            logger.info("Single-arm mode (right only): Left arm actions will be zeroed")
        elif self.arm_mode == ArmMode.LEFT_ONLY:
            logger.info("Single-arm mode (left only): Right arm actions will be zeroed")

    def extract_left_arm(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Extract left arm actions (first 10 dimensions)"""
        return actions[..., self.left_arm_slice]

    def extract_right_arm(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Extract right arm actions (last 10 dimensions)"""
        return actions[..., self.right_arm_slice]

    def extract_pos(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Extract position actions (dual-arm xyz)"""
        if isinstance(actions, torch.Tensor):
            return torch.cat(
                [actions[..., self.left_xyz_slice], actions[..., self.right_xyz_slice]],
                dim=-1,
            )
        else:
            return np.concatenate(
                [actions[..., self.left_xyz_slice], actions[..., self.right_xyz_slice]],
                axis=-1,
            )

    def extract_rot(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Extract rotation actions (dual-arm 6D rotation)"""
        if isinstance(actions, torch.Tensor):
            return torch.cat(
                [actions[..., self.left_rot_slice], actions[..., self.right_rot_slice]],
                dim=-1,
            )
        else:
            return np.concatenate(
                [actions[..., self.left_rot_slice], actions[..., self.right_rot_slice]],
                axis=-1,
            )

    def extract_gripper(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Extract gripper actions (dual-arm gripper)"""
        if isinstance(actions, torch.Tensor):
            return torch.stack(
                [
                    actions[..., self.left_gripper_idx],
                    actions[..., self.right_gripper_idx],
                ],
                dim=-1,
            )
        else:
            return np.stack(
                [
                    actions[..., self.left_gripper_idx],
                    actions[..., self.right_gripper_idx],
                ],
                axis=-1,
            )

    def assemble_actions(
        self,
        pos: Union[np.ndarray, torch.Tensor],
        rot: Union[np.ndarray, torch.Tensor],
        gripper: Union[np.ndarray, torch.Tensor],
    ) -> Union[np.ndarray, torch.Tensor]:
        """
        Assemble full actions from pos/rot/gripper components.

        Args:
            pos: Position (..., 6) - left_xyz(3) + right_xyz(3)
            rot: Rotation (..., 12) - left_6d(6) + right_6d(6)
            gripper: Gripper (..., 2) - left(1) + right(1)

        Returns:
            Full actions (..., 20)
        """
        is_tensor = isinstance(pos, torch.Tensor)
        batch_shape = pos.shape[:-1]

        if is_tensor:
            device = pos.device
            dtype = pos.dtype
            full_actions = torch.zeros(
                *batch_shape, self.action_dim, device=device, dtype=dtype
            )

            # Left arm
            full_actions[..., self.left_xyz_slice] = pos[..., 0:3]
            full_actions[..., self.left_rot_slice] = rot[..., 0:6]
            full_actions[..., self.left_gripper_idx] = gripper[..., 0]

            # Right arm
            full_actions[..., self.right_xyz_slice] = pos[..., 3:6]
            full_actions[..., self.right_rot_slice] = rot[..., 6:12]
            full_actions[..., self.right_gripper_idx] = gripper[..., 1]

            # Combined gripper channel
            full_actions[..., self.gripper_slice] = gripper
        else:
            full_actions = np.zeros((*batch_shape, self.action_dim), dtype=pos.dtype)

            # Left arm
            full_actions[..., self.left_xyz_slice] = pos[..., 0:3]
            full_actions[..., self.left_rot_slice] = rot[..., 0:6]
            full_actions[..., self.left_gripper_idx] = gripper[..., 0]

            # Right arm
            full_actions[..., self.right_xyz_slice] = pos[..., 3:6]
            full_actions[..., self.right_rot_slice] = rot[..., 6:12]
            full_actions[..., self.right_gripper_idx] = gripper[..., 1]

            # Combined gripper channel
            full_actions[..., self.gripper_slice] = gripper

        return full_actions

    def zero_left_arm(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Zero out left arm actions (for single-arm mode)"""
        result = actions.copy() if isinstance(actions, np.ndarray) else actions.clone()
        result[..., self.left_arm_slice] = 0
        return result

    def zero_right_arm(
        self, actions: Union[np.ndarray, torch.Tensor]
    ) -> Union[np.ndarray, torch.Tensor]:
        """Zero out right arm actions (for single-arm mode)"""
        result = actions.copy() if isinstance(actions, np.ndarray) else actions.clone()
        result[..., self.right_arm_slice] = 0
        return result


@dataclass
class PATConfig:
    """
    PAT Unified Configuration Class

    Unified configuration interface supporting three tokenizer backends
    """

    # Tokenizer type
    tokenizer_type: TokenizerType = TokenizerType.UAT

    # Action space configuration
    action_space: ActionSpaceConfig = field(default_factory=ActionSpaceConfig)

    # ========== UAT Specific Configuration ==========
    # Embodiment configuration (multi-robot scenarios)
    embodiment_config: Optional[Dict[str, Any]] = None

    # Token configuration
    n_tokens: int = 16
    n_quantizers: int = 1  # VQ=1, RVQ>=2

    # Latent space
    z_dim: int = 512

    # VQ configuration
    vq_type: Literal["vq", "rvq"] = "vq"
    vq_codebook_size: int = 2048
    vq_commitment_weight: float = 0.25
    vq_decay: float = 0.99
    vq_kmeans_init: bool = True
    vq_threshold_ema_dead_code: int = 2
    vq_quantizer_dropout: float = 0.25

    # Transformer configuration
    encoder_dim: int = 256
    encoder_n_layers: int = 6
    encoder_n_heads: int = 8
    encoder_add_self_attn: bool = False
    encoder_add_causal_mask: bool = False
    encoder_pos_encoding_type: str = "fourier"

    decoder_dim: int = 256
    decoder_n_layers: int = 6
    decoder_n_heads: int = 8
    decoder_add_self_attn: bool = False
    decoder_add_causal_mask: bool = False
    decoder_pos_encoding_type: str = "fourier"
    decoder_cls_size: int = 1

    # ========== OAT Specific Configuration ==========
    # Encoder configuration
    oat_encoder_emb_dim: int = 256
    oat_encoder_head_dim: int = 64
    oat_encoder_depth: int = 6
    oat_encoder_dropout: float = 0.1
    oat_num_registers: int = 16

    # Decoder configuration
    oat_decoder_emb_dim: int = 256
    oat_decoder_head_dim: int = 64
    oat_decoder_depth: int = 6
    oat_decoder_dropout: float = 0.1
    oat_decoder_use_causal: bool = False
    oat_token_dropout_mode: str = "drop"

    # FSQ configuration (if using FSQ quantization)
    oat_fsq_enabled: bool = True
    oat_fsq_levels: List[int] = field(default_factory=lambda: [8, 5, 5, 5])
    oat_fsq_dim: int = 64

    # ========== VQVAE / MultiVQVAE Specific Configuration ==========
    # CNN configuration
    cnn_hidden_size: int = 64
    cnn_output_size: int = 512
    cnn_dropout: float = 0.0

    # VQ configuration
    vqvae_embedding_dim: int = 256
    vqvae_num_embeddings: int = 2048
    vqvae_n_codebooks: int = 4
    vqvae_commitment_cost: float = 0.25
    vqvae_codebook_cost: float = 0.0
    vqvae_codebook_restart_interval: int = 64

    # MultiVQVAE specific
    multivqvae_n_codebooks: Dict[str, int] = field(
        default_factory=lambda: {
            "pos": 6,  # position 6 tokens
            "rot": 3,  # rotation 3 tokens
            "grip": 1,  # gripper 1 token
        }
    )

    # ========== General Configuration ==========
    # Model path (for loading pretrained weights)
    model_path: Optional[str] = None

    # Device
    device: str = "cuda"

    # Precision
    dtype: str = "float32"

    def __post_init__(self):
        # Set default embodiment configuration
        if self.embodiment_config is None:
            self.embodiment_config = {
                "dual_piper_inhouse_v0": {
                    "action_dim": 20,
                    "freq": 30,
                    "duration": 1.0,
                    "description": "20-dim dual-arm action, 30Hz, 1 sec, chunks=30, 3 cameras, center distance 75cm. 1 orbbec + 2 realsense D435i",
                },
            }

        # Validate configuration
        self._validate_config()

    def _validate_config(self):
        """Validate configuration parameters"""
        if self.tokenizer_type == TokenizerType.UAT:
            if self.n_tokens % self.n_quantizers != 0:
                raise ValueError(
                    f"UAT: n_tokens ({self.n_tokens}) must be divisible by n_quantizers ({self.n_quantizers})"
                )
            if self.vq_type == "vq" and self.n_quantizers != 1:
                raise ValueError(
                    f"UAT VQ mode requires n_quantizers=1, current: {self.n_quantizers}"
                )
            if self.vq_type == "rvq" and self.n_quantizers < 2:
                raise ValueError(
                    f"UAT RVQ mode requires n_quantizers>=2, current: {self.n_quantizers}"
                )

        elif self.tokenizer_type == TokenizerType.OAT:
            # OAT: latent horizon must match number of registers
            if self.oat_num_registers < 1:
                raise ValueError(
                    f"OAT: num_registers ({self.oat_num_registers}) must be >= 1"
                )

        elif self.tokenizer_type in [TokenizerType.VQVAE, TokenizerType.MULTIVQVAE]:
            # VQVAE temporal length must be multiple of 8 (3x 2x downsampling)
            if self.action_space.horizon % 8 != 0:
                raise ValueError(
                    f"VQVAE: action_horizon ({self.action_space.horizon}) must be a multiple of 8"
                )

    def to_uat_config(self):
        """Convert to UAT configuration"""
        from .uat.configuration_uat import UATConfig

        return UATConfig(
            embodiment_config=self.embodiment_config,
            n_tokens=self.n_tokens,
            n_quantizers=self.n_quantizers,
            z_dim=self.z_dim,
            vq_type=self.vq_type,
            vq_codebook_size=self.vq_codebook_size,
            vq_commitment_weight=self.vq_commitment_weight,
            vq_decay=self.vq_decay,
            vq_kmeans_init=self.vq_kmeans_init,
            vq_threshold_ema_dead_code=self.vq_threshold_ema_dead_code,
            vq_quantizer_dropout=self.vq_quantizer_dropout,
            encoder_dim=self.encoder_dim,
            encoder_n_layers=self.encoder_n_layers,
            encoder_n_heads=self.encoder_n_heads,
            encoder_add_self_attn=self.encoder_add_self_attn,
            encoder_add_causal_mask=self.encoder_add_causal_mask,
            encoder_pos_encoding_type=self.encoder_pos_encoding_type,
            decoder_dim=self.decoder_dim,
            decoder_n_layers=self.decoder_n_layers,
            decoder_n_heads=self.decoder_n_heads,
            decoder_add_self_attn=self.decoder_add_self_attn,
            decoder_add_causal_mask=self.decoder_add_causal_mask,
            decoder_pos_encoding_type=self.decoder_pos_encoding_type,
            decoder_cls_size=self.decoder_cls_size,
        )

    def to_vqvae_config(self):
        """Convert to VQVAE configuration"""
        return {
            "input_dim": self.action_space.action_dim,
            "embedding_dim": self.vqvae_embedding_dim,
            "cnn_config": {
                "hidden_size": self.cnn_hidden_size,
                "output_size": self.cnn_output_size,
                "dropout": self.cnn_dropout,
            },
            "num_embeddings": self.vqvae_num_embeddings,
            "action_horizon": self.action_space.horizon,
            "n_codebooks": self.vqvae_n_codebooks,
            "commitment_cost": self.vqvae_commitment_cost,
            "codebook_cost": self.vqvae_codebook_cost,
            "codebook_restart_interval": self.vqvae_codebook_restart_interval,
        }

    def to_multivqvae_config(self):
        """Convert to MultiVQVAE configuration"""
        return {
            "input_dim": {
                "pos": self.action_space.pos_dim,
                "rot": self.action_space.rot_dim,
                "grip": self.action_space.grip_dim,
            },
            "embedding_dim": self.vqvae_embedding_dim,
            "cnn_config": {
                "hidden_size": self.cnn_hidden_size,
                "output_size": self.cnn_output_size,
                "dropout": self.cnn_dropout,
            },
            "num_embeddings": self.vqvae_num_embeddings,
            "action_horizon": self.action_space.horizon,
            "n_codebooks": self.multivqvae_n_codebooks,
            "commitment_cost": self.vqvae_commitment_cost,
            "codebook_cost": self.vqvae_codebook_cost,
            "codebook_restart_interval": self.vqvae_codebook_restart_interval,
        }


class BaseTokenizer(ABC):
    """Base Tokenizer Class"""

    def __init__(self, config: PATConfig):
        self.config = config
        self.action_space = config.action_space
        self.device = torch.device(config.device)
        self.dtype = getattr(torch, config.dtype)

    @abstractmethod
    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> Union[List[List[int]], np.ndarray]:
        """Encode actions to tokens"""
        pass

    @abstractmethod
    def decode(
        self,
        tokens: Union[List[List[int]], np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """Decode tokens to actions"""
        pass

    @abstractmethod
    def to(self, device: Union[str, torch.device]) -> "BaseTokenizer":
        """Move to device"""
        pass

    @abstractmethod
    def eval(self) -> "BaseTokenizer":
        """Set to evaluation mode"""
        pass

    def _to_tensor(
        self,
        data: Union[np.ndarray, torch.Tensor, List],
        dtype: Optional[torch.dtype] = None,
    ) -> torch.Tensor:
        """Convert to tensor"""
        if isinstance(data, torch.Tensor):
            return data.to(dtype=dtype or self.dtype, device=self.device)
        elif isinstance(data, np.ndarray):
            return torch.from_numpy(data).to(
                dtype=dtype or self.dtype, device=self.device
            )
        elif isinstance(data, list):
            return torch.tensor(data, dtype=dtype or self.dtype, device=self.device)
        else:
            raise TypeError(f"Unsupported data type：{type(data)}")

    def _to_numpy(self, data: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
        """Convert to numpy"""
        if isinstance(data, torch.Tensor):
            return data.detach().cpu().numpy()
        elif isinstance(data, np.ndarray):
            return data
        else:
            raise TypeError(f"Unsupported data type：{type(data)}")


class UATTokenizer(BaseTokenizer):
    """UAT (ActionCodec) Tokenizer"""

    def __init__(self, config: PATConfig):
        super().__init__(config)
        from .uat.modeling_uat import ActionCodec

        uat_config = config.to_uat_config()
        self.model = ActionCodec(uat_config)

        if config.model_path:
            self._load_weights(config.model_path)

        self.model.to(self.device, dtype=self.dtype)

    def _load_weights(self, path: str):
        """Load pretrained weights"""
        logger.info(f"Loading UAT pretrained weights: {path}")
        if path.endswith(".safetensors"):
            from safetensors.torch import load_file

            state_dict = load_file(path)
        else:
            state_dict = torch.load(path, map_location="cpu", weights_only=True)
        self.model.load_state_dict(state_dict)

    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        padding_mask: Optional[np.ndarray] = None,
        **kwargs,
    ) -> List[List[int]]:
        actions_tensor = self._to_tensor(actions)

        # Single-arm mode processing
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions_tensor = self.action_space.zero_left_arm(actions_tensor)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions_tensor = self.action_space.zero_right_arm(actions_tensor)

        if padding_mask is not None:
            padding_mask = self._to_tensor(padding_mask, dtype=torch.bool)

        tokens = self.model.encode(
            x=actions_tensor,
            embodiment_ids=embodiment_ids,
            padding_mask=padding_mask,
        )
        return tokens

    def decode(
        self,
        tokens: Union[List[List[int]], np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        durations: Optional[Union[List[float], np.ndarray, torch.Tensor]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        if isinstance(tokens, list):
            tokens_tensor = torch.tensor(tokens, dtype=torch.long, device=self.device)
        else:
            tokens_tensor = self._to_tensor(tokens, dtype=torch.long)

        actions, mask = self.model.decode(
            tokens=tokens_tensor,
            embodiment_ids=embodiment_ids,
            durations=durations,
        )

        # Single-arm mode processing: zero out unused arm
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions = self.action_space.zero_left_arm(actions)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions = self.action_space.zero_right_arm(actions)

        return self._to_numpy(actions), self._to_numpy(mask)

    def to(self, device: Union[str, torch.device]) -> "UATTokenizer":
        self.model.to(device)
        self.device = torch.device(device)
        return self

    def eval(self) -> "UATTokenizer":
        self.model.eval()
        return self


class VQVAETokenizer(BaseTokenizer):
    """VQVAE Tokenizer"""

    def __init__(self, config: PATConfig):
        super().__init__(config)
        from .vqvae.models.vqvae import VQVAE

        vqvae_config = config.to_vqvae_config()
        self.model = VQVAE(**vqvae_config)

        if config.model_path:
            self._load_weights(config.model_path)

        self.model.to(self.device, dtype=self.dtype)

    def _load_weights(self, path: str):
        """Load pretrained weights"""
        logger.info(f"Loading VQVAE pretrained weights: {path}")
        if path.endswith(".safetensors"):
            from safetensors.torch import load_file

            state_dict = load_file(path)
        else:
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)
            state_dict = checkpoint.get("module", checkpoint)
        self.model.load_state_dict(state_dict)

    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> np.ndarray:
        actions_tensor = self._to_tensor(actions)

        # Single-arm mode processing
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions_tensor = self.action_space.zero_left_arm(actions_tensor)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions_tensor = self.action_space.zero_right_arm(actions_tensor)

        tokens = self.model.encode(actions_tensor)
        return self._to_numpy(tokens)

    def decode(
        self,
        tokens: Union[List[List[int]], np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        if isinstance(tokens, list):
            tokens_tensor = torch.tensor(tokens, dtype=torch.long, device=self.device)
        else:
            tokens_tensor = self._to_tensor(tokens, dtype=torch.long)

        actions = self.model.decode(tokens_tensor)

        # Single-arm mode processing: zero out unused arm
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions = self.action_space.zero_left_arm(actions)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions = self.action_space.zero_right_arm(actions)

        # Create padding mask (all True, since VQVAE has fixed length)
        batch_size = actions.shape[0]
        mask = np.ones((batch_size, actions.shape[1]), dtype=bool)

        return self._to_numpy(actions), mask

    def to(self, device: Union[str, torch.device]) -> "VQVAETokenizer":
        self.model.to(device)
        self.device = torch.device(device)
        return self

    def eval(self) -> "VQVAETokenizer":
        self.model.eval()
        return self


class OATTokenizer(BaseTokenizer):
    """OAT (Ordered-Action Transformer) Tokenizer"""

    def __init__(self, config: PATConfig):
        super().__init__(config)
        from .oat.encoder import RegisterEncoder
        from .oat.decoder import SinglePassDecoder
        from .oat.fsq import FSQ
        from .oat.modeling_oat import OATModel

        # Create encoder
        # Note: encoder's latent_dim must match FSQ's dim (number of FSQ levels)
        encoder_latent_dim = len(config.oat_fsq_levels)
        encoder = RegisterEncoder(
            sample_dim=config.action_space.action_dim,
            sample_horizon=config.action_space.horizon,
            emb_dim=config.oat_encoder_emb_dim,
            head_dim=config.oat_encoder_head_dim,
            depth=config.oat_encoder_depth,
            pdropout=config.oat_encoder_dropout,
            latent_dim=encoder_latent_dim,
            num_registers=config.oat_num_registers,
        )

        # Create decoder
        # Decoder's latent_dim must match FSQ output dim (len(levels)), latent_proj will project to emb_dim
        decoder_latent_dim = len(config.oat_fsq_levels)
        decoder = SinglePassDecoder(
            sample_dim=config.action_space.action_dim,
            sample_horizon=config.action_space.horizon,
            emb_dim=config.oat_decoder_emb_dim,
            head_dim=config.oat_decoder_head_dim,
            depth=config.oat_decoder_depth,
            pdropout=config.oat_decoder_dropout,
            token_dropout_mode=config.oat_token_dropout_mode,
            use_causal_decoder=config.oat_decoder_use_causal,
            latent_dim=decoder_latent_dim,  # Must match FSQ output dim
            latent_horizon=config.oat_num_registers,
        )

        # Create FSQ quantizer for generation
        fsq_quantizer = FSQ(
            levels=config.oat_fsq_levels,
            drop_quant_p=0.0,  # No dropout during inference
        )

        # Compute codebook size from FSQ levels
        codebook_size = 1
        for level in config.oat_fsq_levels:
            codebook_size *= level

        # Wrap in OATModel with generation support
        self.model = OATModel(
            encoder=encoder,
            decoder=decoder,
            fsq_quantizer=fsq_quantizer,
            latent_dim=decoder_latent_dim,  # This is the FSQ output dim (len(levels))
            num_registers=config.oat_num_registers,
            codebook_size=codebook_size,
            fsq_levels=config.oat_fsq_levels,
        )

        if config.model_path:
            self._load_weights(config.model_path)

        self.model.to(self.device, dtype=self.dtype)

    def _load_weights(self, path: str):
        """Load pretrained weights"""
        logger.info(f"Loading OAT pretrained weights: {path}")

        # Load state dict
        if path.endswith(".safetensors"):
            from safetensors.torch import load_file

            state_dict = load_file(path)
        else:
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)
            state_dict = checkpoint.get("module", checkpoint)

        # Load full state dict into model
        self.model.load_state_dict(state_dict, strict=False)

    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        return_tokens: bool = False,
        use_quantizer: bool = True,
        **kwargs,
    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        """
        Encode actions to latent tokens.

        Args:
            actions: Input actions (batch, horizon, action_dim)
            embodiment_ids: Embodiment IDs (not used for OAT)
            return_tokens: If True, also return discrete tokens
            use_quantizer: If True, apply FSQ quantization (matching original OAT)

        Returns:
            - If return_tokens=False: Quantized latent tokens (batch, num_registers, latent_dim)
            - If return_tokens=True: (quantized_latents, discrete_tokens)
              - quantized_latents: (batch, num_registers, latent_dim) - after FSQ quantization
              - discrete_tokens: (batch, num_registers)
        """
        actions_tensor = self._to_tensor(actions)

        # Single-arm mode processing
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions_tensor = self.action_space.zero_left_arm(actions_tensor)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions_tensor = self.action_space.zero_right_arm(actions_tensor)

        # Encode to latent tokens, optionally return discrete tokens
        # Note: By default (use_quantizer=True), returns quantized latents matching original OAT
        if return_tokens:
            quantized_latents, tokens = self.model.encode(
                actions_tensor, return_tokens=True, use_quantizer=use_quantizer
            )
            return self._to_numpy(quantized_latents), self._to_numpy(tokens)
        else:
            quantized_latents = self.model.encode(
                actions_tensor, return_tokens=False, use_quantizer=use_quantizer
            )
            return self._to_numpy(quantized_latents)

    def decode(
        self,
        latents_or_tokens: Union[np.ndarray, torch.Tensor],
        eval_keep_k: Optional[List[int]] = None,
        is_tokens: bool = False,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Decode latent tokens or discrete tokens to actions.

        Args:
            latents_or_tokens:
                - If is_tokens=False: Latent vectors (batch, num_registers, latent_dim)
                - If is_tokens=True: Discrete token indices (batch, num_registers)
            embodiment_ids: Embodiment IDs (not used for OAT)
            eval_keep_k: Number of tokens to use for each sample in the batch.
                Controls the granularity of decoding (nested dropout).
                - None: Use all tokens (fine granularity)
                - List[int]: Use specified number of tokens for each sample
                - Example: [4, 8] means sample 0 uses 4 tokens, sample 1 uses 8 tokens
            is_tokens: If True, treat input as discrete tokens; otherwise as continuous latents

        Returns:
            Tuple of (reconstructed_actions, padding_mask)
        """
        input_tensor = self._to_tensor(latents_or_tokens)

        # Convert discrete tokens to latents if needed
        if is_tokens or input_tensor.dtype in (torch.int32, torch.int64, torch.long):
            # Input is discrete tokens, convert to latents via FSQ
            tokens_tensor = input_tensor.long()
            fsq = self.model.fsq_quantizer
            latents = fsq.indices_to_embedding(tokens_tensor)
            # latents shape: (B, num_registers, len(fsq_levels))
            # Decoder's latent_proj will handle projection to emb_dim
        else:
            # Input is already continuous latents
            latents = input_tensor
            # the latents still not after fsq.

        # Decode from latent tokens
        actions = self.model.decode(latents, eval_keep_k=eval_keep_k)

        # Single-arm mode processing: zero out unused arm
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions = self.action_space.zero_left_arm(actions)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions = self.action_space.zero_right_arm(actions)

        # Create padding mask (all True, since OAT has fixed length)
        batch_size = actions.shape[0]
        mask = np.ones((batch_size, actions.shape[1]), dtype=bool)

        return self._to_numpy(actions), mask

    def tokenize(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> np.ndarray:
        """
        Tokenize actions to discrete tokens (only returns tokens).

        Args:
            actions: Input actions (batch, horizon, action_dim)
            embodiment_ids: Embodiment IDs (not used for OAT)

        Returns:
            Discrete tokens (batch, num_registers)
        """
        _, tokens = self.encode(
            actions, embodiment_ids=embodiment_ids, return_tokens=True
        )
        return tokens

    def detokenize(
        self,
        tokens: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        eval_keep_k: Optional[List[int]] = None,
        **kwargs,
    ) -> np.ndarray:
        """
        Decode discrete tokens to actions.

        Args:
            tokens: Discrete token indices (batch, num_registers)
            embodiment_ids: Embodiment IDs (not used for OAT)
            eval_keep_k: Number of tokens to use for each sample in the batch.

        Returns:
            Reconstructed actions (batch, horizon, action_dim)
        """
        actions, mask = self.decode(
            tokens,
            embodiment_ids=embodiment_ids,
            eval_keep_k=eval_keep_k,
            is_tokens=True,
        )
        return actions

    def to(self, device: Union[str, torch.device]) -> "OATTokenizer":
        self.model.to(device)
        self.device = torch.device(device)
        return self

    def eval(self) -> "OATTokenizer":
        self.model.eval()
        return self


class MultiVQVAETokenizer(BaseTokenizer):
    """MultiVQVAE Tokenizer (Action decomposition encoding)"""

    def __init__(self, config: PATConfig):
        super().__init__(config)
        from .vqvae.models.multivqvae import MultiVQVAE

        mvqvae_config = config.to_multivqvae_config()
        self.model = MultiVQVAE(**mvqvae_config)

        if config.model_path:
            self._load_weights(config.model_path)

        self.model.to(self.device, dtype=self.dtype)

    def _load_weights(self, path: str):
        """Load pretrained weights"""
        logger.info(f"Loading MultiVQVAE pretrained weights: {path}")
        if path.endswith(".safetensors"):
            from safetensors.torch import load_file

            state_dict = load_file(path)
        else:
            checkpoint = torch.load(path, map_location="cpu", weights_only=True)
            state_dict = checkpoint.get("module", checkpoint)
        self.model.load_state_dict(state_dict)

    def _prepare_multivqvae_input(
        self, actions: torch.Tensor
    ) -> Dict[str, torch.Tensor]:
        """Prepare MultiVQVAE input (decompose into pos/rot/grip)"""
        return {
            "pos": self.action_space.extract_pos(actions),
            "rot": self.action_space.extract_rot(actions),
            "grip": self.action_space.extract_gripper(actions),
        }

    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> np.ndarray:
        actions_tensor = self._to_tensor(actions)

        # Single-arm mode processing
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions_tensor = self.action_space.zero_left_arm(actions_tensor)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions_tensor = self.action_space.zero_right_arm(actions_tensor)

        # Convert to MultiVQVAE input format
        mv_input = self._prepare_multivqvae_input(actions_tensor)

        tokens = self.model.encode(mv_input)
        return self._to_numpy(tokens)

    def decode(
        self,
        tokens: Union[List[List[int]], np.ndarray, torch.Tensor],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        if isinstance(tokens, list):
            tokens_tensor = torch.tensor(tokens, dtype=torch.long, device=self.device)
        else:
            tokens_tensor = self._to_tensor(tokens, dtype=torch.long)

        actions = self.model.decode(tokens_tensor)

        # Single-arm mode processing: zero out unused arm
        if self.action_space.arm_mode == ArmMode.RIGHT_ONLY:
            actions = self.action_space.zero_left_arm(actions)
        elif self.action_space.arm_mode == ArmMode.LEFT_ONLY:
            actions = self.action_space.zero_right_arm(actions)

        # Create padding mask
        batch_size = actions.shape[0]
        mask = np.ones((batch_size, actions.shape[1]), dtype=bool)

        return self._to_numpy(actions), mask

    def to(self, device: Union[str, torch.device]) -> "MultiVQVAETokenizer":
        self.model.to(device)
        self.device = torch.device(device)
        return self

    def eval(self) -> "MultiVQVAETokenizer":
        self.model.eval()
        return self


class PATTokenizer:
    """
    PAT (Policy Action Tokenizer) - Unified Action Tokenizer Interface

    Automatically use different backends by switching tokenizer_type:
    - "uat": Transformer-based, supports multi-embodiment
    - "oat": Register-based transformer encoder
    - "vqvae": CNN-based, efficient encoding
    - "multivqvae": Action decomposition encoding (pos/rot/grip separated)

    Action Space (20 dimensions):
    - Left arm (0-9): xyz(3) + 6d_rot(6) + gripper(1)
    - Right arm (10-19): xyz(3) + 6d_rot(6) + gripper(1)

    Supports single-arm mode: left or right arm zeroed out

    Example:
        >>> # UAT mode
        >>> config = PATConfig(tokenizer_type="uat", n_tokens=16, n_quantizers=1)
        >>> tokenizer = PATTokenizer(config)
        >>> tokens = tokenizer.encode(actions)
        >>> actions_recon, mask = tokenizer.decode(tokens)

        >>> # OAT mode
        >>> config = PATConfig(tokenizer_type="oat", oat_num_registers=16)
        >>> tokenizer = PATTokenizer(config)

        >>> # VQVAE mode
        >>> config = PATConfig(tokenizer_type="vqvae", n_codebooks=4)
        >>> tokenizer = PATTokenizer(config)

        >>> # MultiVQVAE mode
        >>> config = PATConfig(tokenizer_type="multivqvae")
        >>> tokenizer = PATTokenizer(config)

        >>> # Single-arm mode
        >>> config = PATConfig(
        ...     tokenizer_type="uat",
        ...     action_space=ActionSpaceConfig(arm_mode=ArmMode.RIGHT_ONLY)
        ... )
        >>> tokenizer = PATTokenizer(config)
    """

    _tokenizer_classes = {
        TokenizerType.UAT: UATTokenizer,
        TokenizerType.OAT: OATTokenizer,
        TokenizerType.VQVAE: VQVAETokenizer,
        TokenizerType.MULTIVQVAE: MultiVQVAETokenizer,
    }

    def __init__(self, config: PATConfig):
        """
        Initialize PAT Tokenizer

        Args:
            config: PAT configuration object
        """
        self.config = config
        self.action_space = config.action_space
        self.tokenizer_type = config.tokenizer_type

        # Create corresponding tokenizer based on type
        tokenizer_class = self._tokenizer_classes.get(config.tokenizer_type)
        if tokenizer_class is None:
            raise ValueError(
                f"Unsupported tokenizer type: {config.tokenizer_type}. "
                f"Supported types: {list(self._tokenizer_classes.keys())}"
            )

        self._tokenizer = tokenizer_class(config)

        # Internal action normalizer
        self.normalizer = ActionNormalizer(action_dim=self.action_space.action_dim)

        logger.info(f"PAT Tokenizer initialized: {config.tokenizer_type.value}")

    @classmethod
    def from_type(
        cls, tokenizer_type: Union[str, TokenizerType], **kwargs
    ) -> "PATTokenizer":
        """
        Quickly create tokenizer from type string

        Args:
            tokenizer_type: Tokenizer type ("uat", "oat", "vqvae", "multivqvae")
            **kwargs: Override PATConfig parameters

        Returns:
            PATTokenizer instance
        """
        if isinstance(tokenizer_type, str):
            tokenizer_type = TokenizerType(tokenizer_type.lower())

        config = PATConfig(tokenizer_type=tokenizer_type, **kwargs)
        return cls(config)

    @classmethod
    def from_pretrained(
        cls,
        model_path: str,
        tokenizer_type: Optional[Union[str, TokenizerType]] = None,
        **kwargs,
    ) -> "PATTokenizer":
        """
        Load from pretrained weights

        Args:
            model_path: Model weights path (directory containing config.json and model.safetensors)
            tokenizer_type: Tokenizer type (optional, will be auto-inferred from config.json)
            **kwargs: Other PATConfig parameters (will override config.json)

        Returns:
            PATTokenizer instance
        """
        model_path = Path(model_path)

        # Try to load config from config.json
        config_path = model_path / "config.json"
        if config_path.exists():
            logger.info(f"Loading config from {config_path}")
            import json

            with open(config_path, "r") as f:
                config_dict = json.load(f)

            # Update with kwargs (kwargs take precedence)
            config_dict.update(kwargs)

            # Get tokenizer type from config
            if tokenizer_type is None:
                tokenizer_type_str = config_dict.get("tokenizer_type", "uat")
                if isinstance(tokenizer_type_str, str):
                    tokenizer_type = TokenizerType(tokenizer_type_str.lower())
                else:
                    tokenizer_type = tokenizer_type_str

            # Handle action_space config
            action_space_dict = config_dict.get("action_space", {})
            if action_space_dict:
                action_space = ActionSpaceConfig(**action_space_dict)
                config_dict["action_space"] = action_space

            # Handle tokenizer_type - convert string to TokenizerType if needed
            tokenizer_type_str = config_dict.get("tokenizer_type", "uat")
            if isinstance(tokenizer_type_str, str):
                config_dict["tokenizer_type"] = TokenizerType(
                    tokenizer_type_str.lower()
                )

            # Find model file path
            model_file = None
            if (model_path / "model.safetensors").exists():
                model_file = str(model_path / "model.safetensors")
            elif (model_path / "pytorch_model.bin").exists():
                model_file = str(model_path / "pytorch_model.bin")
            config_dict["model_path"] = model_file

            # Create config from dict
            config = PATConfig(**config_dict)
            logger.info(f"Loaded config: tokenizer_type={config.tokenizer_type.value}")
        else:
            # Fallback to inferring from path
            logger.warning(f"Config not found at {config_path}, inferring from path")
            if tokenizer_type is None:
                path_lower = model_path.lower()
                if "uat" in path_lower or "actioncodec" in path_lower:
                    tokenizer_type = TokenizerType.UAT
                elif "oat" in path_lower:
                    tokenizer_type = TokenizerType.OAT
                elif "multivqvae" in path_lower:
                    tokenizer_type = TokenizerType.MULTIVQVAE
                elif "vqvae" in path_lower:
                    tokenizer_type = TokenizerType.VQVAE
                else:
                    raise ValueError(
                        "Cannot infer tokenizer type from path, please specify tokenizer_type explicitly"
                    )

            config = PATConfig(
                tokenizer_type=tokenizer_type, model_path=str(model_path), **kwargs
            )

        # Create tokenizer instance
        tokenizer = cls(config)

        # Load normalizer if available
        normalizer_path = model_path / "normalizer.json"
        if normalizer_path.exists():
            logger.info(f"Loading normalizer from {normalizer_path}")
            tokenizer.load_normalizer(str(normalizer_path))
        elif (model_path / "pat" / "normalizer.json").exists():
            normalizer_path = model_path / "pat" / "normalizer.json"
            logger.info(f"Loading normalizer from {normalizer_path}")
            tokenizer.load_normalizer(str(normalizer_path))
        else:
            logger.warning(
                f"Normalizer config not found at {normalizer_path}, skipping normalizer loading"
            )

        return tokenizer

    def encode(
        self,
        actions: Union[np.ndarray, torch.Tensor, List],
        embodiment_ids: Optional[Union[List[int], int]] = None,
        **kwargs,
    ) -> Union[List[List[int]], np.ndarray]:
        """
        Encode actions to tokens.

        Actions are automatically normalized before encoding if normalizer is fitted.

        Args:
            actions: Action array, shape: (batch, horizon, 20)
            embodiment_ids: Embodiment IDs (only for UAT)
            **kwargs: Other parameters (padding_mask, etc.)

        Returns:
            Token sequence
            - UAT/OAT: List[List[int]], shape (batch, n_tokens)
            - VQVAE/MultiVQVAE: np.ndarray, shape (batch, tokens)
        """
        # Convert to tensor for normalization
        if isinstance(actions, np.ndarray):
            actions_tensor = torch.from_numpy(actions).float()
        elif isinstance(actions, list):
            actions_tensor = torch.tensor(actions, dtype=torch.float32)
        else:
            actions_tensor = actions.float()

        # Auto-normalize if normalizer is fitted
        if self.normalizer.action.is_fitted():
            actions_tensor = self.normalizer.action.normalize(actions_tensor)
            # Convert back to original type
            if isinstance(actions, np.ndarray):
                actions_tensor = actions_tensor.numpy()
            elif isinstance(actions, list):
                actions_tensor = actions_tensor.tolist()
            else:
                pass  # Keep as tensor

        return self._tokenizer.encode(
            actions=actions_tensor, embodiment_ids=embodiment_ids, **kwargs
        )

    def decode(
        self,
        tokens: Optional[Union[List[List[int]], np.ndarray, torch.Tensor]] = None,
        embodiment_ids: Optional[Union[List[int], int]] = None,
        latents: Optional[np.ndarray] = None,
        eval_keep_k: Optional[List[int]] = None,
        **kwargs,
    ) -> Tuple[np.ndarray, np.ndarray]:
        """
        Decode tokens/latents to actions.

        Actions are automatically unnormalized after decoding if normalizer is fitted.

        Args:
            tokens: Token sequence (for VQVAE/UAT)
            embodiment_ids: Embodiment IDs (only for UAT)
            latents: Latent tokens (for OAT)
            eval_keep_k: Number of tokens to use (for OAT nested dropout)
            **kwargs: Other parameters

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                - actions: Reconstructed actions, shape (batch, horizon, 20)
                - mask: Padding mask, shape (batch, horizon)
        """
        # For OAT tokenizer, use latents and eval_keep_k
        if self.tokenizer_type == TokenizerType.OAT:
            actions, mask = self._tokenizer.decode(
                latents_or_tokens=latents if latents is not None else tokens,
                eval_keep_k=eval_keep_k,
                **kwargs,
            )
        else:
            # For VQVAE/UAT tokenizers
            actions, mask = self._tokenizer.decode(
                tokens=tokens, embodiment_ids=embodiment_ids, **kwargs
            )

        # Auto-unnormalize if normalizer is fitted
        if self.normalizer.action.is_fitted():
            actions_tensor = torch.from_numpy(actions).float()
            actions_tensor = self.normalizer.action.unnormalize(actions_tensor)
            actions = actions_tensor.numpy()

        return actions, mask

    def to(self, device: Union[str, torch.device]) -> "PATTokenizer":
        """Move to device"""
        self._tokenizer.to(device)
        return self

    def cuda(self, device_id: int = 0) -> "PATTokenizer":
        """Move to GPU"""
        return self.to(f"cuda:{device_id}")

    def cpu(self) -> "PATTokenizer":
        """Move to CPU"""
        return self.to("cpu")

    def eval(self) -> "PATTokenizer":
        """Set to evaluation mode"""
        self._tokenizer.eval()
        return self

    def train(self) -> "PATTokenizer":
        """Set to training mode"""
        if hasattr(self._tokenizer.model, "train"):
            self._tokenizer.model.train()
        return self

    @property
    def model(self) -> nn.Module:
        """Get underlying model"""
        return self._tokenizer.model

    @property
    def device(self) -> torch.device:
        """Get current device"""
        return self._tokenizer.device

    @property
    def vocab_size(self) -> int:
        """Get vocabulary size"""
        if self.tokenizer_type == TokenizerType.UAT:
            return self.config.vq_codebook_size
        else:
            return self.config.vqvae_num_embeddings

    @property
    def num_tokens(self) -> int:
        """Get number of tokens"""
        if self.tokenizer_type == TokenizerType.UAT:
            return self.config.n_tokens
        elif self.tokenizer_type == TokenizerType.OAT:
            return self.config.oat_num_registers
        elif self.tokenizer_type == TokenizerType.VQVAE:
            # VQVAE: horizon/8 * n_codebooks
            return self.action_space.horizon // 8 * self.config.vqvae_n_codebooks
        else:
            # MultiVQVAE: sum of all codebooks
            return sum(self.config.multivqvae_n_codebooks.values()) * (
                self.action_space.horizon // 8
            )

    def __repr__(self) -> str:
        return (
            f"PATTokenizer(\n"
            f"  type={self.tokenizer_type.value},\n"
            f"  vocab_size={self.vocab_size},\n"
            f"  num_tokens={self.num_tokens},\n"
            f"  action_dim={self.action_space.action_dim},\n"
            f"  horizon={self.action_space.horizon},\n"
            f"  arm_mode={self.action_space.arm_mode.value},\n"
            f"  device={self.device},\n"
            f"  normalizer_fitted={self.normalizer.action.is_fitted()}\n"
            f")"
        )

    def get_normalizer_state(self) -> Dict[str, Any]:
        """Get normalizer state dict for saving."""
        return self.normalizer.state_dict()

    def load_normalizer_state(self, state_dict: Dict[str, Any]) -> None:
        """Load normalizer state dict."""
        self.normalizer.load_state_dict(state_dict)

    def save_normalizer(self, path: str) -> None:
        """Save normalizer state to JSON file."""
        import json

        state = self.get_normalizer_state()
        normalizer_data = {
            "action": {
                "scale": state["action.scale"].tolist(),
                "offset": state["action.offset"].tolist(),
                "fitted": state["action.fitted"].item(),
            }
        }
        with open(path, "w") as f:
            json.dump(normalizer_data, f, indent=2)

    def load_normalizer(self, path: str) -> None:
        """Load normalizer state from JSON file."""
        import json

        with open(path, "r") as f:
            normalizer_data = json.load(f)

        # Get shape from saved data
        scale_data = normalizer_data["action"]["scale"]
        action_dim = len(scale_data)

        # Reinitialize normalizer with correct shape by directly setting buffers
        self.normalizer = ActionNormalizer(action_dim=action_dim)

        # Register new buffers with correct shape
        delattr(self.normalizer.action, "scale")
        delattr(self.normalizer.action, "offset")
        delattr(self.normalizer.action, "fitted")

        self.normalizer.action.register_buffer(
            "scale", torch.ones(action_dim), persistent=True
        )
        self.normalizer.action.register_buffer(
            "offset", torch.zeros(action_dim), persistent=True
        )
        self.normalizer.action.register_buffer(
            "fitted", torch.tensor(False), persistent=True
        )

        # Reconstruct state dict
        state = {
            "action.scale": torch.tensor(scale_data),
            "action.offset": torch.tensor(normalizer_data["action"]["offset"]),
            "action.fitted": torch.tensor(normalizer_data["action"]["fitted"]),
        }
        self.load_normalizer_state(state)


# ========== Helper Functions ==========


def create_uat_tokenizer(
    n_tokens: int = 16,
    n_quantizers: int = 1,
    vq_type: str = "vq",
    codebook_size: int = 2048,
    z_dim: int = 512,
    horizon: int = 24,
    arm_mode: ArmMode = ArmMode.BILATERAL,
    **kwargs,
) -> PATTokenizer:
    """Quickly create UAT Tokenizer"""
    return PATTokenizer.from_type(
        tokenizer_type="uat",
        n_tokens=n_tokens,
        n_quantizers=n_quantizers,
        vq_type=vq_type,
        vq_codebook_size=codebook_size,
        z_dim=z_dim,
        action_space=ActionSpaceConfig(horizon=horizon, arm_mode=arm_mode),
        **kwargs,
    )


def create_vqvae_tokenizer(
    n_codebooks: int = 4,
    embedding_dim: int = 256,
    num_embeddings: int = 2048,
    horizon: int = 24,
    arm_mode: ArmMode = ArmMode.BILATERAL,
    **kwargs,
) -> PATTokenizer:
    """Quickly create VQVAE Tokenizer"""
    return PATTokenizer.from_type(
        tokenizer_type="vqvae",
        vqvae_n_codebooks=n_codebooks,
        vqvae_embedding_dim=embedding_dim,
        vqvae_num_embeddings=num_embeddings,
        action_space=ActionSpaceConfig(horizon=horizon, arm_mode=arm_mode),
        **kwargs,
    )


def create_multivqvae_tokenizer(
    n_codebooks: Dict[str, int] = None,
    embedding_dim: int = 256,
    num_embeddings: int = 2048,
    horizon: int = 24,
    arm_mode: ArmMode = ArmMode.BILATERAL,
    **kwargs,
) -> PATTokenizer:
    """Quickly create MultiVQVAE Tokenizer"""
    if n_codebooks is None:
        n_codebooks = {"pos": 6, "rot": 3, "grip": 1}
    return PATTokenizer.from_type(
        tokenizer_type="multivqvae",
        multivqvae_n_codebooks=n_codebooks,
        vqvae_embedding_dim=embedding_dim,
        vqvae_num_embeddings=num_embeddings,
        action_space=ActionSpaceConfig(horizon=horizon, arm_mode=arm_mode),
        **kwargs,
    )


def create_oat_tokenizer(
    num_registers: int = 16,
    encoder_emb_dim: int = 256,
    encoder_head_dim: int = 64,
    encoder_depth: int = 6,
    encoder_dropout: float = 0.1,
    decoder_emb_dim: int = 256,
    decoder_head_dim: int = 64,
    decoder_depth: int = 6,
    decoder_dropout: float = 0.1,
    decoder_use_causal: bool = False,
    token_dropout_mode: str = "drop",
    horizon: int = 24,
    arm_mode: ArmMode = ArmMode.BILATERAL,
    **kwargs,
) -> PATTokenizer:
    """Quickly create OAT Tokenizer"""
    return PATTokenizer.from_type(
        tokenizer_type="oat",
        oat_num_registers=num_registers,
        oat_encoder_emb_dim=encoder_emb_dim,
        oat_encoder_head_dim=encoder_head_dim,
        oat_encoder_depth=encoder_depth,
        oat_encoder_dropout=encoder_dropout,
        oat_decoder_emb_dim=decoder_emb_dim,
        oat_decoder_head_dim=decoder_head_dim,
        oat_decoder_depth=decoder_depth,
        oat_decoder_dropout=decoder_dropout,
        oat_decoder_use_causal=decoder_use_causal,
        oat_token_dropout_mode=token_dropout_mode,
        action_space=ActionSpaceConfig(horizon=horizon, arm_mode=arm_mode),
        **kwargs,
    )


__all__ = [
    "PATTokenizer",
    "PATConfig",
    "ActionSpaceConfig",
    "TokenizerType",
    "ArmMode",
    "create_uat_tokenizer",
    "create_oat_tokenizer",
    "create_vqvae_tokenizer",
    "create_multivqvae_tokenizer",
]
