web_development/backend/app/services/pattern_detector.py

"""
Chart pattern detection engine.
Detects technical analysis patterns on kline data and returns drawing instructions
for the frontend to overlay on candlestick charts.
"""

import logging
from dataclasses import dataclass, field

import numpy as np

from app.models.schemas import KlineData

logger = logging.getLogger(__name__)


@dataclass
class PatternResult:
    pattern_type: str
    name: str
    direction: str         # "bullish", "bearish", "neutral"
    confidence: float      # 0-100
    start_idx: int
    end_idx: int
    description: str
    draw_lines: list[dict] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "pattern_type": self.pattern_type,
            "name": self.name,
            "direction": self.direction,
            "confidence": round(self.confidence, 1),
            "start_idx": self.start_idx,
            "end_idx": self.end_idx,
            "description": self.description,
            "draw_lines": self.draw_lines,
        }


PATTERN_REGISTRY: dict[str, dict] = {
    "trendline_support": {"name": "支撑趋势线", "description": "连接多个低点形成上升/下降支撑线"},
    "trendline_resistance": {"name": "阻力趋势线", "description": "连接多个高点形成上升/下降阻力线"},
    "ascending_triangle": {"name": "上升三角形", "description": "水平阻力 + 上升支撑，看涨突破形态"},
    "descending_triangle": {"name": "下降三角形", "description": "水平支撑 + 下降阻力，看跌突破形态"},
    "symmetric_triangle": {"name": "对称三角形", "description": "收敛三角形，突破方向待定"},
    "double_top": {"name": "双顶 (M头)", "description": "两个近似高点，看跌反转形态"},
    "double_bottom": {"name": "双底 (W底)", "description": "两个近似低点，看涨反转形态"},
    "head_shoulders": {"name": "头肩顶", "description": "头部高于两肩，看跌反转形态"},
    "inv_head_shoulders": {"name": "头肩底", "description": "头部低于两肩，看涨反转形态"},
    "channel_up": {"name": "上升通道", "description": "平行的上升支撑和阻力"},
    "channel_down": {"name": "下降通道", "description": "平行的下降支撑和阻力"},
    "wedge_rising": {"name": "上升楔形", "description": "收窄的上升通道，看跌"},
    "wedge_falling": {"name": "下降楔形", "description": "收窄的下降通道，看涨"},
    "support_level": {"name": "水平支撑", "description": "价格多次触及的水平支撑位"},
    "resistance_level": {"name": "水平阻力", "description": "价格多次触及的水平阻力位"},
}


def _find_local_extrema(prices: np.ndarray, order: int = 5) -> tuple[list[int], list[int]]:
    highs, lows = [], []
    for i in range(order, len(prices) - order):
        if all(prices[i] >= prices[i - j] for j in range(1, order + 1)) and \
           all(prices[i] >= prices[i + j] for j in range(1, order + 1)):
            highs.append(i)
        if all(prices[i] <= prices[i - j] for j in range(1, order + 1)) and \
           all(prices[i] <= prices[i + j] for j in range(1, order + 1)):
            lows.append(i)
    return highs, lows


def _fit_line(x_indices: list[int], y_values: list[float]) -> tuple[float, float]:
    if len(x_indices) < 2:
        return 0.0, y_values[0] if y_values else 0.0
    x = np.array(x_indices, dtype=float)
    y = np.array(y_values, dtype=float)
    slope, intercept = np.polyfit(x, y, 1)
    return float(slope), float(intercept)


def _make_line(x1: int, y1: float, x2: int, y2: float, color: str, style: str = "solid", label: str = "") -> dict:
    return {"x1": x1, "y1": round(y1, 2), "x2": x2, "y2": round(y2, 2),
            "color": color, "style": style, "label": label}


def detect_trendlines(highs_arr: np.ndarray, lows_arr: np.ndarray, n: int, order: int = 5) -> list[PatternResult]:
    results = []
    high_idxs, low_idxs = _find_local_extrema(highs_arr, order)
    _, low_idxs2 = _find_local_extrema(lows_arr, order)

    if len(low_idxs2) >= 2:
        vals = [float(lows_arr[i]) for i in low_idxs2]
        slope, intercept = _fit_line(low_idxs2, vals)
        y1 = intercept + slope * low_idxs2[0]
        extend_y2 = intercept + slope * (n - 1)
        direction = "bullish" if slope > 0 else "bearish" if slope < 0 else "neutral"
        results.append(PatternResult(
            pattern_type="trendline_support", name="支撑趋势线",
            direction=direction, confidence=65 + min(len(low_idxs2), 5) * 5,
            start_idx=low_idxs2[0], end_idx=n - 1,
            description=f"连接{len(low_idxs2)}个低点，斜率{'上升' if slope > 0 else '下降'}",
            draw_lines=[_make_line(low_idxs2[0], y1, n - 1, extend_y2, "#3fb950", "dashed", "支撑线")],
        ))

    if len(high_idxs) >= 2:
        vals = [float(highs_arr[i]) for i in high_idxs]
        slope, intercept = _fit_line(high_idxs, vals)
        y1 = intercept + slope * high_idxs[0]
        extend_y2 = intercept + slope * (n - 1)
        direction = "bullish" if slope > 0 else "bearish"
        results.append(PatternResult(
            pattern_type="trendline_resistance", name="阻力趋势线",
            direction=direction, confidence=65 + min(len(high_idxs), 5) * 5,
            start_idx=high_idxs[0], end_idx=n - 1,
            description=f"连接{len(high_idxs)}个高点，斜率{'上升' if slope > 0 else '下降'}",
            draw_lines=[_make_line(high_idxs[0], y1, n - 1, extend_y2, "#f85149", "dashed", "阻力线")],
        ))

    return results


def detect_triangles(highs_arr: np.ndarray, lows_arr: np.ndarray, n: int) -> list[PatternResult]:
    results = []
    high_idxs, low_idxs = _find_local_extrema(highs_arr, 5)
    _, low_idxs2 = _find_local_extrema(lows_arr, 5)

    if len(high_idxs) >= 2 and len(low_idxs2) >= 2:
        h_slope, h_int = _fit_line(high_idxs, [float(highs_arr[i]) for i in high_idxs])
        l_slope, l_int = _fit_line(low_idxs2, [float(lows_arr[i]) for i in low_idxs2])

        h_y1 = h_int + h_slope * high_idxs[0]
        h_y2 = h_int + h_slope * high_idxs[-1]
        l_y1 = l_int + l_slope * low_idxs2[0]
        l_y2 = l_int + l_slope * low_idxs2[-1]

        h_range = abs(h_y2 - h_y1)
        l_range = abs(l_y2 - l_y1)
        avg_price = float(np.mean(highs_arr))
        threshold = avg_price * 0.003

        if h_range < threshold and l_slope > 0:
            results.append(PatternResult(
                pattern_type="ascending_triangle", name="上升三角形",
                direction="bullish", confidence=72,
                start_idx=min(high_idxs[0], low_idxs2[0]), end_idx=max(high_idxs[-1], low_idxs2[-1]),
                description="水平阻力线 + 上升支撑线，看涨突破",
                draw_lines=[
                    _make_line(high_idxs[0], h_y1, high_idxs[-1], h_y2, "#f85149", "solid", "水平阻力"),
                    _make_line(low_idxs2[0], l_y1, low_idxs2[-1], l_y2, "#3fb950", "solid", "上升支撑"),
                ],
            ))
        elif l_range < threshold and h_slope < 0:
            results.append(PatternResult(
                pattern_type="descending_triangle", name="下降三角形",
                direction="bearish", confidence=72,
                start_idx=min(high_idxs[0], low_idxs2[0]), end_idx=max(high_idxs[-1], low_idxs2[-1]),
                description="下降阻力线 + 水平支撑线，看跌突破",
                draw_lines=[
                    _make_line(high_idxs[0], h_y1, high_idxs[-1], h_y2, "#f85149", "solid", "下降阻力"),
                    _make_line(low_idxs2[0], l_y1, low_idxs2[-1], l_y2, "#3fb950", "solid", "水平支撑"),
                ],
            ))
        elif h_slope < 0 and l_slope > 0:
            results.append(PatternResult(
                pattern_type="symmetric_triangle", name="对称三角形",
                direction="neutral", confidence=68,
                start_idx=min(high_idxs[0], low_idxs2[0]), end_idx=max(high_idxs[-1], low_idxs2[-1]),
                description="收敛三角形，等待方向突破",
                draw_lines=[
                    _make_line(high_idxs[0], h_y1, high_idxs[-1], h_y2, "#f85149", "solid", "下降阻力"),
                    _make_line(low_idxs2[0], l_y1, low_idxs2[-1], l_y2, "#3fb950", "solid", "上升支撑"),
                ],
            ))
    return results


def detect_double_patterns(highs_arr: np.ndarray, lows_arr: np.ndarray, closes: np.ndarray, n: int) -> list[PatternResult]:
    results = []
    high_idxs, _ = _find_local_extrema(highs_arr, 5)
    _, low_idxs = _find_local_extrema(lows_arr, 5)

    if len(high_idxs) >= 2:
        for i in range(len(high_idxs) - 1):
            h1, h2 = high_idxs[i], high_idxs[i + 1]
            if h2 - h1 < 10:
                continue
            p1, p2 = float(highs_arr[h1]), float(highs_arr[h2])
            avg_p = (p1 + p2) / 2
            if abs(p1 - p2) / avg_p < 0.01:
                neckline = float(np.min(lows_arr[h1:h2 + 1]))
                results.append(PatternResult(
                    pattern_type="double_top", name="双顶 (M头)",
                    direction="bearish", confidence=75,
                    start_idx=h1, end_idx=h2,
                    description=f"两个高点 {p1:.2f}/{p2:.2f}，颈线 {neckline:.2f}",
                    draw_lines=[
                        _make_line(h1, p1, h2, p2, "#f85149", "solid", "双顶连线"),
                        _make_line(h1, neckline, h2, neckline, "#d29922", "dashed", "颈线"),
                    ],
                ))
                break

    if len(low_idxs) >= 2:
        for i in range(len(low_idxs) - 1):
            l1, l2 = low_idxs[i], low_idxs[i + 1]
            if l2 - l1 < 10:
                continue
            p1, p2 = float(lows_arr[l1]), float(lows_arr[l2])
            avg_p = (p1 + p2) / 2
            if abs(p1 - p2) / avg_p < 0.01:
                neckline = float(np.max(highs_arr[l1:l2 + 1]))
                results.append(PatternResult(
                    pattern_type="double_bottom", name="双底 (W底)",
                    direction="bullish", confidence=75,
                    start_idx=l1, end_idx=l2,
                    description=f"两个低点 {p1:.2f}/{p2:.2f}，颈线 {neckline:.2f}",
                    draw_lines=[
                        _make_line(l1, p1, l2, p2, "#3fb950", "solid", "双底连线"),
                        _make_line(l1, neckline, l2, neckline, "#d29922", "dashed", "颈线"),
                    ],
                ))
                break

    return results


def detect_support_resistance(highs_arr: np.ndarray, lows_arr: np.ndarray, n: int) -> list[PatternResult]:
    results = []
    price_range = float(np.max(highs_arr) - np.min(lows_arr))
    tolerance = price_range * 0.008

    high_idxs, _ = _find_local_extrema(highs_arr, 4)
    _, low_idxs = _find_local_extrema(lows_arr, 4)

    resistance_levels = []
    for i in high_idxs:
        p = float(highs_arr[i])
        found = False
        for lvl in resistance_levels:
            if abs(p - lvl["price"]) < tolerance:
                lvl["touches"] += 1
                lvl["indices"].append(i)
                lvl["price"] = (lvl["price"] + p) / 2
                found = True
                break
        if not found:
            resistance_levels.append({"price": p, "touches": 1, "indices": [i]})

    for lvl in resistance_levels:
        if lvl["touches"] >= 2:
            results.append(PatternResult(
                pattern_type="resistance_level", name="水平阻力",
                direction="bearish", confidence=60 + lvl["touches"] * 8,
                start_idx=lvl["indices"][0], end_idx=n - 1,
                description=f"阻力位 {lvl['price']:.2f}，触及 {lvl['touches']} 次",
                draw_lines=[_make_line(lvl["indices"][0], lvl["price"], n - 1, lvl["price"], "#f85149", "dotted", f"阻力 {lvl['price']:.2f}")],
            ))

    support_levels = []
    for i in low_idxs:
        p = float(lows_arr[i])
        found = False
        for lvl in support_levels:
            if abs(p - lvl["price"]) < tolerance:
                lvl["touches"] += 1
                lvl["indices"].append(i)
                lvl["price"] = (lvl["price"] + p) / 2
                found = True
                break
        if not found:
            support_levels.append({"price": p, "touches": 1, "indices": [i]})

    for lvl in support_levels:
        if lvl["touches"] >= 2:
            results.append(PatternResult(
                pattern_type="support_level", name="水平支撑",
                direction="bullish", confidence=60 + lvl["touches"] * 8,
                start_idx=lvl["indices"][0], end_idx=n - 1,
                description=f"支撑位 {lvl['price']:.2f}，触及 {lvl['touches']} 次",
                draw_lines=[_make_line(lvl["indices"][0], lvl["price"], n - 1, lvl["price"], "#3fb950", "dotted", f"支撑 {lvl['price']:.2f}")],
            ))

    return results


def detect_channels(highs_arr: np.ndarray, lows_arr: np.ndarray, n: int) -> list[PatternResult]:
    results = []
    high_idxs, _ = _find_local_extrema(highs_arr, 5)
    _, low_idxs = _find_local_extrema(lows_arr, 5)

    if len(high_idxs) >= 2 and len(low_idxs) >= 2:
        h_slope, h_int = _fit_line(high_idxs, [float(highs_arr[i]) for i in high_idxs])
        l_slope, l_int = _fit_line(low_idxs, [float(lows_arr[i]) for i in low_idxs])

        slope_diff = abs(h_slope - l_slope)
        avg_slope = (abs(h_slope) + abs(l_slope)) / 2

        if avg_slope > 0 and slope_diff / avg_slope < 0.4:
            start = min(high_idxs[0], low_idxs[0])
            end = max(high_idxs[-1], low_idxs[-1])
            h_y1, h_y2 = h_int + h_slope * start, h_int + h_slope * end
            l_y1, l_y2 = l_int + l_slope * start, l_int + l_slope * end

            if h_slope > 0 and l_slope > 0:
                ptype, pname, direction = "channel_up", "上升通道", "bullish"
            elif h_slope < 0 and l_slope < 0:
                ptype, pname, direction = "channel_down", "下降通道", "bearish"
            else:
                return results

            width_start = abs(h_y1 - l_y1)
            width_end = abs(h_y2 - l_y2)
            if width_end < width_start * 0.7:
                ptype = "wedge_rising" if h_slope > 0 else "wedge_falling"
                pname = "上升楔形" if h_slope > 0 else "下降楔形"
                direction = "bearish" if h_slope > 0 else "bullish"

            results.append(PatternResult(
                pattern_type=ptype, name=pname,
                direction=direction, confidence=70,
                start_idx=start, end_idx=end,
                description=f"{pname}，通道宽度 {width_start:.2f} → {width_end:.2f}",
                draw_lines=[
                    _make_line(start, h_y1, end, h_y2, "#f85149", "solid", "上轨"),
                    _make_line(start, l_y1, end, l_y2, "#3fb950", "solid", "下轨"),
                ],
            ))

    return results


def detect_patterns(klines: list[KlineData], enabled_patterns: list[str] | None = None) -> list[dict]:
    if len(klines) < 30:
        return []

    closes = np.array([k.close for k in klines])
    highs = np.array([k.high for k in klines])
    lows = np.array([k.low for k in klines])
    n = len(klines)

    all_patterns: list[PatternResult] = []

    all_types = enabled_patterns or list(PATTERN_REGISTRY.keys())

    need_trendline = any(t.startswith("trendline") for t in all_types)
    need_triangle = any(t.endswith("triangle") for t in all_types)
    need_double = any(t.startswith("double") for t in all_types)
    need_sr = any(t.endswith("level") for t in all_types)
    need_channel = any(t.startswith("channel") or t.startswith("wedge") for t in all_types)

    if need_trendline:
        all_patterns.extend(detect_trendlines(highs, lows, n))
    if need_triangle:
        all_patterns.extend(detect_triangles(highs, lows, n))
    if need_double:
        all_patterns.extend(detect_double_patterns(highs, lows, closes, n))
    if need_sr:
        all_patterns.extend(detect_support_resistance(highs, lows, n))
    if need_channel:
        all_patterns.extend(detect_channels(highs, lows, n))

    filtered = [p for p in all_patterns if p.pattern_type in all_types]
    filtered.sort(key=lambda p: p.confidence, reverse=True)

    return [p.to_dict() for p in filtered]


def get_pattern_registry() -> list[dict]:
    return [{"type": k, **v} for k, v in PATTERN_REGISTRY.items()]