fin_data_unify_plat/backend/app/services/kline/adjustment_service.py

"""
复权计算服务 v2.2
实现前复权 (qfq) 和后复权 (hfq) 计算
"""
import logging
from datetime import datetime, date
from typing import List, Optional, Dict, Tuple
from decimal import Decimal

from sqlalchemy.orm import Session

from app.models.kline import (
    AdjustType, StockKLineItem, StockAdjustFactor
)
from app.repositories.kline.stock_repository import StockKLineRepository

logger = logging.getLogger(__name__)


class AdjustmentService:
    """
    复权计算服务
    
    复权原理：
    - 前复权 (qfq): 以当前价格为基准，向后调整历史价格
      公式: 调整后价格 = 原价格 × 累计复权因子
      
    - 后复权 (hfq): 以上市价格为基准，向前调整后续价格  
      公式: 调整后价格 = 原价格 × 累计复权因子
      
    复权因子来源：分红、配股、拆股等事件
    """
    
    def __init__(self, db: Session):
        self.repository = StockKLineRepository(db)
        self.db = db
    
    def apply_adjustment(
        self,
        symbol: str,
        items: List[StockKLineItem],
        adjust_type: AdjustType,
        factors: Optional[List[StockAdjustFactor]] = None
    ) -> List[StockKLineItem]:
        """
        应用复权计算
        
        Args:
            symbol: 股票代码
            items: K 线数据列表
            adjust_type: 复权类型 (qfq/hfq/none)
            factors: 复权因子列表（可选，不传则从数据库获取）
            
        Returns:
            List[StockKLineItem]: 复权后的 K 线数据
        """
        if adjust_type == AdjustType.NONE or not items:
            return items
        
        # 获取复权因子
        if factors is None:
            start_date = min(item.trade_date for item in items if item.trade_date)
            end_date = max(item.trade_date for item in items if item.trade_date)
            factors = self.repository.get_adjust_factors(symbol, start_date, end_date)
        
        if not factors:
            logger.info(f"{symbol} 无复权因子，返回原始数据")
            return items
        
        # 构建复权因子映射（日期 -> 因子）
        factor_map = self._build_factor_map(factors)
        
        # 应用复权
        if adjust_type == AdjustType.QFQ:
            return self._apply_qfq(items, factor_map, factors)
        elif adjust_type == AdjustType.HFQ:
            return self._apply_hfq(items, factor_map, factors)
        
        return items
    
    def _build_factor_map(
        self,
        factors: List[StockAdjustFactor]
    ) -> Dict[date, StockAdjustFactor]:
        """
        构建复权因子映射
        
        Args:
            factors: 复权因子列表
            
        Returns:
            Dict[date, StockAdjustFactor]: 日期到因子的映射
        """
        factor_map = {}
        for f in factors:
            factor_map[f.ex_date] = f
        return factor_map
    
    def _apply_qfq(
        self,
        items: List[StockKLineItem],
        factor_map: Dict[date, StockAdjustFactor],
        factors: List[StockAdjustFactor]
    ) -> List[StockKLineItem]:
        """
        应用前复权计算
        
        前复权原理：
        - 以最新价格为基准，保持不变
        - 历史价格根据复权因子向前调整
        - 使得价格序列连续，便于技术分析
        
        Args:
            items: K 线数据列表
            factor_map: 复权因子映射
            factors: 复权因子列表
            
        Returns:
            List[StockKLineItem]: 前复权后的 K 线数据
        """
        if not factors:
            return items
        
        # 计算累计复权因子（从最新日期往前累乘）
        # 前复权：越早的数据，累计因子越大
        cumulative_factors = self._calculate_qfq_cumulative_factors(factors)
        
        adjusted_items = []
        for item in items:
            if not item.trade_date:
                adjusted_items.append(item)
                continue
            
            # 查找该日期对应的累计复权因子
            cum_factor = self._get_cumulative_factor_for_date(
                item.trade_date, cumulative_factors, factors
            )
            
            if cum_factor != 1.0:
                adjusted_item = StockKLineItem(
                    symbol=item.symbol,
                    time=item.time,
                    open=self._adjust_price(item.open, cum_factor),
                    high=self._adjust_price(item.high, cum_factor),
                    low=self._adjust_price(item.low, cum_factor),
                    close=self._adjust_price(item.close, cum_factor),
                    volume=item.volume,  # 成交量不调整
                    amount=self._adjust_price(item.amount, cum_factor),
                    trade_date=item.trade_date,
                    is_limit_up=item.is_limit_up,
                    is_limit_down=item.is_limit_down,
                    total_market_cap=self._adjust_price(item.total_market_cap, cum_factor),
                    float_market_cap=self._adjust_price(item.float_market_cap, cum_factor),
                    inst_holding_ratio=item.inst_holding_ratio,
                    trading_days=item.trading_days
                )
                adjusted_items.append(adjusted_item)
            else:
                adjusted_items.append(item)
        
        logger.info(f"前复权计算完成，调整 {len(adjusted_items)} 条数据")
        return adjusted_items
    
    def _apply_hfq(
        self,
        items: List[StockKLineItem],
        factor_map: Dict[date, StockAdjustFactor],
        factors: List[StockAdjustFactor]
    ) -> List[StockKLineItem]:
        """
        应用后复权计算
        
        后复权原理：
        - 以上市价格为基准，保持不变
        - 后续价格根据复权因子向后调整
        - 反映真实的投资收益
        
        Args:
            items: K 线数据列表
            factor_map: 复权因子映射
            factors: 复权因子列表
            
        Returns:
            List[StockKLineItem]: 后复权后的 K 线数据
        """
        if not factors:
            return items
        
        # 计算累计复权因子（从最早日期往后累乘）
        cumulative_factors = self._calculate_hfq_cumulative_factors(factors)
        
        adjusted_items = []
        for item in items:
            if not item.trade_date:
                adjusted_items.append(item)
                continue
            
            # 查找该日期对应的累计复权因子
            cum_factor = self._get_cumulative_factor_for_date_hfq(
                item.trade_date, cumulative_factors, factors
            )
            
            if cum_factor != 1.0:
                adjusted_item = StockKLineItem(
                    symbol=item.symbol,
                    time=item.time,
                    open=self._adjust_price(item.open, cum_factor),
                    high=self._adjust_price(item.high, cum_factor),
                    low=self._adjust_price(item.low, cum_factor),
                    close=self._adjust_price(item.close, cum_factor),
                    volume=item.volume,
                    amount=self._adjust_price(item.amount, cum_factor),
                    trade_date=item.trade_date,
                    is_limit_up=item.is_limit_up,
                    is_limit_down=item.is_limit_down,
                    total_market_cap=self._adjust_price(item.total_market_cap, cum_factor),
                    float_market_cap=self._adjust_price(item.float_market_cap, cum_factor),
                    inst_holding_ratio=item.inst_holding_ratio,
                    trading_days=item.trading_days
                )
                adjusted_items.append(adjusted_item)
            else:
                adjusted_items.append(item)
        
        logger.info(f"后复权计算完成，调整 {len(adjusted_items)} 条数据")
        return adjusted_items
    
    def _calculate_qfq_cumulative_factors(
        self,
        factors: List[StockAdjustFactor]
    ) -> Dict[date, float]:
        """
        计算前复权累计因子
        
        前复权：从最新日期往前累乘
        例如：
        - 2023-06-01: 因子 1.1
        - 2022-06-01: 因子 1.2
        累计因子：
        - 2023-06-01 之后: 1.0
        - 2022-06-01 ~ 2023-06-01: 1.1
        - 2022-06-01 之前: 1.1 * 1.2 = 1.32
        
        Args:
            factors: 复权因子列表（按日期升序排列）
            
        Returns:
            Dict[date, float]: 日期到累计因子的映射
        """
        # 按日期降序排列（从最新到最早）
        sorted_factors = sorted(factors, key=lambda x: x.ex_date, reverse=True)
        
        cumulative = {}
        cum_product = 1.0
        
        for i, f in enumerate(sorted_factors):
            cum_product *= f.adjust_factor
            # 该日期之前的所有数据使用此累计因子
            if i < len(sorted_factors) - 1:
                next_date = sorted_factors[i + 1].ex_date
                cumulative[next_date] = cum_product
            else:
                # 最早的日期之前使用最终累计因子
                cumulative[date(1990, 1, 1)] = cum_product
        
        return cumulative
    
    def _calculate_hfq_cumulative_factors(
        self,
        factors: List[StockAdjustFactor]
    ) -> Dict[date, float]:
        """
        计算后复权累计因子
        
        后复权：从最早日期往后累乘
        例如：
        - 2022-06-01: 因子 1.2
        - 2023-06-01: 因子 1.1
        累计因子：
        - 2022-06-01 之前: 1.0
        - 2022-06-01 ~ 2023-06-01: 1.2
        - 2023-06-01 之后: 1.2 * 1.1 = 1.32
        
        Args:
            factors: 复权因子列表（按日期升序排列）
            
        Returns:
            Dict[date, float]: 日期到累计因子的映射
        """
        # 按日期升序排列（从最早到最新）
        sorted_factors = sorted(factors, key=lambda x: x.ex_date)
        
        cumulative = {}
        cum_product = 1.0
        
        for f in sorted_factors:
            # 该日期及之后的数据使用此累计因子
            cumulative[f.ex_date] = cum_product
            cum_product *= f.adjust_factor
        
        # 最新日期之后使用最终累计因子
        if sorted_factors:
            final_date = date(2100, 12, 31)
            cumulative[final_date] = cum_product
        
        return cumulative
    
    def _get_cumulative_factor_for_date(
        self,
        target_date: date,
        cumulative_factors: Dict[date, float],
        factors: List[StockAdjustFactor]
    ) -> float:
        """
        获取指定日期的前复权累计因子
        
        Args:
            target_date: 目标日期
            cumulative_factors: 累计因子映射
            factors: 复权因子列表
            
        Returns:
            float: 累计复权因子
        """
        # 找到第一个大于目标日期的因子日期
        sorted_dates = sorted(cumulative_factors.keys())
        
        for d in sorted_dates:
            if target_date < d:
                return cumulative_factors[d]
        
        # 如果目标日期大于所有因子日期，因子为 1.0（最新价格不变）
        return 1.0
    
    def _get_cumulative_factor_for_date_hfq(
        self,
        target_date: date,
        cumulative_factors: Dict[date, float],
        factors: List[StockAdjustFactor]
    ) -> float:
        """
        获取指定日期的后复权累计因子
        
        Args:
            target_date: 目标日期
            cumulative_factors: 累计因子映射
            factors: 复权因子列表
            
        Returns:
            float: 累计复权因子
        """
        # 找到第一个小于或等于目标日期的因子日期
        sorted_dates = sorted(cumulative_factors.keys(), reverse=True)
        
        for d in sorted_dates:
            if target_date >= d:
                return cumulative_factors[d]
        
        # 如果目标日期小于所有因子日期，因子为 1.0
        return 1.0
    
    def _adjust_price(
        self,
        price: Optional[float],
        factor: float
    ) -> Optional[float]:
        """
        应用复权因子调整价格
        
        Args:
            price: 原价格
            factor: 复权因子
            
        Returns:
            Optional[float]: 调整后的价格
        """
        if price is None:
            return None
        return round(price * factor, 4)
    
    def calculate_adjust_factor(
        self,
        dividend_ratio: float = 0.0,
        split_ratio: float = 1.0,
       配股比例: float = 0.0,
        配股价格: float = 0.0,
        前收盘价: float = 0.0
    ) -> float:
        """
        计算单次复权因子
        
        复权因子计算公式：
        factor = (前收盘价 - 分红 + 配股比例 × 配股价格) / (前收盘价 × 拆股比例 + 配股比例 × 配股价格)
        
        简化情况：
        - 仅分红: factor = (前收盘价 - 分红) / 前收盘价
        - 仅拆股: factor = 1 / 拆股比例
        
        Args:
            dividend_ratio: 每股分红金额（元）
            split_ratio: 拆股比例（如 2 表示 1 股拆成 2 股）
            配股比例: 配股比例（每 10 股配多少股）
            配股价格: 配股价格
            前收盘价: 除权除息前收盘价
            
        Returns:
            float: 复权因子
        """
        if 前收盘价 <= 0:
            return 1.0
        
        # 简化计算：仅考虑分红和拆股
        if split_ratio > 1:
            # 拆股情况
            return 1.0 / split_ratio
        elif dividend_ratio > 0:
            # 分红情况
            return (前收盘价 - dividend_ratio) / 前收盘价
        else:
            return 1.0
    
    def get_adjust_factors(
        self,
        symbol: str,
        start_date: Optional[date] = None,
        end_date: Optional[date] = None
    ) -> List[StockAdjustFactor]:
        """
        获取股票复权因子
        
        Args:
            symbol: 股票代码
            start_date: 开始日期
            end_date: 结束日期
            
        Returns:
            List[StockAdjustFactor]: 复权因子列表
        """
        return self.repository.get_adjust_factors(symbol, start_date, end_date)


# 服务工厂函数
def get_adjustment_service(db: Session) -> AdjustmentService:
    """获取复权计算服务实例"""
    return AdjustmentService(db)