"""
K线包含关系处理模块
"""

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

logger = logging.getLogger(__name__)


class KLineElement:
    """单个K线元素类"""
    
    def __init__(self, high: float, low: float, open_price: float, 
                 close: float, volume: float, timestamp: pd.Timestamp):
        """
        初始化K线元素
        
        Args:
            high: 最高价
            low: 最低价
            open_price: 开盘价
            close: 收盘价
            volume: 成交量
            timestamp: 时间戳
        """
        self.high = high
        self.low = low
        self.open = open_price
        self.close = close
        self.volume = volume
        self.timestamp = timestamp
        self.direction = 1 if close >= open_price else -1
        
    def __repr__(self):
        return f"KLineElement(H:{self.high}, L:{self.low}, O:{self.open}, C:{self.close})"


class KLine:
    """K线包含关系处理类"""
    
    def __init__(self, data: pd.DataFrame):
        """
        初始化K线处理器
        
        Args:
            data: 包含OHLCV数据的DataFrame
        """
        self.original_data = data.copy()
        self.processed_data = None
        self.kline_elements = []
        self._process_data()
        
    def _process_data(self):
        """处理原始数据为K线元素"""
        self.kline_elements = []
        
        for idx, row in self.original_data.iterrows():
            element = KLineElement(
                high=row['high'],
                low=row['low'],
                open_price=row['open'],
                close=row['close'],
                volume=row['volume'],
                timestamp=idx
            )
            self.kline_elements.append(element)
            
    def has_containment(self, k1: KLineElement, k2: KLineElement) -> bool:
        """
        判断两根K线是否存在包含关系
        
        Args:
            k1: 第一根K线
            k2: 第二根K线
            
        Returns:
            是否存在包含关系
        """
        # K1包含K2：K1的高低点完全包含K2
        k1_contains_k2 = (k1.high >= k2.high and k1.low <= k2.low)
        
        # K2包含K1：K2的高低点完全包含K1
        k2_contains_k1 = (k2.high >= k1.high and k2.low <= k1.low)
        
        return k1_contains_k2 or k2_contains_k1
        
    def merge_contained_klines(self, k1: KLineElement, k2: KLineElement, 
                              direction: int) -> KLineElement:
        """
        合并包含关系的K线
        
        Args:
            k1: 第一根K线
            k2: 第二根K线
            direction: 当前趋势方向（1为上升，-1为下降）
            
        Returns:
            合并后的K线
        """
        if direction > 0:  # 上升趋势中
            # 取两根K线的最高点的最大值，最低点的最大值
            merged_high = max(k1.high, k2.high)
            merged_low = max(k1.low, k2.low)
        else:  # 下降趋势中
            # 取两根K线的最高点的最小值，最低点的最小值
            merged_high = min(k1.high, k2.high)
            merged_low = min(k1.low, k2.low)
            
        # 开盘价和收盘价使用第一根K线的值
        merged_open = k1.open
        merged_close = k2.close
        merged_volume = k1.volume + k2.volume
        merged_timestamp = k2.timestamp  # 使用最后一根K线的时间
        
        return KLineElement(
            high=merged_high,
            low=merged_low,
            open_price=merged_open,
            close=merged_close,
            volume=merged_volume,
            timestamp=merged_timestamp
        )
        
    def handle_containment(self) -> List[KLineElement]:
        """
        处理所有K线的包含关系
        
        Returns:
            处理包含关系后的K线列表
        """
        if len(self.kline_elements) < 2:
            return self.kline_elements.copy()
            
        processed_klines = [self.kline_elements[0]]  # 第一根K线
        
        # 初始方向：根据前两根K线确定
        if len(self.kline_elements) >= 2:
            k1, k2 = self.kline_elements[0], self.kline_elements[1]
            if k2.high > k1.high:
                current_direction = 1  # 上升
            elif k2.high < k1.high:
                current_direction = -1  # 下降
            else:
                current_direction = 1 if k2.low >= k1.low else -1
        else:
            current_direction = 1
            
        i = 1
        while i < len(self.kline_elements):
            current_k = self.kline_elements[i]
            last_processed = processed_klines[-1]
            
            if self.has_containment(last_processed, current_k):
                # 存在包含关系，进行合并
                merged = self.merge_contained_klines(
                    last_processed, current_k, current_direction
                )
                processed_klines[-1] = merged  # 替换最后一个
            else:
                # 不存在包含关系，直接添加
                processed_klines.append(current_k)
                
                # 更新方向
                if current_k.high > last_processed.high:
                    current_direction = 1
                elif current_k.high < last_processed.high:
                    current_direction = -1
                # 如果high相等，保持原方向
                    
            i += 1
            
        logger.info(f"包含关系处理完成：{len(self.kline_elements)} -> {len(processed_klines)}")
        return processed_klines
        
    def to_dataframe(self, processed_klines: Optional[List[KLineElement]] = None) -> pd.DataFrame:
        """
        将处理后的K线转换为DataFrame
        
        Args:
            processed_klines: 处理后的K线列表，如果为None则使用默认处理结果
            
        Returns:
            包含处理后K线的DataFrame
        """
        if processed_klines is None:
            processed_klines = self.handle_containment()
            
        data = []
        for kline in processed_klines:
            data.append({
                'timestamp': kline.timestamp,
                'open': kline.open,
                'high': kline.high,
                'low': kline.low,
                'close': kline.close,
                'volume': kline.volume,
                'direction': kline.direction
            })
            
        df = pd.DataFrame(data)
        if not df.empty:
            df.set_index('timestamp', inplace=True)
            
        return df
        
    def get_processed_data(self) -> pd.DataFrame:
        """
        获取处理包含关系后的数据
        
        Returns:
            处理后的DataFrame
        """
        if self.processed_data is None:
            processed_klines = self.handle_containment()
            self.processed_data = self.to_dataframe(processed_klines)
            
        return self.processed_data
        
    def visualize_containment(self) -> dict:
        """
        生成包含关系可视化信息
        
        Returns:
            包含可视化信息的字典
        """
        original_count = len(self.kline_elements)
        processed_klines = self.handle_containment()
        processed_count = len(processed_klines)
        
        return {
            'original_count': original_count,
            'processed_count': processed_count,
            'merged_count': original_count - processed_count,
            'merge_ratio': (original_count - processed_count) / original_count if original_count > 0 else 0
        }