# -*- coding: utf-8 -*-

"""
通用工具函数模块
提供各种常用的工具函数
"""

import os
import json
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import re
from typing import Any, Dict, List, Union, Optional, Tuple


def format_date(date: Union[datetime, str], format_str: str = '%Y-%m-%d') -> str:
    """
    格式化日期
    
    Args:
        date: 日期对象或日期字符串
        format_str: 格式化字符串
        
    Returns:
        str: 格式化后的日期字符串
    """
    if isinstance(date, str):
        try:
            date_obj = parse_date(date)
            return date_obj.strftime(format_str)
        except Exception:
            return date
    elif isinstance(date, datetime):
        return date.strftime(format_str)
    return str(date)


def parse_date(date_str: str, formats: List[str] = None) -> datetime:
    """
    解析日期字符串
    
    Args:
        date_str: 日期字符串
        formats: 尝试的日期格式列表
        
    Returns:
        datetime: 解析后的日期对象
    """
    if formats is None:
        # 常用的日期格式列表
        formats = [
            '%Y-%m-%d',
            '%Y%m%d',
            '%Y/%m/%d',
            '%Y-%m-%d %H:%M:%S',
            '%Y%m%d %H:%M:%S',
            '%Y/%m/%d %H:%M:%S',
            '%Y-%m-%d %H:%M',
            '%Y%m%d %H:%M',
            '%Y/%m/%d %H:%M'
        ]
    
    for fmt in formats:
        try:
            return datetime.strptime(date_str, fmt)
        except ValueError:
            continue
    
    # 如果所有格式都失败，尝试使用pandas的解析
    try:
        return pd.to_datetime(date_str).to_pydatetime()
    except Exception:
        raise ValueError(f'无法解析日期字符串: {date_str}')


def ensure_directory(directory: str) -> None:
    """
    确保目录存在，如果不存在则创建
    
    Args:
        directory: 目录路径
    """
    if not os.path.exists(directory):
        try:
            os.makedirs(directory, exist_ok=True)
        except Exception as e:
            raise IOError(f'创建目录失败: {directory}. 错误: {str(e)}')


def load_json(file_path: str) -> Dict[str, Any]:
    """
    加载JSON文件
    
    Args:
        file_path: JSON文件路径
        
    Returns:
        dict: 解析后的JSON数据
    """
    if not os.path.exists(file_path):
        raise FileNotFoundError(f'文件不存在: {file_path}')
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f'JSON解析错误: {str(e)}')
    except Exception as e:
        raise IOError(f'读取文件失败: {str(e)}')


def save_json(data: Any, file_path: str, ensure_dir: bool = True, indent: int = 2) -> None:
    """
    保存数据到JSON文件
    
    Args:
        data: 要保存的数据
        file_path: JSON文件路径
        ensure_dir: 是否确保目录存在
        indent: 缩进空格数
    """
    if ensure_dir:
        # 确保目录存在
        dir_path = os.path.dirname(file_path)
        if dir_path:
            ensure_directory(dir_path)
    
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=indent)
    except Exception as e:
        raise IOError(f'保存文件失败: {str(e)}')


def calculate_indicators(
    df: pd.DataFrame,
    indicators: List[str],
    params: Dict[str, Any] = None
) -> pd.DataFrame:
    """
    计算常用技术指标
    
    Args:
        df: 包含OHLCV数据的DataFrame
        indicators: 要计算的指标列表
        params: 指标参数字典
        
    Returns:
        pd.DataFrame: 添加了指标列的DataFrame
    """
    result_df = df.copy()
    params = params or {}
    
    for indicator in indicators:
        if indicator.lower() == 'ma' or indicator.lower().startswith('ma'):
            # 移动平均线
            window = params.get('ma_window', 5)
            if indicator.lower() != 'ma':
                # 尝试从指标名中提取窗口大小，如MA5, MA10等
                match = re.search(r'\d+', indicator)
                if match:
                    window = int(match.group())
            
            result_df[f'MA{window}'] = result_df['close'].rolling(window=window).mean()
        
        elif indicator.lower() == 'ema' or indicator.lower().startswith('ema'):
            # 指数移动平均线
            window = params.get('ema_window', 5)
            if indicator.lower() != 'ema':
                match = re.search(r'\d+', indicator)
                if match:
                    window = int(match.group())
            
            result_df[f'EMA{window}'] = result_df['close'].ewm(span=window, adjust=False).mean()
        
        elif indicator.lower() == 'macd':
            # MACD
            fast_period = params.get('macd_fast', 12)
            slow_period = params.get('macd_slow', 26)
            signal_period = params.get('macd_signal', 9)
            
            # 计算EMA
            ema_fast = result_df['close'].ewm(span=fast_period, adjust=False).mean()
            ema_slow = result_df['close'].ewm(span=slow_period, adjust=False).mean()
            
            # 计算MACD线
            result_df['MACD'] = ema_fast - ema_slow
            # 计算信号线
            result_df['MACD_Signal'] = result_df['MACD'].ewm(span=signal_period, adjust=False).mean()
            # 计算柱状图
            result_df['MACD_Hist'] = result_df['MACD'] - result_df['MACD_Signal']
        
        elif indicator.lower() == 'rsi' or indicator.lower().startswith('rsi'):
            # RSI
            window = params.get('rsi_window', 14)
            if indicator.lower() != 'rsi':
                match = re.search(r'\d+', indicator)
                if match:
                    window = int(match.group())
            
            # 计算价格变化
            delta = result_df['close'].diff()
            gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
            loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
            
            # 计算RS和RSI
            rs = gain / loss
            result_df[f'RSI{window}'] = 100 - (100 / (1 + rs))
        
        elif indicator.lower() == 'kdj' or indicator.lower() == 'stochastic':
            # KDJ指标
            n = params.get('kdj_n', 9)
            m1 = params.get('kdj_m1', 3)
            m2 = params.get('kdj_m2', 3)
            
            # 计算RSV
            low_n = result_df['low'].rolling(window=n).min()
            high_n = result_df['high'].rolling(window=n).max()
            rsv = (result_df['close'] - low_n) / (high_n - low_n) * 100
            
            # 计算K、D、J线
            result_df['K'] = rsv.ewm(com=m1-1, adjust=False).mean()
            result_df['D'] = result_df['K'].ewm(com=m2-1, adjust=False).mean()
            result_df['J'] = 3 * result_df['K'] - 2 * result_df['D']
        
        elif indicator.lower() == 'bollinger' or indicator.lower() == 'bbands':
            # 布林带
            window = params.get('bollinger_window', 20)
            std_dev = params.get('bollinger_std', 2)
            
            result_df['BB_Middle'] = result_df['close'].rolling(window=window).mean()
            result_df['BB_Upper'] = result_df['BB_Middle'] + std_dev * result_df['close'].rolling(window=window).std()
            result_df['BB_Lower'] = result_df['BB_Middle'] - std_dev * result_df['close'].rolling(window=window).std()
        
        elif indicator.lower() == 'volume_ma' or indicator.lower().startswith('vma'):
            # 成交量移动平均线
            window = params.get('volume_ma_window', 5)
            if indicator.lower() != 'volume_ma':
                match = re.search(r'\d+', indicator)
                if match:
                    window = int(match.group())
            
            result_df[f'VMA{window}'] = result_df['volume'].rolling(window=window).mean()
    
    return result_df


def validate_parameters(params: Dict[str, Any], schema: Dict[str, Dict[str, Any]]) -> Tuple[bool, List[str]]:
    """
    验证参数是否符合要求
    
    Args:
        params: 要验证的参数字典
        schema: 参数验证规则，格式为 {param_name: {rule_type: rule_value, ...}}
        
    Returns:
        tuple: (是否验证通过, 错误信息列表)
    """
    errors = []
    
    for param_name, rules in schema.items():
        # 检查参数是否存在且为必填
        if 'required' in rules and rules['required']:
            if param_name not in params or params[param_name] is None:
                errors.append(f'参数 {param_name} 是必填的')
                continue
        
        # 如果参数不存在且不是必填，则跳过后续检查
        if param_name not in params or params[param_name] is None:
            continue
        
        param_value = params[param_name]
        
        # 检查类型
        if 'type' in rules:
            expected_type = rules['type']
            if not isinstance(param_value, expected_type):
                errors.append(f'参数 {param_name} 应为 {expected_type.__name__} 类型，实际为 {type(param_value).__name__}')
        
        # 检查数值范围
        if isinstance(param_value, (int, float)):
            if 'min' in rules and param_value < rules['min']:
                errors.append(f"参数 {param_name} 最小值为 {rules['min']}")
            if 'max' in rules and param_value > rules['max']:
                errors.append(f"参数 {param_name} 最大值为 {rules['max']}")
        
        # 检查字符串长度
        if isinstance(param_value, str):
            if 'min_length' in rules and len(param_value) < rules['min_length']:
                errors.append(f"参数 {param_name} 长度至少为 {rules['min_length']}")
            if 'max_length' in rules and len(param_value) > rules['max_length']:
                errors.append(f"参数 {param_name} 长度最多为 {rules['max_length']}")
            if 'pattern' in rules and not re.match(rules['pattern'], param_value):
                errors.append(f"参数 {param_name} 不符合正则表达式规则: {rules['pattern']}")
        
        # 检查是否在允许的值列表中
        if 'allowed_values' in rules and param_value not in rules['allowed_values']:
            errors.append(f"参数 {param_name} 必须是以下值之一: {rules['allowed_values']}")
    
    return len(errors) == 0, errors


def format_number(value: Union[int, float], precision: int = 2, use_comma: bool = True) -> str:
    """
    格式化数字
    
    Args:
        value: 要格式化的数字
        precision: 小数位数
        use_comma: 是否使用千位分隔符
        
    Returns:
        str: 格式化后的数字字符串
    """
    if value is None:
        return '0'
    
    try:
        # 如果是整数，不显示小数
        if isinstance(value, int):
            if use_comma:
                return f'{value:,}'
            else:
                return str(value)
        
        # 格式化浮点数
        format_str = f'{{:,.{precision}f}}' if use_comma else f'{{:.{precision}f}}'
        return format_str.format(value)
    except Exception:
        return str(value)


def calculate_drawdown(portfolio_values: Union[List[float], pd.Series]) -> pd.Series:
    """
    计算回撤
    
    Args:
        portfolio_values: 投资组合价值序列
        
    Returns:
        pd.Series: 回撤序列
    """
    if isinstance(portfolio_values, list):
        portfolio_values = pd.Series(portfolio_values)
    
    # 计算累计最大值
    cumulative_max = portfolio_values.cummax()
    # 计算回撤
    drawdown = (portfolio_values - cumulative_max) / cumulative_max
    
    return drawdown


def calculate_sharpe_ratio(returns: Union[List[float], pd.Series], risk_free_rate: float = 0.02, annualized: bool = True) -> float:
    """
    计算夏普比率
    
    Args:
        returns: 收益率序列
        risk_free_rate: 无风险利率
        annualized: 是否年化
        
    Returns:
        float: 夏普比率
    """
    if isinstance(returns, list):
        returns = pd.Series(returns)
    
    # 计算超额收益
    excess_returns = returns - (risk_free_rate / 252)  # 假设一年252个交易日
    
    # 计算夏普比率
    mean_excess_return = excess_returns.mean()
    std_excess_return = excess_returns.std()
    
    if std_excess_return == 0:
        return 0
    
    sharpe = mean_excess_return / std_excess_return
    
    # 年化
    if annualized:
        sharpe *= np.sqrt(252)
    
    return sharpe


def chunk_list(lst: List[Any], chunk_size: int) -> List[List[Any]]:
    """
    将列表切分为指定大小的子列表
    
    Args:
        lst: 原始列表
        chunk_size: 每个子列表的大小
        
    Returns:
        list: 子列表的列表
    """
    return [lst[i:i+chunk_size] for i in range(0, len(lst), chunk_size)]


def flatten_list(nested_list: List[List[Any]]) -> List[Any]:
    """
    展平嵌套列表
    
    Args:
        nested_list: 嵌套列表
        
    Returns:
        list: 展平后的列表
    """
    return [item for sublist in nested_list for item in sublist]


def safe_divide(a: float, b: float, default: float = 0.0) -> float:
    """
    安全除法，避免除零错误
    
    Args:
        a: 被除数
        b: 除数
        default: 当除数为零时的默认返回值
        
    Returns:
        float: 除法结果或默认值
    """
    if b == 0:
        return default
    return a / b


def get_file_paths(directory: str, pattern: str = '*') -> List[str]:
    """
    获取目录下所有匹配模式的文件路径
    
    Args:
        directory: 目录路径
        pattern: 文件匹配模式
        
    Returns:
        list: 文件路径列表
    """
    import glob
    
    if not os.path.exists(directory):
        return []
    
    search_pattern = os.path.join(directory, pattern)
    return glob.glob(search_pattern)


def time_function(func):
    """
    装饰器：测量函数执行时间
    
    Args:
        func: 要测量的函数
        
    Returns:
        function: 包装后的函数
    """
    import functools
    
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = datetime.now()
        result = func(*args, **kwargs)
        end_time = datetime.now()
        duration = (end_time - start_time).total_seconds()
        print(f'{func.__name__} 执行时间: {duration:.6f} 秒')
        return result
    
    return wrapper