"""
深入分析cs_rank在有NA值时的行为

问题：cs_rank(expr) -> rank() / count()，当分组内有NA值时：
- rank() 默认跳过NA值
- count() 应该只计算非NA值的数量

但 Polars 的行为可能与预期不同，需要验证。
"""

import polars as pl
import numpy as np


def test_cs_rank_with_na():
    """测试 cs_rank 在有NA值时的行为。"""
    print("=" * 80)
    print("cs_rank 在有NA值时的行为测试")
    print("=" * 80)

    # 测试1：同一截面内有NA值
    print("\n测试1：同一截面内有NA值")
    df = pl.DataFrame(
        {
            "trade_date": ["20250101"] * 5,
            "ts_code": ["001", "002", "003", "004", "005"],
            "value": [1.0, 2.0, 3.0, None, 5.0],
        }
    )
    print("原始数据:")
    print(df)

    result = (
        df.lazy()
        .with_columns(
            [
                pl.col("value").rank().alias("rank"),
                pl.col("value").count().over("trade_date").alias("count"),
            ]
        )
        .with_columns([(pl.col("rank") / pl.col("count")).alias("normalized_rank")])
        .with_columns(
            [
                pl.col("value").rank(method="average").alias("rank_avg"),
                pl.col("value").rank(method="ordinal").alias("rank_ordinal"),
                pl.col("value").rank(method="min").alias("rank_min"),
                pl.col("value").rank(method="max").alias("rank_max"),
                pl.col("value").rank(method="dense").alias("rank_dense"),
            ]
        )
        .collect()
    )
    print("\n排名结果:")
    print(
        result.select(
            [
                "ts_code",
                "value",
                "rank",
                "count",
                "normalized_rank",
                "rank_avg",
                "rank_ordinal",
                "rank_min",
                "rank_dense",
            ]
        )
    )

    # 测试2：验证 count() 的确切行为
    print("\n测试2：验证 count() 的确切行为")
    print("在Polars中，count() 返回分组内的非空值数量")

    # 测试3：模拟 cs_rank(ts_corr(...)) 的情况
    print("\n测试3：模拟 cs_rank(ts_corr(...)) 的嵌套情况")
    df2 = pl.DataFrame(
        {
            "trade_date": ["20250101"] * 5 + ["20250102"] * 5,
            "ts_code": ["001", "002", "003", "004", "005"] * 2,
            "corr_value": [
                0.5,
                0.6,
                None,
                0.8,
                0.9,  # 20250101有NA
                0.3,
                0.4,
                0.5,
                0.6,
                0.7,
            ],  # 20250102无NA
        }
    )
    print("模拟ts_corr输出数据:")
    print(df2)

    result2 = (
        df2.lazy()
        .with_columns(
            [
                pl.col("corr_value").rank().alias("rank"),
                pl.col("corr_value").count().over("trade_date").alias("count"),
            ]
        )
        .with_columns([(pl.col("rank") / pl.col("count")).alias("normalized_rank")])
        .collect()
    )
    print("\ncs_rank结果:")
    print(
        result2.select(
            ["trade_date", "ts_code", "corr_value", "rank", "count", "normalized_rank"]
        )
    )

    # 分析
    print("\n分析:")
    print(
        "如果 count() 只计算非NA值，那么 20250101 的 count 应该是 4，20250102 的 count 应该是 5"
    )
    print("这会导致同一因子在不同日期的排名基准不同")


def test_rolling_corr_boundary():
    """测试 Polars rolling_corr 在边界处的行为。"""
    print("\n" + "=" * 80)
    print("Polars rolling_corr 边界行为测试")
    print("=" * 80)

    # 创建测试数据：不同起始点的相同时间序列
    df1 = pl.DataFrame(
        {
            "trade_date": ["20250101", "20250102", "20250103", "20250104", "20250105"],
            "ts_code": ["001"] * 5,
            "x": [1.0, 2.0, 3.0, 4.0, 5.0],
            "y": [1.0, 2.0, 3.0, 4.0, 5.0],
        }
    ).sort("trade_date")

    df2 = pl.DataFrame(
        {
            "trade_date": [
                "20241229",
                "20241230",
                "20241231",
                "20250101",
                "20250102",
                "20250103",
                "20250104",
                "20250105",
            ],
            "ts_code": ["001"] * 8,
            "x": [-1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            "y": [-1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
        }
    ).sort("trade_date")

    print("\n数据集1（从20250101开始）:")
    print(df1)

    print("\n数据集2（从20241229开始）:")
    print(df2)

    # 计算 rolling_corr(window=3)
    result1 = (
        df1.lazy()
        .with_columns(
            [pl.rolling_corr("x", "y", window_size=3).over("ts_code").alias("corr")]
        )
        .collect()
    )

    result2 = (
        df2.lazy()
        .with_columns(
            [pl.rolling_corr("x", "y", window_size=3).over("ts_code").alias("corr")]
        )
        .collect()
    )

    print("\n数据集1的rolling_corr结果:")
    print(result1.select(["trade_date", "x", "y", "corr"]))

    print("\n数据集2的rolling_corr结果（筛选相同日期）:")
    print(
        result2.filter(
            pl.col("trade_date").is_in(
                ["20250101", "20250102", "20250103", "20250104", "20250105"]
            )
        ).select(["trade_date", "x", "y", "corr"])
    )


if __name__ == "__main__":
    test_cs_rank_with_na()
    test_rolling_corr_boundary()
