import pandas as pd
import numpy as np
import os
import joblib
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.seasonal import seasonal_decompose
from sklearn.metrics import mean_squared_error

class TrendAnalyzer:
    def __init__(self, data_path=None):
        self.model_path = os.path.join(os.path.dirname(__file__), "../final_models/trend_analysis_model.pkl")
        
        if data_path:
            self.data = self._load_data(data_path)
            self.time_series = self._create_time_series()
            self.decomposition = self._decompose_series()
            self.models = self._train_models()
            self._save_model()
        else:
            self._load_model()
    
    def _load_data(self, data_path):
        """加载数据"""
        return pd.read_csv(data_path)
    
    def _create_time_series(self):
        """创建时间序列数据"""
        # 转换时间戳
        self.data['tweet_created'] = pd.to_datetime(self.data['tweet_created'])
        
        # 按小时聚合负面情感数据
        time_series = self.data[
            self.data['airline_sentiment'] == 'negative'
        ].set_index('tweet_created').resample('H').size()
        
        # 重命名列
        time_series = time_series.rename('negative_tweets_count')
        
        # 填充缺失值
        time_series = time_series.fillna(0)
        
        return time_series
    
    def _decompose_series(self):
        """分解时间序列（趋势、季节性、残差）"""
        if len(self.time_series) >= 24 * 7:  # 需要足够的数据进行分解
            decomposition = seasonal_decompose(self.time_series, model='additive', period=24)
            return decomposition
        return None
    
    def _train_models(self):
        """训练ARIMA模型"""
        models = {}
        
        # 为整体趋势训练模型
        try:
            # 简单的ARIMA模型，参数通过AIC选择
            model = ARIMA(self.time_series, order=(1, 1, 1))
            fitted_model = model.fit()
            models['overall'] = fitted_model
            print(f"整体趋势模型训练完成，AIC: {fitted_model.aic:.2f}")
        except Exception as e:
            print(f"训练整体趋势模型失败: {e}")
        
        return models
    
    def _save_model(self):
        """保存模型"""
        os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
        joblib.dump({
            'time_series': self.time_series,
            'decomposition': self.decomposition,
            'models': self.models
        }, self.model_path)
    
    def _load_model(self):
        """加载模型"""
        model_data = joblib.load(self.model_path)
        self.time_series = model_data['time_series']
        self.decomposition = model_data['decomposition']
        self.models = model_data['models']
    
    def predict_future(self, steps=24):
        """预测未来趋势"""
        if 'overall' in self.models:
            forecast = self.models['overall'].forecast(steps=steps)
            return forecast
        return None
    
    def detect_anomalies(self, threshold=3):
        """检测异常点"""
        # 计算移动平均和标准差
        window = 24
        rolling_mean = self.time_series.rolling(window=window).mean()
        rolling_std = self.time_series.rolling(window=window).std()
        
        # 计算Z-score
        z_score = (self.time_series - rolling_mean) / rolling_std
        z_score = z_score.fillna(0)
        
        # 识别异常点
        anomalies = self.time_series[abs(z_score) > threshold]
        
        return anomalies
    
    def get_trend_stats(self):
        """获取趋势统计信息"""
        stats = {
            'total_negative_tweets': self.time_series.sum(),
            'average_per_hour': self.time_series.mean(),
            'max_per_hour': self.time_series.max(),
            'min_per_hour': self.time_series.min(),
            'std_per_hour': self.time_series.std(),
            'period_start': self.time_series.index.min(),
            'period_end': self.time_series.index.max()
        }
        return stats
    
    def get_hourly_pattern(self):
        """获取小时模式"""
        # 计算一天中每个小时的平均投诉量
        hourly_pattern = self.time_series.groupby(self.time_series.index.hour).mean()
        hourly_pattern = hourly_pattern.rename('average_negative_tweets')
        return hourly_pattern
    
    def get_daily_pattern(self):
        """获取周模式"""
        # 计算一周中每天的平均投诉量
        daily_pattern = self.time_series.groupby(self.time_series.index.dayofweek).mean()
        daily_pattern = daily_pattern.rename('average_negative_tweets')
        
        # 重命名索引为星期几
        daily_pattern.index = ['周一', '周二', '周三', '周四', '周五', '周六', '周日']
        
        return daily_pattern
    
    def evaluate_model(self):
        """评估模型性能"""
        if 'overall' not in self.models:
            return None
        
        # 拆分训练集和测试集
        train_size = int(len(self.time_series) * 0.8)
        train, test = self.time_series[:train_size], self.time_series[train_size:]
        
        # 重新训练模型
        model = ARIMA(train, order=(1, 1, 1))
        fitted_model = model.fit()
        
        # 预测测试集
        predictions = fitted_model.forecast(steps=len(test))
        
        # 计算评估指标
        mse = mean_squared_error(test, predictions)
        rmse = np.sqrt(mse)
        
        return {
            'mse': mse,
            'rmse': rmse,
            'test_data': test,
            'predictions': predictions
        }

# 示例用法
if __name__ == "__main__":
    print("开始训练投诉趋势分析模型...")
    data_path = os.path.join(os.path.dirname(__file__), "../data/processed_tweets.csv")
    analyzer = TrendAnalyzer(data_path)
    print("模型训练完成!")
    
    # 获取趋势统计
    stats = analyzer.get_trend_stats()
    print(f"\n趋势统计信息:")
    print(f"总负面推文数: {stats['total_negative_tweets']}")
    print(f"每小时平均: {stats['average_per_hour']:.2f}")
    print(f"最高每小时: {stats['max_per_hour']}")
    print(f"时间范围: {stats['period_start']} 到 {stats['period_end']}")
    
    # 检测异常点
    anomalies = analyzer.detect_anomalies()
    print(f"\n检测到的异常点数量: {len(anomalies)}")
    if len(anomalies) > 0:
        print(f"前5个异常点:")
        print(anomalies.head())
    
    # 获取小时模式
    hourly_pattern = analyzer.get_hourly_pattern()
    print(f"\n小时模式 - 最高投诉时段: {hourly_pattern.idxmax()}:00")
    
    # 预测未来
    forecast = analyzer.predict_future(steps=12)
    if forecast is not None:
        print(f"\n未来12小时预测:")
        print(forecast)
    
    # 评估模型
    evaluation = analyzer.evaluate_model()
    if evaluation:
        print(f"\n模型评估:")
        print(f"RMSE: {evaluation['rmse']:.2f}")
