import math

import pytest
import polars as pl
from numpy.polynomial import Polynomial

from ldnode.loader.datasources.cb_quotations.conv_prem_plan import plan_conv_prem


def _expected_theory_conv_prem(conv_value: float, sample: pl.DataFrame) -> float:
    frame = sample.with_columns(pl.col("conv_prem").log().alias("conv_prem_log"))
    pdf = frame.to_pandas()
    polynomial = Polynomial.fit(
        pdf["conv_value"].astype("float"),
        pdf["conv_prem_log"].astype("float"),
        1,
    ).convert()
    intercept = polynomial.coef[0]
    slope = polynomial.coef[1]
    return math.exp(conv_value * slope + intercept)


def test_plan_conv_prem_excludes_announced_force_redeem_from_regression():
    frame = pl.DataFrame(
        {
            "code": ["A", "B", "C", "FORCE"],
            "conv_value": [80.0, 90.0, 100.0, 85.0],
            "conv_prem": [0.30, 0.24, 0.20, 2.40],
            "close": [120.0, 120.0, 120.0, 120.0],
            "left_years": [3.0, 3.0, 3.0, 3.0],
            "is_call": ["0/15 | 30", "公告不强赎", "未到转股期", "已公告强赎"],
        }
    )
    expected_sample = frame.filter(pl.col("code") != "FORCE").select(
        ["conv_value", "conv_prem"]
    )

    result = plan_conv_prem(frame.lazy()).collect()
    actual = result.filter(pl.col("code") == "A").select("theory_conv_prem").item()
    expected = _expected_theory_conv_prem(80.0, expected_sample)

    assert actual == pytest.approx(expected)
