from datetime import date

import polars as pl
import pytest

from ldnode.calculator.operators.correlation import (
    align_returns,
    correlation_matrix,
    factor_correlation_matrix,
    jaccard_overlap_matrix,
    rolling_correlation_series,
)
from ldnode.calculator.schemas.enumerations import SecurityType


def test_align_returns_uses_inner_join_and_sorts_trade_date() -> None:
    series_map = {
        "beta": pl.DataFrame(
            {
                "trade_date": [date(2024, 1, 3), date(2024, 1, 1), date(2024, 1, 2)],
                "time_return": [0.03, 0.01, 0.02],
            }
        ),
        "alpha": pl.DataFrame(
            {
                "trade_date": [date(2024, 1, 4), date(2024, 1, 2), date(2024, 1, 3)],
                "time_return": [0.04, 0.20, 0.30],
            }
        ),
    }

    aligned = align_returns(series_map)

    assert aligned.to_dict(as_series=False) == {
        "trade_date": [date(2024, 1, 2), date(2024, 1, 3)],
        "beta": [0.02, 0.03],
        "alpha": [0.20, 0.30],
    }


def test_rolling_correlation_series_outputs_expected_pairs_and_null_prefix() -> None:
    aligned_df = pl.DataFrame(
        {
            "trade_date": [date(2024, 1, d) for d in range(1, 6)],
            "A": [1.0, 2.0, 3.0, 4.0, 5.0],
            "B": [2.0, 4.0, 6.0, 8.0, 10.0],
            "C": [5.0, 4.0, 3.0, 2.0, 1.0],
        }
    )

    result = rolling_correlation_series(aligned_df, window_size=3)

    assert result["trade_dates"] == [
        "2024-01-01",
        "2024-01-02",
        "2024-01-03",
        "2024-01-04",
        "2024-01-05",
    ]
    assert result["pairs"]["A+B"][:2] == [None, None]
    assert result["pairs"]["A+B"][2:] == pytest.approx([1.0, 1.0, 1.0])
    assert result["pairs"]["A+C"][:2] == [None, None]
    assert result["pairs"]["A+C"][2:] == pytest.approx([-1.0, -1.0, -1.0])
    assert result["pairs"]["B+C"][:2] == [None, None]
    assert result["pairs"]["B+C"][2:] == pytest.approx([-1.0, -1.0, -1.0])


def test_correlation_matrix_keeps_lower_triangle() -> None:
    aligned_df = pl.DataFrame(
        {
            "trade_date": [date(2024, 1, d) for d in range(1, 5)],
            "A": [1.0, 2.0, 3.0, 4.0],
            "B": [2.0, 4.0, 6.0, 8.0],
            "C": [4.0, 3.0, 2.0, 1.0],
        }
    )

    result = correlation_matrix(aligned_df)

    assert result["labels"] == ["A", "B", "C"]
    assert result["matrix"] == [
        [1.0, None, None],
        [1.0, 1.0, None],
        [-1.0, -1.0, 1.0],
    ]


def test_factor_correlation_matrix_respects_min_samples_threshold() -> None:
    factor_df = pl.DataFrame(
        {
            "f1": [1.0, 2.0, None, None],
            "f2": [1.0, 2.0, None, None],
            "f3": [1.0, 3.0, 5.0, 7.0],
        }
    )

    result = factor_correlation_matrix(factor_df, ["f1", "f2", "f3"])

    assert result["labels"] == ["f1", "f2", "f3"]
    assert result["matrix"] == [
        [1.0, None, None],
        [None, 1.0, None],
        [None, None, 1.0],
    ]


def test_jaccard_overlap_matrix_handles_cross_security_and_average_overlap() -> None:
    res_df_map = {
        "stk_a": pl.DataFrame(
            {
                "trade_date": [date(2024, 1, 1), date(2024, 1, 2)],
                "hold_securities": [["000001", "000002"], ["000002", "000003"]],
            }
        ),
        "stk_b": pl.DataFrame(
            {
                "trade_date": [date(2024, 1, 1), date(2024, 1, 2)],
                "hold_securities": [["000002", "000004"], ["000002", "000003"]],
            }
        ),
        "cb_a": pl.DataFrame(
            {
                "trade_date": [date(2024, 1, 1), date(2024, 1, 2)],
                "hold_securities": [["110001"], ["110002"]],
            }
        ),
    }
    security_type_map = {
        "stk_a": SecurityType.STK,
        "stk_b": SecurityType.STK,
        "cb_a": SecurityType.CB,
    }

    result = jaccard_overlap_matrix(res_df_map, security_type_map)

    assert result["labels"] == ["stk_a", "stk_b", "cb_a"]
    assert result["matrix"] == [
        [1.0, None, None],
        [2.0 / 3.0, 1.0, None],
        [0.0, 0.0, 1.0],
    ]
