"""推文情感分析模型训练和加载模块

实现基于 TF-IDF + LightGBM 的情感分类模型。
"""

from pathlib import Path
from typing import Optional

import numpy as np
import polars as pl
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import LabelEncoder
import lightgbm as lgb
import joblib


class TweetSentimentModel:
    """推文情感分类模型
    
    使用 TF-IDF 特征提取和 LightGBM 分类器。
    """
    
    def __init__(
        self,
        tfidf_vectorizer: Optional[TfidfVectorizer] = None,
        label_encoder: Optional[LabelEncoder] = None,
        airline_encoder: Optional[LabelEncoder] = None,
        classifier: Optional[lgb.LGBMClassifier] = None,
    ):
        """初始化模型
        
        Args:
            tfidf_vectorizer: TF-IDF 向量化器
            label_encoder: 情感标签编码器
            airline_encoder: 航空公司编码器
            classifier: LightGBM 分类器
        """
        self.tfidf_vectorizer = tfidf_vectorizer or TfidfVectorizer(
            max_features=5000,
            ngram_range=(1, 2),
            min_df=2,
            max_df=0.95,
        )
        self.label_encoder = label_encoder or LabelEncoder()
        self.airline_encoder = airline_encoder or LabelEncoder()
        self.classifier = classifier or lgb.LGBMClassifier(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=6,
            random_state=42,
            verbose=-1,
        )
        self._is_fitted = False
    
    def fit(self, texts: np.ndarray, airlines: np.ndarray, sentiments: np.ndarray) -> "TweetSentimentModel":
        """训练模型
        
        Args:
            texts: 推文文本数组
            airlines: 航空公司数组
            sentiments: 情感标签数组
            
        Returns:
            训练好的模型
        """
        # 编码标签
        self.label_encoder.fit(sentiments)
        y = self.label_encoder.transform(sentiments)
        
        # 编码航空公司
        self.airline_encoder.fit(airlines)
        airline_encoded = self.airline_encoder.transform(airlines)
        
        # TF-IDF 特征提取
        X_text = self.tfidf_vectorizer.fit_transform(texts)
        
        # 合并特征
        airline_features = airline_encoded.reshape(-1, 1)
        X = self._combine_features(X_text, airline_features)
        
        # 训练分类器
        self.classifier.fit(X, y)
        
        self._is_fitted = True
        return self
    
    def predict(self, texts: np.ndarray, airlines: np.ndarray) -> np.ndarray:
        """预测情感标签
        
        Args:
            texts: 推文文本数组
            airlines: 航空公司数组
            
        Returns:
            预测的情感标签数组
        """
        if not self._is_fitted:
            raise ValueError("模型尚未训练，请先调用 fit() 方法")
        
        # TF-IDF 特征提取
        X_text = self.tfidf_vectorizer.transform(texts)
        
        # 编码航空公司
        airline_encoded = self.airline_encoder.transform(airlines)
        airline_features = airline_encoded.reshape(-1, 1)
        
        # 合并特征
        X = self._combine_features(X_text, airline_features)
        
        # 预测
        y_pred = self.classifier.predict(X)
        
        # 解码标签
        return self.label_encoder.inverse_transform(y_pred)
    
    def predict_proba(self, texts: np.ndarray, airlines: np.ndarray) -> np.ndarray:
        """预测情感概率
        
        Args:
            texts: 推文文本数组
            airlines: 航空公司数组
            
        Returns:
            预测的概率数组 (n_samples, n_classes)
        """
        if not self._is_fitted:
            raise ValueError("模型尚未训练，请先调用 fit() 方法")
        
        # TF-IDF 特征提取
        X_text = self.tfidf_vectorizer.transform(texts)
        
        # 编码航空公司
        airline_encoded = self.airline_encoder.transform(airlines)
        airline_features = airline_encoded.reshape(-1, 1)
        
        # 合并特征
        X = self._combine_features(X_text, airline_features)
        
        # 预测概率
        return self.classifier.predict_proba(X)
    
    def _combine_features(self, text_features: np.ndarray, airline_features: np.ndarray) -> np.ndarray:
        """合并文本特征和航空公司特征
        
        Args:
            text_features: TF-IDF 文本特征
            airline_features: 航空公司特征
            
        Returns:
            合并后的特征矩阵
        """
        from scipy.sparse import hstack
        return hstack([text_features, airline_features])
    
    def save(self, path: Path) -> None:
        """保存模型
        
        Args:
            path: 保存路径
        """
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        model_data = {
            "tfidf_vectorizer": self.tfidf_vectorizer,
            "label_encoder": self.label_encoder,
            "airline_encoder": self.airline_encoder,
            "classifier": self.classifier,
            "is_fitted": self._is_fitted,
        }
        
        joblib.dump(model_data, path)
    
    @classmethod
    def load(cls, path: Path) -> "TweetSentimentModel":
        """加载模型
        
        Args:
            path: 模型路径
            
        Returns:
            加载的模型
        """
        model_data = joblib.load(path)
        
        model = cls(
            tfidf_vectorizer=model_data["tfidf_vectorizer"],
            label_encoder=model_data["label_encoder"],
            airline_encoder=model_data["airline_encoder"],
            classifier=model_data["classifier"],
        )
        model._is_fitted = model_data["is_fitted"]
        
        return model


def load_model(model_path: Optional[Path] = None) -> TweetSentimentModel:
    """加载预训练模型
    
    Args:
        model_path: 模型路径（可选，默认使用示例模型）
        
    Returns:
        加载的模型
    """
    if model_path is not None and model_path.exists():
        return TweetSentimentModel.load(model_path)
    
    # 创建并返回一个示例模型（使用示例数据训练）
    model = _create_example_model()
    return model


def _create_example_model() -> TweetSentimentModel:
    """创建示例模型（使用示例数据训练）
    
    Returns:
        训练好的示例模型
    """
    # 示例数据
    texts = np.array([
        "@United This is the worst airline ever! My flight was delayed for 5 hours and no one helped!",
        "@Southwest Thank you for the amazing flight! The crew was so helpful and friendly.",
        "@American What is the baggage policy for international flights?",
        "@Delta Terrible service! Lost my luggage and no response from customer support.",
        "@JetBlue Great experience! On time departure and friendly staff.",
        "@United Why is my flight cancelled again? This is unacceptable!",
        "@Southwest Love the free snacks and great customer service!",
        "@American Can you help me with my booking?",
        "@Delta Worst experience ever! Will never fly again!",
        "@JetBlue Thank you for the smooth flight and excellent service!",
    ])
    
    airlines = np.array([
        "United",
        "Southwest",
        "American",
        "Delta",
        "JetBlue",
        "United",
        "Southwest",
        "American",
        "Delta",
        "JetBlue",
    ])
    
    sentiments = np.array([
        "negative",
        "positive",
        "neutral",
        "negative",
        "positive",
        "negative",
        "positive",
        "neutral",
        "negative",
        "positive",
    ])
    
    # 训练模型
    model = TweetSentimentModel()
    model.fit(texts, airlines, sentiments)
    
    return model


if __name__ == "__main__":
    # 示例：加载模型并进行预测
    print("加载模型...")
    model = load_model()
    
    print("\n测试预测...")
    test_texts = np.array([
        "@United This is terrible!",
        "@Southwest Thank you so much!",
        "@American How do I check in?",
    ])
    test_airlines = np.array(["United", "Southwest", "American"])
    
    predictions = model.predict(test_texts, test_airlines)
    probabilities = model.predict_proba(test_texts, test_airlines)
    
    for text, airline, pred, prob in zip(test_texts, test_airlines, predictions, probabilities):
        print(f"\n文本: {text}")
        print(f"航空公司: {airline}")
        print(f"预测: {pred}")
        print(f"概率: {prob}")
