"""
相关性分析算子

提供策略/因子相关性分析所需的纯计算函数，全部基于 Polars 原生实现，无 scipy 依赖。

公开接口：
  - align_returns              : 多策略 time_return 对齐（inner join on trade_date）
  - rolling_correlation_series : 滚动 Pearson 相关系数序列
  - corr_matrix                : 通用相关矩阵（支持 pearson/spearman，下三角）
  - correlation_matrix         : corr_matrix 的 Pearson 别名（策略相关性使用）
  - factor_correlation_matrix  : corr_matrix 的 Spearman 别名（因子分析使用）
  - jaccard_overlap_matrix     : Jaccard+Overlap 双矩阵（下三角）
"""

from itertools import combinations
from typing import Any, Literal

import polars as pl

from ldnode.calculator.schemas.enumerations import SecurityType


def _normalize_corr_value(value: float | None) -> float | None:
    """将 Polars 相关系数结果中的 NaN 统一归一为 None，保持接口输出稳定。"""
    if value is None:
        return None
    return None if value != value else value


def align_returns(series_map: dict[str, pl.DataFrame]) -> pl.DataFrame:
    """
    将多个策略的 time_return 序列按 trade_date inner join 对齐。

    参数
    ----
    series_map : {label: DataFrame(trade_date, time_return)}

    返回
    ----
    宽表 DataFrame，列结构：trade_date, label1, label2, ...（各列为对应策略的 time_return）
    按 trade_date 升序排列。
    """
    labels = list(series_map.keys())
    base_df = series_map[labels[0]].select(
        ["trade_date", pl.col("time_return").alias(labels[0])]
    )
    for label in labels[1:]:
        right = series_map[label].select(
            ["trade_date", pl.col("time_return").alias(label)]
        )
        base_df = base_df.join(right, on="trade_date", how="inner")
    return base_df.sort("trade_date")


def rolling_correlation_series(
    aligned_df: pl.DataFrame,
    window_size: int,
) -> dict[str, Any]:
    """
    计算逐对策略的滚动 Pearson 相关系数序列。

    使用 Polars 原生 rolling_corr：窗口未满时输出 null，方差为 0 时输出 null。

    参数
    ----
    aligned_df  : align_returns 的输出（宽表）
    window_size : 滚动窗口大小（交易日数）

    返回
    ----
    {"trade_dates": ["2024-01-01", ...], "pairs": {"策略A+策略B": [0.8, null, ...]}}
    """
    labels = [c for c in aligned_df.columns if c != "trade_date"]
    trade_dates = [
        d.strftime("%Y-%m-%d") if hasattr(d, "strftime") else str(d)
        for d in aligned_df["trade_date"].to_list()
    ]
    pairs: dict[str, list[float | None]] = {}
    for label_i, label_j in combinations(labels, 2):
        pairs[f"{label_i}+{label_j}"] = (
            aligned_df.select(
                pl.rolling_corr(label_i, label_j, window_size=window_size)
            )
            .to_series()
            .to_list()
        )
    return {"trade_dates": trade_dates, "pairs": pairs}


def corr_matrix(
    df: pl.DataFrame,
    keys: list[str],
    method: Literal["pearson", "spearman"] = "pearson",
    min_samples: int = 0,
) -> dict[str, Any]:
    """
    通用相关矩阵算子（下三角 + 对角线）。

    使用 Polars 原生 pl.corr：方差为 0 时返回 null，直接映射为 Python None。

    参数
    ----
    df          : 包含待计算列的 DataFrame
    keys        : 参与计算的列名列表
    method      : 相关系数方法，"pearson"（默认）或 "spearman"
    min_samples : 最小有效样本数（>0 时对每对列先 drop_nulls，不足则返回 None）
                  因子分析传 3，策略相关性传 0（全量数据，无需过滤）

    返回
    ----
    {"labels": [...], "matrix": [[...], ...]}（上三角为 null）
    """
    n = len(keys)
    raw: list[list[float | None]] = [[None] * n for _ in range(n)]
    for i in range(n):
        for j in range(i + 1):
            if i == j:
                raw[i][j] = 1.0
            else:
                ki, kj = keys[i], keys[j]
                if ki not in df.columns or kj not in df.columns:
                    continue
                sub = df.select([ki, kj]).drop_nulls() if min_samples > 0 else df
                if min_samples > 0 and len(sub) < min_samples:
                    continue
                val = sub.select(pl.corr(ki, kj, method=method)).item()
                raw[i][j] = _normalize_corr_value(val)
                raw[j][i] = None  # 上三角显式置 None（已在初始化时设好，此处仅作说明）
    return {"labels": keys, "matrix": raw}


# ---------------------------------------------------------------------------
# 语义别名：让调用方代码保持可读性
# ---------------------------------------------------------------------------

def correlation_matrix(aligned_df: pl.DataFrame) -> dict[str, Any]:
    """策略全区间 Pearson 相关矩阵。corr_matrix 的 Pearson 别名。"""
    keys = [c for c in aligned_df.columns if c != "trade_date"]
    return corr_matrix(aligned_df, keys, method="pearson")


def factor_correlation_matrix(
    factor_df: pl.DataFrame,
    factor_keys: list[str],
) -> dict[str, Any]:
    """因子 Spearman 相关矩阵。corr_matrix 的 Spearman 别名，启用 min_samples=3 过滤。"""
    return corr_matrix(factor_df, factor_keys, method="spearman", min_samples=3)


def jaccard_overlap_matrix(
    res_df_map: dict[str, pl.DataFrame],
    security_type_map: dict[str, SecurityType],
) -> dict[str, Any]:
    """
    计算 Jaccard + Overlap 持仓重叠度双矩阵（下三角 + 对角线）。

    规则：
    - 对角线 = 1.0
    - 跨品种策略对 = 0.0（不存在共同标的）
    - 同品种：逐调仓日计算后取全区间均值
        Jaccard(A, B) = |A∩B| / |A∪B|
        Overlap(A, B) = |A∩B| / min(|A|, |B|)

    参数
    ----
    res_df_map        : {label: res_df}，res_df 含 trade_date、hold_securities 列
    security_type_map : {label: SecurityType}

    返回
    ----
    {
        "labels"  : [...],
        "jaccard" : [[...], ...],   # 下三角 + 对角线，上三角为 None
        "overlap" : [[...], ...],   # 下三角 + 对角线，上三角为 None
    }
    """
    labels = list(res_df_map.keys())
    n = len(labels)
    jaccard_cache: dict[tuple[int, int], float] = {}
    overlap_cache: dict[tuple[int, int], float] = {}
    for i in range(n):
        for j in range(i + 1):
            if i == j:
                jaccard_cache[(i, j)] = 1.0
                overlap_cache[(i, j)] = 1.0
                continue
            label_i, label_j = labels[i], labels[j]
            if security_type_map[label_i] != security_type_map[label_j]:
                jaccard_cache[(i, j)] = 0.0
                overlap_cache[(i, j)] = 0.0
                continue
            holdings_i = res_df_map[label_i].select(
                ["trade_date", pl.col("hold_securities").alias("hs_i")]
            )
            holdings_j = res_df_map[label_j].select(
                ["trade_date", pl.col("hold_securities").alias("hs_j")]
            )
            merged = holdings_i.join(holdings_j, on="trade_date", how="inner")
            if merged.is_empty():
                jaccard_cache[(i, j)] = 0.0
                overlap_cache[(i, j)] = 0.0
                continue
            jaccard_values: list[float] = []
            overlap_values: list[float] = []
            for row in merged.iter_rows(named=True):
                set_i = set(row["hs_i"]) if row["hs_i"] is not None else set()
                set_j = set(row["hs_j"]) if row["hs_j"] is not None else set()
                intersection = len(set_i & set_j)
                union = set_i | set_j
                jaccard_values.append(intersection / len(union) if union else 0.0)
                min_size = min(len(set_i), len(set_j))
                overlap_values.append(intersection / min_size if min_size > 0 else 0.0)
            jaccard_cache[(i, j)] = (
                sum(jaccard_values) / len(jaccard_values) if jaccard_values else 0.0
            )
            overlap_cache[(i, j)] = (
                sum(overlap_values) / len(overlap_values) if overlap_values else 0.0
            )

    jaccard_matrix: list[list[float | None]] = []
    overlap_matrix: list[list[float | None]] = []
    for row_idx in range(n):
        jrow: list[float | None] = []
        orow: list[float | None] = []
        for col_idx in range(n):
            if row_idx >= col_idx:
                jrow.append(float(jaccard_cache[(row_idx, col_idx)]))
                orow.append(float(overlap_cache[(row_idx, col_idx)]))
            else:
                jrow.append(None)
                orow.append(None)
        jaccard_matrix.append(jrow)
        overlap_matrix.append(orow)
    return {"labels": labels, "jaccard": jaccard_matrix, "overlap": overlap_matrix}
