from transformers.configuration_utils import PretrainedConfig
from typing import Dict, Any


from transformers.configuration_utils import PretrainedConfig
from typing import Dict, Any


class TimeRCDConfig(PretrainedConfig):
    """
    Configuration class for Time_RCD model.

    This is the configuration class to store the configuration of a [`Time_RCD`] model. It is used to
    instantiate a Time_RCD model according to the specified arguments, defining the model architecture.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs.
    Read the documentation from [`PretrainedConfig`] for more information.

    Args:
        d_model (`int`, *optional*, defaults to 512):
            Dimension of model hidden states.
        d_proj (`int`, *optional*, defaults to 256):
            Dimension of projection layer.
        patch_size (`int`, *optional*, defaults to 4):
            Size of time series patches.
        num_layers (`int`, *optional*, defaults to 8):
            Number of transformer layers.
        num_heads (`int`, *optional*, defaults to 8):
            Number of attention heads.
        d_ff_dropout (`float`, *optional*, defaults to 0.1):
            Dropout rate for feed-forward networks.
        use_rope (`bool`, *optional*, defaults to True):
            Whether to use Rotary Position Embedding.
        activation (`str`, *optional*, defaults to "gelu"):
            Activation function name.
        num_features (`int`, *optional*, defaults to 1):
            Number of input features in the time series.
        dropout (`float`, *optional*, defaults to 0.1):
            Dropout rate for the model.
        max_seq_len (`int`, *optional*, defaults to 512):
            Maximum sequence length.
        win_size (`int`, *optional*, defaults to 5000):
            Window size for inference.
        batch_size (`int`, *optional*, defaults to 64):
            Default batch size for inference.
    """

    model_type = "time_rcd"

    def __init__(
        self,
        d_model: int = 512,
        d_proj: int = 256,
        patch_size: int = 4,  # Your specific configuration
        num_layers: int = 8,
        num_heads: int = 8,
        d_ff_dropout: float = 0.1,
        use_rope: bool = True,
        activation: str = "gelu",
        num_features: int = 1,
        dropout: float = 0.1,
        max_seq_len: int = 512,
        win_size: int = 5000,
        batch_size: int = 64,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.d_model = d_model
        self.d_proj = d_proj
        self.patch_size = patch_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.d_ff_dropout = d_ff_dropout
        self.use_rope = use_rope
        self.activation = activation
        self.num_features = num_features
        self.dropout = dropout
        self.max_seq_len = max_seq_len
        self.win_size = win_size
        self.batch_size = batch_size

    @classmethod
    def from_pretrained_config(cls, original_config_dict: Dict[str, Any]):
        """Convert from your original configuration format."""
        return cls(
            d_model=original_config_dict.get("ts_config", {}).get("d_model", 512),
            d_proj=original_config_dict.get("ts_config", {}).get("d_proj", 256),
            patch_size=original_config_dict.get("ts_config", {}).get("patch_size", 16),
            num_layers=original_config_dict.get("ts_config", {}).get("num_layers", 8),
            num_heads=original_config_dict.get("ts_config", {}).get("num_heads", 8),
            d_ff_dropout=original_config_dict.get("ts_config", {}).get("d_ff_dropout", 0.1),
            use_rope=original_config_dict.get("ts_config", {}).get("use_rope", True),
            activation=original_config_dict.get("ts_config", {}).get("activation", "gelu"),
            num_features=original_config_dict.get("ts_config", {}).get("num_features", 1),
            dropout=original_config_dict.get("dropout", 0.1),
            max_seq_len=original_config_dict.get("max_seq_len", 512),
            win_size=original_config_dict.get("win_size", 5000),
            batch_size=original_config_dict.get("batch_size", 64),
        )


# Backward compatibility alias
AnomalyCLIPConfig = TimeRCDConfig

# Register config with AutoConfig when using trust_remote_code
try:
    from transformers import AutoConfig
    AutoConfig.register("time_rcd", TimeRCDConfig)
except Exception:
    pass  # Silently fail if already registered or in restricted environment