from __future__ import annotations

from transformers import PretrainedConfig


class EfficientRAEConfig(PretrainedConfig):
    """Configuration for the spatial compression modules around a frozen RAE."""

    model_type = "efficient_rae"

    def __init__(
        self,
        image_size: int = 256,
        patch_size: int = 16,
        encoder_hidden_size: int = 1024,
        pool_size: int = 2,
        hidden_size: int = 1024,
        num_attention_heads: int = 16,
        mlp_ratio: float = 4.0,
        pooler_num_hidden_layers: int = 16,
        unpooler_num_hidden_layers: int = 4,
        pooler_depth_attention: bool = True,
        unpooler_depth_attention: bool = True,
        depth_attention_stride: int = 4,
        rms_norm_eps: float = 1e-6,
        layer_norm_eps: float = 1e-6,
        dropout: float = 0.1,
        attention_dropout: float = 0.1,
        semantic_num_queries: int | None = None,
        mask_ratio_min: float = -0.1,
        mask_ratio_max: float = 0.5,
        noise_t_max: float = 1.0,
        checkpoint_format_version: int = 2,
        **kwargs,
    ):
        # Legacy adapter checkpoints included RAE source metadata. Ignore it:
        # the RAE is now loaded independently and injected into the model.
        kwargs.pop("base_model_name_or_path", None)
        kwargs.pop("base_model_revision", None)
        super().__init__(**kwargs)
        self.image_size = image_size
        self.patch_size = patch_size
        self.encoder_hidden_size = encoder_hidden_size
        self.pool_size = pool_size
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.mlp_ratio = mlp_ratio
        self.pooler_num_hidden_layers = pooler_num_hidden_layers
        self.unpooler_num_hidden_layers = unpooler_num_hidden_layers
        self.pooler_depth_attention = pooler_depth_attention
        self.unpooler_depth_attention = unpooler_depth_attention
        self.depth_attention_stride = depth_attention_stride
        self.rms_norm_eps = rms_norm_eps
        self.layer_norm_eps = layer_norm_eps
        self.dropout = dropout
        self.attention_dropout = attention_dropout
        self.semantic_num_queries = (
            semantic_num_queries
            if semantic_num_queries is not None
            else self.encoder_grid_size**2 + 5
        )
        self.mask_ratio_min = mask_ratio_min
        self.mask_ratio_max = mask_ratio_max
        self.noise_t_max = noise_t_max
        self.checkpoint_format_version = checkpoint_format_version
        self._validate()

    @property
    def encoder_grid_size(self) -> int:
        return self.image_size // self.patch_size

    @property
    def compact_grid_size(self) -> int:
        return self.encoder_grid_size // self.pool_size

    @property
    def intermediate_size(self) -> int:
        return int(self.hidden_size * self.mlp_ratio)

    def _validate(self) -> None:
        if self.image_size < 1 or self.patch_size < 1:
            raise ValueError("image_size and patch_size must be positive")
        if self.image_size % self.patch_size:
            raise ValueError("image_size must be divisible by patch_size")
        if self.pool_size < 1 or self.encoder_grid_size % self.pool_size:
            raise ValueError("encoder grid size must be divisible by pool_size")
        if self.hidden_size < 1 or self.hidden_size % self.num_attention_heads:
            raise ValueError("hidden_size must be divisible by num_attention_heads")
        head_size = self.hidden_size // self.num_attention_heads
        if head_size % 4:
            raise ValueError("attention head size must be divisible by 4 for 2D RoPE")
        if self.mlp_ratio <= 0:
            raise ValueError("mlp_ratio must be positive")
        if self.pooler_num_hidden_layers < 1 or self.unpooler_num_hidden_layers < 1:
            raise ValueError("pooler and unpooler must each have at least one layer")
        if self.depth_attention_stride < 1:
            raise ValueError("depth_attention_stride must be positive")
        if self.rms_norm_eps <= 0 or self.layer_norm_eps <= 0:
            raise ValueError("normalization epsilons must be positive")
        if not 0.0 <= self.dropout < 1.0:
            raise ValueError("dropout must be in [0, 1)")
        if not 0.0 <= self.attention_dropout < 1.0:
            raise ValueError("attention_dropout must be in [0, 1)")
        if self.encoder_hidden_size % self.num_attention_heads:
            raise ValueError(
                "encoder_hidden_size must be divisible by num_attention_heads"
            )
        if self.semantic_num_queries < 1:
            raise ValueError("semantic_num_queries must be positive")
        if self.mask_ratio_min > self.mask_ratio_max:
            raise ValueError("mask_ratio_min must not exceed mask_ratio_max")
        if self.mask_ratio_max > 1.0:
            raise ValueError("mask_ratio_max must not exceed 1")
        if not 0.0 <= self.noise_t_max <= 1.0:
            raise ValueError("noise_t_max must be in [0, 1]")
        if self.checkpoint_format_version != 2:
            raise ValueError("only checkpoint_format_version=2 is supported")


EfficientRAEConfig.register_for_auto_class()