"""航空公司情感分析模型训练模块"""

from pathlib import Path
import joblib
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestClassifier
from sklearn.impute import SimpleImputer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, f1_score, recall_score
from sklearn.model_selection import train_test_split, TimeSeriesSplit
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder, StandardScaler
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.svm import SVC
from xgboost import XGBClassifier
from sklearn.neural_network import MLPClassifier

from src.data import DataProcessor, get_feature_columns

MODELS_DIR = Path("models")
MODEL_PATH = MODELS_DIR / "sentiment_model.pkl"


def get_pipeline(model_type: str = "rf") -> Pipeline:
    """构建sklearn处理流水线
    
    1. 数值特征 -> 缺失填充 (均值) -> 标准化
    2. 类别特征 -> 缺失填充 (众数) -> OneHot编码
    3. 文本特征 -> TF-IDF向量化 (n-gram: 1-2, 最大特征: 5000)
    4. 模型 -> 多种算法对比
    """
    numeric_features, categorical_features = get_feature_columns()
    
    # 数值处理管道
    numeric_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="mean")),
            ("scaler", StandardScaler()),
        ]
    )

    # 类别处理管道
    categorical_transformer = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore")),
        ]
    )

    # 文本处理管道
    text_transformer = TfidfVectorizer(
        max_features=5000,
        ngram_range=(1, 2),
        stop_words='english'
    )

    # 组合预处理
    preprocessor = ColumnTransformer(
        transformers=[
            ("num", numeric_transformer, numeric_features),
            ("cat", categorical_transformer, categorical_features),
            ("text", text_transformer, "cleaned_text"),
        ]
    )

    # 选择模型（多算法对比）
    if model_type == "lr":
        clf = LogisticRegression(random_state=42, max_iter=1000)
    elif model_type == "rf":
        clf = RandomForestClassifier(n_estimators=100, max_depth=10, random_state=42)
    elif model_type == "svm":
        clf = SVC(kernel='linear', probability=True, random_state=42)
    elif model_type == "xgb":
        clf = XGBClassifier(random_state=42, eval_metric='mlogloss')
    elif model_type == "mlp":
        clf = MLPClassifier(hidden_layer_sizes=(100, 50), max_iter=1000, random_state=42)
    else:
        clf = RandomForestClassifier(random_state=42)

    return Pipeline(steps=[("preprocessor", preprocessor), ("classifier", clf)])


def evaluate_model(model, X_test, y_test, model_name: str):
    """评估模型性能"""
    y_pred = model.predict(X_test)
    
    # 计算指标（侧重召回率）
    macro_f1 = f1_score(y_test, y_pred, average='macro')
    recall_negative = recall_score(y_test, y_pred, labels=['negative'], average='macro', zero_division=0)
    
    print(f"[{model_name}]")
    print(f"  Macro-F1: {macro_f1:.4f}")
    print(f"  负面情感召回率: {recall_negative:.4f}")
    
    return macro_f1, recall_negative


def train() -> None:
    """执行训练流程 - 只训练最佳模型"""
    print(">>> 1. 数据准备与时序划分")
    processor = DataProcessor("data/Tweets.csv")
    df = processor.load_data()
    
    # 数据预处理
    df_processed = processor.extract_features(df)
    
    # 准备训练数据
    X = df_processed[['cleaned_text'] + get_feature_columns()[0] + get_feature_columns()[1]]
    y = df_processed['airline_sentiment']
    
    # 时序划分训练测试集（确保验证集时间晚于训练集）
    df_sorted = df_processed.sort_values('tweet_created')
    split_idx = int(0.8 * len(df_sorted))
    
    X_train = X.iloc[:split_idx]
    X_test = X.iloc[split_idx:]
    y_train = y.iloc[:split_idx]
    y_test = y.iloc[split_idx:]
    
    print(f"训练集大小: {X_train.shape}, 测试集大小: {X_test.shape}")
    print(f"时间范围: {df_sorted['tweet_created'].min()} 到 {df_sorted['tweet_created'].max()}")

    print("\n>>> 2. 训练最佳模型 (LogisticRegression)")
    
    # 直接训练最佳模型
    pipe = get_pipeline("lr")
    pipe.fit(X_train, y_train)
    
    macro_f1, recall_negative = evaluate_model(pipe, X_test, y_test, "LogisticRegression")
    
    print("\n>>> 3. 模型详细评估")
    
    # 详细分类报告
    y_pred = pipe.predict(X_test)
    print("\n分类报告:")
    print(classification_report(y_test, y_pred))
    
    print("\n>>> 4. 误差分析")
    
    # 计算错误样本
    error_mask = y_pred != y_test
    error_count = error_mask.sum()
    
    # 假阴性分析（漏判负面）
    fn_mask = (y_test == 'negative') & (y_pred != 'negative')
    fn_count = fn_mask.sum()
    
    print(f"总计错误样本数: {error_count}")
    print(f"假阴性错误（漏判负面）: {fn_count}")
    
    print("\n>>> 5. 保存模型")
    
    # 确保模型目录存在
    MODELS_DIR.mkdir(exist_ok=True)
    
    # 保存模型
    joblib.dump(pipe, MODEL_PATH)
    
    # 保存特征名称（用于预测时特征对齐）
    feature_names = {
        'numeric_features': get_feature_columns()[0],
        'categorical_features': get_feature_columns()[1],
        'text_feature': 'cleaned_text'
    }
    joblib.dump(feature_names, MODELS_DIR / "feature_names.pkl")
    
    print(f"最佳模型 (LogisticRegression) 已保存至 {MODEL_PATH}")
    print(f"最终性能: Macro-F1 = {macro_f1:.4f}")
    print(f"负面情感召回率 = {recall_negative:.4f}")
    
    if macro_f1 >= 0.80:
        print("✅ 达到目标: Macro-F1 ≥ 0.80")
    else:
        print("❌ 未达到目标: Macro-F1 < 0.80")


if __name__ == "__main__":
    train()