"""
test_unit_cases.py
==================

Unit tests for ``LogisticPredictor.predict``. Each test encodes a
specific clinical archetype (subtype, regimen, demographics) and
asserts the probability falls in the expected band.

Tests fall in three groups:

1.  **High-probability archetypes** (>=60% ypN0) -- aggressive biology
    + modern regimen (KEYNOTE-522 TNBC, dual-HER2 in HER2-enriched).
2.  **Low-probability archetypes** (<=20%) -- luminal-like elderly low
    Ki67, NET-only regimens, low-genomic-risk luminals.
3.  **Edge cases / input validation** -- invalid subtypes, ages, etc.

API contract assumed (from Agent B's module / data dictionary):

    predict(subtype, ki67_pct, age, cN_stage, regimen,
            has_pertuzumab=False, has_carboplatin=False,
            has_pembrolizumab=False, genomic_risk='unknown')
        -> {'probability': float, 'risk_band': str, 'confidence': str, ...}

Run:  pytest tests/test_unit_cases.py
"""

from __future__ import annotations

import math
import sys
from pathlib import Path

import pytest

THIS_DIR = Path(__file__).resolve().parent
ROOT = THIS_DIR.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

try:
    from nivel1_calculadora.predictor import LogisticPredictor
except Exception:                                                # noqa: BLE001
    pytest.skip("nivel1_calculadora.predictor not yet available; skipping unit tests.",
                allow_module_level=True)


# ---------------------------------------------------------------------------
# Helper
# ---------------------------------------------------------------------------
@pytest.fixture(scope="module")
def predictor() -> LogisticPredictor:
    return LogisticPredictor()


def _call(p: LogisticPredictor, **kwargs):
    """Convenience wrapper with default flags."""
    defaults = dict(
        has_pertuzumab=False,
        has_carboplatin=False,
        has_pembrolizumab=False,
        genomic_risk="unknown",
    )
    defaults.update(kwargs)
    return p.predict(**defaults)


# ---------------------------------------------------------------------------
# 1) High-probability archetypes
# ---------------------------------------------------------------------------
def test_TNBC_pembro_young_cN1_should_be_high(predictor):
    out = _call(predictor, subtype="TNBC", ki67_pct=70, age=42,
                cN_stage="cN1", regimen="AC_T",
                has_carboplatin=True, has_pembrolizumab=True)
    assert out["probability"] >= 0.60, \
        f"Expected >=60% (KEYNOTE-522 archetype), got {out['probability']:.2%}"
    assert out["risk_band"] == "high"


def test_HER2_enriched_pertu_should_be_very_high(predictor):
    out = _call(predictor, subtype="HR-/HER2+", ki67_pct=60, age=50,
                cN_stage="cN1", regimen="TCH", has_pertuzumab=True)
    assert out["probability"] >= 0.70, \
        f"Expected >=70% (HER2-enriched + pertuzumab), got {out['probability']:.2%}"
    assert out["risk_band"] == "high"


def test_HER2_enriched_no_pertu_still_high(predictor):
    """OR=40 dominates: even without pertuzumab, HR-/HER2+ should be high."""
    out = _call(predictor, subtype="HR-/HER2+", ki67_pct=50, age=55,
                cN_stage="cN1", regimen="TCH")
    assert out["probability"] >= 0.60


def test_triple_positive_cN2_pertu_should_be_medium_high(predictor):
    out = _call(predictor, subtype="HR+/HER2+", ki67_pct=40, age=55,
                cN_stage="cN2", regimen="TCH", has_pertuzumab=True)
    assert 0.35 <= out["probability"] <= 0.70, \
        f"Expected 35-70% (triple+ cN2 pertu), got {out['probability']:.2%}"


# ---------------------------------------------------------------------------
# 2) Low-probability archetypes
# ---------------------------------------------------------------------------
def test_luminal_old_low_ki67_cN1_should_be_low(predictor):
    out = _call(predictor, subtype="HR+/HER2-", ki67_pct=10, age=68,
                cN_stage="cN1", regimen="AC_T")
    assert out["probability"] <= 0.15, \
        f"Expected <=15% (elderly luminal low Ki67), got {out['probability']:.2%}"
    assert out["risk_band"] == "low"


def test_luminal_low_genomic_risk_should_be_very_low(predictor):
    out = _call(predictor, subtype="HR+/HER2-", ki67_pct=15, age=60,
                cN_stage="cN1", regimen="AC_T", genomic_risk="low")
    assert out["probability"] <= 0.10


def test_luminal_high_ki67_younger_still_moderate(predictor):
    """High Ki67 in HR+/HER2- gives a 3x OR boost; should clear 15%."""
    out = _call(predictor, subtype="HR+/HER2-", ki67_pct=35, age=45,
                cN_stage="cN1", regimen="AC_T")
    assert out["probability"] >= 0.15


# ---------------------------------------------------------------------------
# 3) Subtype ordering invariants (biology check)
# ---------------------------------------------------------------------------
def test_HER2_enriched_higher_than_TNBC_same_inputs(predictor):
    her2 = _call(predictor, subtype="HR-/HER2+", ki67_pct=40,
                 age=50, cN_stage="cN1", regimen="TCH")
    tnbc = _call(predictor, subtype="TNBC", ki67_pct=40,
                 age=50, cN_stage="cN1", regimen="AC_T")
    assert her2["probability"] > tnbc["probability"], \
        "HR-/HER2+ should always dominate TNBC on raw subtype OR."


def test_TNBC_higher_than_luminal_same_inputs(predictor):
    tnbc = _call(predictor, subtype="TNBC", ki67_pct=40,
                 age=50, cN_stage="cN1", regimen="AC_T")
    lum = _call(predictor, subtype="HR+/HER2-", ki67_pct=40,
                age=50, cN_stage="cN1", regimen="AC_T")
    assert tnbc["probability"] > lum["probability"]


def test_cN2_lower_than_cN1_same_inputs(predictor):
    cn1 = _call(predictor, subtype="TNBC", ki67_pct=60,
                age=50, cN_stage="cN1", regimen="AC_T")
    cn2 = _call(predictor, subtype="TNBC", ki67_pct=60,
                age=50, cN_stage="cN2", regimen="AC_T")
    assert cn1["probability"] > cn2["probability"]


def test_pertuzumab_increases_HER2pos_probability(predictor):
    no_p = _call(predictor, subtype="HR-/HER2+", ki67_pct=40,
                 age=50, cN_stage="cN1", regimen="TCH")
    yes_p = _call(predictor, subtype="HR-/HER2+", ki67_pct=40,
                  age=50, cN_stage="cN1", regimen="TCH",
                  has_pertuzumab=True)
    assert yes_p["probability"] > no_p["probability"]


def test_pembrolizumab_no_effect_outside_TNBC(predictor):
    """Per the model, pembrolizumab OR only applies in TNBC."""
    no_p = _call(predictor, subtype="HR+/HER2-", ki67_pct=20,
                 age=50, cN_stage="cN1", regimen="AC_T")
    yes_p = _call(predictor, subtype="HR+/HER2-", ki67_pct=20,
                  age=50, cN_stage="cN1", regimen="AC_T",
                  has_pembrolizumab=True)
    assert math.isclose(yes_p["probability"], no_p["probability"], abs_tol=1e-9)


# ---------------------------------------------------------------------------
# 4) Risk-band coverage
# ---------------------------------------------------------------------------
def test_risk_band_low_returned_for_low_probability(predictor):
    out = _call(predictor, subtype="HR+/HER2-", ki67_pct=10,
                age=70, cN_stage="cN2", regimen="AC_T",
                genomic_risk="low")
    assert out["risk_band"] == "low"


def test_risk_band_high_returned_for_high_probability(predictor):
    out = _call(predictor, subtype="HR-/HER2+", ki67_pct=60,
                age=45, cN_stage="cN1", regimen="TCH",
                has_pertuzumab=True)
    assert out["risk_band"] == "high"


# ---------------------------------------------------------------------------
# 5) Edge cases / input validation
# ---------------------------------------------------------------------------
def test_invalid_subtype_raises(predictor):
    with pytest.raises(ValueError):
        _call(predictor, subtype="ER+", ki67_pct=20,
              age=50, cN_stage="cN1", regimen="AC_T")


def test_invalid_cN_raises(predictor):
    with pytest.raises(ValueError):
        _call(predictor, subtype="TNBC", ki67_pct=60,
              age=50, cN_stage="cN4", regimen="AC_T")


def test_negative_age_raises(predictor):
    with pytest.raises(ValueError):
        _call(predictor, subtype="TNBC", ki67_pct=60,
              age=-5, cN_stage="cN1", regimen="AC_T")


def test_unrealistic_age_raises(predictor):
    with pytest.raises(ValueError):
        _call(predictor, subtype="TNBC", ki67_pct=60,
              age=150, cN_stage="cN1", regimen="AC_T")


def test_invalid_genomic_risk_raises(predictor):
    with pytest.raises(ValueError):
        _call(predictor, subtype="HR+/HER2-", ki67_pct=20,
              age=50, cN_stage="cN1", regimen="AC_T",
              genomic_risk="medium")


def test_probability_is_in_unit_interval(predictor):
    """Stress test: probability must always be in [0, 1]."""
    for s in ("HR-/HER2+", "HR+/HER2+", "TNBC", "HR+/HER2-"):
        for age in (30, 50, 70):
            for cn in ("cN1", "cN2", "cN3"):
                out = _call(predictor, subtype=s, ki67_pct=30,
                            age=age, cN_stage=cn, regimen="AC_T",
                            has_pertuzumab=True,
                            has_carboplatin=True,
                            has_pembrolizumab=True)
                assert 0.0 <= out["probability"] <= 1.0


def test_returns_required_keys(predictor):
    out = _call(predictor, subtype="TNBC", ki67_pct=60,
                age=50, cN_stage="cN1", regimen="AC_T")
    for key in ("probability", "risk_band", "confidence"):
        assert key in out, f"missing required key: {key}"
