from pydantic import Field, model_validator, PrivateAttr, ConfigDict
from typing import List, Dict, Any, Optional, Union

from ...core.base.hyperparameters_base import ModelHyperparameters

class BSMModelHyperparameters(ModelHyperparameters):
    """
    Hyperparameters for the BSM (Buyer Seller Messaging) model training,
    extending the base ModelHyperparameters.
    
    Fields are organized into three tiers:
    1. Tier 1: Essential User Inputs - fields that users must explicitly provide
    2. Tier 2: System Inputs with Defaults - fields with reasonable defaults that can be overridden
    3. Tier 3: Derived Fields - fields calculated from other fields (private attributes with properties)
    """
    # ===== Essential User Inputs (Tier 1) =====
    # These are fields that users must explicitly provide
    
    # For BERT model configuration
    tokenizer: str = Field(
        description="Tokenizer name or path (e.g., from Hugging Face)")
    
    # For text field specification
    text_name: str = Field(
        description="Name of the primary text field to be processed")
        
    # ===== System Inputs with Defaults (Tier 2) =====
    # These are fields with reasonable defaults that users can override
    
    # Override model_class from base
    model_class: str = Field(
        default="multimodal_bert", 
        description="Model class identifier")

    # Trainer and Optimization parameters
    # For optimizer
    lr_decay: float = Field(
        default=0.05, 
        description="Learning rate decay")
    
    adam_epsilon: float = Field(
        default=1e-08, 
        description="Epsilon for Adam optimizer")
    
    momentum: float = Field(
        default=0.9, 
        description="Momentum for SGD optimizer (if SGD is chosen)")
    
    # For scheduler
    run_scheduler: bool = Field(
        default=True, 
        description="Run learning rate scheduler flag")
    
    val_check_interval: float = Field(
        default=0.25, 
        description="Validation check interval during training (float for fraction of epoch, int for steps)")
    
    # For trainer
    warmup_steps: int = Field(
        default=300, 
        gt=0, 
        le=1000, 
        description="Warmup steps for learning rate scheduler")
    
    weight_decay: float = Field(
        default=0.0, 
        description="Weight decay for optimizer (L2 penalty)")
    
    gradient_clip_val: float = Field(
        default=1.0, 
        description="Value for gradient clipping to prevent exploding gradients")
    
    fp16: bool = Field(
        default=False, 
        description="Enable 16-bit mixed precision training (requires compatible hardware)")
    
    # Early stopping and Checkpointing parameters
    early_stop_metric: str = Field(
        default='val_loss', 
        description="Metric for early stopping")
    
    early_stop_patience: int = Field(
        default=3, 
        gt=0, 
        le=10, 
        description="Patience for early stopping")
    
    load_ckpt: bool = Field(
        default=False, 
        description="Load checkpoint flag")

    # Text Preprocessing and Tokenization parameters
    text_field_overwrite: bool = Field(
        default=False, 
        description="Overwrite text field if it exists (e.g. during feature engineering)")
    
    # For chunking long texts
    chunk_trancate: bool = Field(
        default=True, 
        description="Chunk truncation flag for long texts") # Typo 'trancate' kept as per original
    
    max_total_chunks: int = Field(
        default=3, 
        description="Maximum total chunks for processing long texts")
    
    # For tokenizer settings
    max_sen_len: int = Field(
        default=512, 
        description="Maximum sentence length for tokenizer")
    
    fixed_tokenizer_length: bool = Field(
        default=True, 
        description="Use fixed tokenizer length")
    
    text_input_ids_key: str = Field(
        default='input_ids', 
        description="Key name for input_ids from tokenizer output")
    
    text_attention_mask_key: str = Field(
        default='attention_mask', 
        description="Key name for attention_mask from tokenizer output")
    
    # Model structure parameters
    # For Convolutional layers
    num_channels: List[int] = Field(
        default=[100, 100], 
        description="Number of channels for convolutional layers")
    
    num_layers: int = Field(
        default=2, 
        description="Number of layers in the model (e.g., BiLSTM, Transformer encoders)")
    
    dropout_keep: float = Field(
        default=0.1, 
        description="Dropout keep probability")
    
    kernel_size: List[int] = Field(
        default=[3, 5, 7], 
        description="Kernel sizes for convolutional layers")
    
    is_embeddings_trainable: bool = Field(
        default=True, 
        description="Trainable embeddings flag")
    
    # For BERT fine-tuning
    pretrained_embedding: bool = Field(
        default=True, 
        description="Use pretrained embeddings")
    
    reinit_layers: int = Field(
        default=2, 
        description="Number of layers to reinitialize from pretrained model")
    
    reinit_pooler: bool = Field(
        default=True, 
        description="Reinitialize pooler layer flag")
    
    # For Multimodal BERT
    hidden_common_dim: int = Field(
        default=100, 
        description="Common hidden dimension for multimodal model")
        
    # ===== Derived Fields (Tier 3) =====
    # These are fields calculated from other fields
    
    _model_config_dict: Optional[Dict[str, Any]] = PrivateAttr(default=None)
    _tokenizer_config: Optional[Dict[str, Any]] = PrivateAttr(default=None)
    
    # Explicitly define the model_config
    model_config = ConfigDict(
        arbitrary_types_allowed=True,
        validate_assignment=True,
        extra='forbid',
        protected_namespaces=()
    )

    @property
    def model_config_dict(self) -> Dict[str, Any]:
        """Get complete model configuration dictionary derived from hyperparameters."""
        if self._model_config_dict is None:
            self._model_config_dict = {
                "hidden_common_dim": self.hidden_common_dim,
                "num_layers": self.num_layers,
                "dropout": self.dropout_keep,
                "num_channels": self.num_channels,
                "kernel_size": self.kernel_size,
                "trainable_embeddings": self.is_embeddings_trainable,
                "pretrained": self.pretrained_embedding,
                "reinit_layers": self.reinit_layers,
                "reinit_pooler": self.reinit_pooler,
            }
        return self._model_config_dict
    
    @property
    def tokenizer_config(self) -> Dict[str, Any]:
        """Get tokenizer configuration dictionary derived from hyperparameters."""
        if self._tokenizer_config is None:
            self._tokenizer_config = {
                "name": self.tokenizer,
                "max_length": self.max_sen_len,
                "fixed_length": self.fixed_tokenizer_length,
                "text_field": self.text_name,
                "input_ids_key": self.text_input_ids_key,
                "attention_mask_key": self.text_attention_mask_key
            }
        return self._tokenizer_config

    @model_validator(mode='after')
    def validate_bsm_hyperparameters(self) -> 'BSMModelHyperparameters':
        """Validate BSM model-specific hyperparameters and initialize derived fields."""
        # Call the base model validator first to initialize its derived fields
        super().validate_dimensions()
        
        # Initialize derived fields
        self._model_config_dict = None
        self._tokenizer_config = None
        
        # Perform BSM-specific validations
        if len(self.num_channels) != self.num_layers:
            raise ValueError(f"Length of num_channels ({len(self.num_channels)}) must match num_layers ({self.num_layers})")
            
        return self
    
    def get_public_init_fields(self) -> Dict[str, Any]:
        """
        Override get_public_init_fields to include BSM-specific derived fields.
        Gets a dictionary of public fields suitable for initializing a child config.
        """
        # Get fields from parent class
        base_fields = super().get_public_init_fields()
        
        # Add derived fields that should be exposed
        derived_fields = {
            # If you need to expose any derived fields, add them here
        }
        
        # Combine (derived fields take precedence if overlap)
        return {**base_fields, **derived_fields}
        
    def get_trainer_config(self) -> Dict[str, Any]:
        """
        Get trainer configuration dictionary for PyTorch Lightning.
        This combines various trainer-related settings.
        
        Returns:
            Dict[str, Any]: Configuration dictionary for trainer
        """
        return {
            "max_epochs": self.max_epochs,
            "gpus": self.device if self.device >= 0 else 0,
            "gradient_clip_val": self.gradient_clip_val,
            "val_check_interval": self.val_check_interval,
            "precision": 16 if self.fp16 else 32,
            "early_stop_metric": self.early_stop_metric,
            "early_stop_patience": self.early_stop_patience,
        }
