import math

import pandas as pd

from rarelens_ml.features import RAW_COLUMNS, build


def raw(**overrides: list) -> pd.DataFrame:
    base = {
        "impact": ["HIGH", "LOW", None],
        "consequence": ["stop_gained", "synonymous_variant", None],
        "cadd_phred": ["35", "2.1", "-"],
        "am_pathogenicity": ["0.98", None, "-"],
    }
    base.update(overrides)
    return pd.DataFrame(base, index=[10, 11, 12])


def test_raw_columns_are_the_serving_contract() -> None:
    assert RAW_COLUMNS == ["impact", "consequence", "cadd_phred", "am_pathogenicity"]


def test_allele_frequency_is_not_a_feature() -> None:
    """It dominated the model and the ranking already scores it, auditably and only once.

    Keeping it here also meant learning ACMG's own frequency-based benign rule from labels that
    rule produced, which is most of why the headline AUROC looked so good.
    """
    assert "gnomad_af" not in RAW_COLUMNS
    assert "gnomad_af" not in build(raw(gnomad_af=[0.0, 0.5, None])).columns


def test_build_ranks_impact_and_coerces_numbers() -> None:
    out = build(raw())
    assert out["impact_rank"].tolist() == [3, 1, 0]
    assert out["cadd_phred"].iloc[0] == 35.0
    assert math.isnan(out["cadd_phred"].iloc[2])  # VEP writes "-" for missing
    assert math.isnan(out["am_pathogenicity"].iloc[1])


def test_build_keeps_the_input_index() -> None:
    assert build(raw()).index.tolist() == [10, 11, 12]


def test_build_makes_consequence_categorical() -> None:
    assert isinstance(build(raw())["consequence"].dtype, pd.CategoricalDtype)