"""分型级信号：绕开笔/线段/中枢的确认链，直接用分型 + 背驰做预设转折点。

动机：笔确认滞后 9~10 根、线段 101~121 根、买卖点 16~21 根，全部超过 alpha 半衰期（4~6 根）。
而分型只需右侧 KLC 完成即可确认，滞后通常 1~3 根，是唯一来得及的结构。
"""
from __future__ import annotations

import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from chanlun import TF_DF
from chanlun.core.ChanEnum import Chan_FX_TYPE


@dataclass
class FxSignal:
    """一个分型转折信号。confirm_idx 是实时可交易时刻。"""

    direction: int          # +1 底分型（潜在买） / -1 顶分型（潜在卖）
    fx_idx: int             # 分型极值所在K线
    confirm_idx: int        # 右侧KLC完成、分型可被确认的K线
    lag: int                # confirm_idx - fx_idx
    price: float            # 分型极值价
    confirm_price: float    # 确认时刻收盘价
    seg_macd_area: float    # 本段（前一反向分型 -> 本分型）的 MACD 面积
    prev_seg_macd_area: float   # 前一个同向段的 MACD 面积
    prev_extreme: float     # 前一个同向分型的极值价
    is_divergence: bool     # 是否背驰：价格创新极值但力度衰减
    ratio: float            # 面积比 本段/前段，越小背驰越强


def extract_fx_signals(chan: TF_DF, df: pd.DataFrame) -> list[FxSignal]:
    """从已建好的缠论结构里抽取分型信号，并就地算好背驰。"""
    idx_of = {t: i for i, t in enumerate(df["date"].dt.strftime("%Y-%m-%d %H:%M:%S"))}
    n = len(df)
    closes = df["close"].to_numpy(dtype=float)
    macdhist = (
        df["macdhist"].to_numpy(dtype=float)
        if "macdhist" in df.columns
        else np.zeros(n)
    )
    if "macdhist" not in df.columns and hasattr(chan, "dataframe"):
        if "macdhist" in chan.dataframe.columns:
            macdhist = chan.dataframe["macdhist"].to_numpy(dtype=float)

    # 按时间收集已成型的分型
    raw = []
    for klc in chan.klc_list:
        if klc.fx not in (Chan_FX_TYPE.TOP, Chan_FX_TYPE.BOTTOM):
            continue
        if klc.next is None or klc.next.end_klu is None:
            continue
        e_key = str(klc.end_time)
        c_key = str(klc.next.end_klu.time)
        if e_key not in idx_of or c_key not in idx_of:
            continue
        fx_idx = idx_of[e_key]
        confirm_idx = idx_of[c_key]
        if confirm_idx <= fx_idx:
            continue
        d = 1 if klc.fx == Chan_FX_TYPE.BOTTOM else -1
        raw.append({
            "d": d,
            "fx_idx": fx_idx,
            "confirm_idx": confirm_idx,
            "price": float(klc.low if d == 1 else klc.high),
        })

    raw.sort(key=lambda r: r["fx_idx"])

    def area(lo: int, hi: int, sign: int) -> float:
        """区间内顺方向的 MACD 柱面积。sign=-1 取负柱（下跌段），+1 取正柱。"""
        if hi <= lo:
            return 0.0
        seg = macdhist[lo : hi + 1]
        vals = seg[seg < 0] if sign < 0 else seg[seg > 0]
        return float(np.abs(vals).sum())

    signals: list[FxSignal] = []
    for i, r in enumerate(raw):
        d = r["d"]
        # 本段起点 = 上一个反向分型；前一同向段 = 再往前一组
        prev_opp = None
        prev_same = None
        prev_opp2 = None
        for j in range(i - 1, -1, -1):
            if prev_opp is None and raw[j]["d"] == -d:
                prev_opp = raw[j]
                continue
            if prev_opp is not None and prev_same is None and raw[j]["d"] == d:
                prev_same = raw[j]
                continue
            if prev_same is not None and prev_opp2 is None and raw[j]["d"] == -d:
                prev_opp2 = raw[j]
                break
        if prev_opp is None:
            continue

        sign = -1 if d == 1 else 1  # 底分型前是下跌段，取负柱
        cur_area = area(prev_opp["fx_idx"], r["fx_idx"], sign)
        prev_area = (
            area(prev_opp2["fx_idx"], prev_same["fx_idx"], sign)
            if (prev_same is not None and prev_opp2 is not None)
            else 0.0
        )

        # 背驰：价格创新极值（底更低 / 顶更高）但力度反而衰减
        new_extreme = False
        prev_extreme = np.nan
        if prev_same is not None:
            prev_extreme = prev_same["price"]
            new_extreme = (
                r["price"] <= prev_extreme if d == 1 else r["price"] >= prev_extreme
            )
        ratio = cur_area / prev_area if prev_area > 0 else np.nan
        is_div = bool(new_extreme and prev_area > 0 and cur_area < prev_area)

        signals.append(
            FxSignal(
                direction=d,
                fx_idx=r["fx_idx"],
                confirm_idx=r["confirm_idx"],
                lag=r["confirm_idx"] - r["fx_idx"],
                price=r["price"],
                confirm_price=float(closes[r["confirm_idx"]]),
                seg_macd_area=cur_area,
                prev_seg_macd_area=prev_area,
                prev_extreme=float(prev_extreme) if prev_extreme == prev_extreme else np.nan,
                is_divergence=is_div,
                ratio=float(ratio) if ratio == ratio else np.nan,
            )
        )
    return signals


def signals_to_frame(signals: list[FxSignal]) -> pd.DataFrame:
    return pd.DataFrame([s.__dict__ for s in signals])


def add_forward_returns(
    sig: pd.DataFrame, df: pd.DataFrame, horizons=(3, 5, 10, 20, 40)
) -> pd.DataFrame:
    """以 confirm_idx 收盘价入场的方向调整收益。"""
    closes = df["close"].to_numpy(dtype=float)
    n = len(df)
    out = sig.copy()
    for h in horizons:
        vals = []
        for i, d in zip(out["confirm_idx"], out["direction"]):
            j = int(i) + h
            vals.append(d * (closes[j] - closes[int(i)]) / closes[int(i)] if j < n else np.nan)
        out[f"ret_{h}"] = vals
    return out