tests/test_risk_models/test_prostate_mortality_model.py

"""Tests for the Prostate Cancer Mortality prediction model.

Web calculator available at: https://prostate.predict.cam/tool

NOTE: This implementation is based on the Stata reference code from the published
model. There are discrepancies with the current web calculator, particularly for
cases with comorbidities (charlson=1), suggesting the web calculator may use an
updated version or different calibration. Our implementation matches the Stata
code exactly.
"""

import pytest

from sentinel.risk_models.prostate_mortality import ProstateMortalityRiskModel
from sentinel.user_input import (
    Anthropometrics,
    ClinicalTests,
    Demographics,
    Ethnicity,
    Lifestyle,
    PersonalMedicalHistory,
    ProstateCancerTreatment,
    PSATest,
    Sex,
    SmokingHistory,
    SmokingStatus,
    UserInput,
)


def create_test_user(
    age: int,
    psa: float,
    grade_group: int,
    t_stage: int,
    charlson: int = 0,
    treatment: ProstateCancerTreatment = ProstateCancerTreatment.CONSERVATIVE,
    ethnicity: Ethnicity = Ethnicity.WHITE,
) -> UserInput:
    """Create a test user with specified prostate cancer parameters.

    Args:
        age: Patient age in years.
        psa: PSA value in ng/mL.
        grade_group: Histological grade group (1-5).
        t_stage: Clinical T stage (1-4).
        charlson: Charlson comorbidity score (0-1).
        treatment: Primary treatment received.
        ethnicity: Patient ethnicity.

    Returns:
        UserInput instance configured for prostate mortality testing.
    """
    return UserInput(
        demographics=Demographics(
            age_years=age,
            sex=Sex.MALE,
            ethnicity=ethnicity,
            anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
        ),
        lifestyle=Lifestyle(
            smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
        ),
        personal_medical_history=PersonalMedicalHistory(
            prostate_cancer_grade_group=grade_group,
            prostate_cancer_t_stage=t_stage,
            charlson_comorbidity_score=charlson,
            prostate_cancer_treatment=treatment,
        ),
        family_history=[],
        clinical_tests=ClinicalTests(
            psa=PSATest(value_ng_ml=psa),
        ),
    )


GROUND_TRUTH_CASES = [
    {
        "name": "low_risk_young",
        "age": 55,
        "psa": 5.0,
        "grade_group": 1,
        "t_stage": 1,
        "charlson": 0,
        "treatment": ProstateCancerTreatment.CONSERVATIVE,
        "expected_pcsm_15yr": 7.0,
        "expected_npcm_15yr": 8.0,
        "expected_overall_15yr": 15.0,
    },
    {
        "name": "medium_risk_example",
        "age": 65,
        "psa": 11.0,
        "grade_group": 4,
        "t_stage": 2,
        "charlson": 0,
        "treatment": ProstateCancerTreatment.CONSERVATIVE,
        "expected_pcsm_15yr": 19.0,
        "expected_npcm_15yr": 26.0,
        "expected_overall_15yr": 45.0,
    },
    {
        "name": "moderate_risk",
        "age": 60,
        "psa": 8.0,
        "grade_group": 3,
        "t_stage": 2,
        "charlson": 0,
        "treatment": ProstateCancerTreatment.CONSERVATIVE,
        "expected_pcsm_15yr": 15.0,
        "expected_npcm_15yr": 15.0,
        "expected_overall_15yr": 30.0,
    },
]


class TestProstateMortalityRiskModel:
    """Test suite for ProstateMortalityRiskModel."""

    def setup_method(self) -> None:
        """Set up test fixtures."""
        self.model = ProstateMortalityRiskModel()

    def test_metadata(self) -> None:
        """Test model metadata including name, cancer type, and references."""
        assert self.model.name == "prostate_mortality"
        assert self.model.cancer_type() == "prostate"
        assert "Predict Prostate" in self.model.description()
        assert "PCSM" in self.model.description()
        assert "mortality" in self.model.interpretation().lower()
        assert len(self.model.references()) > 0
        assert any("predict" in ref.lower() for ref in self.model.references())

    def test_absolute_risk_basic(self) -> None:
        """Test basic absolute risk calculation returns valid percentages."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
                charlson_comorbidity_score=0,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=5.0),
            ),
        )
        risks = self.model.absolute_risk(user, years=15)
        assert risks["pcsm"] > 0
        assert risks["npcm"] > 0
        assert risks["overall"] > 0
        assert risks["pcsm"] < 100
        assert risks["npcm"] < 100
        assert risks["overall"] <= 100

    @pytest.mark.parametrize("case", GROUND_TRUTH_CASES, ids=lambda case: case["name"])
    def test_ground_truth_cases(self, case) -> None:
        """Test model predictions against validated web calculator results.

        Args:
            case: Test case dictionary with patient parameters and expected results.
        """
        if case["expected_pcsm_15yr"] is None:
            pytest.skip("TODO: Fill in expected values from web calculator")

        user = create_test_user(
            age=case["age"],
            psa=case["psa"],
            grade_group=case["grade_group"],
            t_stage=case["t_stage"],
            charlson=case["charlson"],
            treatment=case["treatment"],
            ethnicity=case.get("ethnicity", Ethnicity.WHITE),
        )

        risks = self.model.absolute_risk(user, years=15)
        assert risks["pcsm"] == pytest.approx(case["expected_pcsm_15yr"], abs=4.0)
        assert risks["npcm"] == pytest.approx(case["expected_npcm_15yr"], abs=4.0)
        assert risks["overall"] == pytest.approx(case["expected_overall_15yr"], abs=4.0)

    def test_compute_score_with_male_user_input(self) -> None:
        """Test compute_score returns formatted string for valid male user."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
                charlson_comorbidity_score=0,
                prostate_cancer_treatment=ProstateCancerTreatment.CONSERVATIVE,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=5.0),
            ),
        )

        score = self.model.compute_score(user)
        assert "PCSM" in score
        assert "NPCM" in score
        assert "Overall" in score
        assert "%" in score

    def test_compute_score_rejects_female_user(self) -> None:
        """Test model correctly rejects female patients."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.FEMALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=165.0, weight_kg=70.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=5.0),
            ),
        )

        score = self.model.compute_score(user)
        assert "N/A" in score
        assert "male" in score.lower()

    def test_validation_errors(self) -> None:
        """Test validation errors for missing required fields."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(),
            family_history=[],
            clinical_tests=ClinicalTests(),
        )

        score = self.model.compute_score(user)
        assert "N/A" in score
        assert "Invalid" in score

    def test_age_out_of_range(self) -> None:
        """Test age outside validated range raises error."""
        user = UserInput(
            demographics=Demographics(
                age_years=30,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=5.0),
            ),
        )

        score = self.model.compute_score(user)
        assert "N/A" in score

    def test_psa_validation(self) -> None:
        """Test PSA value validation at boundary values."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(),
        )

        score = self.model.compute_score(user)
        assert "N/A" in score

    def test_competing_risks_consistency(self) -> None:
        """Test competing risks sum correctly to overall mortality."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=3,
                prostate_cancer_t_stage=2,
                charlson_comorbidity_score=1,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=10.0),
            ),
        )

        risks = self.model.absolute_risk(user, years=15)
        total = risks["pcsm"] + risks["npcm"]
        assert total == pytest.approx(risks["overall"], abs=0.5)

    def test_different_time_horizons(self) -> None:
        """Test model predictions for different time horizons."""
        user = UserInput(
            demographics=Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            lifestyle=Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=2,
                prostate_cancer_t_stage=1,
            ),
            family_history=[],
            clinical_tests=ClinicalTests(
                psa=PSATest(value_ng_ml=5.0),
            ),
        )

        risk_5yr = self.model.absolute_risk(user, years=5)
        risk_10yr = self.model.absolute_risk(user, years=10)
        risk_15yr = self.model.absolute_risk(user, years=15)

        assert risk_5yr["pcsm"] < risk_10yr["pcsm"] < risk_15yr["pcsm"]
        assert risk_5yr["npcm"] < risk_10yr["npcm"] < risk_15yr["npcm"]
        assert risk_5yr["overall"] < risk_10yr["overall"] < risk_15yr["overall"]

    def test_treatment_effect_radical_vs_conservative(self) -> None:
        """Test radical treatment reduces PCSM compared to conservative treatment."""
        base_input = {
            "demographics": Demographics(
                age_years=65,
                sex=Sex.MALE,
                ethnicity=Ethnicity.WHITE,
                anthropometrics=Anthropometrics(height_cm=175.0, weight_kg=80.0),
            ),
            "lifestyle": Lifestyle(
                smoking=SmokingHistory(status=SmokingStatus.NEVER, pack_years=0.0),
            ),
            "family_history": [],
            "clinical_tests": ClinicalTests(
                psa=PSATest(value_ng_ml=10.0),
            ),
        }

        user_conservative = UserInput(
            **base_input,
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=3,
                prostate_cancer_t_stage=2,
                charlson_comorbidity_score=0,
                prostate_cancer_treatment=ProstateCancerTreatment.CONSERVATIVE,
            ),
        )

        user_radical = UserInput(
            **base_input,
            personal_medical_history=PersonalMedicalHistory(
                prostate_cancer_grade_group=3,
                prostate_cancer_t_stage=2,
                charlson_comorbidity_score=0,
                prostate_cancer_treatment=ProstateCancerTreatment.RADICAL,
            ),
        )

        risk_conservative = self.model.absolute_risk(user_conservative, years=15)
        risk_radical = self.model.absolute_risk(user_radical, years=15)

        assert risk_radical["pcsm"] < risk_conservative["pcsm"]