Chan/chanlun/pipeline/builders/kline.py

"""TF_DF builder mixin — 由 split_tfdf_builders 自动生成，逻辑与原 TF_DF 一致。"""
from __future__ import annotations

from datetime import timedelta
from decimal import Decimal

import numpy as np
import pandas as pd
import talib.abstract as ta
from pandas import DataFrame
from technical.util import resample_to_interval

from chanlun.core.ChanBI import ChanBI
from chanlun.core.ChanBIZS import ChanBIZS
from chanlun.core.ChanBSP import ChanBSP
from chanlun.core.ChanEnum import (
    Chan_BI_DIR,
    Chan_BSP_DIR,
    Chan_BSP_TYPE,
    Chan_FX_TYPE,
    Chan_K_DIR,
    Chan_KLC_FX,
    Chan_KLC_STATE,
    Chan_KLINE_DIR,
    Chan_KLU_PATTERN,
    Chan_PRICE_TREND,
    Chan_SEG_DIR,
    Chan_ZS_DIR,
)
from chanlun.core.ChanKLC import ChanKLC
from chanlun.core.ChanKLU import ChanKLU
from chanlun.core.ChanSBI import ChanSBI
from chanlun.core.ChanSEG import ChanSEG
from chanlun.core.ChanZS import ChanZS, ChanZS_Big
from chanlun.indicators.ChanMACD import ChanMACD


class KlineBuilderMixin:
	def get_klu_state(self, dataframe):
		klc_list = self.get_klc_list(self.get_klu_list(dataframe))
		bi_list = self.cal_bi_list(klc_list)
		klu_state_list = []
		klc_index = 0
		for index in range(0, len(dataframe)):
			if klc_index == len(klc_list):
				klc_index = len(klc_list) - 1
			klc = klc_list[klc_index]
			if klc.end_klu and klc.end_klu.idx == index:
				if klc.klc_state == Chan_KLC_STATE.S10:
					klu_state_list.append("10")
					#print(klc.end_time, klc.klc_fx_type)
				elif klc.klc_state == Chan_KLC_STATE.S_10:
					klu_state_list.append("-10")
					#print(klc.end_time, klc.klc_fx_type)
				elif klc.klc_state == Chan_KLC_STATE.S11:
					klu_state_list.append("11")
					#print(klc.end_time, klc.klc_fx_type)
				elif klc.klc_state == Chan_KLC_STATE.S_11:
					klu_state_list.append("-11")
					#print(klc.end_time, klc.klc_fx_type)
				else:
					klu_state_list.append("00")
				klc_index += 1
			else:
				klu_state_list.append("00")
		print(klu_state_list[:20])
		return klu_state_list


	def check_fx1(self, klc):
		if klc.pre and klc.next:
			if klc.high > klc.pre.high and klc.high > klc.next.high and klc.low > klc.pre.low and klc.low > klc.next.low:
				if klc.pre.pre and klc.next.next:
					if klc.high > klc.pre.pre.high and klc.high > klc.next.next.high:
						#if (klc.close > klc.ema52 or klc.next.close > klc.next.ema52) and klc.macd > 0:
						klc.set_fx(Chan_FX_TYPE.TOP)
						#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "TOP")
						return Chan_FX_TYPE.TOP
			elif klc.low < klc.pre.low and klc.low < klc.next.low and klc.high < klc.pre.high and klc.high < klc.next.high:
				#if (klc.close < klc.ema52 or klc.next.close < klc.next.ema52) and klc.macd < 0:
				if klc.pre.pre and klc.next.next:
					if klc.low < klc.pre.pre.low and klc.low < klc.next.next.low:
						klc.set_fx(Chan_FX_TYPE.BOTTOM)
						#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "BOTTOM")
						return Chan_FX_TYPE.BOTTOM
		return Chan_FX_TYPE.UNKNOWN

	def check_fx(self, klc):
		if klc.pre and klc.next:
			if klc.high > klc.pre.high and klc.high > klc.next.high and klc.low > klc.pre.low and klc.low > klc.next.low:
				#if (klc.close > klc.ema52 or klc.next.close > klc.next.ema52) and klc.macd > 0:
				klc.set_fx(Chan_FX_TYPE.TOP)
				#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "TOP")
				return Chan_FX_TYPE.TOP
			elif klc.low < klc.pre.low and klc.low < klc.next.low and klc.high < klc.pre.high and klc.high < klc.next.high:
				#if (klc.close < klc.ema52 or klc.next.close < klc.next.ema52) and klc.macd < 0:
				klc.set_fx(Chan_FX_TYPE.BOTTOM)
				#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "BOTTOM")
				return Chan_FX_TYPE.BOTTOM
		return Chan_FX_TYPE.UNKNOWN

	def check_fx2(self, klc):
		if klc.pre and klc.next:
			if klc.high > klc.pre.close and klc.close > klc.next.close and klc.close > klc.pre.close and klc.close > klc.next.close:
				#if (klc.close > klc.ema52 or klc.next.close > klc.next.ema52) and klc.macd > 0:
				klc.set_fx(Chan_FX_TYPE.TOP)
				#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "TOP")
				return Chan_FX_TYPE.TOP
			elif klc.low < klc.pre.close and klc.close < klc.next.close and klc.close < klc.pre.close and klc.close < klc.next.close:
				#if (klc.close < klc.ema52 or klc.next.close < klc.next.ema52) and klc.macd < 0:
				klc.set_fx(Chan_FX_TYPE.BOTTOM)
				#print(klc.start_time, klc.end_time,klc.next.start_time, klc.next.end_time, klc.macd, klc.state, klc.fx, "BOTTOM")
				return Chan_FX_TYPE.BOTTOM
		return Chan_FX_TYPE.UNKNOWN

	def check_fx_pattern(self, klc):
		klu_list = klc.pre.klu_list + klc.klu_list + klc.next.klu_list
		
		self.cal_klu_pattern(klu_list)
		p = ""
		for klu in klu_list:
			p += klu.to_string()
		#print(p)

	def cal_volume_ratio(self, dataframe, window=10):
		df = dataframe.copy()
		# 计算过去N根K线的平均成交量
		df['avg_volume'] = df['volume'].rolling(window=window).mean()
		# 计算量比
		df['volume_ratio'] = df['volume'] / df['avg_volume']
		# 填充缺失值（前N根K线）
		df['volume_ratio'] = df['volume_ratio'].fillna(1.0)
		return df['volume_ratio']

	def cal_kl_data(self, dataframe:DataFrame):
		fields = "time,open,high,low,close,volume"
		klu_list = []
		last_klu = None
		for i in range(0, len(dataframe)):
			item = dataframe.iloc[i]
			date = item['date']
			o = item['open']
			h = item['high']
			l = item['low']
			c = item['close']
			v = item['volume']
			# time_obj = date.fromtimestamp(date)
			# date = date + timedelta(hours=8)
			time_str = date.strftime('%Y-%m-%d %H:%M:%S')
			item_data = [
				time_str,
				o,
				h,
				l,
				c,
				v
			]
			# klu =  KLU(self.create_item_dict(item_data, GetColumnNameFromFieldList(fields)))
			klu = ChanKLU(time_str, o, h, l, c, v)
			# print(klu.time, klu.open, klu.high, klu.low, klu.close, klu.volume)	
			klu.set_idx(i)
			klu_list.append(klu)
			if last_klu:
				last_klu.set_next(klu)
				klu.set_pre(last_klu)
			last_klu = klu
			if 'macd' in item:
				klu.set_indicators(item)
		return klu_list

	def get_kl_data(self, dataframe:DataFrame):
		return self.cal_kl_data(dataframe)

	def get_klc_list(self, klu_list):
		klc_list = []
		last_klu = None
		macd = ChanMACD(klu_list)
		klu_list = macd.cal_macd_state()
		ema_up_list = []
		ema_down_list = []
		ema_up_count = 0
		ema_down_count = 0
		last_klu = None
		for klu in klu_list:
			ema = klu.ema52
			last_ema = last_klu.ema52 if last_klu else 0
			if klu.close >= ema:
				ema_up_count += 1
			elif klu.close < ema:
				ema_down_count += 1
			if last_klu and last_klu.close >= last_ema and klu.close < ema:
				ema_up_list.append(ema_up_count)
				#print(last_klu.time, ema_up_count, "UP END")
				ema_up_count = 0
			elif last_klu and last_klu.close < last_ema and klu.close >= ema:
				ema_down_list.append(ema_down_count)
				#print(last_klu.time, ema_down_count, "DOWN END")
				ema_down_count = 0
			if len(klc_list) > 0:
				last_klc = klc_list[-1]
				if klu.exception:
					ddir = Chan_KLINE_DIR.DOWN
					if last_klc.high < klu.high:
						ddir = Chan_KLINE_DIR.UP
					klc = ChanKLC(klu, index=len(klc_list), ddir=ddir)
					klc.high = klu.close if klu.close > klu.open else klu.open
					klc.low = klu.open if klu.close > klu.open else klu.close
					klc_list.append(klc)
					last_klc.set_next(klc)
					klc.set_pre(last_klc)
					last_klc.set_end_klu(last_klu)
					klc.set_pre_fx()
					#print(klu.time, klu.high, klu.low, klu.close, klu.open, klu.exception)
				else:
					included = last_klc.check_klu_included(klu)
					if not included:
						ddir = Chan_KLINE_DIR.DOWN
						if last_klc.high < klu.high:
							ddir = Chan_KLINE_DIR.UP
						klc = ChanKLC(klu, index=len(klc_list), ddir=ddir)
						klc_list.append(klc)
						last_klc.set_next(klc)
						klc.set_pre(last_klc)
						last_klc.set_end_klu(last_klu)
						klc.set_pre_fx()
					else:
						last_klc.add_klu(klu)
			else:
				ddir = Chan_KLINE_DIR.UP
				if klu.open > klu.close:
					ddir = Chan_KLINE_DIR.DOWN     
				klc = ChanKLC(klu, 0, ddir)
				klc_list.append(klc)
			last_klu = klu
		klc_list = self.cal_trend(klc_list)
		#print(ema52_up_list, ema52_down_list)
		return klc_list


	def get_klu_list(self, dataframe):
		klu_list = self.get_kl_data(dataframe)
		#klu_list = self.cal_klu_pattern(klu_list)
		return klu_list

	def cal_klu_pattern(self, klu_list):
		"""
		计算裸K的pattern - 识别反转形态
		"""
		if not klu_list or len(klu_list) < 3:
			return klu_list
			
		for i, klu in enumerate(klu_list):
			# 单根K线反转模式识别
			self._detect_single_reversal_pattern(klu)
			
			# 双根K线形态识别
			if i >= 1:
				self._detect_double_pattern(klu_list[i-1], klu)
			
			# 三根K线形态识别
			if i >= 2:
				self._detect_triple_pattern(klu_list[i-2], klu_list[i-1], klu)
			
			#if klu.pattern != Chan_KLU_PATTERN.UNKNOWN:
				#print(klu.time, klu.pattern, klu.lower_shadow_ratio, klu.upper_shadow_ratio, klu.body_ratio, klu.lower_shadow_ratio/klu.body_ratio, klu.upper_shadow_ratio/klu.body_ratio)
		return klu_list
	

	def _detect_single_reversal_pattern(self, klu):
		"""检测单根K线反转模式"""
		body = abs(klu.close - klu.open)
		upper_shadow = klu.high - max(klu.close, klu.open)
		lower_shadow = min(klu.close, klu.open) - klu.low
		total_range = klu.high - klu.low
		
		# 避免除零
		if total_range == 0:
			return
		
		body_ratio = body / total_range
		upper_ratio = upper_shadow / total_range
		lower_ratio = lower_shadow / total_range
		#print(klu.time, upper_ratio, lower_ratio, body_ratio, upper_ratio/body_ratio, lower_ratio/body_ratio)
		# 避免body_ratio为0时的除零错误
		if body_ratio == 0:
			return
		# 锤子线/上吊线 - 反转信号
		if lower_ratio / body_ratio >= 2:
			# 锤子线：底部反转，需要前面一段
			if klu.close > klu.open and klu.pre:
				klu.set_pattern(Chan_KLU_PATTERN.HAMMER)  # 底部反转
			# 上吊线：顶部反转，需要前一根是上涨趋势
			elif klu.close < klu.open and klu.pre:
				klu.set_pattern(Chan_KLU_PATTERN.HANGING_MAN)  # 顶部反转
		
		# 倒锤子线/射击之星 - 反转信号
		elif upper_ratio / body_ratio >= 2:
			# 倒锤子线：底部反转，需要前一根是下跌趋势
			if klu.close > klu.open and klu.pre:
				klu.set_pattern(Chan_KLU_PATTERN.INVERTED_HAMMER)  # 底部反转
			# 射击之星：顶部反转，需要前一根是上涨趋势
			elif klu.close < klu.open and klu.pre:
				klu.set_pattern(Chan_KLU_PATTERN.SHOOTING_STAR)  # 顶部反转
		
		# 十字星 - 反转信号
		elif body_ratio <= 0.1:
			if upper_ratio > 0.4 and lower_ratio > 0.4:
				klu.set_pattern(Chan_KLU_PATTERN.LONG_LEGGED_DOJI)  # 强烈反转信号
			elif upper_ratio > 0.4 and lower_ratio <= 0.1:
				# 墓碑十字星：顶部反转，需要前一根是上涨趋势
				if klu.pre and klu.pre.close > klu.pre.open:
					klu.set_pattern(Chan_KLU_PATTERN.GRAVESTONE_DOJI)  # 顶部反转
			elif lower_ratio > 0.4 and upper_ratio <= 0.1:
				# 蜻蜓十字星：底部反转，需要前一根是下跌趋势
				if klu.pre and klu.pre.close < klu.pre.open:
					klu.set_pattern(Chan_KLU_PATTERN.DRAGONFLY_DOJI)  # 底部反转
			else:
				klu.set_pattern(Chan_KLU_PATTERN.DOJI)  # 一般反转信号
	

	def _detect_double_pattern(self, prev_klu, curr_klu):
		"""检测两根K线形成的形态
		包括：吞没形态(看涨/看跌)、乌云盖顶、曙光初现
		"""
		# 如果前一根K线已经有形态，不再识别双K线形态
		if prev_klu.pattern != Chan_KLU_PATTERN.UNKNOWN:
			return
		
		# 计算K线实体
		prev_body = abs(prev_klu.close - prev_klu.open)
		curr_body = abs(curr_klu.close - curr_klu.open)
		
		# 判断K线颜色（阴阳）
		prev_bullish = prev_klu.close > prev_klu.open
		curr_bullish = curr_klu.close > curr_klu.open
		
		# 检查是否存在长期趋势（至少需要5根K线的趋势）
		def check_long_trend(klu, bullish_trend=True, min_bars=5):
			"""检查是否存在长期趋势
			bullish_trend=True: 检查上涨趋势
			bullish_trend=False: 检查下跌趋势
			min_bars: 最少需要多少根K线形成趋势
			"""
			if not klu or not klu.pre:
				return False
			return True
				
			# 使用EMA指标判断长期趋势
			if klu.ema52 > 0:
				if bullish_trend and klu.close < klu.ema52:
					return False
				if not bullish_trend and klu.close > klu.ema52:
					return False
			
			# 检查连续的K线方向
			count = 0
			current = klu.pre
			
			while current and count < min_bars:
				if not current.pre:
					break
					
				if bullish_trend:
					# 上涨趋势：当前收盘价高于前一根收盘价
					if current.close <= current.pre.close:
						break
				else:
					# 下跌趋势：当前收盘价低于前一根收盘价
					if current.close >= current.pre.close:
						break
						
				count += 1
				current = current.pre
				
			return count >= min_bars
		
		# 1. 看涨吞没形态：前阴后阳，后者完全吞没前者
		# 要求前面有明显的下跌趋势
		if not prev_bullish and curr_bullish and \
		   abs(curr_klu.open - prev_klu.close) < 10 and \
		   curr_klu.close > prev_klu.open and \
		   check_long_trend(prev_klu, bullish_trend=False, min_bars=5):
			curr_klu.set_pattern(Chan_KLU_PATTERN.BULLISH_ENGULFING)
			return
		
		# 2. 看跌吞没形态：前阳后阴，后者完全吞没前者
		# 要求前面有明显的上涨趋势
		if prev_bullish and not curr_bullish and \
		   abs(curr_klu.open - prev_klu.close) < 10 and \
		   curr_klu.close < prev_klu.open and \
		   check_long_trend(prev_klu, bullish_trend=True, min_bars=5):
			curr_klu.set_pattern(Chan_KLU_PATTERN.BEARISH_ENGULFING)
			return
		
		# 3. 乌云盖顶：前阳后阴，后者开盘价高于前者最高价，收盘价在前者实体中部以下
		# 要求前面有明显的上涨趋势
		if prev_bullish and not curr_bullish and \
		   curr_klu.open > prev_klu.high and \
		   curr_klu.close < (prev_klu.open + prev_klu.close) / 2 and \
		   curr_klu.close > prev_klu.open and \
		   check_long_trend(prev_klu, bullish_trend=True, min_bars=5):
			curr_klu.set_pattern(Chan_KLU_PATTERN.DARK_CLOUD_COVER)
			return
		
		# 4. 曙光初现：前阴后阳，后者开盘价低于前者最低价，收盘价在前者实体中部以上
		# 要求前面有明显的下跌趋势
		if not prev_bullish and curr_bullish and \
		   curr_klu.open < prev_klu.low and \
		   curr_klu.close > (prev_klu.open + prev_klu.close) / 2 and \
		   curr_klu.close < prev_klu.open and \
		   check_long_trend(prev_klu, bullish_trend=False, min_bars=5):
			curr_klu.set_pattern(Chan_KLU_PATTERN.PIERCING_LINE)
			return
		
		# 平顶和平底移至三根K线形态中判断


	def _detect_triple_pattern(self, first_klu, second_klu, third_klu):
		"""检测三根K线形成的形态
		包括：早晨之星、黄昏之星、平顶、平底
		"""
		# 如果前两根K线已经有形态，不再识别三K线形态
		if first_klu.pattern != Chan_KLU_PATTERN.UNKNOWN or \
		   second_klu.pattern != Chan_KLU_PATTERN.UNKNOWN:
			return
		
		# 判断K线颜色（阴阳）
		first_bullish = first_klu.close > first_klu.open
		second_bullish = second_klu.close > second_klu.open
		third_bullish = third_klu.close > third_klu.open
		
		# 计算实体大小
		first_body = abs(first_klu.close - first_klu.open)
		second_body = abs(second_klu.close - second_klu.open)
		third_body = abs(third_klu.close - third_klu.open)
		
		# 检查是否存在长期趋势（至少需要5根K线的趋势）
		def check_long_trend(klu, bullish_trend=True, min_bars=5):
			"""检查是否存在长期趋势
			bullish_trend=True: 检查上涨趋势
			bullish_trend=False: 检查下跌趋势
			min_bars: 最少需要多少根K线形成趋势
			"""
			if not klu or not klu.pre:
				return False
				
			# 使用EMA指标判断长期趋势
			if klu.ema52 > 0:
				if bullish_trend and klu.close < klu.ema52:
					return False
				if not bullish_trend and klu.close > klu.ema52:
					return False
			
			# 检查连续的K线方向
			count = 0
			current = klu.pre
			
			while current and count < min_bars:
				if not current.pre:
					break
					
				if bullish_trend:
					# 上涨趋势：当前收盘价高于前一根收盘价
					if current.close <= current.pre.close:
						break
				else:
					# 下跌趋势：当前收盘价低于前一根收盘价
					if current.close >= current.pre.close:
						break
						
				count += 1
				current = current.pre
				
			return count >= min_bars
		
		# 1. 早晨之星：第一根阴线，第二根十字星或小实体，第三根阳线
		# 要求前面有明显的下跌趋势
		if not first_bullish and third_bullish and \
		   second_body < first_body * 0.3 and \
		   third_body > first_body * 0.5 and \
		   max(second_klu.open, second_klu.close) < first_klu.close and \
		   min(second_klu.open, second_klu.close) < third_klu.open and \
		   third_klu.close > (first_klu.open + first_klu.close) / 2 and \
		   check_long_trend(first_klu, bullish_trend=False, min_bars=7):
			third_klu.set_pattern(Chan_KLU_PATTERN.MORNING_STAR)
			return
		
		# 2. 黄昏之星：第一根阳线，第二根十字星或小实体，第三根阴线
		# 要求前面有明显的上涨趋势
		if first_bullish and not third_bullish and \
		   second_body < first_body * 0.3 and \
		   third_body > first_body * 0.5 and \
		   min(second_klu.open, second_klu.close) > first_klu.close and \
		   max(second_klu.open, second_klu.close) > third_klu.open and \
		   third_klu.close < (first_klu.open + first_klu.close) / 2 and \
		   check_long_trend(first_klu, bullish_trend=True, min_bars=7):
			third_klu.set_pattern(Chan_KLU_PATTERN.EVENING_STAR)
			return
			
		# 3. 平顶：三根K线的最高点几乎相同（上升趋势中更有意义）
		# 要求前面有明显的上涨趋势
		if (abs(first_klu.high - second_klu.high) / first_klu.high < 0.0002 and 
			abs(second_klu.high - third_klu.high) / second_klu.high < 0.0002 and
			check_long_trend(first_klu, bullish_trend=True, min_bars=7)):
			# 额外确认：价格接近阻力位或关键技术指标
			is_near_resistance = False
			
			# 检查是否接近EMA52阻力位
			if first_klu.ema52 > 0:
				resistance_level = first_klu.ema52
				if abs(first_klu.high - resistance_level) / resistance_level < 0.01:
					is_near_resistance = True
			
			# 检查是否有成交量确认（成交量减少表示上涨动能减弱）
			volume_confirmation = False
			if (first_klu.volume > 0 and second_klu.volume > 0 and third_klu.volume > 0 and
				third_klu.volume < second_klu.volume and second_klu.volume < first_klu.volume):
				volume_confirmation = True
				
			if is_near_resistance or volume_confirmation:
				third_klu.set_pattern(Chan_KLU_PATTERN.TWEEZER_TOP)
				return
		
		# 4. 平底：三根K线的最低点几乎相同（下降趋势中更有意义）
		# 要求前面有明显的下跌趋势
		if (abs(first_klu.low - second_klu.low) / first_klu.low < 0.0002 and 
			abs(second_klu.low - third_klu.low) / second_klu.low < 0.0002 and
			check_long_trend(first_klu, bullish_trend=False, min_bars=7)):
			# 额外确认：价格接近支撑位或关键技术指标
			is_near_support = False
			
			# 检查是否接近EMA52支撑位
			if first_klu.ema52 > 0:
				support_level = first_klu.ema52
				if abs(first_klu.low - support_level) / support_level < 0.01:
					is_near_support = True
			
			# 检查是否有成交量确认（成交量减少表示下跌动能减弱）
			volume_confirmation = False
			if (first_klu.volume > 0 and second_klu.volume > 0 and third_klu.volume > 0 and
				third_klu.volume < second_klu.volume and second_klu.volume < first_klu.volume):
				volume_confirmation = True
				
			if is_near_support or volume_confirmation:
				third_klu.set_pattern(Chan_KLU_PATTERN.TWEEZER_BOTTOM)
				return