JUKUAN/Lib/future/kline_reconstruction.py

# K线复原工具
# 用于从交易记录CSV文件中提取开仓记录，获取对应的K线数据并绘制包含均线的K线图

from jqdata import *
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from datetime import datetime, timedelta, date
import re
import os
import warnings
warnings.filterwarnings('ignore')

# 中文字体设置（虽然图片内文字用英文，但保留设置以防需要）
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
plt.rcParams['axes.unicode_minus'] = False


def _get_current_directory():
    """
    获取当前文件所在目录，兼容 Jupyter notebook 环境
    
    返回:
        str: 当前目录路径
    """
    try:
        # 在普通 Python 脚本中，使用 __file__
        current_dir = os.path.dirname(os.path.abspath(__file__))
    except NameError:
        # 在 Jupyter notebook 环境中，__file__ 不存在，使用当前工作目录
        current_dir = os.getcwd()
        # 如果当前目录不是 future 目录，尝试查找
        if not os.path.exists(os.path.join(current_dir, 'transaction.csv')):
            # 尝试查找 future 目录
            if 'future' not in current_dir:
                # 尝试向上查找 future 目录
                parent_dir = os.path.dirname(current_dir)
                future_dir = os.path.join(parent_dir, 'future')
                if os.path.exists(os.path.join(future_dir, 'transaction.csv')):
                    current_dir = future_dir
    return current_dir


def read_and_filter_open_positions(csv_path):
    """
    读取CSV文件并筛选出开仓记录
    
    参数:
        csv_path (str): CSV文件路径
    
    返回:
        pandas.DataFrame: 包含开仓记录的DataFrame
    """
    try:
        df = pd.read_csv(csv_path, encoding='utf-8-sig')
        
        # 筛选交易类型第一个字符为"开"的行
        open_positions = df[df['交易类型'].str[0] == '开'].copy()
        
        print(f"从CSV文件中读取到 {len(df)} 条记录")
        print(f"筛选出 {len(open_positions)} 条开仓记录")
        
        return open_positions
    except Exception as e:
        print(f"读取CSV文件时出错: {str(e)}")
        return pd.DataFrame()


def extract_contract_code_and_date(row):
    """
    从标的列提取合约编号，从日期列提取日期
    
    参数:
        row (pandas.Series): DataFrame的一行数据
    
    返回:
        tuple: (contract_code, trade_date) 或 (None, None) 如果提取失败
    """
    try:
        # 提取合约编号：从"标的"列中提取括号内的内容
        target_str = str(row['标的'])
        match = re.search(r'\(([^)]+)\)', target_str)
        if match:
            contract_code = match.group(1)
        else:
            print(f"无法从标的 '{target_str}' 中提取合约编号")
            return None, None
        
        # 提取日期
        date_str = str(row['日期'])
        try:
            trade_date = datetime.strptime(date_str, '%Y-%m-%d').date()
        except:
            print(f"日期格式错误: {date_str}")
            return None, None
        
        return contract_code, trade_date
    except Exception as e:
        print(f"提取合约编号和日期时出错: {str(e)}")
        return None, None


def calculate_trade_days_range(trade_date, days_before=60, days_after=10):
    """
    计算交易日范围：往前days_before个交易日，往后days_after个交易日
    
    参数:
        trade_date (date): 开仓日期
        days_before (int): 往前交易日数量，默认60
        days_after (int): 往后交易日数量，默认10
    
    返回:
        tuple: (start_date, end_date) 或 (None, None) 如果计算失败
    """
    try:
        # 往前找：从trade_date往前找days_before个交易日
        # get_trade_days(end_date=trade_date, count=n) 返回包括trade_date在内的n个交易日
        # 所以需要count=days_before+1，第一个就是days_before个交易日前的日期
        trade_days_before = get_trade_days(end_date=trade_date, count=days_before + 1)
        if len(trade_days_before) < days_before + 1:
            print(f"无法获取足够的往前交易日，只获取到 {len(trade_days_before)} 个")
            return None, None
        
        start_date = trade_days_before[0].date()
        
        # 往后找：从trade_date往后找days_after个交易日
        # get_trade_days(start_date=trade_date, count=n) 返回包括trade_date在内的n个交易日
        # 所以需要count=days_after+1，最后一个就是days_after个交易日后的日期
        trade_days_after = get_trade_days(start_date=trade_date, count=days_after + 1)
        if len(trade_days_after) < days_after + 1:
            print(f"无法获取足够的往后交易日，只获取到 {len(trade_days_after)} 个")
            return None, None
        
        end_date = trade_days_after[-1].date()
        
        return start_date, end_date
    except Exception as e:
        print(f"计算交易日范围时出错: {str(e)}")
        return None, None


def get_kline_data(contract_code, start_date, end_date):
    """
    获取指定合约在时间范围内的K线数据
    
    参数:
        contract_code (str): 合约编号，如 'JD2502.XDCE'
        start_date (date): 开始日期
        end_date (date): 结束日期
    
    返回:
        pandas.DataFrame: 包含OHLC数据的DataFrame，如果获取失败返回None
    """
    try:
        # 使用get_price获取K线数据
        price_data = get_price(
            contract_code,
            start_date=start_date,
            end_date=end_date,
            frequency='1d',
            fields=['open', 'close', 'high', 'low']
        )
        
        if price_data is None or len(price_data) == 0:
            print(f"未获取到 {contract_code} 在 {start_date} 至 {end_date} 的数据")
            return None
        
        return price_data
    except Exception as e:
        print(f"获取K线数据时出错: {str(e)}")
        return None


def calculate_moving_averages(data):
    """
    计算5K, 10K, 20K, 30K均线
    
    参数:
        data (pandas.DataFrame): 包含close列的DataFrame
    
    返回:
        pandas.DataFrame: 添加了均线列的DataFrame
    """
    data = data.copy()
    
    # 计算均线
    data['ma5'] = data['close'].rolling(window=5).mean()
    data['ma10'] = data['close'].rolling(window=10).mean()
    data['ma20'] = data['close'].rolling(window=20).mean()
    data['ma30'] = data['close'].rolling(window=30).mean()
    
    return data


def filter_data_with_ma(data):
    """
    过滤掉任何一条均线为空的日期
    
    参数:
        data (pandas.DataFrame): 包含均线列的DataFrame
    
    返回:
        pandas.DataFrame: 过滤后的DataFrame
    """
    # 过滤掉任何一条均线为空的日期
    filtered_data = data.dropna(subset=['ma5', 'ma10', 'ma20', 'ma30'])
    
    return filtered_data


def plot_kline_chart(data, contract_code, trade_date, save_path):
    """
    绘制K线图（包含均线和开仓日期标注）
    
    参数:
        data (pandas.DataFrame): 包含OHLC和均线数据的DataFrame
        contract_code (str): 合约编号
        trade_date (date): 开仓日期
        save_path (str): 保存路径
    """
    try:
        # 创建图表
        fig, ax = plt.subplots(figsize=(16, 10))
        
        # 准备数据
        dates = data.index
        opens = data['open']
        highs = data['high']
        lows = data['low']
        closes = data['close']
        
        # 找到开仓日期在数据中的位置
        trade_date_idx = None
        for i, date_idx in enumerate(dates):
            if isinstance(date_idx, date):
                if date_idx == trade_date:
                    trade_date_idx = i
                    break
            elif isinstance(date_idx, datetime):
                if date_idx.date() == trade_date:
                    trade_date_idx = i
                    break
        
        # 绘制K线
        for i in range(len(data)):
            date_idx = dates[i]
            open_price = opens.iloc[i]
            high_price = highs.iloc[i]
            low_price = lows.iloc[i]
            close_price = closes.iloc[i]
            
            # K线颜色：红涨绿跌
            color = 'red' if close_price > open_price else 'green'
            edge_color = 'darkred' if close_price > open_price else 'darkgreen'
            
            # 绘制影线（最高价到最低价的竖线）
            ax.plot([i, i], [low_price, high_price], color='black', linewidth=1)
            
            # 绘制实体（开盘价到收盘价的矩形）
            body_height = abs(close_price - open_price)
            if body_height == 0:
                body_height = 0.01  # 避免高度为0
            bottom = min(open_price, close_price)
            
            # 使用矩形绘制K线实体
            rect = patches.Rectangle((i-0.4, bottom), 0.8, body_height,
                                   linewidth=1, edgecolor=edge_color, 
                                   facecolor=color, alpha=0.8)
            ax.add_patch(rect)
        
        # 绘制均线
        ax.plot(range(len(data)), data['ma5'], label='MA5', color='blue', linewidth=1.5, alpha=0.8)
        ax.plot(range(len(data)), data['ma10'], label='MA10', color='orange', linewidth=1.5, alpha=0.8)
        ax.plot(range(len(data)), data['ma20'], label='MA20', color='purple', linewidth=1.5, alpha=0.8)
        ax.plot(range(len(data)), data['ma30'], label='MA30', color='brown', linewidth=1.5, alpha=0.8)
        
        # 标注开仓日期位置
        if trade_date_idx is not None:
            trade_price = closes.iloc[trade_date_idx]
            ax.plot(trade_date_idx, trade_price, marker='*', markersize=15, 
                   color='yellow', markeredgecolor='black', markeredgewidth=1.5,
                   label='Open Position', zorder=10)
            # 添加垂直线
            ax.axvline(x=trade_date_idx, color='yellow', linestyle='--', 
                      linewidth=2, alpha=0.7, zorder=5)
        
        # 设置图表标题和标签（使用英文）
        contract_simple = contract_code.split('.')[0]  # 提取合约编号的简约部分
        ax.set_title(f'{contract_simple} ({contract_code}) K-Line Chart\n'
                   f'Period: {dates[0].strftime("%Y-%m-%d")} to {dates[-1].strftime("%Y-%m-%d")} '
                   f'({len(data)} bars)', 
                   fontsize=14, fontweight='bold', pad=20)
        
        ax.set_xlabel('Time', fontsize=12)
        ax.set_ylabel('Price', fontsize=12)
        ax.grid(True, alpha=0.3)
        ax.legend(loc='upper left', fontsize=10)
        
        # 设置x轴标签
        step = max(1, len(data) // 10)  # 显示约10个时间标签
        tick_positions = range(0, len(data), step)
        tick_labels = []
        for pos in tick_positions:
            date_val = dates[pos]
            if isinstance(date_val, date):
                tick_labels.append(date_val.strftime('%Y-%m-%d'))
            elif isinstance(date_val, datetime):
                tick_labels.append(date_val.strftime('%Y-%m-%d'))
            else:
                tick_labels.append(str(date_val))
        
        ax.set_xticks(tick_positions)
        ax.set_xticklabels(tick_labels, rotation=45, ha='right')
        
        # 添加统计信息（使用英文）
        max_price = highs.max()
        min_price = lows.min()
        latest_close = closes.iloc[-1]
        first_close = closes.iloc[0]
        total_change = (latest_close - first_close) / first_close * 100
        
        stats_text = (f'High: {max_price:.2f}\n'
                    f'Low: {min_price:.2f}\n'
                    f'Latest Close: {latest_close:.2f}\n'
                    f'Total Change: {total_change:+.2f}%')
        
        ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
               verticalalignment='top', bbox=dict(boxstyle='round', 
               facecolor='wheat', alpha=0.8), fontsize=10)
        
        # 调整布局并保存
        plt.tight_layout()
        plt.savefig(save_path, dpi=150, bbox_inches='tight')
        plt.close()
        
        print(f"K线图已保存到: {save_path}")
        
    except Exception as e:
        print(f"绘制K线图时出错: {str(e)}")
        raise


def reconstruct_kline_from_transactions(csv_path=None, output_dir=None):
    """
    主函数：从交易记录中复原K线图
    
    参数:
        csv_path (str): CSV文件路径，默认为 'Lib/future/transaction.csv'
        output_dir (str): 输出目录，默认为 'Lib/future/K'
    """
    # 设置默认路径
    if csv_path is None:
        # 获取当前文件所在目录
        # 在 Jupyter notebook 中，__file__ 不存在，使用当前工作目录
        try:
            current_dir = os.path.dirname(os.path.abspath(__file__))
        except NameError:
            # 在 Jupyter notebook 环境中，使用当前工作目录
            current_dir = os.getcwd()
            # 如果当前目录不是 future 目录，尝试查找
            if not os.path.exists(os.path.join(current_dir, 'transaction.csv')):
                # 尝试查找 future 目录
                if 'future' in current_dir:
                    pass  # 已经在 future 目录中
                else:
                    # 尝试向上查找 future 目录
                    parent_dir = os.path.dirname(current_dir)
                    future_dir = os.path.join(parent_dir, 'future')
                    if os.path.exists(os.path.join(future_dir, 'transaction.csv')):
                        current_dir = future_dir
        csv_path = os.path.join(current_dir, 'transaction.csv')
    
    if output_dir is None:
        # 获取当前文件所在目录
        try:
            current_dir = os.path.dirname(os.path.abspath(__file__))
        except NameError:
            # 在 Jupyter notebook 环境中，使用当前工作目录
            current_dir = os.getcwd()
            # 如果当前目录不是 future 目录，尝试查找
            if not os.path.exists(os.path.join(current_dir, 'transaction.csv')):
                # 尝试查找 future 目录
                if 'future' in current_dir:
                    pass  # 已经在 future 目录中
                else:
                    # 尝试向上查找 future 目录
                    parent_dir = os.path.dirname(current_dir)
                    future_dir = os.path.join(parent_dir, 'future')
                    if os.path.exists(os.path.join(future_dir, 'transaction.csv')):
                        current_dir = future_dir
        output_dir = os.path.join(current_dir, 'K')
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    print(f"输出目录: {output_dir}")
    
    # 1. 读取和筛选开仓记录
    print("\n=== 步骤1: 读取和筛选开仓记录 ===")
    open_positions = read_and_filter_open_positions(csv_path)
    
    if len(open_positions) == 0:
        print("未找到开仓记录，退出")
        return
    
    # 2. 处理每条开仓记录
    print(f"\n=== 步骤2: 处理 {len(open_positions)} 条开仓记录 ===")
    success_count = 0
    fail_count = 0
    
    for idx, row in open_positions.iterrows():
        print(f"\n--- 处理第 {idx + 1}/{len(open_positions)} 条记录 ---")
        
        try:
            # 提取合约编号和日期
            contract_code, trade_date = extract_contract_code_and_date(row)
            if contract_code is None or trade_date is None:
                print(f"跳过：无法提取合约编号或日期")
                fail_count += 1
                continue
            
            print(f"合约编号: {contract_code}, 开仓日期: {trade_date}")
            
            # 计算交易日范围
            start_date, end_date = calculate_trade_days_range(trade_date, days_before=60, days_after=10)
            if start_date is None or end_date is None:
                print(f"跳过：无法计算交易日范围")
                fail_count += 1
                continue
            
            print(f"数据范围: {start_date} 至 {end_date}")
            
            # 获取K线数据
            kline_data = get_kline_data(contract_code, start_date, end_date)
            if kline_data is None or len(kline_data) == 0:
                print(f"跳过：无法获取K线数据")
                fail_count += 1
                continue
            
            print(f"获取到 {len(kline_data)} 条K线数据")
            
            # 计算均线
            kline_data = calculate_moving_averages(kline_data)
            
            # 过滤数据
            filtered_data = filter_data_with_ma(kline_data)
            if len(filtered_data) == 0:
                print(f"跳过：过滤后无有效数据")
                fail_count += 1
                continue
            
            print(f"过滤后剩余 {len(filtered_data)} 条有效数据")
            
            # 生成文件名
            contract_simple = contract_code.split('.')[0]  # 提取合约编号的简约部分
            filename = f"{contract_simple}_{trade_date.strftime('%Y%m%d')}.png"
            save_path = os.path.join(output_dir, filename)
            
            # 绘制K线图
            plot_kline_chart(filtered_data, contract_code, trade_date, save_path)
            
            success_count += 1
            print(f"✓ 成功处理")
            
        except Exception as e:
            print(f"✗ 处理时出错: {str(e)}")
            fail_count += 1
            continue
    
    # 输出统计信息
    print(f"\n=== 处理完成 ===")
    print(f"成功: {success_count} 条")
    print(f"失败: {fail_count} 条")
    print(f"总计: {len(open_positions)} 条")


# 使用示例
if __name__ == "__main__":
    print("=" * 60)
    print("K线复原工具")
    print("=" * 60)
    
    reconstruct_kline_from_transactions()
    
    print("\n=== 完成 ===")