check_ss_strategy_stock/src/trading_calculator.py

from datetime import timedelta, date
from collections import defaultdict
import logging

# 配置日志记录
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(filename)s - %(message)s')

class TradingCalculator:
    """
    根据每日持仓信息计算交易日期和买卖点。
    """
    def __init__(self, pdf_files_info):
        """
        初始化 TradingCalculator。

        Args:
            pdf_files_info (list): 包含字典的列表，每个字典代表一个 PDF 文件信息。
                                   格式: [{'record_date': date, 'codes': [code1, ...], 'path': str}, ...]
                                   需要确保列表按 record_date 升序排列。
        """
        self.pdf_files_info = pdf_files_info
        # 记录日期到持仓代码集合的映射
        self.holdings_by_record_date = {info['record_date']: set(info['codes']) for info in pdf_files_info}
        # 原始记录日期列表，已排序
        self.record_dates = sorted(self.holdings_by_record_date.keys())
        print(f"初始化 TradingCalculator，共有 {len(self.record_dates)} 个记录日期。")

    def determine_buy_sell_dates(self):
        """
        根据记录日期和每日持仓，计算每个股票的买入卖出日期。

        Returns:
            list: 包含交易记录的列表。
                  格式: [{'code': str, 'buy_date': date or str, 'sell_date': date or str}, ...]
                  买入或卖出日期可能为字符串 "未知 (范围外)"。
        """
        if not self.record_dates:
            logging.warning("没有记录日期，无法计算买卖点。")
            return []

        # 按股票代码收集其出现的所有记录日期
        stock_record_dates = defaultdict(list)
        for record_date in self.record_dates:
            holdings = self.holdings_by_record_date.get(record_date, set())
            for code in holdings:
                stock_record_dates[code].append(record_date)

        transactions = []
        # record_dates 已经是排序好的完整日期列表
        all_available_dates = self.record_dates

        first_overall_record_date = all_available_dates[0]
        last_overall_record_date = all_available_dates[-1]

        for code, dates in stock_record_dates.items():
            if not dates:
                continue

            # dates 已经是 stock_record_dates 中按 key (record_date) 排序后append的，所以是排序好的
            # dates.sort() # 确保日期有序，但这里应该已经是排序好的

            # 识别连续的持有日期段（基于记录日期）
            segments = []
            current_segment = [dates[0]]

            for i in range(1, len(dates)):
                current_date = dates[i]
                last_date_in_segment = current_segment[-1]

                # 检查当前记录日期是否是上一记录日期的下一个日期（在 all_available_dates 列表中）
                try:
                    last_date_index = all_available_dates.index(last_date_in_segment)
                    # 确保索引+1在列表范围内，并且下一个日期就是当前日期
                    if last_date_index + 1 < len(all_available_dates) and all_available_dates[last_date_index + 1] == current_date:
                        # 日期连续
                        current_segment.append(current_date)
                    else:
                        # 日期不连续，当前段结束，开始新段
                        segments.append(current_segment)
                        current_segment = [current_date]
                except ValueError:
                     # This should not happen if dates come from self.record_dates
                    logging.error(f"日期 {last_date_in_segment} 未在 record_dates 中找到，代码 {code}")
                    # Treat as interruption if date is unexpected
                    segments.append(current_segment)
                    current_segment = [current_date]


            # 添加最后一个段
            if current_segment:
                segments.append(current_segment)

            # 根据连续段确定买卖日期
            # 买入日期是持有段开始记录日期的前一个记录日期（作为交易日）
            # 卖出日期是持有段结束记录日期的前一个记录日期（作为交易日），除非是最后一天，则卖出日期未知
            # PRD 提到交易日是上一个记录日，持仓是基于当天记录日。
            # 记录日 i 的持仓对应交易日 i-1
            # 持仓段 [记录日 i, 记录日 j] 对应的交易日段为 [记录日 i-1, 记录日 j-1]
            # 买入发生在 记录日 i-1
            # 卖出发生在 记录日 j

            for segment in segments:
                segment_start_record_date = segment[0] # 持有段的第一个记录日期
                segment_end_record_date = segment[-1] # 持有段的最后一个记录日期

                buy_date = None
                sell_date = None

                # 确定买入日期 (上一个记录日对应的交易日)
                if segment_start_record_date == first_overall_record_date:
                    # 如果持有从第一个记录日开始，买入日期未知
                    logging.info(f"{code} 在第一个记录日 {first_overall_record_date} 持有，买入日期未知 (范围外)。")
                    buy_date = None
                else:
                    # 找到 segment_start_record_date 在 all_available_dates 中的前一个记录日期
                    try:
                        start_date_index = all_available_dates.index(segment_start_record_date)
                        if start_date_index > 0:
                            # 买入日期是持有段开始记录日期的前一个记录日期
                            buy_date = all_available_dates[start_date_index - 1]
                        else:
                            # 这应该被上面的 if 覆盖，作为备用
                            logging.warning(f"{code} 的开始记录日期 {segment_start_record_date} 是记录日列表的第一个日期，买入日期未知。")
                            buy_date = None
                    except ValueError:
                        logging.error(f"开始记录日期 {segment_start_record_date} 未在 record_dates 中找到，代码 {code}")
                        buy_date = "错误" # Should not happen

                # 确定卖出日期
                if segment_end_record_date == last_overall_record_date:
                    # 如果持有到最后一个记录日结束，卖出日期未知
                    logging.info(f"{code} 在最后一个记录日 {last_overall_record_date} 仍持有，卖出日期未知 (范围外)。")
                    sell_date = None
                else:
                    # 卖出日期是持有段的最后一个记录日期本身 (因为PRD说卖出是当天)
                    sell_date = segment_end_record_date
                    logging.info(f"{code} 在 {sell_date} 卖出。")

                # 如果买入日期是日期对象，卖出日期也是日期对象，才添加交易
                # 如果任何一个未知，则不添加具体的交易记录，只记录日志
                if isinstance(buy_date, date) or isinstance(sell_date, date):
                     transactions.append({'code': code, 'buy_date': buy_date, 'sell_date': sell_date})
                     logging.info(f"{code} 的交易 ({buy_date} -> {sell_date}) 已添加。")
                else:
                    logging.info(f"{code} 的交易 ({buy_date} -> {sell_date}) 因日期范围外而忽略。")


        print(f"获得了 {len(transactions)} 笔买卖交易 (含范围外未确定日期的)。")
        # 过滤掉未知日期的交易，只返回有确定买卖日期的交易
        final_transactions = [t for t in transactions if isinstance(t['buy_date'], date) and isinstance(t['sell_date'], date)]
        partial_transactions = [t for t in transactions if not isinstance(t['buy_date'], date) or not isinstance(t['sell_date'], date)]
        print(f"确定了 {len(final_transactions)} 笔有明确买卖日期的交易。")
        print(f"忽略了 {len(partial_transactions)} 笔有范围外未确定日期的交易。")
        return final_transactions