import os
import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional, Tuple
import datetime
import json
import pickle
from sklearn.preprocessing import MinMaxScaler


class DataProcessor:
    """数据处理类，用于处理和存储加密货币数据"""
    
    def __init__(self, storage_path: str = "./data"):
        """
        初始化数据处理器
        
        Args:
            storage_path: 数据存储路径
        """
        self.storage_path = storage_path
        os.makedirs(storage_path, exist_ok=True)
        self.scalers = {}  # 用于存储每个特征的缩放器
    
    def preprocess_market_data(self, symbol: str, data: pd.DataFrame) -> pd.DataFrame:
        """
        预处理市场数据
        
        Args:
            symbol: 交易对符号，例如 'BTCUSDT'
            data: 原始市场数据
            
        Returns:
            预处理后的数据
        """
        # 确保数据类型正确
        data = data.copy()
        
        # 将时间戳设置为索引
        if 'timestamp' in data.columns:
            data.set_index('timestamp', inplace=True)
        
        # 移除不必要的列
        columns_to_keep = ['open', 'high', 'low', 'close', 'volume']
        data = data[columns_to_keep]
        
        # 检查并处理缺失值
        data.bfill(inplace=True)  # 前向填充
        data.bfill(inplace=True)  # 后向填充
        
        # 添加技术指标
        data = self.add_technical_indicators(data)
        
        return data
    
    def add_technical_indicators(self, data: pd.DataFrame) -> pd.DataFrame:
        """
        添加技术指标
        
        Args:
            data: 原始价格数据
            
        Returns:
            添加了技术指标的数据
        """
        df = data.copy()
        
        # 计算移动平均线
        df['MA5'] = df['close'].rolling(window=5).mean()
        df['MA10'] = df['close'].rolling(window=10).mean()
        df['MA20'] = df['close'].rolling(window=20).mean()
        df['MA50'] = df['close'].rolling(window=50).mean()
        
        # 计算相对强弱指标(RSI)
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # 计算MACD
        exp1 = df['close'].ewm(span=12, adjust=False).mean()
        exp2 = df['close'].ewm(span=26, adjust=False).mean()
        macd = exp1 - exp2
        signal = macd.ewm(span=9, adjust=False).mean()
        
        df['MACD'] = macd
        df['MACD_Signal'] = signal
        df['MACD_Hist'] = macd - signal
        
        # 计算布林带
        df['MA20_std'] = df['close'].rolling(window=20).std()
        df['Bollinger_Upper'] = df['MA20'] + (df['MA20_std'] * 2)
        df['Bollinger_Lower'] = df['MA20'] - (df['MA20_std'] * 2)
        
        # 计算ATR (Average True Range)
        high_low = df['high'] - df['low']
        high_close = abs(df['high'] - df['close'].shift())
        low_close = abs(df['low'] - df['close'].shift())
        
        ranges = pd.concat([high_low, high_close, low_close], axis=1)
        true_range = ranges.max(axis=1)
        df['ATR'] = true_range.rolling(14).mean()
        
        # 填充计算指标产生的NaN值
        df.bfill(inplace=True)
        
        return df
    
    def normalize_data(self, data: pd.DataFrame, symbol: str) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """
        归一化数据
        
        Args:
            data: 预处理后的数据
            symbol: 交易对符号
            
        Returns:
            归一化后的数据和缩放器
        """
        df = data.copy()
        scalers = {}
        
        # 对每列进行归一化
        for column in df.columns:
            scaler = MinMaxScaler(feature_range=(0, 1))
            df[column] = scaler.fit_transform(df[[column]])
            scalers[column] = scaler
        
        # 保存缩放器以便后续使用
        self.scalers[symbol] = scalers
        
        return df, scalers
    
    def prepare_model_input(self, data: pd.DataFrame, window_size: int = 60) -> Tuple[np.ndarray, np.ndarray]:
        """
        准备模型输入数据（时间序列窗口）
        
        Args:
            data: 归一化后的数据
            window_size: 时间窗口大小
            
        Returns:
            X和y数据
        """
        X, y = [], []
        
        for i in range(len(data) - window_size):
            X.append(data.iloc[i:(i + window_size)].values)
            y.append(data.iloc[i + window_size]['close'])
        
        return np.array(X), np.array(y)
    
    def save_data(self, symbol: str, data: pd.DataFrame, data_type: str = 'raw') -> str:
        """
        保存数据到文件
        
        Args:
            symbol: 交易对符号
            data: 要保存的数据
            data_type: 数据类型，例如 'raw', 'processed', 'normalized'
            
        Returns:
            保存文件的路径
        """
        # 创建目录（如果不存在）
        symbol_dir = os.path.join(self.storage_path, symbol)
        os.makedirs(symbol_dir, exist_ok=True)
        
        # 生成文件名
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"{symbol}_{data_type}_{timestamp}.csv"
        file_path = os.path.join(symbol_dir, filename)
        
        # 保存数据
        data.to_csv(file_path)
        
        return file_path
    
    def save_scalers(self, symbol: str) -> str:
        """
        保存缩放器到文件
        
        Args:
            symbol: 交易对符号
            
        Returns:
            保存文件的路径
        """
        if symbol not in self.scalers:
            raise ValueError(f"没有找到{symbol}的缩放器")
        
        # 创建目录（如果不存在）
        symbol_dir = os.path.join(self.storage_path, symbol)
        os.makedirs(symbol_dir, exist_ok=True)
        
        # 生成文件名
        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"{symbol}_scalers_{timestamp}.pkl"
        file_path = os.path.join(symbol_dir, filename)
        
        # 保存缩放器
        with open(file_path, 'wb') as f:
            pickle.dump(self.scalers[symbol], f)
        
        return file_path
    
    def load_data(self, file_path: str) -> pd.DataFrame:
        """
        从文件加载数据
        
        Args:
            file_path: 文件路径
            
        Returns:
            加载的数据
        """
        return pd.read_csv(file_path, index_col=0, parse_dates=True)
    
    def load_scalers(self, file_path: str) -> Dict[str, Any]:
        """
        从文件加载缩放器
        
        Args:
            file_path: 文件路径
            
        Returns:
            加载的缩放器
        """
        with open(file_path, 'rb') as f:
            return pickle.load(f)
    
    def get_latest_data_file(self, symbol: str, data_type: str = 'processed') -> Optional[str]:
        """
        获取最新的数据文件路径
        
        Args:
            symbol: 交易对符号
            data_type: 数据类型
            
        Returns:
            最新的数据文件路径，如果不存在则返回None
        """
        symbol_dir = os.path.join(self.storage_path, symbol)
        
        if not os.path.exists(symbol_dir):
            return None
        
        # 查找匹配的文件
        files = [f for f in os.listdir(symbol_dir) if f.startswith(f"{symbol}_{data_type}") and f.endswith('.csv')]
        
        if not files:
            return None
        
        # 按文件名（包含时间戳）排序
        files.sort(reverse=True)
        
        return os.path.join(symbol_dir, files[0])
    
    def get_latest_scaler_file(self, symbol: str) -> Optional[str]:
        """
        获取最新的缩放器文件路径
        
        Args:
            symbol: 交易对符号
            
        Returns:
            最新的缩放器文件路径，如果不存在则返回None
        """
        symbol_dir = os.path.join(self.storage_path, symbol)
        
        if not os.path.exists(symbol_dir):
            return None
        
        # 查找匹配的文件
        files = [f for f in os.listdir(symbol_dir) if f.startswith(f"{symbol}_scalers") and f.endswith('.pkl')]
        
        if not files:
            return None
        
        # 按文件名（包含时间戳）排序
        files.sort(reverse=True)
        
        return os.path.join(symbol_dir, files[0]) 