from typing import Any, List, Optional
from transformers.configuration_utils import PretrainedConfig
from transformers import Qwen2_5_VLProcessor, AutoProcessor
from transformers.models.auto.configuration_auto import CONFIG_MAPPING


class MonkeyOCRv2ViTAEEncoderConfig(PretrainedConfig):

    model_type: str = "monkeyocrv2_vitae_encoder"

    def __init__(
        self,
        num_channels: int = 3,
        patch_size: int = 32,               # preprocessor tile size in pixels (= encoder stride for 1:1 mapping)
        temporal_patch_size: int = 1,
        # ViTAEv2-S parameters (see Table 1 in paper):
        stage_dims: List[int] = None,       # [64, 128, 256, 512]  token_dims
        stage_depths: List[int] = None,     # [2, 2, 8, 2]   NC_depth (official ViTAEv2-S)
        stage_heads: List[int] = None,      # [1, 2, 4, 8]   NC_heads
        downsample_ratios: List[int] = None,  # [4, 2, 2, 2]  RC downsample ratio per stage
        kernel_sizes: List[int] = None,       # [7, 3, 3, 3]  RC kernel size per stage
        rc_tokens_type: List[str] = None,     # ['window', 'window', 'transformer', 'transformer']
        nc_tokens_type: List[str] = None,     # ['window', 'window', 'transformer', 'transformer']
        nc_groups: List[int] = None,        # [1, 32, 64, 128]  NC_group
        rc_groups: List[int] = None,        # [1, 16, 32, 64]   RC PCM group (RC0→RC3)
        rc_heads: List[int] = None,         # [1, 1, 2, 4]      RC attn heads (RC0→RC3)
        rc_embed_dims: List[int] = None,    # [64, 64, 128, 256] PRM/PCM intermediate dim per RC
        prm_embed_dim: int = 64,            # legacy fallback for RC0 PRM dim
        window_size: int = 7,               # windowed-attn window size (stages 1-2 + RC1-RC2)
        mlp_ratio: float = 4.0,
        hidden_size: int = 1024,            # LLM projection output dim
        rms_norm_eps: float = 1e-5,
        use_bias: bool = False,
        attn_implementation: str = "sdpa",  # "eager"/"sdpa"/"flash_attention_2"
        initializer_range: float = 0.02,
        init_merger_std: float = 0.02,
        is_causal: bool = False,
        post_norm: bool = True,
        gradient_checkpointing: bool = False,
        **kwargs: Any,
    ):
        super().__init__(**kwargs)
        self.num_channels = num_channels
        self.patch_size = patch_size
        self.temporal_patch_size = temporal_patch_size
        self.stage_dims = stage_dims if stage_dims is not None else [64, 128, 256, 512]
        self.stage_depths = stage_depths if stage_depths is not None else [2, 2, 8, 2]
        self.stage_heads = stage_heads if stage_heads is not None else [1, 2, 4, 8]
        self.downsample_ratios = downsample_ratios if downsample_ratios is not None else [4, 2, 2, 2]
        self.kernel_sizes = kernel_sizes if kernel_sizes is not None else [7, 3, 3, 3]
        self.rc_tokens_type = rc_tokens_type if rc_tokens_type is not None else ["window", "window", "transformer", "transformer"]
        self.nc_tokens_type = nc_tokens_type if nc_tokens_type is not None else ["window", "window", "transformer", "transformer"]
        self.nc_groups = nc_groups if nc_groups is not None else [1, 32, 64, 128]
        self.rc_groups = rc_groups if rc_groups is not None else [1, 16, 32, 64]
        self.rc_heads = rc_heads if rc_heads is not None else [1, 1, 2, 4]
        self.rc_embed_dims = rc_embed_dims if rc_embed_dims is not None else [64, 64, 128, 256]
        self.prm_embed_dim = prm_embed_dim
        self.window_size = window_size
        self.mlp_ratio = mlp_ratio
        self.hidden_size = hidden_size
        self.rms_norm_eps = rms_norm_eps
        self.use_bias = use_bias
        self.attn_implementation = attn_implementation
        self.initializer_range = initializer_range
        self.init_merger_std = init_merger_std
        self.is_causal = is_causal
        self.post_norm = post_norm
        self.gradient_checkpointing = gradient_checkpointing


class MonkeyOCRv2ViTAEProcessor(Qwen2_5_VLProcessor):
    attributes = ["image_processor", "tokenizer"]

    def __init__(self, image_processor=None, tokenizer=None, chat_template=None,
                 encoder_stride=None, **kwargs):
        super().__init__(image_processor, tokenizer, chat_template=chat_template)
        self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
        self.image_token_id = 151655 if not hasattr(tokenizer, "image_token_id") else tokenizer.image_token_id
        # encoder_stride: total conv downsampling applied to each preprocessor tile.
        # With encoder_stride == preprocessor.patch_size, stage4 outputs 1 token per tile.
        pp_size = getattr(image_processor, 'patch_size', 32) if image_processor is not None else 32
        self.encoder_stride = encoder_stride if encoder_stride is not None else pp_size


AutoProcessor.register("monkeyocrv2_vitae", MonkeyOCRv2ViTAEProcessor)
CONFIG_MAPPING.register("monkeyocrv2_vitae_encoder", MonkeyOCRv2ViTAEEncoderConfig)