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],
        "gnomad_af": [None, "0.12", 0.001],
        "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", "gnomad_af", "cadd_phred", "am_pathogenicity"]


def test_build_ranks_impact_and_coerces_numbers() -> None:
    out = build(raw())
    assert out["impact_rank"].tolist() == [3, 1, 0]
    assert out["gnomad_af"].tolist() == [0.0, 0.12, 0.001]  # missing AF means absent from gnomAD
    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)