"""
航空公司情感分析机器学习模块
实现传统机器学习模型，用于客户流失预测
"""

import pandas as pd
import numpy as np
import re
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.preprocessing import LabelEncoder, OneHotEncoder
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
from sklearn.pipeline import Pipeline
from sklearn.compose import ColumnTransformer
import xgboost as xgb
import joblib
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
import warnings
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False

class FeatureEngineering:
    """特征工程类"""
    
    def __init__(self):
        self.text_vectorizer = None
        self.label_encoders = {}
        self.onehot_encoder = None
        
    def extract_text_features(self, text_series):
        """提取文本统计特征"""
        features = []
        
        for text in text_series:
            if pd.isna(text):
                features.append([0, 0, 0, 0, 0, 0])
                continue
                
            # 文本长度
            text_length = len(text)
            
            # 词数
            word_count = len(text.split())
            
            # 标点比例
            punctuation_count = len(re.findall(r'[!"#$%&\'()*+,-./:;<=>?@[\\]^_`{|}~]', text))
            punctuation_ratio = punctuation_count / text_length if text_length > 0 else 0
            
            # 大写字母比例
            uppercase_count = len(re.findall(r'[A-Z]', text))
            uppercase_ratio = uppercase_count / text_length if text_length > 0 else 0
            
            # 社交媒体特征
            mention_count = len(re.findall(r'@\w+', text))
            hashtag_count = len(re.findall(r'#\w+', text))
            url_count = len(re.findall(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', text))
            
            features.append([text_length, word_count, punctuation_ratio, uppercase_ratio, mention_count, hashtag_count, url_count])
        
        return pd.DataFrame(features, columns=[
            'text_length', 'word_count', 'punctuation_ratio', 'uppercase_ratio', 
            'mention_count', 'hashtag_count', 'url_count'
        ])
    
    def extract_temporal_features(self, datetime_series):
        """提取时间特征"""
        features = []
        
        for dt_str in datetime_series:
            if pd.isna(dt_str):
                features.append([0, 0])
                continue
                
            try:
                dt = pd.to_datetime(dt_str)
                hour = dt.hour
                weekday = dt.weekday()
                features.append([hour, weekday])
            except:
                features.append([0, 0])
        
        return pd.DataFrame(features, columns=['hour', 'weekday'])
    
    def fit_transform(self, df):
        """拟合并转换特征"""
        
        # 文本特征 - TF-IDF
        self.text_vectorizer = TfidfVectorizer(
            ngram_range=(1, 2),
            max_features=5000,
            stop_words='english'
        )
        tfidf_features = self.text_vectorizer.fit_transform(df['text'].fillna(''))
        
        # 文本统计特征
        text_stat_features = self.extract_text_features(df['text'])
        
        # 结构化特征编码
        # 航空公司 - one-hot编码
        airline_encoded = pd.get_dummies(df['airline'], prefix='airline')
        
        # 负面原因 - 多标签特征
        negative_reason_encoded = pd.get_dummies(df['negativereason'].fillna('No_Reason'), prefix='neg_reason')
        
        # 时间特征
        temporal_features = self.extract_temporal_features(df['tweet_created'])
        
        # 合并所有特征
        feature_list = [
            tfidf_features,
            text_stat_features,
            airline_encoded,
            negative_reason_encoded,
            temporal_features,
            df[['retweet_count', 'airline_sentiment_confidence', 'negativereason_confidence']].fillna(0)
        ]
        
        # 转换为稀疏矩阵并合并
        from scipy.sparse import hstack
        sparse_features = hstack([feature_list[0]] + [pd.DataFrame(f).values for f in feature_list[1:]])
        
        return sparse_features
    
    def transform(self, df):
        """转换新数据"""
        if self.text_vectorizer is None:
            raise ValueError("必须先调用fit_transform方法")
            
        tfidf_features = self.text_vectorizer.transform(df['text'].fillna(''))
        text_stat_features = self.extract_text_features(df['text'])
        
        airline_encoded = pd.get_dummies(df['airline'], prefix='airline')
        # 确保列对齐
        expected_airline_cols = [f'airline_{airline}' for airline in self.text_vectorizer.get_feature_names_out()[:6]]
        for col in expected_airline_cols:
            if col not in airline_encoded.columns:
                airline_encoded[col] = 0
        airline_encoded = airline_encoded[expected_airline_cols]
        
        negative_reason_encoded = pd.get_dummies(df['negativereason'].fillna('No_Reason'), prefix='neg_reason')
        temporal_features = self.extract_temporal_features(df['tweet_created'])
        
        feature_list = [
            tfidf_features,
            text_stat_features,
            airline_encoded,
            negative_reason_encoded,
            temporal_features,
            df[['retweet_count', 'airline_sentiment_confidence', 'negativereason_confidence']].fillna(0)
        ]
        
        from scipy.sparse import hstack
        return hstack([feature_list[0]] + [pd.DataFrame(f).values for f in feature_list[1:]])

class MachineLearningPipeline:
    """机器学习管道类"""
    
    def __init__(self):
        self.feature_engineer = FeatureEngineering()
        self.models = {}
        self.results = {}
        self.target_encoder = LabelEncoder()
        
    def prepare_data(self, df):
        """准备训练数据"""
        # 业务关注：判断客户是否会流失（负面情感）
        # 将情感分类转换为二分类问题：负面 vs 非负面
        df = df.copy()
        df['target'] = (df['airline_sentiment'] == 'negative').astype(int)
        
        # 去重处理：基于文本内容去重，避免数据泄露
        df_unique = df.drop_duplicates(subset=['text'])
        print(f'去重后样本数: {len(df_unique)} (原样本数: {len(df)})')
        
        # 先划分数据，避免数据泄露
        train_df, test_df = train_test_split(
            df_unique, test_size=0.2, random_state=42, stratify=df_unique['target']
        )
        
        # 在训练集上拟合特征工程
        X_train = self.feature_engineer.fit_transform(train_df)
        y_train = self.target_encoder.fit_transform(train_df['target'])
        
        # 在测试集上转换特征
        X_test = self.feature_engineer.transform(test_df)
        y_test = self.target_encoder.transform(test_df['target'])
        
        return X_train, X_test, y_train, y_test, train_df, test_df
    
    def train_models(self, df, test_size=0.2, random_state=42):
        """训练所有模型"""
        X_train, X_test, y_train, y_test, train_df, test_df = self.prepare_data(df)
        
        # 定义模型
        model_configs = {
            'LogisticRegression': {
                'model': LogisticRegression(random_state=random_state, max_iter=1000),
                'params': {'C': 1.0}
            },
            'RandomForest': {
                'model': RandomForestClassifier(random_state=random_state, n_estimators=100),
                'params': {}
            },
            'SVM': {
                'model': SVC(kernel='linear', random_state=random_state, probability=True),
                'params': {'C': 1.0}
            },
            'XGBoost': {
                'model': xgb.XGBClassifier(random_state=random_state, n_estimators=100),
                'params': {}
            }
        }
        
        # 训练和评估每个模型
        for name, config in model_configs.items():
            print(f"训练模型: {name}")
            
            model = config['model']
            model.fit(X_train, y_train)
            
            # 预测
            y_pred = model.predict(X_test)
            y_pred_proba = model.predict_proba(X_test)[:, 1]
            
            # 计算指标
            accuracy = accuracy_score(y_test, y_pred)
            macro_f1 = f1_score(y_test, y_pred, average='macro')
            
            # 重点关注假阴性（漏判负面情感）
            tn, fp, fn, tp = confusion_matrix(y_test, y_pred).ravel()
            false_negative_rate = fn / (fn + tp) if (fn + tp) > 0 else 0
            
            self.models[name] = model
            self.results[name] = {
                'accuracy': accuracy,
                'macro_f1': macro_f1,
                'false_negative_rate': false_negative_rate,
                'confusion_matrix': (tn, fp, fn, tp),
                'classification_report': classification_report(y_test, y_pred, output_dict=True)
            }
            
            print(f"{name} - 准确率: {accuracy:.4f}, Macro-F1: {macro_f1:.4f}, 假阴性率: {false_negative_rate:.4f}")
        
        return self.results
    
    def cross_validation(self, df, cv=5):
        """交叉验证"""
        X, y, _ = self.prepare_data(df)
        
        cv_results = {}
        for name, model in self.models.items():
            cv_scores = cross_val_score(model, X, y, cv=cv, scoring='f1_macro')
            cv_results[name] = {
                'mean_f1': cv_scores.mean(),
                'std_f1': cv_scores.std(),
                'scores': cv_scores
            }
        
        return cv_results
    
    def save_models(self, filepath="models/"):
        """保存模型"""
        import os
        os.makedirs(filepath, exist_ok=True)
        
        for name, model in self.models.items():
            joblib.dump(model, f"{filepath}{name}_model.pkl")
        
        # 保存特征工程器
        joblib.dump(self.feature_engineer, f"{filepath}feature_engineer.pkl")
        joblib.dump(self.target_encoder, f"{filepath}target_encoder.pkl")
        
        print(f"模型已保存到: {filepath}")
    
    def load_models(self, filepath="models/"):
        """加载模型"""
        import os
        for name in ['LogisticRegression', 'RandomForest', 'SVM', 'XGBoost']:
            model_path = f"{filepath}{name}_model.pkl"
            if os.path.exists(model_path):
                self.models[name] = joblib.load(model_path)
        
        self.feature_engineer = joblib.load(f"{filepath}feature_engineer.pkl")
        self.target_encoder = joblib.load(f"{filepath}target_encoder.pkl")
        
        print("模型加载完成")
    
    def predict_single(self, text, airline, tweet_created, retweet_count=0, 
                      sentiment_confidence=0, negative_reason=None, negative_reason_confidence=0):
        """单条预测"""
        # 创建单条数据
        data = {
            'text': text,
            'airline': airline,
            'tweet_created': tweet_created,
            'retweet_count': retweet_count,
            'airline_sentiment_confidence': sentiment_confidence,
            'negativereason': negative_reason,
            'negativereason_confidence': negative_reason_confidence
        }
        
        df_single = pd.DataFrame([data])
        X = self.feature_engineer.transform(df_single)
        
        predictions = {}
        for name, model in self.models.items():
            proba = model.predict_proba(X)[0, 1]  # 负面情感概率
            predictions[name] = {
                'probability': proba,
                'prediction': '负面' if proba > 0.5 else '非负面'
            }
        
        return predictions

def create_visualizations(results):
    """创建可视化图表"""
    fig, axes = plt.subplots(2, 2, figsize=(15, 12))
    
    # 模型性能对比
    models = list(results.keys())
    accuracies = [results[model]['accuracy'] for model in models]
    f1_scores = [results[model]['macro_f1'] for model in models]
    
    axes[0, 0].bar(models, accuracies, color=['skyblue', 'lightgreen', 'lightcoral', 'gold'])
    axes[0, 0].set_title('模型准确率对比')
    axes[0, 0].set_ylabel('准确率')
    
    axes[0, 1].bar(models, f1_scores, color=['skyblue', 'lightgreen', 'lightcoral', 'gold'])
    axes[0, 1].set_title('模型Macro-F1对比')
    axes[0, 1].set_ylabel('Macro-F1')
    
    # 假阴性率对比
    fn_rates = [results[model]['false_negative_rate'] for model in models]
    axes[1, 0].bar(models, fn_rates, color=['red', 'orange', 'yellow', 'green'])
    axes[1, 0].set_title('假阴性率对比 (越低越好)')
    axes[1, 0].set_ylabel('假阴性率')
    
    # 混淆矩阵热力图（使用最佳模型）
    best_model = max(results.items(), key=lambda x: x[1]['macro_f1'])[0]
    tn, fp, fn, tp = results[best_model]['confusion_matrix']
    cm = np.array([[tn, fp], [fn, tp]])
    
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=axes[1, 1])
    axes[1, 1].set_title(f'{best_model} 混淆矩阵')
    axes[1, 1].set_xlabel('预测标签')
    axes[1, 1].set_ylabel('真实标签')
    
    plt.tight_layout()
    return fig