from datetime import date

import polars as pl
import pytest

from ldnode.calculator.plans.polars import (
    log_factor_weight_stats,
    planning_screen_weight,
    planning_weight_column,
)
from ldnode.calculator.schemas import HoldRangeType, HoldWeight


def _factor_hold_method():
    return {
        "hold_range_type": HoldRangeType.NUM,
        "hold_range_value": {"start": 1, "end": 6},
        "hold_threshold": 0,
        "hold_weight": HoldWeight.FACTOR,
        "max_hold_weight": 1,
        "hold_weight_factor": "factor_alpha",
    }


def test_screen_factor_weight_returns_original_factor_values():
    frame = pl.DataFrame(
        {
            "trade_date": [date(2026, 6, 8)] * 6,
            "factor_alpha": [1.0, 2.0, 3.0, 0.0, 0.0, 0.0],
        }
    )

    result = planning_screen_weight(frame.lazy(), _factor_hold_method()).collect()

    assert result["_weight"].to_list() == [1.0, 2.0, 3.0, 0.0, 0.0, 0.0]


def test_vector_factor_weight_uses_positive_denominator_and_ignores_zero():
    frame = pl.DataFrame(
        {
            "trade_date": [date(2026, 6, 8)] * 6,
            "code": ["A", "B", "C", "D", "E", "F"],
            "rank": [1, 2, 3, 4, 5, 6],
            "factor_alpha": [1.0, 2.0, 3.0, 0.0, 0.0, 0.0],
        }
    )
    parameters = {"hold_method": _factor_hold_method()}

    result = planning_weight_column(frame.lazy(), parameters).collect()

    assert result["weight"].to_list() == pytest.approx(
        [1.0 / 6.0, 2.0 / 6.0, 3.0 / 6.0, 0.0, 0.0, 0.0]
    )


def test_factor_weight_stats_blocks_all_invalid_values():
    frame = pl.DataFrame(
        {
            "trade_date": [date(2026, 6, 8)] * 3,
            "factor_alpha": [0.0, 0.0, 0.0],
        }
    )

    with pytest.raises(ValueError, match="FACTOR 权重因子全部无效"):
        log_factor_weight_stats(frame, "unit_test", "factor_alpha")


def test_factor_weight_stats_blocks_null_values():
    frame = pl.DataFrame(
        {
            "trade_date": [date(2026, 6, 8)] * 3,
            "factor_alpha": [1.0, None, 2.0],
        }
    )

    with pytest.raises(ValueError, match="FACTOR 权重因子存在空值"):
        log_factor_weight_stats(frame, "unit_test", "factor_alpha")


def test_factor_weight_stats_blocks_negative_values():
    frame = pl.DataFrame(
        {
            "trade_date": [date(2026, 6, 8)] * 3,
            "factor_alpha": [1.0, -1.0, 2.0],
        }
    )

    with pytest.raises(ValueError, match="FACTOR 权重因子存在负值"):
        log_factor_weight_stats(frame, "unit_test", "factor_alpha")
