"""
批量识别对称三角形 - 从 pkl 文件读取 OHLCV 数据
每个 pkl 文件包含 108 个股票 × N 个交易日的矩阵
"""

import os
import sys
import pickle
import numpy as np
import pandas as pd

# 让脚本能找到 src/ 下的模块
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..", "src"))

from sym_triangle import detect_sym_triangle, line_y, pivots_fractal, fit_boundary_line

# ============================================================================
# 【可调参数区】
# ============================================================================

# --- 数据源 ---
DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "data")

# --- 窗口大小 ---
WINDOW = 400

# --- 枢轴点检测 ---
PIVOT_K = 20

# --- 边界线拟合 ---
BOUNDARY_N_SEGMENTS = 2
BOUNDARY_SOURCE = "full"

# --- 斜率约束 ---
UPPER_SLOPE_MAX = 0.10
LOWER_SLOPE_MIN = -0.10

# --- 触碰判定 ---
TOUCH_TOL = 0.10
TOUCH_LOSS_MAX = 0.10

# --- 收敛要求 ---
SHRINK_RATIO = 0.8

# --- 突破判定 ---
BREAK_TOL = 0.001
VOL_WINDOW = 20
VOL_K = 1.3
FALSE_BREAK_M = 5

# --- 输出控制 ---
PRINT_DEBUG = False  # 批量时关闭调试输出
SAVE_ALL_CHARTS = False  # True=保存所有股票图，False=只保存识别到的


# ============================================================================
# pkl 数据加载
# ============================================================================

class FakeModule:
    """空壳模块，绕过 model 依赖"""
    ndarray = np.ndarray


def load_pkl(pkl_path: str) -> dict:
    """加载 pkl 文件，返回字典 {mtx, dtes, tkrs, tkrs_name, ...}"""
    # 注入空壳模块
    sys.modules['model'] = FakeModule()
    sys.modules['model.index_info'] = FakeModule()
    
    with open(pkl_path, 'rb') as f:
        data = pickle.load(f)
    return data


def load_ohlcv_from_pkl(data_dir: str) -> tuple:
    """
    从 pkl 文件加载 OHLCV 数据
    
    Returns:
        open_mtx, high_mtx, low_mtx, close_mtx, volume_mtx: shape=(n_stocks, n_days)
        dates: shape=(n_days,) 真实日期 (如 20050104)
        tkrs: shape=(n_stocks,) 股票代码
        tkrs_name: shape=(n_stocks,) 股票名称
    """
    open_data = load_pkl(os.path.join(data_dir, "open.pkl"))
    high_data = load_pkl(os.path.join(data_dir, "high.pkl"))
    low_data = load_pkl(os.path.join(data_dir, "low.pkl"))
    close_data = load_pkl(os.path.join(data_dir, "close.pkl"))
    volume_data = load_pkl(os.path.join(data_dir, "volume.pkl"))
    
    # 使用 close 的元数据
    dates = close_data["dtes"]
    tkrs = close_data["tkrs"]
    tkrs_name = close_data["tkrs_name"]
    
    return (
        open_data["mtx"],
        high_data["mtx"],
        low_data["mtx"],
        close_data["mtx"],
        volume_data["mtx"],
        dates,
        tkrs,
        tkrs_name,
    )


def get_stock_df(
    stock_idx: int,
    open_mtx: np.ndarray,
    high_mtx: np.ndarray,
    low_mtx: np.ndarray,
    close_mtx: np.ndarray,
    volume_mtx: np.ndarray,
    dates: np.ndarray,
) -> pd.DataFrame:
    """提取单个股票的 DataFrame"""
    df = pd.DataFrame({
        "date": dates,
        "open": open_mtx[stock_idx, :],
        "high": high_mtx[stock_idx, :],
        "low": low_mtx[stock_idx, :],
        "close": close_mtx[stock_idx, :],
        "volume": volume_mtx[stock_idx, :],
    })
    # 过滤掉 NaN/0 值
    df = df.replace(0, np.nan).dropna().reset_index(drop=True)
    return df


# ============================================================================
# 绘图
# ============================================================================

def plot_sym_triangle(df: pd.DataFrame, res, stock_id: str, out_path: str) -> None:
    import matplotlib.pyplot as plt

    close = df["close"].to_numpy(dtype=float)
    x = np.arange(len(df), dtype=float)
    dates = df["date"].to_numpy()
    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)
    
    start_date = dates[start] if len(dates) > start else start
    end_date = dates[end] if len(dates) > end else end
    
    plt.title(
        f"[{stock_id}] sym_triangle: "
        f"range={start_date}-{end_date}, "
        f"slope=({a_u:.4f},{a_l:.4f}), "
        f"width_ratio={res.width_ratio:.2f}, "
        f"touches=({res.touches_upper},{res.touches_lower})"
    )
    
    if len(dates) > 0:
        step = max(1, len(dates) // 8)
        idx = np.arange(0, len(dates), step)
        plt.xticks(idx, dates[idx], rotation=45, ha="right")
    
    plt.legend()
    plt.tight_layout()
    plt.savefig(out_path, dpi=150)
    plt.close()


# ============================================================================
# 主流程
# ============================================================================

def main() -> None:
    print("=" * 60)
    print("Symmetric Triangle Batch Detection - from pkl files")
    print("=" * 60)
    
    # 1. 加载数据
    print("\n[1] Loading OHLCV pkl files...")
    open_mtx, high_mtx, low_mtx, close_mtx, volume_mtx, dates, tkrs, tkrs_name = load_ohlcv_from_pkl(DATA_DIR)
    n_stocks, n_days = close_mtx.shape
    print(f"    Stocks: {n_stocks}")
    print(f"    Days: {n_days}")
    print(f"    Date range: {dates[0]} ~ {dates[-1]}")
    
    # 2. 准备输出目录
    outputs_dir = os.path.join(os.path.dirname(__file__), "..", "..", "outputs", "sym_triangles")
    os.makedirs(outputs_dir, exist_ok=True)
    
    # 3. 遍历所有股票
    print(f"\n[2] Scanning {n_stocks} stocks...")
    detected = []
    
    for i in range(n_stocks):
        stock_code = tkrs[i]
        stock_name = tkrs_name[i]
        stock_id = f"{stock_code}"  # 使用真实股票代码
        
        # 提取单只股票数据
        df = get_stock_df(i, open_mtx, high_mtx, low_mtx, close_mtx, volume_mtx, dates)
        
        if len(df) < WINDOW:
            if PRINT_DEBUG:
                print(f"  [{stock_id}] 跳过: 数据不足 ({len(df)} < {WINDOW})")
            continue
        
        # 运行检测
        res = detect_sym_triangle(
            df,
            window=WINDOW,
            pivot_k=PIVOT_K,
            touch_tol=TOUCH_TOL,
            touch_loss_max=TOUCH_LOSS_MAX,
            shrink_ratio=SHRINK_RATIO,
            break_tol=BREAK_TOL,
            vol_window=VOL_WINDOW,
            vol_k=VOL_K,
            false_break_m=FALSE_BREAK_M,
            upper_slope_max=UPPER_SLOPE_MAX,
            lower_slope_min=LOWER_SLOPE_MIN,
            boundary_fit=True,
            boundary_n_segments=BOUNDARY_N_SEGMENTS,
            boundary_source=BOUNDARY_SOURCE,
        )
        
        if res is not None:
            detected.append((i, stock_id, res, df))
            out_path = os.path.join(outputs_dir, f"{stock_id}.png")
            plot_sym_triangle(df, res, stock_id, out_path)
            print(f"  [OK] {stock_id} -> {out_path}")
        elif PRINT_DEBUG:
            print(f"  [--] {stock_id} not detected")
    
    # 4. 汇总结果
    print("\n" + "=" * 60)
    print(f"Scan completed! {len(detected)}/{n_stocks} stocks have symmetric triangles")
    print("=" * 60)
    
    if detected:
        print("\nDetected stocks:")
        for i, stock_id, res, _ in detected:
            name = tkrs_name[i]
            print(f"  - {stock_id} ({name}): slope=({res.upper_coef[0]:.4f},{res.lower_coef[0]:.4f}), "
                  f"width_ratio={res.width_ratio:.2f}, breakout={res.breakout}")
        print(f"\nCharts saved to: {outputs_dir}")
    
    # 5. Save summary CSV
    if detected:
        summary_path = os.path.join(outputs_dir, "summary.csv")
        summary_data = []
        for i, stock_id, res, _ in detected:
            summary_data.append({
                "stock_idx": i,
                "stock_code": stock_id,
                "stock_name": tkrs_name[i],
                "start_date": dates[res.start],
                "end_date": dates[res.end],
                "upper_slope": res.upper_coef[0],
                "lower_slope": res.lower_coef[0],
                "width_ratio": res.width_ratio,
                "touches_upper": res.touches_upper,
                "touches_lower": res.touches_lower,
                "breakout": res.breakout,
            })
        pd.DataFrame(summary_data).to_csv(summary_path, index=False, encoding="utf-8-sig")
        print(f"Summary saved: {summary_path}")


if __name__ == "__main__":
    main()
