import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.metrics import classification_report, confusion_matrix
import joblib
import os

class PropagationRiskAssessor:
    def __init__(self, data_path=None):
        self.model_path = os.path.join(os.path.dirname(__file__), "../final_models/propagation_risk_model.pkl")
        self.label_encoder = LabelEncoder()
        self.scaler = StandardScaler()
        
        if data_path:
            self.data = self._load_data(data_path)
            self._preprocess_data()
            self._extract_basic_features()
            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['tweet_created'] = pd.to_datetime(self.data['tweet_created'])
        
        # 按时间排序
        self.data = self.data.sort_values('tweet_created')
        
        # 计算基本传播特征
        self.data['hour_of_day'] = self.data['tweet_created'].dt.hour
        self.data['day_of_week'] = self.data['tweet_created'].dt.dayofweek
        self.data['is_weekend'] = (self.data['day_of_week'] >= 5).astype(int)
    
    def _extract_basic_features(self):
        """提取基本特征"""
        # 计算每条推文的传播潜力
        self.data['propagation_potential'] = self.data['retweet_count'] * self.data['airline_sentiment_confidence']
        return self.data
    
    def _create_risk_labels(self):
        """创建传播风险标签"""
        # 处理缺失值
        self.data['retweet_count'] = self.data['retweet_count'].fillna(0)
        self.data['airline_sentiment_confidence'] = self.data['airline_sentiment_confidence'].fillna(0.5)
        
        # 基于转发数和增长率创建风险标签
        self.data['risk_score'] = (
            self.data['retweet_count'] * 0.7 +
            np.where(self.data['airline_sentiment'] == 'negative', 50, 0) * 0.3
        )
        
        # 创建分类标签
        self.data['risk_label'] = pd.cut(
            self.data['risk_score'],
            bins=[0, 10, 50, float('inf')],
            labels=['低风险', '中风险', '高风险']
        )
    
    def _train_model(self):
        """训练随机森林分类模型"""
        # 创建风险标签
        self._create_risk_labels()
        
        # 定义特征
        features = [
            'retweet_count',
            'airline_sentiment_confidence',
            'hour_of_day',
            'day_of_week',
            'is_weekend'
        ]
        
        # 处理情感特征
        self.data['sentiment_numeric'] = self.label_encoder.fit_transform(self.data['airline_sentiment'])
        features.append('sentiment_numeric')
        
        # 处理所有特征的缺失值
        for feature in features:
            if self.data[feature].isnull().any():
                if self.data[feature].dtype in ['int64', 'float64']:
                    self.data[feature] = self.data[feature].fillna(self.data[feature].mean())
                else:
                    self.data[feature] = self.data[feature].fillna(self.data[feature].mode()[0])
        
        X = self.data[features]
        y = self.data['risk_label']
        
        # 处理目标变量中的NaN值
        valid_indices = ~y.isnull()
        X = X[valid_indices]
        y = y[valid_indices]
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
        
        # 标准化特征
        self.scaler.fit(X_train)
        X_train_scaled = self.scaler.transform(X_train)
        X_test_scaled = self.scaler.transform(X_test)
        
        # 训练随机森林模型
        model = RandomForestClassifier(
            n_estimators=100,
            max_depth=8,
            random_state=42,
            class_weight='balanced'
        )
        model.fit(X_train_scaled, y_train)
        
        # 评估模型
        y_pred = model.predict(X_test_scaled)
        print("传播风险评估模型性能:")
        print(classification_report(y_test, y_pred))
        print("混淆矩阵:")
        print(confusion_matrix(y_test, y_pred))
        
        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,
            'scaler': self.scaler
        }, 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.scaler = model_data['scaler']
    
    def assess_risk(self, data):
        """评估传播风险"""
        if isinstance(data, pd.DataFrame):
            df = data.copy()
        else:
            df = pd.DataFrame(data)
        
        # 预处理数据
        if 'tweet_created' in df:
            df['tweet_created'] = pd.to_datetime(df['tweet_created'])
            df['hour_of_day'] = df['tweet_created'].dt.hour
            df['day_of_week'] = df['tweet_created'].dt.dayofweek
            df['is_weekend'] = (df['day_of_week'] >= 5).astype(int)
        else:
            # 使用默认值
            df['hour_of_day'] = 12
            df['day_of_week'] = 0
            df['is_weekend'] = 0
        
        # 处理情感特征
        df['sentiment_numeric'] = self.label_encoder.transform(df['airline_sentiment'])
        
        # 选择特征
        features = [
            'retweet_count',
            'airline_sentiment_confidence',
            'hour_of_day',
            'day_of_week',
            'is_weekend',
            'sentiment_numeric'
        ]
        
        # 标准化特征
        X_scaled = self.scaler.transform(df[features])
        
        # 预测风险等级
        risk_prediction = self.model.predict(X_scaled)
        risk_probabilities = self.model.predict_proba(X_scaled)
        
        return {
            'risk_level': risk_prediction,
            'risk_probabilities': risk_probabilities
        }
    


# 示例用法
if __name__ == "__main__":
    print("开始训练传播风险评估模型...")
    data_path = os.path.join(os.path.dirname(__file__), "../data/processed_tweets.csv")
    print(f"加载数据: {data_path}")
    assessor = PropagationRiskAssessor(data_path)
    print("模型训练完成!")
    
    # 测试单条数据
    print("\n测试单条数据:")
    test_data = {
        'retweet_count': [25],
        'airline_sentiment': ['negative'],
        'airline_sentiment_confidence': [0.95],
        'tweet_created': ['2025-01-14 14:30:00']
    }
    
    result = assessor.assess_risk(test_data)
    print(f"传播风险等级: {result['risk_level'][0]}")
    print(f"风险概率: {result['risk_probabilities'][0]}")
    
    # 测试低风险数据
    print("\n测试低风险数据:")
    low_risk_data = {
        'retweet_count': [1],
        'airline_sentiment': ['positive'],
        'airline_sentiment_confidence': [0.8],
        'tweet_created': ['2025-01-14 10:00:00']
    }
    
    result = assessor.assess_risk(low_risk_data)
    print(f"传播风险等级: {result['risk_level'][0]}")
    print(f"风险概率: {result['risk_probabilities'][0]}")
