app/age_align.py

"""Align optional patient age to age-banded OHIP premiums / add-ons.

Many Schedule codes carry an explicit age window in their description
(e.g. ``ages 0 to 15``, ``2 - 11 years of age``, ``aged 65 years and above``,
``under 2 years``). When the clinician supplies an optional age (in months),
we:

1. Steer premium retrieval with age-appropriate keywords.
2. Drop age-banded premiums that fall outside the patient's age.
3. Leave non-age-banded premiums (after-hours, immunization, chronic-disease,
   discharge, …) untouched.

When age is omitted, no age filtering is applied — physicians stay unblocked.
Age alone is clinical signal, not an identifier (unlike DOB / health card).
"""
from __future__ import annotations

import math
import re
from collections.abc import Callable
from dataclasses import dataclass
from re import Match

# Sentinel for open-ended upper bounds (e.g. 65+).
_INF_MONTHS = 200 * 12


@dataclass(frozen=True)
class AgeAlignment:
    age_months: int
    label: str
    keywords: str


@dataclass(frozen=True)
class AgeWindow:
    """Inclusive age window in whole months."""

    lo_months: int
    hi_months: int  # inclusive


def align_patient_age(age_months: int | None) -> AgeAlignment | None:
    """Build steering keywords / label for a patient age in months."""
    if age_months is None or age_months < 0:
        return None
    years = age_months / 12.0
    if age_months < 24:
        label = f"{age_months} months old"
        keywords = (
            "under 2 years of age infant toddler paediatric age premium "
            "newborn child ages 0 to 15"
        )
    elif years < 12:
        y = int(years)
        label = f"{y} years old"
        keywords = (
            f"{y} years of age child paediatric 2 - 11 years of age "
            "ages 0 to 15 aged 19 years and below"
        )
    elif years < 18:
        y = int(years)
        label = f"{y} years old"
        keywords = (
            f"{y} years of age adolescent 12 - 17 years of age "
            "ages 0 to 15 aged 19 years and below less than age 22"
        )
    elif years < 65:
        y = int(years)
        label = f"{y} years old"
        keywords = (
            f"adult age {y} years 18 to 64 inclusive aged 19 years and below"
            if y <= 19
            else f"adult age {y} years 18 to 64 inclusive"
        )
    else:
        y = int(years)
        label = f"{y} years old"
        keywords = f"adult {y} years of age and older aged 65 years and above"
    return AgeAlignment(age_months=age_months, label=label, keywords=keywords)


def premium_age_window(doc: dict) -> AgeWindow | None:
    """Extract an age window from a code's description, or None if not age-banded."""
    text = " ".join(
        [
            doc.get("description_text") or "",
            doc.get("description") or "",
            doc.get("rules_and_constraints") or "",
        ]
    )
    if not text.strip():
        return None
    return _parse_age_window(text)


def premium_applies_to_age(doc: dict, age_months: int) -> bool:
    """True if this premium is compatible with the patient's age.

    Non-age-banded premiums always apply. Age-banded ones must contain
    ``age_months`` in their inclusive window.
    """
    window = premium_age_window(doc)
    if window is None:
        return True
    return window.lo_months <= age_months <= window.hi_months


def filter_premiums_for_age(
    premiums: list[dict], age_months: int
) -> list[dict]:
    """Drop age-banded premiums outside the patient's age."""
    return [p for p in premiums if premium_applies_to_age(p, age_months)]


def months_from_years(years: float) -> int:
    """Convert a whole/fractional year value to whole months (floor)."""
    if years < 0:
        return 0
    return int(math.floor(years * 12))


# --- description parsers -----------------------------------------------------

def _year_hi(years_inclusive: int) -> int:
    """Last month index still inside an inclusive year upper bound."""
    return (years_inclusive + 1) * 12 - 1


def _parse_age_window(text: str) -> AgeWindow | None:
    # Skip exception clauses ("except for patients under 4 years…") — those
    # are carve-outs, not eligibility windows.
    scrubbed = re.sub(
        r"except\s+for\s+patients?\s+under\s+\d+\s+years?[^.]*",
        " ",
        text,
        flags=re.I,
    )
    patterns: list[tuple[re.Pattern[str], Callable[[Match[str]], AgeWindow]]] = [
        (
            re.compile(
                r"ages?\s+(\d+)\s*(?:to|-|–)\s*(\d+)(?:\s+inclusive)?",
                re.I,
            ),
            lambda m: AgeWindow(int(m.group(1)) * 12, _year_hi(int(m.group(2)))),
        ),
        (
            re.compile(
                r"(\d+)\s*[-–]\s*(\d+)\s*years?\s+of\s+age",
                re.I,
            ),
            lambda m: AgeWindow(int(m.group(1)) * 12, _year_hi(int(m.group(2)))),
        ),
        (
            re.compile(r"aged\s+(\d+)\s+years?\s+and\s+below", re.I),
            lambda m: AgeWindow(0, _year_hi(int(m.group(1)))),
        ),
        (
            re.compile(r"aged\s+(\d+)\s+years?\s+and\s+above", re.I),
            lambda m: AgeWindow(int(m.group(1)) * 12, _INF_MONTHS),
        ),
        (
            re.compile(r"(\d+)\s+years?\s+of\s+age\s+and\s+older", re.I),
            lambda m: AgeWindow(int(m.group(1)) * 12, _INF_MONTHS),
        ),
        (
            re.compile(r"less than age\s+(\d+)", re.I),
            lambda m: AgeWindow(0, int(m.group(1)) * 12 - 1),
        ),
        (
            re.compile(r"under\s+two\s+years?", re.I),
            lambda _m: AgeWindow(0, 23),
        ),
        (
            re.compile(r"under\s+(\d+)\s+years?", re.I),
            lambda m: AgeWindow(0, int(m.group(1)) * 12 - 1),
        ),
        (
            re.compile(r"\bnewborn\b|\binfant\b", re.I),
            lambda _m: AgeWindow(0, 11),
        ),
    ]
    for pattern, builder in patterns:
        m = pattern.search(scrubbed)
        if not m:
            continue
        # Newborn/infant fallback only when no numeric age cue exists.
        if pattern.pattern.startswith(r"\bnewborn"):
            if re.search(
                r"ages?\s+\d+|years?\s+of\s+age|aged\s+\d+|under\s+\d+|under\s+two",
                scrubbed,
                re.I,
            ):
                continue
        try:
            return builder(m)
        except (TypeError, ValueError):
            continue
    return None