"""Runnable example shipped inside the export (copied next to the loader shim as ``example.py``).

    pip install -r requirements.txt
    python example.py            # uses this directory as the checkpoint

Prints one answer of each question type. Needs a GPU with room for the checkpoint; pass a directory
as the first argument to score a different export.
"""

import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

from carejev_omni_loader import load_carejev_omni  # noqa: E402  (shipped next to this file)

STATE = (
    "PATIENT: 68-year-old male. Race: white.\n"
    "STUDY ENTRY: hospital day 2; primary disease category: acute respiratory failure or multiple "
    "organ system failure with sepsis (ARF/MOSF).\n"
    "COMORBIDITIES: 2 recorded (diabetes).\n"
    "SEVERITY: SUPPORT coma score 0; APACHE III physiology score 61; SUPPORT physiology score 33.9.\n"
    "FUNCTION: 1 activities of daily living impaired (0 = independent, 7 = fully dependent).\n"
    "DAY-3 PHYSIOLOGY: mean arterial pressure 72 mmHg; heart rate 104 /min; respiratory rate 28 /min; "
    "temperature 38.1 °C; white cell count 15.2 x10^9/L; creatinine 2.1 mg/dL; sodium 139 mmol/L; "
    "arterial pH 7.33; blood urea nitrogen 44 mg/dL; urine output 900 mL/day."
)

LOS_LEVELS = [
    "Less than 1 day",
    "1 day",
    "2 days",
    "3 days",
    "4 days",
    "5 days",
    "6 days",
    "7 days",
    "8 to 14 days",
    "More than 14 days",
]


def main() -> None:
    path = sys.argv[1] if len(sys.argv) > 1 else str(HERE)
    clf = load_carejev_omni(path)

    # noul: a yes/no question. The answer is P(yes); option order never changes it.
    noul = clf.predict(
        state=STATE,
        question="Will this patient die before being discharged from this hospitalization?",
        options=["Yes", "No"],
    )
    print("noul   P(yes) =", round(noul["probabilities"]["Yes"], 3))

    # choice: one of k options. Probabilities are averaged over rotations of a canonical order.
    choice = clf.predict(
        state=STATE,
        question="Which single problem best explains this presentation?",
        options=["Sepsis", "Heart failure", "Pulmonary embolism", "Diabetic ketoacidosis"],
    )
    print("choice      ->", choice["prediction"], "| confidence", round(choice["confidence"], 3))

    # score: ordered levels. `score` is the expected level, i.e. sum(index * probability).
    score = clf.predict(
        state=STATE,
        question="How long will this patient's hospital stay be in total?",
        options=LOS_LEVELS,
    )
    top = max(score["probabilities"], key=score["probabilities"].get)
    expected = sum(i * p for i, p in enumerate(score["probabilities"].values()))
    print(f"score       -> {top} | expected level {expected:.2f} of {len(LOS_LEVELS) - 1}")

    # Images, heart sounds and other media go in the same call:
    #   clf.predict(state=..., question=..., options=[...], media="ecg.png", modality="image")
    #   clf.predict(state=..., question=..., options=["Yes", "No"], media="pcg.wav", modality="audio")
    print("\nResearch use only. Not a medical device. See README.md.")


if __name__ == "__main__":
    main()