Configuring Age and Sex-Specific Reference Ranges in LIMS

from __future__ import annotations

import datetime as dt
from enum import Enum

from pydantic import BaseModel, Field, model_validator


class Sex(str, Enum):
    male = "male"
    female = "female"
    other = "other"
    unknown = "unknown"
    all = "ALL"  # sex-independent analyte or explicit fallback cohort


class CohortInterval(BaseModel):
    analyte_loinc: str
    method_version: str
    age_low_days: int = Field(ge=0)
    age_high_days: int | None = None  # None => open-ended upper bound
    sex: Sex
    lower: float
    upper: float
    critical_low: float | None = None
    critical_high: float | None = None
    priority: int = 0  # higher wins; e.g. gestational cohort > chronological
    version: str

    @model_validator(mode="after")
    def _check_bounds(self) -> "CohortInterval":
        if self.age_high_days is not None and self.age_high_days <= self.age_low_days:
            raise ValueError(
                f"{self.analyte_loinc}: age range [{self.age_low_days}, "
                f"{self.age_high_days}) is empty or inverted"
            )
        if self.upper <= self.lower:
            raise ValueError(
                f"{self.analyte_loinc}: reference bound upper<=lower"
            )
        return self

def age_in_days(birth_date: dt.date, collection_date: dt.date) -> int:
    """Deterministic, leap-year-safe age in whole days at collection."""
    if collection_date < birth_date:
        raise ValueError("collection_date precedes birth_date")
    return (collection_date - birth_date).days


def normalize_sex(raw: str | None) -> Sex:
    """Map an inbound demographic value to the administrative-gender enum.

    Missing or unrecognized values become `unknown` — never a silent
    default to male/female, which would fabricate a cohort assignment.
    """
    if not raw:
        return Sex.unknown
    try:
        return Sex(raw.strip().lower())
    except ValueError:
        return Sex.unknown

from typing import Protocol


class IntervalStore(Protocol):
    async def fetch(
        self, loinc: str, specimen: str, method_version: str
    ) -> list[CohortInterval]: ...


class ReferenceRangeResolver:
    def __init__(self, store: IntervalStore) -> None:
        self._store = store

    async def resolve(
        self,
        *,
        loinc: str,
        specimen: str,
        method_version: str,
        age_days: int,
        sex: Sex,
    ) -> tuple[CohortInterval | None, bool]:
        cohorts = await self._store.fetch(loinc, specimen, method_version)
        ordered = sorted(cohorts, key=lambda c: c.priority, reverse=True)

        for c in ordered:
            upper = c.age_high_days if c.age_high_days is not None else float("inf")
            if c.age_low_days <= age_days < upper and c.sex in (sex, Sex.all):
                return c, False  # exact demographic match, no fallback

        fallback = next((c for c in ordered if c.sex is Sex.all), None)
        return fallback, fallback is not None

class RangeResolutionAudit(BaseModel):
    trace_id: str
    analyte_loinc: str
    specimen: str
    method_version: str
    age_days: int
    sex: Sex
    matched_version: str | None
    lower: float | None
    upper: float | None
    fallback_applied: bool
    resolved_at: dt.datetime = Field(
        default_factory=lambda: dt.datetime.now(dt.timezone.utc)
    )


def build_audit(
    *, trace_id: str, specimen: str, age_days: int, sex: Sex,
    interval: CohortInterval | None, loinc: str, method_version: str,
    fallback_applied: bool,
) -> RangeResolutionAudit:
    return RangeResolutionAudit(
        trace_id=trace_id,
        analyte_loinc=loinc,
        specimen=specimen,
        method_version=method_version,
        age_days=age_days,
        sex=sex,
        matched_version=interval.version if interval else None,
        lower=interval.lower if interval else None,
        upper=interval.upper if interval else None,
        fallback_applied=fallback_applied,
    )

import pytest


class StubStore:
    def __init__(self, rows: list[CohortInterval]) -> None:
        self._rows = rows

    async def fetch(self, loinc, specimen, method_version):
        return self._rows


PED = CohortInterval(analyte_loinc="2324-2", method_version="v3",
                     age_low_days=0, age_high_days=1825, sex=Sex.all,
                     lower=145, upper=420, priority=0, version="glu-ped-1")
ADULT = CohortInterval(analyte_loinc="2324-2", method_version="v3",
                       age_low_days=1825, age_high_days=None, sex=Sex.all,
                       lower=110, upper=240, priority=0, version="glu-adt-1")


@pytest.mark.asyncio
async def test_upper_boundary_is_exclusive():
    resolver = ReferenceRangeResolver(StubStore([PED, ADULT]))
    interval, fb = await resolver.resolve(
        loinc="2324-2", specimen="serum", method_version="v3",
        age_days=1825, sex=Sex.male,
    )
    assert interval.version == "glu-adt-1"  # 1825 days -> adult, not pediatric
    assert fb is False


def test_audit_round_trips():
    audit = build_audit(trace_id="OBR-21-9001", specimen="serum",
                        age_days=1825, sex=Sex.male, interval=ADULT,
                        loinc="2324-2", method_version="v3",
                        fallback_applied=False)
    assert RangeResolutionAudit.model_validate_json(
        audit.model_dump_json()
    ) == audit