"""
Model architectures cho Aspect-Based Sentiment Analysis
Hỗ trợ nhiều architectures: Transformer-based, CNN, LSTM, và hybrid models
"""
import torch
import os
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
    RobertaPreTrainedModel, RobertaModel, 
    BertPreTrainedModel, BertModel,
    XLMRobertaPreTrainedModel, XLMRobertaModel,
    BartPreTrainedModel, BartModel, BartForSequenceClassification,
    T5PreTrainedModel, T5EncoderModel,
    AutoConfig, AutoModel, AutoTokenizer,
    PreTrainedModel
)
from transformers.modeling_outputs import SequenceClassifierOutput
from typing import Optional


class BaseABSA(PreTrainedModel):
    """Base class cho tất cả ABSA models"""
    def __init__(self, config):
        super().__init__(config)
        self.num_aspects = config.num_aspects
        self.num_sentiments = config.num_sentiments

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
        raise NotImplementedError

    def get_sentiment_classifiers(self, hidden_size):
        """Create sentiment classifiers cho từng aspect"""
        return nn.ModuleList([
            nn.Linear(hidden_size, self.num_sentiments + 1)  # +1 cho "none"
            for _ in range(self.num_aspects)
        ])


# ========== Transformer-based Models ==========

class TransformerForABSA(RobertaPreTrainedModel):
    """RoBERTa-based model (cho PhoBERT, ViSoBERT, RoBERTa-GRU)"""
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)
        self.roberta = RobertaModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.hidden_size, config.num_sentiments + 1)
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
        # RoBERTa-based models don't use token_type_ids, so we ignore it
        kwargs.pop('token_type_ids', None)
        # Filter kwargs to only include valid arguments for RobertaModel
        model_kwargs = {
            k: v for k, v in kwargs.items() 
            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
                    'output_attentions', 'output_hidden_states']
        }
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
        pooled = self.dropout(outputs.pooler_output)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            B, A, _ = all_logits.size()
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
        
        # T5 returns BaseModelOutput, which has hidden_states
        # But we need to handle it properly
        hidden_states = getattr(outputs, 'hidden_states', None)
        attentions = getattr(outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    def save_pretrained(self, save_directory: str, **kwargs):
        # Ensure directory exists
        os.makedirs(save_directory, exist_ok=True)
        
        # Save backbone
        self.roberta.save_pretrained(save_directory, **kwargs)
        
        # Update and save config with custom attributes
        config = self.roberta.config
        config.num_aspects = len(self.sentiment_classifiers)
        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1  # -1 vì không tính lớp "none"
        # Auto map để AutoModel tự động load đúng class
        # models.py sẽ được upload vào root của repo
        config.auto_map = {
            "AutoModel": "models.TransformerForABSA",
            "AutoModelForSequenceClassification": "models.TransformerForABSA"
        }
        # Lưu thêm thông tin vào config để dễ dàng load lại
        if not hasattr(config, 'custom_model_type'):
            config.custom_model_type = 'TransformerForABSA'
        config.save_pretrained(save_directory, **kwargs)
        
        # Save full state_dict (bao gồm cả sentiment_classifiers)
        sd = kwargs.get("state_dict", None) or self.state_dict()
        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        
        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
        if num_aspects is None:
            num_aspects = getattr(config, 'num_aspects', None)
            if num_aspects is None:
                raise ValueError("num_aspects must be provided or present in config")
        
        if num_sentiments is None:
            num_sentiments = getattr(config, 'num_sentiments', None)
            if num_sentiments is None:
                raise ValueError("num_sentiments must be provided or present in config")
        
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        
        model = cls(config)
        
        # Load backbone weights
        model.roberta = RobertaModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
        )
        
        # Load full state_dict nếu có (bao gồm sentiment_classifiers)
        try:
            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
            if os.path.exists(state_dict_path):
                state_dict = torch.load(state_dict_path, map_location="cpu")
                model.load_state_dict(state_dict, strict=False)
            elif "state_dict" in kwargs:
                model.load_state_dict(kwargs["state_dict"], strict=False)
        except Exception as e:
            print(f"⚠ Warning: Could not load full state_dict: {e}")
        
        return model


class BERTForABSA(BertPreTrainedModel):
    """BERT-based model (cho mBERT)"""
    def __init__(self, config):
        super().__init__(config)
        self.bert = BertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.hidden_size, config.num_sentiments + 1)
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, token_type_ids=None, **kwargs):
        # BERT models can use token_type_ids, but for single sentence tasks, it's usually all zeros
        # Filter kwargs to only include valid arguments for BertModel
        model_kwargs = {
            k: v for k, v in kwargs.items() 
            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
                    'output_attentions', 'output_hidden_states']
        }
        # BERT expects token_type_ids, but if not provided, it will default to all zeros
        if token_type_ids is not None:
            model_kwargs['token_type_ids'] = token_type_ids
        
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.bert(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
        pooled = self.dropout(outputs.pooler_output)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
        
        # T5 returns BaseModelOutput, which has hidden_states
        # But we need to handle it properly
        hidden_states = getattr(outputs, 'hidden_states', None)
        attentions = getattr(outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    def save_pretrained(self, save_directory: str, **kwargs):
        """Save model with custom attributes"""
        os.makedirs(save_directory, exist_ok=True)
        self.bert.save_pretrained(save_directory, **kwargs)
        config = self.bert.config
        config.num_aspects = len(self.sentiment_classifiers)
        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
        config.auto_map = {
            "AutoModel": "models.BERTForABSA",
            "AutoModelForSequenceClassification": "models.BERTForABSA"
        }
        if not hasattr(config, 'custom_model_type'):
            config.custom_model_type = 'BERTForABSA'
        config.save_pretrained(save_directory, **kwargs)
        sd = kwargs.get("state_dict", None) or self.state_dict()
        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        
        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
        if num_aspects is None:
            num_aspects = getattr(config, 'num_aspects', None)
            if num_aspects is None:
                raise ValueError("num_aspects must be provided or present in config")
        
        if num_sentiments is None:
            num_sentiments = getattr(config, 'num_sentiments', None)
            if num_sentiments is None:
                raise ValueError("num_sentiments must be provided or present in config")
        
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        model = cls(config)
        model.bert = BertModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
        )
        
        # Load full state_dict if available
        try:
            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
            if os.path.exists(state_dict_path):
                state_dict = torch.load(state_dict_path, map_location="cpu")
                model.load_state_dict(state_dict, strict=False)
            elif "state_dict" in kwargs:
                model.load_state_dict(kwargs["state_dict"], strict=False)
        except Exception as e:
            print(f"⚠ Warning: Could not load full state_dict: {e}")
        
        return model


class XLMRobertaForABSA(XLMRobertaPreTrainedModel):
    """XLM-RoBERTa-based model"""
    def __init__(self, config):
        super().__init__(config)
        self.roberta = XLMRobertaModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.hidden_size, config.num_sentiments + 1)
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict)
        pooled = self.dropout(outputs.pooler_output)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
        
        # T5 returns BaseModelOutput, which has hidden_states
        # But we need to handle it properly
        hidden_states = getattr(outputs, 'hidden_states', None)
        attentions = getattr(outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int, num_sentiments: int, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        model = cls(config)
        model.roberta = XLMRobertaModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config",)},
        )
        return model


class RoBERTaGRUForABSA(RobertaPreTrainedModel):
    """RoBERTa + GRU hybrid model"""
    base_model_prefix = "roberta"

    def __init__(self, config):
        super().__init__(config)
        self.roberta = RobertaModel(config)
        self.gru = nn.GRU(
            config.hidden_size, config.hidden_size, 
            num_layers=2, batch_first=True, bidirectional=True, dropout=0.2
        )
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.hidden_size * 2, config.num_sentiments + 1)  # *2 vì bidirectional
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.roberta(input_ids, attention_mask=attention_mask, return_dict=return_dict)
        
        # Use last_hidden_state thay vì pooler_output
        sequence_output = outputs.last_hidden_state  # [B, L, H]
        
        # GRU layer
        gru_out, _ = self.gru(sequence_output)  # [B, L, H*2]
        # Take last timestep
        pooled = self.dropout(gru_out[:, -1, :])  # [B, H*2]
        
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
        
        # T5 returns BaseModelOutput, which has hidden_states
        # But we need to handle it properly
        hidden_states = getattr(outputs, 'hidden_states', None)
        attentions = getattr(outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int, num_sentiments: int, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        model = cls(config)
        model.roberta = RobertaModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config",)},
        )
        return model


class BartForABSA(BartPreTrainedModel):
    """BART-based model (cho BartPho)"""
    def __init__(self, config):
        super().__init__(config)
        self.model = BartModel(config)
        self.dropout = nn.Dropout(config.dropout)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.d_model, config.num_sentiments + 1)
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
        # BART models don't use token_type_ids, so we ignore it
        kwargs.pop('token_type_ids', None)
        # Filter kwargs to only include valid arguments for BartModel
        # Remove training-specific arguments that BartModel doesn't accept
        model_kwargs = {
            k: v for k, v in kwargs.items() 
            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
                    'output_attentions', 'output_hidden_states']
        }
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        
        # IMPORTANT: For BART, we need to access encoder output directly
        # BartModel.forward() returns decoder output in last_hidden_state
        # We need to call encoder separately to get encoder hidden states
        # Only call encoder once (don't call full model.forward() to avoid double computation)
        encoder_outputs = self.model.get_encoder()(
            input_ids, 
            attention_mask=attention_mask, 
            return_dict=True,
            **{k: v for k, v in model_kwargs.items()}
        )
        sequence_output = encoder_outputs.last_hidden_state  # [B, L, H] - encoder output
        
        # Mean pooling with attention mask (weighted mean to avoid padding tokens)
        if attention_mask is not None:
            # Expand attention mask to match sequence_output dimensions
            attention_mask_expanded = attention_mask.unsqueeze(-1).expand(sequence_output.size()).float()
            # Sum over sequence length, divide by number of non-padding tokens
            sum_embeddings = torch.sum(sequence_output * attention_mask_expanded, dim=1)
            sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
            pooled = sum_embeddings / sum_mask  # [B, H]
        else:
            pooled = sequence_output.mean(dim=1)  # [B, H]
        
        pooled = self.dropout(pooled)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + (encoder_outputs.hidden_states, encoder_outputs.attentions)) if loss is not None else (all_logits,)
        
        # Use encoder outputs for hidden_states and attentions
        hidden_states = getattr(encoder_outputs, 'hidden_states', None)
        attentions = getattr(encoder_outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    def save_pretrained(self, save_directory: str, **kwargs):
        """Save model with custom attributes"""
        os.makedirs(save_directory, exist_ok=True)
        self.model.save_pretrained(save_directory, **kwargs)
        config = self.model.config
        config.num_aspects = len(self.sentiment_classifiers)
        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
        config.auto_map = {
            "AutoModel": "models.BartForABSA",
            "AutoModelForSequenceClassification": "models.BartForABSA"
        }
        if not hasattr(config, 'custom_model_type'):
            config.custom_model_type = 'BartForABSA'
        config.save_pretrained(save_directory, **kwargs)
        sd = kwargs.get("state_dict", None) or self.state_dict()
        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        
        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
        if num_aspects is None:
            num_aspects = getattr(config, 'num_aspects', None)
            if num_aspects is None:
                raise ValueError("num_aspects must be provided or present in config")
        
        if num_sentiments is None:
            num_sentiments = getattr(config, 'num_sentiments', None)
            if num_sentiments is None:
                raise ValueError("num_sentiments must be provided or present in config")
        
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        model = cls(config)
        model.model = BartModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
        )
        
        # Load full state_dict if available
        try:
            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
            if os.path.exists(state_dict_path):
                state_dict = torch.load(state_dict_path, map_location="cpu")
                model.load_state_dict(state_dict, strict=False)
            elif "state_dict" in kwargs:
                model.load_state_dict(kwargs["state_dict"], strict=False)
        except Exception as e:
            print(f"⚠ Warning: Could not load full state_dict: {e}")
        
        return model


class T5ForABSA(T5PreTrainedModel):
    """T5-based model (cho ViT5) - sử dụng encoder only"""
    def __init__(self, config):
        super().__init__(config)
        self.encoder = T5EncoderModel(config)
        self.dropout = nn.Dropout(config.dropout_rate)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(config.d_model, config.num_sentiments + 1)
            for _ in range(config.num_aspects)
        ])
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None, labels=None, return_dict=None, **kwargs):
        # T5 models don't use token_type_ids, so we ignore it
        kwargs.pop('token_type_ids', None)
        # Filter kwargs to only include valid arguments for T5EncoderModel
        model_kwargs = {
            k: v for k, v in kwargs.items() 
            if k in ['position_ids', 'head_mask', 'inputs_embeds', 
                    'output_attentions', 'output_hidden_states']
        }
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        outputs = self.encoder(input_ids, attention_mask=attention_mask, return_dict=return_dict, **model_kwargs)
        
        # Mean pooling with attention mask (weighted mean to avoid padding tokens)
        sequence_output = outputs.last_hidden_state  # [B, L, H]
        if attention_mask is not None:
            # Expand attention mask to match sequence_output dimensions
            attention_mask_expanded = attention_mask.unsqueeze(-1).expand(sequence_output.size()).float()
            # Sum over sequence length, divide by number of non-padding tokens
            sum_embeddings = torch.sum(sequence_output * attention_mask_expanded, dim=1)
            sum_mask = torch.clamp(attention_mask_expanded.sum(dim=1), min=1e-9)
            pooled = sum_embeddings / sum_mask  # [B, H]
        else:
            pooled = sequence_output.mean(dim=1)  # [B, H]
        
        pooled = self.dropout(pooled)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if not return_dict:
            return ((loss, all_logits) + outputs[2:]) if loss is not None else (all_logits,) + outputs[2:]
        
        # T5 returns BaseModelOutput, which has hidden_states
        # But we need to handle it properly
        hidden_states = getattr(outputs, 'hidden_states', None)
        attentions = getattr(outputs, 'attentions', None)
        
        return SequenceClassifierOutput(
            loss=loss, logits=all_logits,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    def save_pretrained(self, save_directory: str, **kwargs):
        """Save model with custom attributes"""
        os.makedirs(save_directory, exist_ok=True)
        self.encoder.save_pretrained(save_directory, **kwargs)
        config = self.encoder.config
        config.num_aspects = len(self.sentiment_classifiers)
        config.num_sentiments = self.sentiment_classifiers[0].out_features - 1
        config.auto_map = {
            "AutoModel": "models.T5ForABSA",
            "AutoModelForSequenceClassification": "models.T5ForABSA"
        }
        if not hasattr(config, 'custom_model_type'):
            config.custom_model_type = 'T5ForABSA'
        config.save_pretrained(save_directory, **kwargs)
        sd = kwargs.get("state_dict", None) or self.state_dict()
        torch.save(sd, os.path.join(save_directory, "pytorch_model.bin"))

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, num_aspects: int = None, num_sentiments: int = None, **kwargs):
        config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
        
        # Nếu num_aspects và num_sentiments không được truyền vào, đọc từ config
        if num_aspects is None:
            num_aspects = getattr(config, 'num_aspects', None)
            if num_aspects is None:
                raise ValueError("num_aspects must be provided or present in config")
        
        if num_sentiments is None:
            num_sentiments = getattr(config, 'num_sentiments', None)
            if num_sentiments is None:
                raise ValueError("num_sentiments must be provided or present in config")
        
        config.num_aspects = num_aspects
        config.num_sentiments = num_sentiments
        model = cls(config)
        model.encoder = T5EncoderModel.from_pretrained(
            pretrained_model_name_or_path, config=config,
            **{k: v for k, v in kwargs.items() if k not in ("config", "state_dict")},
        )
        
        # Load full state_dict if available
        try:
            state_dict_path = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
            if os.path.exists(state_dict_path):
                state_dict = torch.load(state_dict_path, map_location="cpu")
                model.load_state_dict(state_dict, strict=False)
            elif "state_dict" in kwargs:
                model.load_state_dict(kwargs["state_dict"], strict=False)
        except Exception as e:
            print(f"⚠ Warning: Could not load full state_dict: {e}")
        
        return model


# ========== Non-Transformer Models ==========

class TextCNNForABSA(nn.Module):
    """TextCNN model - không dùng transformers"""
    def __init__(self, vocab_size, embed_dim, num_filters, filter_sizes, num_aspects, num_sentiments, max_length=256):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.convs = nn.ModuleList([
            nn.Conv1d(embed_dim, num_filters, kernel_size=fs)
            for fs in filter_sizes
        ])
        self.dropout = nn.Dropout(0.5)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(len(filter_sizes) * num_filters, num_sentiments + 1)
            for _ in range(num_aspects)
        ])

    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True):
        # input_ids: [B, L]
        x = self.embedding(input_ids)  # [B, L, E]
        x = x.permute(0, 2, 1)  # [B, E, L]
        
        conv_outputs = []
        for conv in self.convs:
            conv_out = F.relu(conv(x))  # [B, F, L']
            pooled = F.max_pool1d(conv_out, kernel_size=conv_out.size(2))  # [B, F, 1]
            conv_outputs.append(pooled.squeeze(2))  # [B, F]
        
        x = torch.cat(conv_outputs, dim=1)  # [B, F*len(filter_sizes)]
        x = self.dropout(x)
        
        all_logits = torch.stack([cls(x) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if return_dict:
            return SequenceClassifierOutput(
                loss=loss, logits=all_logits,
                hidden_states=None, attentions=None
            )
        return (loss, all_logits) if loss is not None else (all_logits,)


class BiLSTMForABSA(nn.Module):
    """BiLSTM model - không dùng transformers"""
    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, num_aspects, num_sentiments, dropout=0.3):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.lstm = nn.LSTM(
            embed_dim, hidden_dim, num_layers,
            batch_first=True, bidirectional=True, dropout=dropout
        )
        self.dropout = nn.Dropout(dropout)
        self.sentiment_classifiers = nn.ModuleList([
            nn.Linear(hidden_dim * 2, num_sentiments + 1)  # *2 vì bidirectional
            for _ in range(num_aspects)
        ])

    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True):
        x = self.embedding(input_ids)  # [B, L, E]
        lstm_out, (h_n, c_n) = self.lstm(x)  # [B, L, H*2]
        
        # Use last non-padding hidden state instead of always using last timestep
        # This is important because padding tokens can be at the end
        if attention_mask is not None:
            # Find the last non-padding token for each sequence
            # attention_mask: [B, L] where 1 = real token, 0 = padding
            seq_lengths = attention_mask.sum(dim=1) - 1  # -1 for 0-indexing
            # Ensure seq_lengths are within valid range
            seq_lengths = torch.clamp(seq_lengths, min=0, max=lstm_out.size(1) - 1)
            # Get last hidden state for each sequence: [B, H*2]
            batch_size = lstm_out.size(0)
            pooled = lstm_out[torch.arange(batch_size, device=lstm_out.device), seq_lengths, :]
        else:
            # Fallback: use last timestep if no attention mask
            pooled = lstm_out[:, -1, :]  # [B, H*2]
        
        pooled = self.dropout(pooled)
        all_logits = torch.stack([cls(pooled) for cls in self.sentiment_classifiers], dim=1)

        loss = None
        if labels is not None:
            logits_flat = all_logits.view(-1, all_logits.size(-1))
            targets_flat = labels.view(-1)
            loss = nn.CrossEntropyLoss()(logits_flat, targets_flat)

        if return_dict:
            return SequenceClassifierOutput(
                loss=loss, logits=all_logits,
                hidden_states=None, attentions=None
            )
        return (loss, all_logits) if loss is not None else (all_logits,)


# ========== Model Factory ==========

def get_model_class(model_name: str):
    """Factory function để lấy model class dựa trên model name"""
    model_name_lower = model_name.lower()
    
    # RoBERTa-GRU (check first to avoid confusion)
    if 'roberta' in model_name_lower and ('gru' in model_name_lower or 'roberta-base-gru' in model_name_lower):
        return RoBERTaGRUForABSA
    
    # Roberta-based (PhoBERT v1/v2, ViSoBERT)
    if 'phobert' in model_name_lower or 'visobert' in model_name_lower or 'roberta' in model_name_lower:
        return TransformerForABSA
    
    # XLM-RoBERTa
    elif 'xlm-roberta' in model_name_lower or 'xlm_roberta' in model_name_lower:
        return XLMRobertaForABSA
    
    # BERT
    elif 'bert' in model_name_lower and 'roberta' not in model_name_lower:
        return BERTForABSA
    
    # BART
    elif 'bart' in model_name_lower:
        return BartForABSA
    
    # T5
    elif 't5' in model_name_lower or 'vit5' in model_name_lower:
        return T5ForABSA
    
    # TextCNN
    elif 'textcnn' in model_name_lower or 'cnn' in model_name_lower:
        return TextCNNForABSA
    
    # BiLSTM
    elif 'bilstm' in model_name_lower or 'lstm' in model_name_lower:
        return BiLSTMForABSA
    
    # Default: try Roberta
    else:
        return TransformerForABSA


def create_model(model_name: str, num_aspects: int, num_sentiments: int, vocab_size=None, **kwargs):
    """
    Create model instance dựa trên model name
    
    Args:
        model_name: Tên model hoặc model ID từ Hugging Face
        num_aspects: Số lượng aspects
        num_sentiments: Số lượng sentiment classes
        vocab_size: Vocabulary size (chỉ cần cho TextCNN/BiLSTM)
        **kwargs: Additional arguments
    """
    model_class = get_model_class(model_name)
    
    # RoBERTa-GRU cần base model riêng
    if model_class == RoBERTaGRUForABSA:
        # Use roberta-base as base model for RoBERTa-GRU
        base_model_name = 'roberta-base'
        return model_class.from_pretrained(
            base_model_name,
            num_aspects=num_aspects,
            num_sentiments=num_sentiments,
            trust_remote_code=True,
            **kwargs
        )
    
    # Non-transformer models
    if model_class in [TextCNNForABSA, BiLSTMForABSA]:
        if vocab_size is None:
            raise ValueError(f"vocab_size is required for {model_class.__name__}")
        
        if model_class == TextCNNForABSA:
            return TextCNNForABSA(
                vocab_size=vocab_size,
                embed_dim=kwargs.get('embed_dim', 300),
                num_filters=kwargs.get('num_filters', 100),
                filter_sizes=kwargs.get('filter_sizes', [3, 4, 5]),
                num_aspects=num_aspects,
                num_sentiments=num_sentiments,
                max_length=kwargs.get('max_length', 256)
            )
        elif model_class == BiLSTMForABSA:
            return BiLSTMForABSA(
                vocab_size=vocab_size,
                embed_dim=kwargs.get('embed_dim', 300),
                hidden_dim=kwargs.get('hidden_dim', 256),
                num_layers=kwargs.get('num_layers', 2),
                num_aspects=num_aspects,
                num_sentiments=num_sentiments,
                dropout=kwargs.get('dropout', 0.3)
            )
    
    # Transformer models
    else:
        return model_class.from_pretrained(
            model_name,
            num_aspects=num_aspects,
            num_sentiments=num_sentiments,
            trust_remote_code=True,
            **kwargs
        )