"""
因子分析核心计算逻辑（stk/cb/etf 共用）

主要职责：
1. 识别调仓日并计算各期持仓内因子值
2. 计算持有期收益（close 价格法）
3. 计算每期 Rank IC（Spearman 相关）
4. 汇总统计指标、滚动 IC 均值、累计 IC
5. 计算因子相关矩阵
"""

import re
from typing import Any, Generator

import polars as pl

from ldnode import logger
from ldnode.calculator.handlers.models import TaskResponse
from ldnode.calculator.handlers.utils import Timer, response
from ldnode.calculator.operators.correlation import _normalize_corr_value
from ldnode.calculator.operators import factor_correlation_matrix


_HOLD_SECURITY_CODE_PATTERN = re.compile(r"\(([^()]+)\)$")


def _normalize_security_code(security: str) -> str:
    """兼容 `名称(代码)` 与纯 `代码` 两种持仓标识格式，统一提取为纯代码。"""
    matched = _HOLD_SECURITY_CODE_PATTERN.search(security)
    return matched.group(1) if matched else security


def _normalize_holdings(holdings: list[str] | None) -> list[str]:
    """批量规范化持仓标识，供因子分析与相关矩阵计算统一使用。"""
    if not holdings:
        return []
    return [_normalize_security_code(security) for security in holdings]


def _build_factor_analysis_fact_table(
    res_df: pl.DataFrame,
    screen_df: pl.DataFrame,
    factor_keys: list[str],
) -> tuple[list, pl.DataFrame]:
    """
    构造因子分析统一事实表，避免按调仓期重复切片与 join。

    返回:
        valid_dates: 可计算 IC 的调仓日列表（最后一期无 next date，已剔除）
        fact_df: 列结构
            rebalance_date / next_trade_date / code / close_ti / close_tnext /
            holding_return / factor_keys...
    """
    all_rebalance_dates = _identify_rebalance_dates(res_df)
    if len(all_rebalance_dates) < 2:
        return [], pl.DataFrame()

    valid_dates = all_rebalance_dates[:-1]
    next_dates = all_rebalance_dates[1:]
    rebalance_pair_df = pl.DataFrame(
        {
            "rebalance_date": valid_dates,
            "next_trade_date": next_dates,
        }
    )

    # 使用全排名截面（rank 不为 null）作为 IC 计算范围，
    # 替代原来仅使用持仓标的的方式，使 IC 反映因子对全候选池的预测力（学术标准）
    ranked_securities_df = (
        screen_df
        .filter(
            pl.col("trade_date").is_in(valid_dates),
            pl.col("rank").is_not_null(),
        )
        .select(
            pl.col("trade_date").alias("rebalance_date"),
            pl.col("code").cast(pl.String),
        )
        .unique()
    )

    if ranked_securities_df.is_empty():
        return valid_dates, pl.DataFrame()

    current_factor_df = screen_df.select(
        [
            pl.col("trade_date").alias("rebalance_date"),
            pl.col("code").cast(pl.String),  # screen_df.code 为 cat，需与 normalized_holdings_df.code(str) 对齐
            pl.col("close").alias("close_ti"),
            *factor_keys,
        ]
    )
    next_close_df = screen_df.select(
        [
            pl.col("trade_date").alias("next_trade_date"),
            pl.col("code").cast(pl.String),  # 同上
            pl.col("close").alias("close_tnext"),
        ]
    )

    fact_df = (
        rebalance_pair_df
        .join(ranked_securities_df, on="rebalance_date", how="inner")
        .join(current_factor_df, on=["rebalance_date", "code"], how="inner")
        .join(next_close_df, on=["next_trade_date", "code"], how="inner")
        .with_columns(
            ((pl.col("close_tnext") / pl.col("close_ti")) - 1.0).alias("holding_return")
        )
        .sort(["rebalance_date", "code"])
    )
    return valid_dates, fact_df


def _compute_factor_analysis(
    res_df: pl.DataFrame,
    screen_df: pl.DataFrame,
    factor_keys: list[str],
    rolling_period: int,
    task_id: str,
    log_prefix: str,
) -> Generator[TaskResponse[Any], None, None]:
    """
    因子分析核心计算，生成三帧输出。

    Args:
        res_df: 回测结果 DataFrame，含 trade_date / hold_securities 等列
        screen_df: 完整截面 DataFrame，含 code / trade_date / close / 各因子列
        factor_keys: 需要分析的因子名称列表
        rolling_period: 滚动 IC 均值窗口期数
        task_id: 任务 ID（用于日志）
        log_prefix: 日志前缀（如 "股票因子分析"）
    """
    with Timer(f"{log_prefix} 识别调仓日. 任务ID: {task_id}", logger=logger.event_backtest):
        rebalance_dates, ic_by_factor = _calc_ic_series(res_df, screen_df, factor_keys)

    with Timer(f"{log_prefix} 汇总统计指标. 任务ID: {task_id}", logger=logger.event_backtest):
        ic_stats = _calc_ic_stats(ic_by_factor, factor_keys)
        ic_series_data = _calc_ic_series_data(
            rebalance_dates, ic_by_factor, factor_keys, rolling_period
        )
        corr_matrix = _calc_factor_correlation(screen_df, res_df, factor_keys)

    # 帧1：各因子 IC 统计指标
    yield response(
        data_type="factor_ic_stats",
        data={"factors": ic_stats},
    )

    # 帧2：IC 时序数据（用于图表）
    yield response(
        data_type="factor_ic_series",
        data={
            "rebalance_dates": [d.strftime("%Y-%m-%d") for d in rebalance_dates],
            "factors": ic_series_data,
        },
    )

    # 帧3：因子相关矩阵（下三角）
    yield response(
        data_type="factor_correlation_matrix",
        data=corr_matrix,
    )


def _identify_rebalance_dates(res_df: pl.DataFrame) -> list:
    """
    识别调仓日：hold_securities 从 T-1 到 T 发生变化（或第一天）的日期列表。

    res_df 需含 trade_date（date 类型或字符串）、hold_securities（List[str]）列。
    返回已排序的日期对象列表。
    """
    # 按 trade_date 升序排列
    df = res_df.sort("trade_date")
    dates = df["trade_date"].to_list()
    holdings = df["hold_securities"].to_list()

    rebalance_dates = []
    prev_holdings = None
    for d, h in zip(dates, holdings):
        normalized_holdings = _normalize_holdings(h)
        h_set = set(normalized_holdings)
        prev_set = set(_normalize_holdings(prev_holdings))
        # 第一天直接算调仓日；或持仓集合发生变化时
        if prev_holdings is None or h_set != prev_set:
            rebalance_dates.append(d)
        prev_holdings = normalized_holdings

    return rebalance_dates


def _calc_ic_series(
    res_df: pl.DataFrame,
    screen_df: pl.DataFrame,
    factor_keys: list[str],
) -> tuple[list, dict[str, list[float | None]]]:
    """
    计算各调仓期的 Rank IC（Spearman 相关）。

    返回:
        rebalance_dates: 有效的调仓日列表（最后一个调仓期因无 T_next 而被剔除）
        ic_by_factor: {factor_key: [ic_T1, ic_T2, ...]}
    """
    valid_dates, fact_df = _build_factor_analysis_fact_table(res_df, screen_df, factor_keys)
    if not valid_dates:
        return [], {k: [] for k in factor_keys}

    ic_by_factor: dict[str, list[float | None]] = {k: [] for k in factor_keys}
    fact_by_date = {
        key[0]: frame
        for key, frame in fact_df.partition_by("rebalance_date", as_dict=True).items()
    }

    for t_i in valid_dates:
        merged = fact_by_date.get(t_i)
        if merged is None or len(merged) < 3:
            # 无持仓，IC 设为 None
            for k in factor_keys:
                ic_by_factor[k].append(None)
            continue

        holding_returns = merged["holding_return"].to_list()

        # 对每个因子计算 Spearman IC
        for k in factor_keys:
            if k not in merged.columns:
                ic_by_factor[k].append(None)
                continue
            factor_vals = merged[k].to_list()

            # 过滤 None / NaN
            pairs = [
                (f, r)
                for f, r in zip(factor_vals, holding_returns)
                if f is not None and r is not None
                and f == f and r == r  # NaN 自比较为 False
            ]
            if len(pairs) < 3:
                ic_by_factor[k].append(None)
                continue

            fv, rv = zip(*pairs)
            # Polars Spearman：方差为 0 时返回 null → Python None
            ic_df = pl.DataFrame({"f": list(fv), "r": list(rv)})
            ic_by_factor[k].append(
                _normalize_corr_value(
                    ic_df.select(pl.corr("f", "r", method="spearman")).item()
                )
            )

    return valid_dates, ic_by_factor


def _calc_ic_stats(
    ic_by_factor: dict[str, list[float | None]],
    factor_keys: list[str],
) -> list[dict]:
    """汇总每个因子的 IC 统计指标"""
    result = []
    for k in factor_keys:
        ic_list = ic_by_factor.get(k, [])
        # 只保留非 None 的有效 IC 值
        valid_ic = [v for v in ic_list if v is not None]

        if not valid_ic:
            result.append({
                "factor": k,
                "current_ic": None,
                "ic_mean": None,
                "ic_std": None,
                "ir": None,
                "ic_win_rate": None,
                "ic_max": None,
                "ic_min": None,
                "ic_skew": None,
                "ic_kurt": None,
            })
            continue

        ic_arr = pl.Series(valid_ic)
        ic_mean = float(ic_arr.mean())
        ic_std = float(ic_arr.std(ddof=1)) if len(valid_ic) > 1 else None
        ir = (ic_mean / ic_std) if ic_std else None
        _skew = ic_arr.skew()
        _kurt = ic_arr.kurtosis()

        result.append({
            "factor": k,
            # current_ic = 最后一个有效 IC（对应最后一个有数据的调仓期）
            "current_ic": valid_ic[-1],
            "ic_mean": ic_mean,
            "ic_std": ic_std,
            "ir": ir,
            "ic_win_rate": float(sum(1 for v in valid_ic if v > 0) / len(valid_ic)),
            "ic_max": float(ic_arr.max()),
            "ic_min": float(ic_arr.min()),
            "ic_skew": float(_skew) if _skew is not None else None,
            "ic_kurt": float(_kurt) if _kurt is not None else None,
        })

    return result


def _calc_ic_series_data(
    rebalance_dates: list,
    ic_by_factor: dict[str, list[float | None]],
    factor_keys: list[str],
    rolling_period: int,
) -> dict[str, dict]:
    """
    计算每个因子的 IC 时序数据：
    - ic: 原始 IC 序列
    - rolling_ic_mean: 滚动均值（前 rolling_period-1 期为 None）
    - cumulative_ic: 累计 IC（前缀和）
    """
    result = {}
    n = len(rebalance_dates)

    for k in factor_keys:
        ic_list = ic_by_factor.get(k, [None] * n)

        # 滚动 IC 均值
        rolling = []
        for i in range(n):
            window = ic_list[max(0, i - rolling_period + 1): i + 1]
            if len(window) < rolling_period:
                rolling.append(None)
            else:
                # None 视为 0 参与均值，分母固定为 rolling_period（而非 valid count）
                # 边界：rolling_period=1 时窗口始终满足，None IC 输出 0.0（非 None）
                total = sum(v if v is not None else 0.0 for v in window)
                rolling.append(total / rolling_period)

        # 累计 IC（前缀和，None 值贡献 0）
        cumulative = []
        running = 0.0
        for v in ic_list:
            running += v if v is not None else 0.0
            cumulative.append(running)

        result[k] = {
            "ic": ic_list,
            "rolling_ic_mean": rolling,
            "cumulative_ic": cumulative,
        }

    return result


def _calc_factor_correlation(
    screen_df: pl.DataFrame,
    res_df: pl.DataFrame,
    factor_keys: list[str],
) -> dict:
    """
    计算因子 Spearman 相关矩阵（基于全时段持仓内因子值）。

    方法：收集所有调仓日持仓内的因子值，rank 后计算 Pearson 相关矩阵。
    返回下三角矩阵（i > j 有值，i == j 为 1.0，i < j 为 None）。
    """
    n = len(factor_keys)

    if n == 0:
        return {"labels": [], "matrix": []}

    if n == 1:
        return {"labels": factor_keys, "matrix": [[1.0]]}

    _, fact_df = _build_factor_analysis_fact_table(res_df, screen_df, factor_keys)
    if fact_df.is_empty():
        # 无数据，返回对角矩阵
        matrix = _build_lower_triangle_matrix(
            [[1.0 if i == j else None for j in range(n)] for i in range(n)]
        )
        return {"labels": factor_keys, "matrix": matrix}

    return factor_correlation_matrix(fact_df.select(factor_keys), factor_keys)


def _build_lower_triangle_matrix(raw: list[list]) -> list[list]:
    """保持下三角矩阵格式（i < j 设为 None）"""
    n = len(raw)
    result = []
    for i in range(n):
        row = []
        for j in range(n):
            if j > i:
                row.append(None)
            else:
                row.append(raw[i][j])
        result.append(row)
    return result
