"""
策略相关性分析 Handler（Issue #5）

跨品种 handler，不区分单一品种。

功能：
  1. 对多个策略（stk / cb / etf）各执行一次事件驱动回测（record_history=False）
  2. 对齐各策略 time_return 序列（按 trade_date inner join）
  3. 计算滚动 Pearson 相关系数序列
  4. 计算全区间 Pearson 相关矩阵（下三角 + 对角线）
  5. 计算 Jaccard 持仓重叠度矩阵（同品种策略按日均值，跨品种固定为 0）
  6. 以 inline JSON dict 形式 yield 三帧结果（与因子分析接口对齐）

输出帧格式：
  - correlation_matrix : {"labels": [...], "matrix": [[...], ...]}（下三角+对角线有值，上三角 null）
  - correlation_series : {"trade_dates": [...], "pairs": {"策略A+策略B": [...]}}
  - overlap_matrix     : {"labels": [...], "jaccard": [[...], ...], "overlap": [[...], ...]}（下三角+对角线）
"""

from typing import Any, Generator

import polars as pl

from ldnode import logger
from ldnode.calculator.operators import (
    align_returns,
    cb_event_backtest,
    correlation_matrix,
    etf_event_backtest,
    jaccard_overlap_matrix,
    rolling_correlation_series,
    stk_event_backtest,
)
from ldnode.calculator.plans.polars import maybe_resolve_joins
from ldnode.calculator.schemas import CB_BUILTIN_FACTORS, ETF_BUILTIN_FACTORS, STK_BUILTIN_FACTORS
from ldnode.calculator.schemas.enumerations import SecurityType
from ldnode.loader.files import CB_DATA, ETF_DATA, STK_DATA

from .models import TaskResponse
from .models.strategy_correlation import StrategyCorrelationParameters
from .utils import Timer, error_response, response, wrap_handler


# ---------------------------------------------------------------------------
# 内部辅助（仅保留品种路由，计算逻辑已迁至 operators/correlation.py）
# ---------------------------------------------------------------------------

_SCOPE_CONFIG: dict[SecurityType, tuple[str, list[str]]] = {
    SecurityType.STK: ("stk", STK_BUILTIN_FACTORS),
    SecurityType.CB:  ("cb",  CB_BUILTIN_FACTORS),
    SecurityType.ETF: ("etf", ETF_BUILTIN_FACTORS),
}

_DATA_FILE = {
    SecurityType.STK: STK_DATA,
    SecurityType.CB:  CB_DATA,
    SecurityType.ETF: ETF_DATA,
}

_BACKTEST_FN = {
    SecurityType.STK: stk_event_backtest,
    SecurityType.CB:  cb_event_backtest,
    SecurityType.ETF: etf_event_backtest,
}


def _run_backtest(entry: dict, merged_config: dict) -> pl.DataFrame:
    """
    根据品种类型调用对应回测函数，返回 res_df。

    返回值含列：trade_date、time_return、hold_securities 等。
    record_history=False 以节省内存；不需要 return_screen_df。
    """
    sec_type = SecurityType(entry["security_type"])

    if sec_type not in _BACKTEST_FN:
        raise ValueError(f"不支持的品种类型: {sec_type}")

    data_file = _DATA_FILE[sec_type]
    _, default_builtins = _SCOPE_CONFIG[sec_type]
    backtest_fn = _BACKTEST_FN[sec_type]

    lf = data_file.output_lazyframe()
    lf, builtin_factors = maybe_resolve_joins(
        lf, sec_type.value, {"strategy_config": merged_config}, default_builtins,
    )
    logger.calculator.info(
        f"[Correlation] 并表完成. label={entry.get('label')} "
        f"scope={sec_type.name} builtin_factors={len(builtin_factors)}"
    )

    res_df, _ = backtest_fn(
        lf=lf,
        parameters=merged_config,
        record_history=False,
        builtin_factors=builtin_factors,
    )
    return res_df



# ---------------------------------------------------------------------------
# Handler 主体
# ---------------------------------------------------------------------------

@wrap_handler("策略相关性分析", is_generator=True)
def handle_strategy_correlation_analysis(
    parameters: StrategyCorrelationParameters,
) -> Generator[TaskResponse[Any]]:
    """响应策略相关性分析任务（Issue #5）

    遍历多个策略（stk/cb/etf），依次执行回测，然后计算：
      - 滚动 Pearson 相关系数序列（逐对策略）
      - 全区间 Pearson 相关矩阵（N×N 下三角）
      - Jaccard 持仓重叠度矩阵（N×N 下三角）
    """
    task_id = parameters["task_id"]
    global_settings = parameters["global_settings"]
    strategies = parameters["strategies"]
    rolling_period = parameters["rolling_period"]

    # ------------------------------------------------------------------
    # 1. 入参校验
    # ------------------------------------------------------------------
    if len(strategies) < 2:
        yield error_response("strategies 至少需要 2 个策略")
        return

    labels = [s["label"] for s in strategies]
    if len(labels) != len(set(labels)):
        yield error_response("strategies 中存在重复的 label，每个策略的 label 必须唯一")
        return

    # ------------------------------------------------------------------
    # 2. 逐策略执行回测
    # ------------------------------------------------------------------
    # res_df_map    : {label: res_df}，保存完整回测结果（含 hold_securities）
    # series_map    : {label: DataFrame(trade_date, time_return)}，仅用于对齐计算
    # sec_type_map  : {label: security_type}，用于 Jaccard 跨品种判断
    res_df_map: dict[str, pl.DataFrame] = {}
    series_map: dict[str, pl.DataFrame] = {}
    sec_type_map: dict[str, SecurityType] = {}

    for entry in strategies:
        label = entry["label"]
        sec_type = SecurityType(entry["security_type"])
        sec_type_map[label] = sec_type

        # 合并全局配置到 core_config
        merged_config = dict(entry["core_config"])
        merged_config["date_range"] = {
            "start_date": global_settings["start_date"],
            "end_date": global_settings["end_date"],
        }
        merged_config["capital"] = global_settings["initial_capital"]
        merged_config["commission"] = global_settings["commission"]
        merged_config["benchmark"] = global_settings["benchmark"]

        with Timer(
            f"策略相关性分析-策略回测结束. label={label} task_id={task_id}",
            logger=logger.event_backtest,
        ):
            res_df = _run_backtest(entry, merged_config)

        res_df_map[label] = res_df
        # 仅保留对齐所需的两列
        series_map[label] = res_df.select(["trade_date", "time_return"])

        # 诊断日志：time_return 统计，用于排查全 null 问题
        tr = res_df["time_return"]
        logger.event_backtest.info(
            f"策略相关性分析-time_return统计. label={label} "
            f"count={tr.len()} null={tr.null_count()} "
            f"std={tr.std():.6f} min={tr.min():.6f} max={tr.max():.6f}"
        )

    # ------------------------------------------------------------------
    # 3 & 4. 对齐 time_return 序列（inner join on trade_date）
    # ------------------------------------------------------------------
    aligned_df = align_returns(series_map)

    # ------------------------------------------------------------------
    # 5. 计算滚动 Pearson 相关系数序列
    # ------------------------------------------------------------------
    with Timer(
        f"策略相关性分析-滚动相关系数计算结束. task_id={task_id}",
        logger=logger.event_backtest,
    ):
        correlation_series_data = rolling_correlation_series(aligned_df, rolling_period)

    yield response(
        data_type="correlation_series",
        data=correlation_series_data,
    )

    # ------------------------------------------------------------------
    # 6. 计算全区间 Pearson 相关矩阵
    # ------------------------------------------------------------------
    with Timer(
        f"策略相关性分析-相关矩阵计算结束. task_id={task_id}",
        logger=logger.event_backtest,
    ):
        correlation_matrix_data = correlation_matrix(aligned_df)

    yield response(
        data_type="correlation_matrix",
        data=correlation_matrix_data,
    )

    # ------------------------------------------------------------------
    # 7. 计算 Jaccard 持仓重叠度矩阵
    # ------------------------------------------------------------------
    with Timer(
        f"策略相关性分析-Jaccard矩阵计算结束. task_id={task_id}",
        logger=logger.event_backtest,
    ):
        overlap_matrix_data = jaccard_overlap_matrix(res_df_map, sec_type_map)

    yield response(
        data_type="overlap_matrix",
        data=overlap_matrix_data,
    )
