import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, MinMaxScaler
import lightgbm as lgb
import joblib
import os

class CrisisIndexCalculator:
    def __init__(self, data_path=None):
        self.model_path = os.path.join(os.path.dirname(__file__), "../final_models/crisis_index_model.pkl")
        self.label_encoder = LabelEncoder()
        self.scaler = MinMaxScaler(feature_range=(0, 100))
        self.reason_encoder = LabelEncoder()
        
        if data_path:
            self.data = self._load_data(data_path)
            self._preprocess_data()
            self.model = self._train_model()
            self._save_model()
        else:
            self._load_model()
    
    def _load_data(self, data_path):
        """加载数据"""
        return pd.read_csv(data_path)
    
    def _preprocess_data(self):
        """预处理数据，准备特征"""
        # 处理情感极性
        self.data['sentiment_numeric'] = self.label_encoder.fit_transform(self.data['airline_sentiment'])
        
        # 处理负面原因（填充缺失值）
        self.data['negativereason'] = self.data['negativereason'].fillna('Unknown')
        self.data['negativereason_numeric'] = self.reason_encoder.fit_transform(self.data['negativereason'])
        
        # 填充置信度缺失值
        self.data['airline_sentiment_confidence'] = self.data['airline_sentiment_confidence'].fillna(0.5)
        self.data['negativereason_confidence'] = self.data['negativereason_confidence'].fillna(0.0)
        
        # 定义特征和目标
        self.features = [
            'airline_sentiment_confidence',
            'negativereason_confidence',
            'sentiment_numeric',
            'negativereason_numeric',
            'retweet_count'
        ]
        
        # 创建目标变量：基于情感和置信度的组合分数
        self.data['crisis_score'] = self._create_crisis_score()
        
    def _create_crisis_score(self):
        """创建危机分数作为模型训练目标"""
        # 基础分数：负面情感权重调整，避免分数过于集中
        base_score = np.where(self.data['airline_sentiment'] == 'negative', 50, 
                              np.where(self.data['airline_sentiment'] == 'neutral', 20, 5))
        
        # 情感置信度调整（降低权重）
        sentiment_adj = self.data['airline_sentiment_confidence'] * 20
        
        # 负面原因置信度调整（降低权重，仅对负面情感有效）
        reason_adj = np.where(self.data['airline_sentiment'] == 'negative', 
                            self.data['negativereason_confidence'] * 15, 0)
        
        # 转发数影响（对数转换，降低权重）
        retweet_adj = np.log1p(self.data['retweet_count']) * 3
        
        # 总分数
        total_score = base_score + sentiment_adj + reason_adj + retweet_adj
        
        # 归一化到0-100
        return np.clip(total_score, 0, 100)
    
    def _train_model(self):
        """训练LightGBM回归模型"""
        X = self.data[self.features]
        y = self.data['crisis_score']
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        # 定义LightGBM模型
        model = lgb.LGBMRegressor(
            objective='regression',
            n_estimators=100,
            max_depth=6,
            learning_rate=0.1,
            subsample=0.8,
            colsample_bytree=0.8,
            random_state=42
        )
        
        # 训练模型
        model.fit(X_train, y_train)
        
        # 评估模型
        train_score = model.score(X_train, y_train)
        test_score = model.score(X_test, y_test)
        print(f"Crisis Index Model - Train R²: {train_score:.4f}, Test R²: {test_score:.4f}")
        
        return model
    
    def _save_model(self):
        """保存模型和编码器"""
        os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
        joblib.dump({
            'model': self.model,
            'label_encoder': self.label_encoder,
            'reason_encoder': self.reason_encoder,
            'features': self.features
        }, self.model_path)
    
    def _load_model(self):
        """加载预训练模型"""
        model_data = joblib.load(self.model_path)
        self.model = model_data['model']
        self.label_encoder = model_data['label_encoder']
        self.reason_encoder = model_data['reason_encoder']
        self.features = model_data['features']
    
    def calculate_crisis_index(self, data):
        """计算给定数据的危机程度"""
        # 预处理输入数据
        processed_data = self._preprocess_input(data)
        
        # 预测危机指数
        crisis_index = self.model.predict(processed_data[self.features])
        
        # 确保在0-100范围内
        return np.clip(crisis_index, 0, 100)
    
    def _preprocess_input(self, data):
        """预处理输入数据"""
        if isinstance(data, pd.DataFrame):
            df = data.copy()
        else:
            df = pd.DataFrame(data)
        
        # 处理情感极性
        df['sentiment_numeric'] = self.label_encoder.transform(df['airline_sentiment'])
        
        # 处理负面原因
        df['negativereason'] = df['negativereason'].fillna('Unknown')
        # 处理未知的负面原因
        known_reasons = set(self.reason_encoder.classes_)
        df['negativereason'] = df['negativereason'].apply(
            lambda x: x if x in known_reasons else 'Unknown'
        )
        df['negativereason_numeric'] = self.reason_encoder.transform(df['negativereason'])
        
        # 填充置信度缺失值
        df['airline_sentiment_confidence'] = df['airline_sentiment_confidence'].fillna(0.5)
        df['negativereason_confidence'] = df['negativereason_confidence'].fillna(0.0)
        
        # 确保转发数存在
        if 'retweet_count' not in df:
            df['retweet_count'] = 0
        
        return df
    
    def get_risk_level(self, crisis_index):
        """根据危机指数返回风险等级"""
        if crisis_index >= 70:
            return "高风险"
        elif crisis_index >= 40:
            return "中风险"
        else:
            return "低风险"

# 示例用法
if __name__ == "__main__":
    print("开始训练危机指数计算模型...")
    # 训练模型
    data_path = os.path.join(os.path.dirname(__file__), "../data/processed_tweets.csv")
    print(f"加载数据: {data_path}")
    calculator = CrisisIndexCalculator(data_path)
    print("模型训练完成!")
    
    # 测试单条数据
    print("\n测试单条数据:")
    test_data = {
        'airline_sentiment': ['negative'],
        'airline_sentiment_confidence': [0.98],
        'negativereason': ['Late Flight'],
        'negativereason_confidence': [0.85],
        'retweet_count': [15]
    }
    
    index = calculator.calculate_crisis_index(test_data)
    risk_level = calculator.get_risk_level(index[0])
    print(f"危机指数: {index[0]:.2f}, 风险等级: {risk_level}")
    
    # 测试中性数据
    print("\n测试中性数据:")
    neutral_data = {
        'airline_sentiment': ['neutral'],
        'airline_sentiment_confidence': [0.75],
        'negativereason': [None],
        'negativereason_confidence': [0.0],
        'retweet_count': [2]
    }
    
    index = calculator.calculate_crisis_index(neutral_data)
    risk_level = calculator.get_risk_level(index[0])
    print(f"危机指数: {index[0]:.2f}, 风险等级: {risk_level}")
