from datetime import date
from types import SimpleNamespace

import polars as pl
import pytest

from ldnode.calculator.handlers.stk_factor_analysis_handler import (
    handle_stk_factor_analysis,
)
from ldnode.calculator.handlers.strategy_correlation_handler import (
    handle_strategy_correlation_analysis,
)
from ldnode.calculator.schemas.enumerations import SecurityType


def _build_strategy_res_df(time_returns: list[float], holdings: list[list[str]]) -> pl.DataFrame:
    return pl.DataFrame(
        {
            "trade_date": [date(2024, 1, d) for d in range(1, len(time_returns) + 1)],
            "time_return": time_returns,
            "hold_securities": holdings,
        }
    )


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


def _build_factor_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", "CCC", "DDD", "AAA", "CCC", "DDD"],
            "close": [10.0, 11.0, 12.0, 13.0, 15.0, 17.0, 14.0, 16.0, 18.0],
            "factor_a": [1.0, 2.0, 3.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
            "factor_b": [3.0, 2.0, 1.0, 1.0, 0.0, -1.0, -2.0, -3.0, -4.0],
        }
    )


def test_strategy_correlation_handler_yields_three_business_frames_and_fee(monkeypatch) -> None:
    import ldnode.calculator.handlers.strategy_correlation_handler as strategy_handler

    res_df_map = {
        "alpha": _build_strategy_res_df(
            [0.01, 0.02, 0.03],
            [["000001", "000002"], ["000001", "000003"], ["000001", "000003"]],
        ),
        "beta": _build_strategy_res_df(
            [0.02, 0.04, 0.06],
            [["000002", "000004"], ["000001", "000003"], ["000001", "000003"]],
        ),
    }

    monkeypatch.setattr(
        strategy_handler,
        "_run_backtest",
        lambda entry, merged_config: res_df_map[entry["label"]],
    )

    results = list(
        handle_strategy_correlation_analysis(
            {
                "task_id": "task-strategy-handler",
                "global_settings": {
                    "start_date": "2024-01-01",
                    "end_date": "2024-01-31",
                    "initial_capital": 1_000_000.0,
                    "commission": 0.001,
                    "benchmark": "000300.SH",
                },
                "strategies": [
                    {
                        "security_type": SecurityType.STK,
                        "label": "alpha",
                        "core_config": {},
                    },
                    {
                        "security_type": SecurityType.STK,
                        "label": "beta",
                        "core_config": {},
                    },
                ],
                "rolling_period": 2,
            }
        )
    )

    assert [item["data_type"] for item in results] == [
        "correlation_series",
        "correlation_matrix",
        "overlap_matrix",
        "fee",
    ]
    assert results[0]["data"]["pairs"]["alpha+beta"][0] is None
    assert results[0]["data"]["pairs"]["alpha+beta"][1:] == pytest.approx([1.0, 1.0])
    assert results[1]["data"]["matrix"] == [[1.0, None], [1.0, 1.0]]
    assert results[2]["data"]["matrix"] == [[1.0, None], [7.0 / 9.0, 1.0]]


def test_stk_factor_analysis_handler_returns_error_when_rank_factors_empty() -> None:
    results = list(
        handle_stk_factor_analysis(
            {
                "task_id": "task-empty-rank-factors",
                "strategy_config": {"rank_factors": []},
                "rolling_period": 3,
            }
        )
    )

    assert [item["data_type"] for item in results] == ["error", "fee"]
    assert results[0]["data"] == "rank_factors 为空，无法进行因子分析"


def test_stk_factor_analysis_handler_yields_factor_frames_and_fee(monkeypatch) -> None:
    import ldnode.calculator.handlers.stk_factor_analysis_handler as stk_handler

    monkeypatch.setattr(
        stk_handler,
        "STK_DATA",
        SimpleNamespace(output_lazyframe=lambda: None),
    )
    monkeypatch.setattr(
        stk_handler,
        "stk_event_backtest",
        lambda **kwargs: (_build_factor_res_df(), None, _build_factor_screen_df()),
    )

    results = list(
        handle_stk_factor_analysis(
            {
                "task_id": "task-stk-factor",
                "strategy_config": {
                    "rank_factors": [
                        {"factor": "factor_a"},
                        {"factor": "factor_b"},
                    ]
                },
                "rolling_period": 2,
            }
        )
    )

    assert [item["data_type"] for item in results] == [
        "factor_ic_stats",
        "factor_ic_series",
        "factor_correlation_matrix",
        "fee",
    ]
    assert [item["factor"] for item in results[0]["data"]["factors"]] == [
        "factor_a",
        "factor_b",
    ]
    assert results[1]["data"]["rebalance_dates"] == ["2024-01-01"]
    assert results[2]["data"]["labels"] == ["factor_a", "factor_b"]
