"""Transformer-based sentiment model for Myanmar text."""

import logging
from typing import Any, Dict, Optional

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForSequenceClassification,
    AutoTokenizer,
)

from .base_model import BaseModel

logger = logging.getLogger(__name__)


class TransformerSentimentModel(BaseModel):
    """Transformer-based sentiment classification model."""
    
    def __init__(
        self,
        model_name: str = "bert-base-multilingual-cased",
        num_labels: int = 4,
        dropout: float = 0.1,
        freeze_encoder: bool = False,
    ):
        """
        Args:
            model_name: Pretrained model name
            num_labels: Number of sentiment labels
            dropout: Dropout rate
            freeze_encoder: Whether to freeze encoder weights
        """
        super().__init__()
        
        self.model_name = model_name
        self.num_labels = num_labels
        
        # Load pretrained config
        self.config = AutoConfig.from_pretrained(model_name)
        
        # Load pretrained model
        self.transformer = AutoModel.from_pretrained(model_name)
        
        # Classification head
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(self.config.hidden_size, num_labels)
        
        # Freeze encoder if requested
        if freeze_encoder:
            for param in self.transformer.parameters():
                param.requires_grad = False
        
        self.to(self.device)
    
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass."""
        outputs = self.transformer(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        
        # Use [CLS] token representation
        pooled_output = outputs.last_hidden_state[:, 0, :]
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        
        return logits
    
    def predict(
        self,
        texts: list,
        tokenizer,
        batch_size: int = 16,
    ) -> Dict[str, Any]:
        """Make predictions on texts."""
        self.eval()
        
        all_probs = []
        
        with torch.no_grad():
            for i in range(0, len(texts), batch_size):
                batch_texts = texts[i:i + batch_size]
                
                encoding = tokenizer(
                    batch_texts,
                    padding=True,
                    truncation=True,
                    max_length=512,
                    return_tensors="pt",
                )
                
                input_ids = encoding["input_ids"].to(self.device)
                attention_mask = encoding["attention_mask"].to(self.device)
                
                logits = self.forward(input_ids, attention_mask)
                probs = torch.softmax(logits, dim=-1)
                
                all_probs.append(probs.cpu().numpy())
        
        import numpy as np
        all_probs = np.vstack(all_probs)
        
        sentiment_labels = ["negative", "neutral", "positive", "sarcastic"]
        
        predictions = []
        for i, probs in enumerate(all_probs):
            pred_idx = probs.argmax()
            predictions.append({
                "text": texts[i],
                "sentiment": sentiment_labels[pred_idx],
                "confidence": probs[pred_idx],
                "probabilities": {
                    label: probs[j] for j, label in enumerate(sentiment_labels)
                },
            })
        
        return {"predictions": predictions}
    
    def extract_features(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Extract hidden features."""
        outputs = self.transformer(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        return outputs.last_hidden_state


def load_pretrained_model(
    model_path: str,
    num_labels: int = 4,
) -> TransformerSentimentModel:
    """Load a pretrained model from path or HuggingFace."""
    # Check if it's a HuggingFace model
    if "/" in model_path:
        return TransformerSentimentModel(
            model_name=model_path,
            num_labels=num_labels,
        )
    
    # Load from local checkpoint
    model = TransformerSentimentModel(num_labels=num_labels)
    checkpoint = torch.load(model_path, map_location="cpu")
    
    if "model_state_dict" in checkpoint:
        model.load_state_dict(checkpoint["model_state_dict"])
    elif "model" in checkpoint:
        model.transformer = checkpoint["model"]
    
    return model


if __name__ == "__main__":
    print("Testing TransformerSentimentModel...")
    
    model = TransformerSentimentModel(
        model_name="bert-base-multilingual-cased",
        num_labels=4,
    )
    
    print(f"Total parameters: {model.get_num_parameters():,}")
    print(f"Trainable parameters: {model.get_num_trainable_parameters():,}")