Chan/chanlun/analysis/wyckoff/events.py

"""威科夫阶段与事件（启发式）。"""
from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd


def _bar_time(df: pd.DataFrame, i: int):
	row = df.iloc[i]
	if "date" in df.columns and pd.notna(row["date"]):
		return row["date"]
	if "timestamp" in df.columns:
		return row["timestamp"]
	return i


def _avg_vol(df: pd.DataFrame, i: int, win: int = 20) -> float:
	a = max(0, i - win + 1)
	v = df["volume"].astype(float).iloc[a : i + 1]
	m = float(v.mean()) if len(v) else 0.0
	return m if m > 0 else 1.0


def detect_bias_and_events(
	df: pd.DataFrame,
	tr: Dict[str, Any],
) -> Tuple[str, List[Dict[str, Any]], Dict[str, Any]]:
	"""
	返回 bias、events、volume_confirm。

	Spring/UTAD 相对「结构高低」判定：取区间内次低/次高（剔除单根极值），
	避免箱体把假破低点吃进 lo 后永远刺不破、从而无 C 阶段。
	"""
	hi = float(tr["high"])
	lo = float(tr["low"])
	mid = float(tr["mid"])
	tol = float(tr.get("tol") or (hi - lo) * 0.05)
	s = int(tr["abs_start_idx"])
	e = int(tr["abs_end_idx"])
	events: List[Dict[str, Any]] = []

	# 结构边界：用次低/次高作假破参照（至少 8 根才启用）
	seg = df.iloc[s : e + 1]
	event_lo, event_hi = lo, hi
	if len(seg) >= 8:
		lows = seg["low"].astype(float)
		highs = seg["high"].astype(float)
		# nsmallest(2) 的较大者 = 次低；nlargest(2) 的较小者 = 次高
		event_lo = float(lows.nsmallest(min(2, len(lows))).iloc[-1])
		event_hi = float(highs.nlargest(min(2, len(highs))).iloc[-1])
		# 勿比公布箱沿更「松」：结构带应在箱内
		event_lo = max(event_lo, lo)
		event_hi = min(event_hi, hi)
		# 若次低仍等于极值（多根同价），略抬参照便于识别收回
		if abs(event_lo - lo) < 1e-12:
			event_lo = lo + max(tol * 0.35, (hi - lo) * 0.02)
		if abs(event_hi - hi) < 1e-12:
			event_hi = hi - max(tol * 0.35, (hi - lo) * 0.02)

	# 扫描区间内及之后（含 tail_reserve）
	scan_end = int(tr.get("abs_scan_end_idx", min(len(df) - 1, e + 15)))
	scan_end = min(len(df) - 1, max(scan_end, e))
	spring = None
	utad = None
	sos = None
	sod = None  # sign of weakness / distribution breakdown
	lps = None
	lpsy = None

	for i in range(s + 2, scan_end + 1):
		row = df.iloc[i]
		low = float(row["low"])
		high = float(row["high"])
		close = float(row["close"])
		vol = float(row["volume"]) if "volume" in df.columns else 0.0
		avg_v = _avg_vol(df, i)
		ratio = vol / avg_v if avg_v else 0.0

		# Spring: pierce below structural support then close back
		if spring is None and low < event_lo - tol * 0.35 and close >= event_lo - tol * 0.35:
			vol_ok = ratio <= 1.35 or (i + 1 <= scan_end and float(df.iloc[min(i + 1, scan_end)]["volume"]) / avg_v < 1.2)
			spring = {
				"type": "Spring",
				"time": _bar_time(df, i),
				"price": low,
				"note": "假破下沿后收回",
				"volume_ratio": round(ratio, 3),
				"volume_ok": bool(vol_ok),
				"idx": i,
			}

		# UTAD: pierce above structural resistance then close back
		if utad is None and high > event_hi + tol * 0.35 and close <= event_hi + tol * 0.35:
			vol_ok = ratio >= 0.8
			utad = {
				"type": "UTAD",
				"time": _bar_time(df, i),
				"price": high,
				"note": "假破上沿后跌回",
				"volume_ratio": round(ratio, 3),
				"volume_ok": bool(vol_ok),
				"idx": i,
			}

		# SOS: close above high with volume
		if sos is None and close > hi + tol * 0.15:
			vol_ok = ratio >= 1.15
			sos = {
				"type": "SOS",
				"time": _bar_time(df, i),
				"price": close,
				"note": "放量上破交易区间",
				"volume_ratio": round(ratio, 3),
				"volume_ok": bool(vol_ok),
				"idx": i,
			}

		# SOW / breakdown
		if sod is None and close < lo - tol * 0.15:
			vol_ok = ratio >= 1.15
			sod = {
				"type": "SOW",
				"time": _bar_time(df, i),
				"price": close,
				"note": "放量下破交易区间",
				"volume_ratio": round(ratio, 3),
				"volume_ok": bool(vol_ok),
				"idx": i,
			}

	# LPS after SOS: pullback that holds above mid/high-band with lighter volume
	if sos is not None:
		si = int(sos["idx"])
		for i in range(si + 1, min(len(df), si + 25)):
			row = df.iloc[i]
			low = float(row["low"])
			close = float(row["close"])
			vol = float(row["volume"]) if "volume" in df.columns else 0.0
			avg_v = _avg_vol(df, i)
			ratio = vol / avg_v if avg_v else 0.0
			if low >= mid - tol and close >= hi - tol * 2:
				vol_ok = ratio <= 1.05
				lps = {
					"type": "LPS",
					"time": _bar_time(df, i),
					"price": low,
					"note": "突破后缩量回踩不破",
					"volume_ratio": round(ratio, 3),
					"volume_ok": bool(vol_ok),
					"idx": i,
				}
				break

	if sod is not None:
		si = int(sod["idx"])
		for i in range(si + 1, min(len(df), si + 25)):
			row = df.iloc[i]
			high = float(row["high"])
			close = float(row["close"])
			vol = float(row["volume"]) if "volume" in df.columns else 0.0
			avg_v = _avg_vol(df, i)
			ratio = vol / avg_v if avg_v else 0.0
			if high <= mid + tol and close <= lo + tol * 2:
				vol_ok = ratio <= 1.05
				lpsy = {
					"type": "LPSY",
					"time": _bar_time(df, i),
					"price": high,
					"note": "下跌突破后缩量反抽不过",
					"volume_ratio": round(ratio, 3),
					"volume_ok": bool(vol_ok),
					"idx": i,
				}
				break

	# 冲突清理：已判定吸筹且有 SOS 时，丢弃更早的 UTAD（避免阶段/图面误导）
	# 派发且有 SOW 时，丢弃更晚才合理的 Spring 假信号同理在偏置后再滤
	keep = []
	for ev in (spring, sos, lps, utad, sod, lpsy):
		if not ev:
			continue
		keep.append(ev)

	# bias（先算）
	last_c = float(df["close"].iloc[-1])
	bias = "unknown"
	if sos and (not sod or int(sos.get("idx", 0)) >= int(sod.get("idx", 0))):
		bias = "accumulation"
	elif sod and (not sos or int(sod.get("idx", 0)) > int(sos.get("idx", 0))):
		bias = "distribution"
	elif spring and not utad:
		bias = "accumulation"
	elif utad and not spring:
		bias = "distribution"
	elif last_c >= mid:
		bias = "accumulation"
	else:
		bias = "distribution"

	filtered = []
	for ev in keep:
		if bias == "accumulation" and ev["type"] == "UTAD" and sos and int(ev["idx"]) <= int(sos["idx"]):
			continue
		if bias == "distribution" and ev["type"] == "Spring" and sod and int(ev["idx"]) <= int(sod["idx"]):
			continue
		filtered.append(ev)

	events = [{k: v for k, v in ev.items() if k != "idx"} for ev in filtered]

	avg_volume = float(df["volume"].astype(float).iloc[max(0, e - 20) : e + 1].mean()) if "volume" in df.columns else 0.0
	volume_confirm = {
		"avg_volume": avg_volume,
		"event_checks": {ev["type"]: {"volume_ok": ev.get("volume_ok"), "volume_ratio": ev.get("volume_ratio")} for ev in events},
	}
	return bias, events, volume_confirm


def build_phases(
	df: pd.DataFrame,
	tr: Dict[str, Any],
	bias: str,
	events: List[Dict[str, Any]],
	min_bars: int = 3,
) -> List[Dict[str, Any]]:
	"""
	按威科夫事件锚点切分 A–E（启发式）。

	吸筹：A停止 → B筑底 → C测试(Spring) → D拉升(SOS…LPS) → E离开
	派发：A停止 → B筑顶 → C测试(UTAD) → D派发(SOW…LPSY) → E离开

	无 Spring/UTAD 时：若已有 SOS/SOW，用突破前末次沿带测试补 C；仍无则省略 C。
	"""
	s = int(tr["abs_start_idx"])
	e = int(tr["abs_end_idx"])
	hi = float(tr["high"])
	lo = float(tr["low"])
	n_last = len(df) - 1
	min_span = max(2, min_bars - 1)
	range_len = max(1, e - s)

	def _match_idx(t) -> Optional[int]:
		if t is None:
			return None
		lo = max(0, s - 2)
		hi = min(len(df), e + 40)
		for i in range(lo, hi):
			if _bar_time(df, i) == t:
				return i
		try:
			tt = pd.Timestamp(t)
			sample = None
			if "date" in df.columns and len(df):
				sample = df["date"].iloc[min(s, n_last)]
			if sample is not None and getattr(sample, "tzinfo", None) is not None and tt.tzinfo is None:
				tt = tt.tz_localize(sample.tzinfo)
			for i in range(lo, hi):
				bt = _bar_time(df, i)
				try:
					if abs((pd.Timestamp(bt) - tt).total_seconds()) <= 1:
						return i
				except Exception:
					continue
		except Exception:
			pass
		return None

	event_idx: Dict[str, int] = {}
	for ev in events:
		idx = _match_idx(ev.get("time"))
		if idx is not None:
			event_idx[str(ev.get("type"))] = idx

	accum = bias != "distribution"
	if accum:
		c_ev = event_idx.get("Spring")
		d_ev = event_idx.get("SOS")
		d_tail = event_idx.get("LPS") or d_ev
	else:
		c_ev = event_idx.get("UTAD")
		d_ev = event_idx.get("SOW")
		d_tail = event_idx.get("LPSY") or d_ev

	# 有 D 无明确测试事件时：用突破前最后一次触及下/上沿作为 C（次级测试）
	if c_ev is None and d_ev is not None:
		band = lo + (hi - lo) * 0.28 if accum else hi - (hi - lo) * 0.28
		for i in range(int(d_ev) - 1, s + 1, -1):
			row = df.iloc[i]
			if accum and float(row["low"]) <= band:
				c_ev = i
				break
			if not accum and float(row["high"]) >= band:
				c_ev = i
				break

	def _lab(phase: str) -> str:
		if accum:
			m = {"A": "A停止下跌", "B": "B筑底", "C": "C测试", "D": "D拉升", "E": "E离开"}
		else:
			m = {"A": "A停止上涨", "B": "B筑顶", "C": "C测试", "D": "D派发", "E": "E离开"}
		return m.get(phase, phase)

	a_end = s + max(min_bars, range_len // 5)

	c_start = c_end = None
	if c_ev is not None:
		c_start = max(s, int(c_ev) - 1)
		c_end = min(n_last, int(c_ev) + 1)

	if d_ev is not None:
		d_start = int(d_ev)
		d_end = min(n_last, max(int(d_tail or d_ev), d_start) + max(min_bars, range_len // 8))
		if d_tail is not None:
			d_end = max(d_end, min(n_last, int(d_tail) + 1))
	else:
		d_start = d_end = None

	if c_start is not None:
		b_end = max(a_end + 1, c_start)
	elif d_start is not None:
		b_end = max(a_end + 1, d_start)
	else:
		b_end = max(a_end + 1, e)

	if d_end is not None:
		e_start = min(n_last, d_end)
		e_end = n_last
	else:
		e_start = e_end = None

	raw = [("A", s, a_end), ("B", a_end, b_end)]
	if c_start is not None and c_end is not None:
		raw.append(("C", c_start, c_end))
	if d_start is not None and d_end is not None:
		raw.append(("D", d_start, d_end))
	if e_start is not None and e_end is not None and e_end > e_start:
		raw.append(("E", e_start, e_end))

	phases: List[Dict[str, Any]] = []
	cursor = s
	for phase, _a, _b in raw:
		if cursor >= n_last:
			break
		a = max(int(_a), cursor)
		b = int(max(int(_b), a))
		need = 1 if phase == "C" else min_span
		if b < a + need:
			b = min(n_last, a + need)
		b = int(np.clip(b, a, n_last))
		if b < a:
			continue
		if phases and phases[-1].get("_a") == a and phases[-1].get("_b") == b:
			continue
		phases.append(
			{
				"phase": phase,
				"label": _lab(phase),
				"start_time": _bar_time(df, a),
				"end_time": _bar_time(df, b),
				"_a": a,
				"_b": b,
			}
		)
		cursor = b
	for p in phases:
		p.pop("_a", None)
		p.pop("_b", None)
	return phases