# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Configuration module for the PI0 Policy.

This module defines the `PI0Config` class, which handles the configuration parameters
for the PI0 Vision-Language-Action Flow Model. It includes settings for the model architecture,
optimization, scheduling, and data processing.
"""

import logging
from dataclasses import dataclass, field
from typing import Literal

from opentau.configs.policies import PreTrainedConfig
from opentau.configs.types import FeatureType, NormalizationMode, PolicyFeature
from opentau.optim.optimizers import AdamWConfig
from opentau.optim.schedulers import (
    CosineDecayWithWarmupSchedulerConfig,
    LRSchedulerConfig,
)


@PreTrainedConfig.register_subclass("pi0")
@dataclass
class PI0Config(PreTrainedConfig):
    """Configuration class for the PI0 Policy.

    This class defines the configuration parameters for the PI0 model, including
    input/output structure, model architecture, training settings, and preprocessing.

    Args:
        n_obs_steps: Number of observation steps to use. Defaults to 1.
        chunk_size: Size of the action chunk. The upper bound for n_action_steps. Defaults to 50.
        n_action_steps: Number of action steps to predict. Defaults to 50.
        safety_buffer: Safety buffer size. Defaults to 0.
        normalization_mapping: Mapping of feature names to normalization modes.
            Defaults to identity for visual features and mean-std for state and action.
        max_state_dim: Maximum dimension for state vectors. Shorter vectors are padded. Defaults to 32.
        max_action_dim: Maximum dimension for action vectors. Shorter vectors are padded. Defaults to 32.
        resize_imgs_with_padding: Target size (height, width) for image resizing with padding.
            Defaults to (224, 224).
        empty_cameras: Number of empty camera inputs to add. Used for specific adaptations like
            Aloha simulation. Defaults to 0.
        tokenizer_max_length: Maximum length for tokenizer. Defaults to 48.
        proj_width: Width of the projection layer. Defaults to 1024.
        dropout: Dropout rate. Defaults to 0.1.
        num_steps: Number of flow matching steps for decoding. Defaults to 10.
        advantage_threshold: Advantage binarization threshold for AWR. Defaults to 0.0.
        advantage: Advantage conditioning mode. One of "ignore", "on", "use".
            "use" uses values from dataset, "ignore" disables conditioning,
            "on" sets advantage to True (for expert demos). Defaults to "use".
        init_strategy: Initialization strategy. One of "no_init", "full_he_init", "expert_only_he_init".
            Defaults to "full_he_init".
        use_cache: Whether to use KV cache during inference. Defaults to True.
        attention_implementation: Attention implementation to use ("eager" or "fa2"). Defaults to "eager".
        freeze_vision_encoder: Whether to freeze the vision encoder during fine-tuning. Defaults to True.
        train_expert_only: Whether to train only the expert module. Defaults to False.
        train_state_proj: Whether to train the state projection layer. Defaults to True.
        optimizer_lr: Learning rate for the optimizer. Defaults to 2.5e-5.
        optimizer_betas: Beta parameters for AdamW optimizer. Defaults to (0.9, 0.95).
        optimizer_eps: Epsilon parameter for AdamW optimizer. Defaults to 1e-8.
        optimizer_weight_decay: Weight decay for AdamW optimizer. Defaults to 1e-10.
        scheduler_warmup_steps: Number of warmup steps for the scheduler. Defaults to 1_000.
        scheduler_decay_steps: Number of decay steps for the scheduler. Defaults to 30_000.
        scheduler_decay_lr: Target learning rate after decay. Defaults to 2.5e-6.
    """

    # Input / output structure.
    n_obs_steps: int = 1
    chunk_size: int = 50
    n_action_steps: int = 50
    safety_buffer: int = 0

    normalization_mapping: dict[str, NormalizationMode] = field(
        default_factory=lambda: {
            "VISUAL": NormalizationMode.IDENTITY,
            "STATE": NormalizationMode.MEAN_STD,
            "ACTION": NormalizationMode.MEAN_STD,
        }
    )

    # Shorter state and action vectors will be padded
    max_state_dim: int = 32
    max_action_dim: int = 32

    # Image preprocessing
    resize_imgs_with_padding: tuple[int, int] = (224, 224)

    # Add empty images. Used by pi0_aloha_sim which adds the empty
    # left and right wrist cameras in addition to the top camera.
    empty_cameras: int = 0

    # Tokenizer
    tokenizer_max_length: int = 48

    # Projector
    proj_width: int = 1024

    # Dropout
    dropout: float = 0.1

    # Decoding
    num_steps: int = 10

    # advantage binarization threshold for AWR
    advantage_threshold: float = 0.0

    # When set to "use", the advantage values provided in the dataset will be used.
    # When set to "ignore", no advantage conditioning will be applied.
    # When set to "on", the advantage will always be True.
    # This should only be "on" when training on expert demonstrations or interventions.
    advantage: Literal["ignore", "on", "use"] = "use"

    # Initialization strategy
    init_strategy: Literal["no_init", "full_he_init", "expert_only_he_init"] = "full_he_init"

    # Attention utils
    use_cache: bool = True
    attention_implementation: str = "eager"  # or fa2

    # Finetuning settings
    freeze_vision_encoder: bool = True
    train_expert_only: bool = False
    train_state_proj: bool = True

    # Training presets
    optimizer_lr: float = 2.5e-5
    optimizer_betas: tuple[float, float] = (0.9, 0.95)
    optimizer_eps: float = 1e-8
    optimizer_weight_decay: float = 1e-10

    scheduler_warmup_steps: int = 1_000
    scheduler_decay_steps: int = 30_000
    scheduler_decay_lr: float = 2.5e-6

    # TODO: Add EMA

    def __post_init__(self):
        """Post-initialization validation."""
        super().__post_init__()

        # TODO(Steven): Validate device and amp? in all policy configs?
        """Input validation (not exhaustive)."""
        if self.n_action_steps > self.chunk_size:
            raise ValueError(
                f"The chunk size is the upper bound for the number of action steps per model invocation. Got "
                f"{self.n_action_steps} for `n_action_steps` and {self.chunk_size} for `chunk_size`."
            )
        if self.n_obs_steps != 1:
            raise ValueError(
                f"Multiple observation steps not handled yet. Got `nobs_steps={self.n_obs_steps}`"
            )

        assert self.init_strategy in ["no_init", "full_he_init", "expert_only_he_init"], (
            f"Invalid init strategy: {self.init_strategy} must be one of ['no_init', 'full_he_init', 'expert_only_he_init']"
        )

        if self.init_strategy == "expert_only_he_init" and self.pretrained_path == "lerobot/pi0":
            raise ValueError(
                "You cannot load pretrained PI0 model when init_strategy is 'expert_only_he_init' due to differences in PaliGemma tokenizer vocab sizes."
            )

        if self.pretrained_path is not None and self.pretrained_path != "lerobot/pi0":
            logging.info("Setting init_strategy to 'no_init' because we are resuming from a checkpoint.")
            self.init_strategy = "no_init"

    def validate_features(self) -> None:
        """Validates the features and adds empty cameras if configured.

        This method checks feature configurations and dynamically adds empty camera inputs
        to `self.input_features` based on the `empty_cameras` parameter.
        """
        # TODO: implement value error
        # if not self.image_features and not self.env_state_feature:
        #     raise ValueError("You must provide at least one image or the environment state among the inputs.")

        for i in range(self.empty_cameras):
            key = f"observation.images.empty_camera_{i}"
            empty_camera = PolicyFeature(
                type=FeatureType.VISUAL,
                shape=(3, 480, 640),
            )
            self.input_features[key] = empty_camera

    def get_optimizer_preset(self) -> AdamWConfig:
        """Returns the default optimizer configuration.

        Returns:
            AdamWConfig: The optimizer configuration with default parameters.
        """
        return AdamWConfig(
            lr=self.optimizer_lr,
            betas=self.optimizer_betas,
            eps=self.optimizer_eps,
            weight_decay=self.optimizer_weight_decay,
        )

    def get_scheduler_preset(self) -> LRSchedulerConfig:
        """Returns the default scheduler configuration.

        Returns:
            CosineDecayWithWarmupSchedulerConfig: The scheduler configuration with default parameters.
        """
        return CosineDecayWithWarmupSchedulerConfig(
            peak_lr=self.optimizer_lr,
            decay_lr=self.scheduler_decay_lr,
            num_warmup_steps=self.scheduler_warmup_steps,
            num_decay_steps=self.scheduler_decay_steps,
        )

    @property
    def observation_delta_indices(self) -> None:
        """Indices for observation deltas.

        Returns:
            None: As observation deltas are not used.
        """
        return None

    @property
    def action_delta_indices(self) -> list[int]:
        """Indices for action deltas.

        Returns:
            list[int]: A list of indices corresponding to the chunk size.
        """
        return list(range(self.chunk_size))

    @property
    def reward_delta_indices(self) -> None:
        """Indices for reward deltas.

        Returns:
            None: As reward deltas are not used.
        """
        return None
