from __future__ import annotations

from dataclasses import dataclass
from typing import List, Literal, Optional, Tuple

import numpy as np
import pandas as pd


@dataclass
class SymTriangleResult:
    start: int
    end: int
    upper_coef: Tuple[float, float]  # (a_u, b_u)
    lower_coef: Tuple[float, float]  # (a_l, b_l)
    width_start: float
    width_end: float
    width_ratio: float
    touches_upper: int
    touches_lower: int
    apex_x: float
    breakout: Literal["up", "down", "none"]
    breakout_idx: Optional[int]
    volume_confirmed: Optional[bool]
    false_breakout: Optional[bool]


def pivots_fractal(
    high: np.ndarray, low: np.ndarray, k: int = 3
) -> Tuple[np.ndarray, np.ndarray]:
    """左右窗口分形：返回 pivot_high_idx, pivot_low_idx。"""
    n = len(high)
    ph: List[int] = []
    pl: List[int] = []
    for i in range(k, n - k):
        if high[i] == np.max(high[i - k : i + k + 1]):
            ph.append(i)
        if low[i] == np.min(low[i - k : i + k + 1]):
            pl.append(i)
    return np.array(ph, dtype=int), np.array(pl, dtype=int)


def fit_line(x: np.ndarray, y: np.ndarray) -> Tuple[float, float]:
    """拟合 y = a*x + b。"""
    a, b = np.polyfit(x, y, deg=1)
    return float(a), float(b)


def fit_boundary_line(
    x: np.ndarray, y: np.ndarray, mode: str = "upper", n_segments: int = 3
) -> Tuple[float, float]:
    """
    边界线拟合（分段取极值）：把时间轴分成 n_segments 段，每段取一个极值点。
    - mode="upper": 每段取最高点（让线连接各段峰值）
    - mode="lower": 每段取最低点（让线连接各段谷底）
    这样可以确保拟合的点在时间上分散，更接近手动画线。
    """
    if len(x) < 2:
        return fit_line(x, y)
    
    n_segments = min(n_segments, len(x))
    if n_segments < 2:
        n_segments = 2
    
    # 按 x 排序
    sort_idx = np.argsort(x)
    x_sorted = x[sort_idx]
    y_sorted = y[sort_idx]
    
    # 分段取极值
    segment_size = len(x_sorted) // n_segments
    boundary_x = []
    boundary_y = []
    
    for i in range(n_segments):
        start = i * segment_size
        if i == n_segments - 1:
            end = len(x_sorted)  # 最后一段包含剩余所有点
        else:
            end = (i + 1) * segment_size
        
        if start >= end:
            continue
        
        seg_x = x_sorted[start:end]
        seg_y = y_sorted[start:end]
        
        if mode == "upper":
            idx = np.argmax(seg_y)
        else:
            idx = np.argmin(seg_y)
        
        boundary_x.append(seg_x[idx])
        boundary_y.append(seg_y[idx])
    
    if len(boundary_x) < 2:
        return fit_line(x, y)
    
    return fit_line(np.array(boundary_x), np.array(boundary_y))


def line_y(a: float, b: float, x: np.ndarray) -> np.ndarray:
    return a * x + b


def detect_sym_triangle(
    df: pd.DataFrame,
    window: int = 120,
    pivot_k: int = 3,
    touch_tol: float = 0.006,
    touch_loss_max: float = 0.015,
    shrink_ratio: float = 0.6,
    break_tol: float = 0.003,
    vol_window: int = 20,
    vol_k: float = 1.5,
    false_break_m: int = 5,
    upper_slope_max: float = 0.0,
    lower_slope_min: float = 0.0,
    boundary_fit: bool = True,
    boundary_n_segments: int = 3,
    boundary_source: str = "pivots",
) -> Optional[SymTriangleResult]:
    """
    只检测“最近一个窗口”是否存在对称三角形，并给出突破/确认/假突破信息。
    - 实时：false_breakout 只能返回 None（因为需要未来 m 根确认）
    - 回测：如果 df 包含突破后的 m 根数据，才会输出 false_breakout True/False
    """
    required = {"open", "high", "low", "close"}
    if not required.issubset(df.columns):
        raise ValueError(f"df must contain columns: {sorted(required)}")

    n = len(df)
    if n < max(window, 2 * pivot_k + 5):
        return None

    high = df["high"].to_numpy(dtype=float)
    low = df["low"].to_numpy(dtype=float)
    close = df["close"].to_numpy(dtype=float)
    has_volume = "volume" in df.columns and df["volume"].notna().any()
    volume = df["volume"].to_numpy(dtype=float) if has_volume else None

    ph_idx, pl_idx = pivots_fractal(high, low, k=pivot_k)

    end = n - 1
    start = max(0, end - window + 1)
    x_all = np.arange(n, dtype=float)

    ph_in = ph_idx[(ph_idx >= start) & (ph_idx <= end)]
    pl_in = pl_idx[(pl_idx >= start) & (pl_idx <= end)]
    if len(ph_in) < 2 or len(pl_in) < 2:
        return None

    # 拟合两条线（边界拟合：让线贴着峰/谷）
    if boundary_fit:
        if boundary_source == "full":
            x_upper = x_all[start : end + 1]
            y_upper = high[start : end + 1]
            x_lower = x_all[start : end + 1]
            y_lower = low[start : end + 1]
        else:
            x_upper = x_all[ph_in]
            y_upper = high[ph_in]
            x_lower = x_all[pl_in]
            y_lower = low[pl_in]
        a_u, b_u = fit_boundary_line(x_upper, y_upper, mode="upper", n_segments=boundary_n_segments)
        a_l, b_l = fit_boundary_line(x_lower, y_lower, mode="lower", n_segments=boundary_n_segments)
    else:
        a_u, b_u = fit_line(x_all[ph_in], high[ph_in])
        a_l, b_l = fit_line(x_all[pl_in], low[pl_in])

    # 斜率：对称三角形硬条件（close-only 可放宽上/下沿斜率）
    if not (a_u <= upper_slope_max and a_l >= lower_slope_min):
        return None

    # 宽度收敛
    upper_start = float(line_y(a_u, b_u, np.array([start]))[0])
    lower_start = float(line_y(a_l, b_l, np.array([start]))[0])
    upper_end = float(line_y(a_u, b_u, np.array([end]))[0])
    lower_end = float(line_y(a_l, b_l, np.array([end]))[0])
    width_start = upper_start - lower_start
    width_end = upper_end - lower_end
    if width_start <= 0 or width_end <= 0:
        return None
    width_ratio = width_end / width_start
    if width_ratio > shrink_ratio:
        return None

    # 触碰程度（用 pivot 点贴线程度判定）
    ph_dist = np.abs(high[ph_in] - line_y(a_u, b_u, x_all[ph_in])) / np.maximum(
        line_y(a_u, b_u, x_all[ph_in]), 1e-9
    )
    pl_dist = np.abs(low[pl_in] - line_y(a_l, b_l, x_all[pl_in])) / np.maximum(
        line_y(a_l, b_l, x_all[pl_in]), 1e-9
    )
    touches_upper = int((ph_dist <= touch_tol).sum())
    touches_lower = int((pl_dist <= touch_tol).sum())
    # 用损失函数替代“>=2触碰”的硬条件
    loss_upper = float(np.mean(ph_dist)) if len(ph_dist) else float("inf")
    loss_lower = float(np.mean(pl_dist)) if len(pl_dist) else float("inf")
    if loss_upper > touch_loss_max or loss_lower > touch_loss_max:
        return None

    # apex（两线交点）仅做合理性输出（可扩展为硬条件）
    denom = a_u - a_l
    apex_x = float((b_l - b_u) / denom) if abs(denom) > 1e-12 else float("inf")

    # === 突破判定（用最后一根 close 做实时判断）===
    upper_last = upper_end
    lower_last = lower_end
    breakout: Literal["up", "down", "none"] = "none"
    breakout_idx: Optional[int] = None
    if close[end] > upper_last * (1 + break_tol):
        breakout = "up"
        breakout_idx = end
    elif close[end] < lower_last * (1 - break_tol):
        breakout = "down"
        breakout_idx = end

    # 成交量确认（可选）
    volume_confirmed: Optional[bool] = None
    if breakout != "none" and has_volume and volume is not None:
        vol_ma = pd.Series(volume).rolling(vol_window).mean().to_numpy()
        if np.isfinite(vol_ma[end]) and vol_ma[end] > 0:
            volume_confirmed = bool(volume[end] > vol_ma[end] * vol_k)
        else:
            volume_confirmed = None

    # 假突破（需要未来 m 根数据，实时无法得知）
    false_breakout: Optional[bool] = None
    if breakout != "none" and breakout_idx is not None:
        # 如果数据不足以观察未来 m 根，则返回 None
        if breakout_idx + false_break_m < n:
            false_breakout = False
            for t in range(breakout_idx + 1, breakout_idx + false_break_m + 1):
                upper_t = float(line_y(a_u, b_u, np.array([t]))[0])
                lower_t = float(line_y(a_l, b_l, np.array([t]))[0])
                if breakout == "up" and close[t] < upper_t:
                    false_breakout = True
                    break
                if breakout == "down" and close[t] > lower_t:
                    false_breakout = True
                    break
        else:
            false_breakout = None

    return SymTriangleResult(
        start=start,
        end=end,
        upper_coef=(a_u, b_u),
        lower_coef=(a_l, b_l),
        width_start=float(width_start),
        width_end=float(width_end),
        width_ratio=float(width_ratio),
        touches_upper=touches_upper,
        touches_lower=touches_lower,
        apex_x=apex_x,
        breakout=breakout,
        breakout_idx=breakout_idx,
        volume_confirmed=volume_confirmed,
        false_breakout=false_breakout,
    )


def plot_sym_triangle(df: pd.DataFrame, res: SymTriangleResult) -> None:
    """简单可视化验算（需要 matplotlib）。"""
    import matplotlib.pyplot as plt

    close = df["close"].to_numpy(dtype=float)
    x = np.arange(len(df), dtype=float)
    a_u, b_u = res.upper_coef
    a_l, b_l = res.lower_coef

    start, end = res.start, res.end
    xw = np.arange(start, end + 1, dtype=float)
    upper = line_y(a_u, b_u, xw)
    lower = line_y(a_l, b_l, xw)

    plt.figure(figsize=(12, 5))
    plt.plot(x, close, linewidth=1.2, label="close")
    plt.plot(xw, upper, linewidth=2, label="upper")
    plt.plot(xw, lower, linewidth=2, label="lower")
    plt.axvline(end, color="gray", linestyle="--", linewidth=1)
    plt.title(
        "sym_triangle: "
        f"width_ratio={res.width_ratio:.2f}, "
        f"touches=({res.touches_upper},{res.touches_lower}), "
        f"breakout={res.breakout}, "
        f"vol_ok={res.volume_confirmed}, "
        f"false={res.false_breakout}"
    )
    plt.legend()
    plt.show()
