from enum import Enum
from typing import Optional
from transformers import (
    PretrainedConfig,
    AutoConfig,
    SiglipTextConfig, 
    SiglipVisionConfig,
    CLIPTextConfig,
    CLIPVisionConfig
)


LOGIT_SCALE_CLIP_PRETRAINED = 4.605170249938965


class RetLayerStrategy(str, Enum):
    VIT_B = 'vit_b'
    VIT_L = 'vit_l'
    VIT_H = 'vit_h'


RET_LAYER_STRATEGY_MAPPING = {
    RetLayerStrategy.VIT_B: (2, 6, 10),
    RetLayerStrategy.VIT_L: (2, 17, 22),
    RetLayerStrategy.VIT_H: (3, 24, 30)
}


class Ret2Config(PretrainedConfig):
    model_type = 'ret2'

    def __init__(
        self,
        text_name_or_path: Optional[str] = None,
        text_frozen: bool = True,
        text_layer_strategy: str = RetLayerStrategy.VIT_B,
        text_global_feats_proj_size: int = 0,
        vision_name_or_path: Optional[str] = None,
        vision_frozen: bool = True,
        vision_layer_strategy: str = RetLayerStrategy.VIT_L,
        vision_global_feats_proj_size: int = 0,
        tokenizer_max_length: int = None,
        w_final_size: int = 0,
        logit_scale_init_value: float = LOGIT_SCALE_CLIP_PRETRAINED,
        logit_scale_clamp: float = 100.0,
        hidden_size: int = 1024,
        dropout_p: float = 0.05,
        attention_dropout: float = 0.05,
        input_gate_bias: float = 0.0,
        forget_gate_bias: float = 0.0,             
        **kwargs
    ):
        super().__init__(**kwargs)

        self.text_name_or_path = text_name_or_path
        self.text_frozen = text_frozen
        self.text_layer_strategy = text_layer_strategy
        self.text_global_feats_proj_size = text_global_feats_proj_size
        self.vision_name_or_path = vision_name_or_path
        self.vision_frozen = vision_frozen
        self.vision_layer_strategy = vision_layer_strategy
        self.vision_global_feats_proj_size = vision_global_feats_proj_size
        self.tokenizer_max_length = tokenizer_max_length
        self.w_final_size = w_final_size
        self.logit_scale_init_value = logit_scale_init_value
        self.logit_scale_clamp = logit_scale_clamp
        self.hidden_size = hidden_size
        self.dropout_p = dropout_p
        self.attention_dropout = attention_dropout
        self.input_gate_bias = input_gate_bias
        self.forget_gate_bias = forget_gate_bias

        self._load_sub_configs()

    def _load_sub_configs(self):
        """Load text and vision configs from their respective paths."""
        if self.text_name_or_path is None:
            self.text_config = None
        else:
            if 'siglip' in self.text_name_or_path.lower():
                auto_cls = SiglipTextConfig
            elif 'clip-vit-' in self.text_name_or_path.lower():
                auto_cls = CLIPTextConfig
            else:
                auto_cls = AutoConfig
            self.text_config = auto_cls.from_pretrained(self.text_name_or_path)

        if self.vision_name_or_path is None:
            self.vision_config = None
        else:
            if 'siglip' in self.vision_name_or_path.lower():
                auto_cls = SiglipVisionConfig
            elif 'clip-vit-' in self.vision_name_or_path.lower():
                auto_cls = CLIPVisionConfig
            else:
                auto_cls = AutoConfig
            self.vision_config = auto_cls.from_pretrained(self.vision_name_or_path)

    def to_dict(self):
        """
        Override to_dict to exclude text_config and vision_config from serialization.
        """
        ret = super().to_dict()
        
        ret.pop('text_config', None)
        ret.pop('vision_config', None)
        
        return ret

    @classmethod
    def from_dict(cls, config_dict, **kwargs):
        """
        Override from_dict to ensure configs are loaded after deserialization.
        """
        config = super().from_dict(config_dict, **kwargs)
        if isinstance(config, tuple):
            config[0]._load_sub_configs()
        else:
            config._load_sub_configs()
        
        return config
    
    @property
    def text_layer_idxs(self):
        return RET_LAYER_STRATEGY_MAPPING[self.text_layer_strategy]
    
    @property
    def vision_layer_idxs(self):
        return RET_LAYER_STRATEGY_MAPPING[self.vision_layer_strategy]