"""
数据处理模块：负责K线数据的清理、格式化和预处理
"""

import pandas as pd
import numpy as np
from typing import Optional, Tuple, List
import logging

logger = logging.getLogger(__name__)


class DataProcessor:
    """数据处理器：负责K线数据的处理和验证"""
    
    @staticmethod
    def validate_klines(df: pd.DataFrame) -> bool:
        """
        验证K线数据的完整性和正确性
        
        Args:
            df: K线数据DataFrame
            
        Returns:
            是否通过验证
        """
        required_columns = ['open', 'high', 'low', 'close', 'volume']
        
        # 检查必需列是否存在
        if not all(col in df.columns for col in required_columns):
            logger.error("缺少必需的列")
            return False
            
        # 检查数据是否为空
        if df.empty:
            logger.error("数据为空")
            return False
            
        # 检查价格关系是否正确
        invalid_rows = (
            (df['high'] < df['low']) |
            (df['high'] < df['open']) |
            (df['high'] < df['close']) |
            (df['low'] > df['open']) |
            (df['low'] > df['close'])
        )
        
        if invalid_rows.any():
            logger.warning(f"发现 {invalid_rows.sum()} 行无效的价格关系")
            
        # 检查是否有NaN值
        if df[required_columns].isnull().any().any():
            logger.warning("数据中包含NaN值")
            
        return True
        
    @staticmethod
    def clean_klines(df: pd.DataFrame) -> pd.DataFrame:
        """
        清理K线数据
        
        Args:
            df: 原始K线数据
            
        Returns:
            清理后的K线数据
        """
        df_clean = df.copy()
        
        # 移除NaN值
        df_clean = df_clean.dropna()
        
        # 修正无效的价格关系
        # 如果high < max(open, close)，则设置high = max(open, close, low)
        df_clean['high'] = np.maximum.reduce([
            df_clean['high'], 
            df_clean['open'], 
            df_clean['close'],
            df_clean['low']
        ])
        
        # 如果low > min(open, close)，则设置low = min(open, close, high)
        df_clean['low'] = np.minimum.reduce([
            df_clean['low'], 
            df_clean['open'], 
            df_clean['close'],
            df_clean['high']
        ])
        
        # 确保volume非负
        df_clean['volume'] = np.maximum(df_clean['volume'], 0)
        
        # 按时间排序
        df_clean = df_clean.sort_index()
        
        logger.info(f"数据清理完成，剩余 {len(df_clean)} 条记录")
        return df_clean
        
    @staticmethod
    def add_technical_indicators(df: pd.DataFrame) -> pd.DataFrame:
        """
        添加技术指标
        
        Args:
            df: K线数据
            
        Returns:
            包含技术指标的数据
        """
        df_with_indicators = df.copy()
        
        # 添加价格范围
        df_with_indicators['range'] = df_with_indicators['high'] - df_with_indicators['low']
        
        # 添加实体大小
        df_with_indicators['body'] = abs(df_with_indicators['close'] - df_with_indicators['open'])
        
        # 添加上影线长度
        df_with_indicators['upper_shadow'] = df_with_indicators['high'] - np.maximum(
            df_with_indicators['open'], 
            df_with_indicators['close']
        )
        
        # 添加下影线长度
        df_with_indicators['lower_shadow'] = np.minimum(
            df_with_indicators['open'], 
            df_with_indicators['close']
        ) - df_with_indicators['low']
        
        # 添加K线方向
        df_with_indicators['direction'] = np.where(
            df_with_indicators['close'] > df_with_indicators['open'], 1, -1
        )
        
        return df_with_indicators
        
    @staticmethod
    def resample_klines(df: pd.DataFrame, new_timeframe: str) -> pd.DataFrame:
        """
        重采样K线数据到新的时间周期
        
        Args:
            df: 原始K线数据
            new_timeframe: 新的时间周期，如'4H', '1D'
            
        Returns:
            重采样后的K线数据
        """
        try:
            # 重采样规则
            agg_dict = {
                'open': 'first',
                'high': 'max',
                'low': 'min',
                'close': 'last',
                'volume': 'sum'
            }
            
            # 执行重采样
            resampled = df.resample(new_timeframe).agg(agg_dict)
            
            # 移除空值
            resampled = resampled.dropna()
            
            logger.info(f"重采样到 {new_timeframe}，得到 {len(resampled)} 条记录")
            return resampled
            
        except Exception as e:
            logger.error(f"重采样失败: {e}")
            raise
            
    @staticmethod
    def calculate_returns(df: pd.DataFrame) -> pd.DataFrame:
        """
        计算收益率
        
        Args:
            df: K线数据
            
        Returns:
            包含收益率的数据
        """
        df_with_returns = df.copy()
        
        # 计算收盘价收益率
        df_with_returns['returns'] = df_with_returns['close'].pct_change()
        
        # 计算对数收益率
        df_with_returns['log_returns'] = np.log(df_with_returns['close'] / df_with_returns['close'].shift(1))
        
        return df_with_returns
        
    @staticmethod
    def get_data_summary(df: pd.DataFrame) -> dict:
        """
        获取数据摘要信息
        
        Args:
            df: K线数据
            
        Returns:
            数据摘要字典
        """
        summary = {
            'total_records': len(df),
            'date_range': {
                'start': df.index.min(),
                'end': df.index.max()
            },
            'price_range': {
                'min': df['low'].min(),
                'max': df['high'].max()
            },
            'volume_stats': {
                'total': df['volume'].sum(),
                'avg': df['volume'].mean(),
                'max': df['volume'].max()
            }
        }
        
        if len(df) > 0:
            summary['latest_price'] = df['close'].iloc[-1]
            
        return summary