from datetime import date
import math

import polars as pl
import pytest

from ldnode.calculator.handlers._factor_analysis_core import (
    _build_factor_analysis_fact_table,
    _calc_factor_correlation,
    _calc_ic_series,
    _calc_ic_series_data,
    _calc_ic_stats,
    _identify_rebalance_dates,
    _normalize_holdings,
    _normalize_security_code,
)


def _build_res_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "trade_date": [date(2024, 1, 1), date(2024, 1, 2), date(2024, 1, 3)],
            "hold_securities": [
                ["Alpha(AAA)", "Beta(BBB)", "Gamma(CCC)"],
                ["Alpha(AAA)", "Gamma(CCC)", "Delta(DDD)"],
                ["Alpha(AAA)", "Gamma(CCC)", "Delta(DDD)"],
            ],
        }
    )


def _build_screen_df() -> pl.DataFrame:
    return pl.DataFrame(
        {
            "trade_date": [
                date(2024, 1, 1),
                date(2024, 1, 1),
                date(2024, 1, 1),
                date(2024, 1, 2),
                date(2024, 1, 2),
                date(2024, 1, 2),
                date(2024, 1, 3),
                date(2024, 1, 3),
                date(2024, 1, 3),
            ],
            "code": ["AAA", "BBB", "CCC", "AAA", "BBB", "CCC", "AAA", "CCC", "DDD"],
            "close": [10.0, 10.0, 10.0, 11.0, 12.0, 13.0, 12.0, 14.0, 16.0],
            "factor_a": [1.0, 2.0, 3.0, 5.0, 4.0, 6.0, 7.0, 8.0, 9.0],
            "factor_b": [3.0, 2.0, 1.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            "factor_constant": [1.0] * 9,
        }
    )


def test_identify_rebalance_dates_only_marks_first_day_and_holdings_change() -> None:
    res_df = _build_res_df()

    result = _identify_rebalance_dates(res_df)

    assert result == [date(2024, 1, 1), date(2024, 1, 2)]


def test_normalize_security_code_supports_plain_code_and_name_code() -> None:
    assert _normalize_security_code("123013.SZ") == "123013.SZ"
    assert _normalize_security_code("横河转债(123013.SZ)") == "123013.SZ"
    assert _normalize_holdings(["横河转债(123013.SZ)", "113576.SH"]) == [
        "123013.SZ",
        "113576.SH",
    ]


def test_calc_ic_series_returns_expected_rank_ic_values() -> None:
    res_df = _build_res_df()
    screen_df = _build_screen_df()

    rebalance_dates, ic_by_factor = _calc_ic_series(
        res_df,
        screen_df,
        ["factor_a", "factor_b", "factor_constant"],
    )

    assert rebalance_dates == [date(2024, 1, 1)]
    assert ic_by_factor["factor_a"] == pytest.approx([1.0])
    assert ic_by_factor["factor_b"] == pytest.approx([-1.0])
    assert ic_by_factor["factor_constant"] == [None]


def test_build_factor_analysis_fact_table_reuses_normalized_holdings_and_computes_returns() -> None:
    valid_dates, fact_df = _build_factor_analysis_fact_table(
        _build_res_df(),
        _build_screen_df(),
        ["factor_a", "factor_b"],
    )

    assert valid_dates == [date(2024, 1, 1)]
    assert fact_df.select("code").to_series().to_list() == ["AAA", "BBB", "CCC"]
    assert fact_df.select("holding_return").to_series().to_list() == pytest.approx(
        [0.1, 0.2, 0.3]
    )
    assert fact_df.columns == [
        "rebalance_date",
        "next_trade_date",
        "code",
        "close_ti",
        "factor_a",
        "factor_b",
        "close_tnext",
        "holding_return",
    ]


def test_calc_ic_stats_uses_last_valid_ic_and_ignores_none() -> None:
    result = _calc_ic_stats(
        {
            "factor_a": [None, 0.5, -0.5],
            "factor_b": [None, None],
        },
        ["factor_a", "factor_b"],
    )

    factor_a = next(item for item in result if item["factor"] == "factor_a")
    factor_b = next(item for item in result if item["factor"] == "factor_b")

    assert factor_a["current_ic"] == -0.5
    assert factor_a["ic_mean"] == 0.0
    assert math.isclose(factor_a["ic_std"], math.sqrt(0.5), rel_tol=1e-9)
    assert factor_a["ir"] == 0.0
    assert factor_a["ic_win_rate"] == 0.5
    assert factor_a["ic_max"] == 0.5
    assert factor_a["ic_min"] == -0.5
    assert factor_b == {
        "factor": "factor_b",
        "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,
    }


def test_calc_ic_series_data_builds_rolling_mean_and_cumulative_sum() -> None:
    result = _calc_ic_series_data(
        [date(2024, 1, 1), date(2024, 1, 2), date(2024, 1, 3)],
        {"factor_a": [1.0, None, -0.5]},
        ["factor_a"],
        rolling_period=2,
    )

    assert result == {
        "factor_a": {
            "ic": [1.0, None, -0.5],
            "rolling_ic_mean": [None, 1.0, -0.5],
            "cumulative_ic": [1.0, 1.0, 0.5],
        }
    }


def test_calc_factor_correlation_collects_rebalance_holdings_and_returns_lower_triangle() -> None:
    result = _calc_factor_correlation(
        _build_screen_df(),
        _build_res_df(),
        ["factor_a", "factor_b", "factor_constant"],
    )

    assert result["labels"] == ["factor_a", "factor_b", "factor_constant"]
    assert result["matrix"][0] == [1.0, None, None]
    assert result["matrix"][1][0] < 0
    assert result["matrix"][1][1] == 1.0
    assert result["matrix"][2] == [None, None, 1.0]
