from pathlib import Path

import numpy as np
import pandas as pd
import pytest

from rarelens_ml.train import label, read_vep_tab

HEADER = [
    "Uploaded_variation", "Location", "Allele", "Consequence", "IMPACT", "SYMBOL",
    "gnomADe_AF", "CLIN_SIG", "CADD_PHRED", "am_pathogenicity",
]


def write_vep_tab(path: Path, rows: list[list[str]]) -> Path:
    lines = [
        "## ENSEMBL VARIANT EFFECT PREDICTOR v113.0",
        "## Column descriptions:",
        "#" + "\t".join(HEADER),
        *("\t".join(r) for r in rows),
    ]
    path.write_text("\n".join(lines) + "\n")
    return path


@pytest.mark.parametrize(
    ("clin_sig", "expected"),
    [
        # VEP writes lowercase, comma-separated terms from co-located ClinVar records.
        ("pathogenic", 1),
        ("pathogenic,likely_pathogenic", 1),
        ("likely_benign", 0),
        ("benign,likely_benign", 0),
        # ClinVar VCF CLNSIG spelling must keep working too.
        ("Pathogenic/Likely_pathogenic", 1),
        ("Benign", 0),
        ("uncertain_significance", None),
        ("pathogenic,benign", None),  # conflicting evidence is not a label
        ("-", None),
        ("", None),
        (np.nan, None),
    ],
)
def test_label(clin_sig: object, expected: int | None) -> None:
    assert label(clin_sig) == expected


def test_read_vep_tab_uses_the_hash_header_and_keeps_dashes(tmp_path: Path) -> None:
    tsv = write_vep_tab(
        tmp_path / "x.vep.tsv",
        [["22_1_A_G", "22:1", "G", "missense_variant", "MODERATE", "TBX1", "-", "pathogenic", "28", "0.9"]],
    )
    df = read_vep_tab(tsv)
    assert list(df.columns) == HEADER
    assert df.loc[0, "gnomADe_AF"] == "-"
    assert df.loc[0, "CLIN_SIG"] == "pathogenic"


def test_load_returns_raw_serving_columns_and_labels(tmp_path: Path) -> None:
    from rarelens_ml.features import RAW_COLUMNS
    from rarelens_ml.train import load

    tsv = write_vep_tab(
        tmp_path / "x.vep.tsv",
        [
            ["a", "22:1", "G", "missense_variant", "MODERATE", "TBX1", "0.0001", "pathogenic", "28", "0.9"],
            ["b", "22:2", "A", "synonymous_variant", "LOW", "CHEK2", "0.12", "benign", "3", "-"],
            ["c", "22:3", "T", "intron_variant", "MODIFIER", "CHEK2", "0.3", "uncertain_significance", "1", "-"],
        ],
    )
    X, y = load(str(tsv))
    assert list(X.columns) == RAW_COLUMNS
    assert y.tolist() == [1, 0]  # the VUS row is dropped


def test_logged_model_returns_probabilities_from_raw_columns(tmp_path: Path) -> None:
    """The registered model must take the raw columns serving sends and return P(pathogenic)."""
    import mlflow

    from rarelens_ml.train import fit, log_and_register

    rng = np.random.default_rng(0)
    n = 400
    impact = rng.choice(["HIGH", "MODERATE", "LOW", "MODIFIER"], n)
    y = pd.Series(((impact == "HIGH") | (rng.random(n) < 0.1)).astype(int))
    X = pd.DataFrame({
        "impact": impact,
        "consequence": rng.choice(["stop_gained", "missense_variant", "intron_variant"], n),
        "gnomad_af": rng.random(n).round(4).astype(str),   # strings, as read from the DB
        "cadd_phred": (rng.random(n) * 40).round(1).astype(str),
        "am_pathogenicity": "-",
    })

    mlflow.set_tracking_uri(f"sqlite:///{tmp_path}/mlflow.db")
    mlflow.set_experiment("test")
    clf = fit(X, y)
    version = log_and_register(clf, model_name="rarelens-test", alias="production")

    model = mlflow.pyfunc.load_model("models:/rarelens-test@production")
    scores = np.asarray(model.predict(X.head(50)))
    assert version == "1"
    assert scores.shape == (50,)
    assert ((scores >= 0) & (scores <= 1)).all()
    assert not set(np.unique(scores)) <= {0.0, 1.0}, "got class labels, expected probabilities"