"""缠论买卖点信号有效性评估：事件研究（event study）。

核心口径约定：
  - 入场时刻一律取 bsp.sure_time（笔被确认的那根K线收盘），而非分型时间 end_time。
    分型时间在当时是不可知的，用它回测等于开了未来函数。
  - BUY 视为做多，SELL 视为做空，收益按方向调整后统一为正=盈利。
"""
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_BSP_DIR

DEFAULT_HORIZONS = (1, 3, 5, 10, 20, 40)


@dataclass
class BspEvent:
    """一个可交易的买卖点事件。"""

    bsp_type: str
    direction: int  # +1 做多 / -1 做空
    fx_time: pd.Timestamp  # 分型时间（信号形态出现）
    entry_time: pd.Timestamp  # 确认时间（可交易）
    entry_idx: int
    entry_price: float
    lag_bars: int  # 确认滞后了多少根K线


def build_bi_zs(chan: TF_DF, zs_source: str) -> list:
    """构造笔中枢列表。

    seg  —— 在每个线段内部找笔中枢，中枢必须等所属线段成形，确认慢一层。
    pure —— 直接在扁平笔序列上滚动，不依赖线段，确认更快。
    """
    if zs_source == "seg":
        return chan.cal_bi_zs(chan.seg_list)
    if zs_source == "pure":
        return chan.cal_bi_zs_list_pure(chan.bi_list)
    raise ValueError(f"未知的中枢来源: {zs_source}")


def run_pipeline(df: pd.DataFrame, tf: str, zs_source: str = "seg") -> tuple[TF_DF, list]:
    """跑完整缠论 pipeline，返回引擎与买卖点列表。"""
    chan = TF_DF(df, 1, tf)
    bi_zs_list = build_bi_zs(chan, zs_source)
    bsp_list = chan.find_all_bsp(chan.bi_list, bi_zs_list) if bi_zs_list else []
    return chan, bsp_list


def _time_index_map(df: pd.DataFrame) -> dict[str, int]:
    """K线收盘时间 -> 行号。缠论内部把时间存成无时区字符串，按字符串对齐最稳。"""
    keys = df["date"].dt.strftime("%Y-%m-%d %H:%M:%S")
    return {k: i for i, k in enumerate(keys)}


def to_events(bsp_list: list, df: pd.DataFrame) -> list[BspEvent]:
    """把 ChanBSP 转成以确认时刻为准的可交易事件。"""
    idx_map = _time_index_map(df)
    closes = df["close"].to_numpy(dtype=float)
    events: list[BspEvent] = []

    for bsp in bsp_list:
        if not bsp.is_sure or bsp.sure_time is None:
            continue
        entry_key, fx_key = str(bsp.sure_time), str(bsp.end_time)
        if entry_key not in idx_map or fx_key not in idx_map:
            continue
        entry_idx = idx_map[entry_key]
        fx_idx = idx_map[fx_key]
        entry_ts = pd.Timestamp(entry_key)
        fx_ts = pd.Timestamp(fx_key)
        events.append(
            BspEvent(
                bsp_type=str(bsp.type).replace("Chan_BSP_TYPE.", ""),
                direction=1 if bsp.dir == Chan_BSP_DIR.BUY else -1,
                fx_time=fx_ts,
                entry_time=entry_ts,
                entry_idx=entry_idx,
                entry_price=float(closes[entry_idx]),
                lag_bars=entry_idx - fx_idx,
            )
        )
    return events


def forward_returns(
    events: list[BspEvent], df: pd.DataFrame, horizons=DEFAULT_HORIZONS
) -> pd.DataFrame:
    """计算每个事件在各持有期的方向调整收益，以及 MFE/MAE。"""
    closes = df["close"].to_numpy(dtype=float)
    highs = df["high"].to_numpy(dtype=float)
    lows = df["low"].to_numpy(dtype=float)
    n = len(df)
    rows = []

    for ev in events:
        row = {
            "bsp_type": ev.bsp_type,
            "direction": ev.direction,
            "side": "LONG" if ev.direction == 1 else "SHORT",
            "fx_time": ev.fx_time,
            "entry_time": ev.entry_time,
            "entry_idx": ev.entry_idx,
            "entry_price": ev.entry_price,
            "lag_bars": ev.lag_bars,
        }
        for h in horizons:
            j = ev.entry_idx + h
            if j >= n:
                row[f"ret_{h}"] = np.nan
                row[f"mfe_{h}"] = np.nan
                row[f"mae_{h}"] = np.nan
                continue
            seg = slice(ev.entry_idx + 1, j + 1)
            row[f"ret_{h}"] = ev.direction * (closes[j] - ev.entry_price) / ev.entry_price
            if ev.direction == 1:
                best, worst = highs[seg].max(), lows[seg].min()
            else:
                best, worst = lows[seg].min(), highs[seg].max()
            row[f"mfe_{h}"] = ev.direction * (best - ev.entry_price) / ev.entry_price
            row[f"mae_{h}"] = ev.direction * (worst - ev.entry_price) / ev.entry_price
        rows.append(row)

    return pd.DataFrame(rows)


def baseline_stats(df: pd.DataFrame, horizons=DEFAULT_HORIZONS) -> pd.DataFrame:
    """基准：全样本每根K线无条件持有的收益分布（多头视角）。"""
    closes = df["close"].to_numpy(dtype=float)
    rows = []
    for h in horizons:
        fwd = (closes[h:] - closes[:-h]) / closes[:-h]
        rows.append(
            {
                "horizon": h,
                "base_mean_long": fwd.mean(),
                "base_median_long": np.median(fwd),
                "base_winrate_long": (fwd > 0).mean(),
                "base_std": fwd.std(ddof=1),
            }
        )
    return pd.DataFrame(rows)


def _tstat(x: np.ndarray) -> float:
    if len(x) < 2:
        return np.nan
    sd = x.std(ddof=1)
    return np.nan if sd == 0 else float(x.mean() / (sd / np.sqrt(len(x))))


def summarize(
    fwd: pd.DataFrame, df: pd.DataFrame, horizons=DEFAULT_HORIZONS, by_type: bool = True
) -> pd.DataFrame:
    """汇总各类买卖点在各持有期的表现，并给出对基准的超额。"""
    base = baseline_stats(df, horizons).set_index("horizon")
    groups: list[tuple[str, pd.DataFrame]] = [("ALL", fwd)]
    if by_type:
        groups += [("ALL_LONG", fwd[fwd.direction == 1]), ("ALL_SHORT", fwd[fwd.direction == -1])]
        groups += [(t, g) for t, g in fwd.groupby("bsp_type")]

    rows = []
    for name, g in groups:
        if g.empty:
            continue
        for h in horizons:
            r = g[f"ret_{h}"].dropna().to_numpy()
            if len(r) == 0:
                continue
            # 基准需按方向调整：做空的无条件期望是多头期望的相反数
            dirs = g.loc[g[f"ret_{h}"].notna(), "direction"].to_numpy()
            base_mean = float(np.mean(dirs) * base.loc[h, "base_mean_long"])
            rows.append(
                {
                    "group": name,
                    "horizon": h,
                    "n": len(r),
                    "mean": r.mean(),
                    "median": np.median(r),
                    "winrate": (r > 0).mean(),
                    "excess": r.mean() - base_mean,
                    "tstat": _tstat(r),
                    "mfe": g[f"mfe_{h}"].dropna().mean(),
                    "mae": g[f"mae_{h}"].dropna().mean(),
                }
            )
    return pd.DataFrame(rows)


def fmt_pct(x: float) -> str:
    return "n/a" if pd.isna(x) else f"{x * 100:+.2f}%"