Recipe 12.8: Python Implementation Example

pip install boto3

import json
import logging
import math
import random
import statistics
import uuid
from collections import defaultdict
from datetime import date, datetime, timedelta, timezone
from decimal import Decimal
from statistics import mean, median, stdev

import boto3
from botocore.config import Config

# Structured logging. In production, ship JSON-formatted records
# to CloudWatch Logs Insights for cross-call investigation.
# Trajectory data is PHI in the strongest sense (longitudinal,
# disease-specific, often genetically suggestive). Log structural
# metadata only (run_id, patient_id_hash, disease_name,
# model_version, runtime_ms), never raw observation values, never
# treatment timeline content, never per-patient posterior samples.
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# Adaptive retry handles throttling from S3, DynamoDB, EventBridge,
# CloudWatch, HealthLake, and SageMaker. The nightly trajectory
# pipeline touches every cohort patient, so retries should be
# quick and capped. A stuck dependency must not balloon a
# multi-hour training window into a multi-day incident.
BOTO3_RETRY_CONFIG = Config(
    retries={"max_attempts": 5, "mode": "adaptive"})

# Module-level clients. Reused across function calls within the
# pipeline so each call does not pay the connection cost. The
# demo wires up MockS3 / MockTable / MockEventBus / MockCloudWatch /
# MockHealthLake via run_demo() and never touches these real
# handles; they are staged here so production wiring is a one-line
# swap. boto3 client construction is lazy (no network call until
# first use), so the unused handles are free at import.
REGION = "us-east-1"
s3_client          = boto3.client("s3",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)
dynamodb           = boto3.resource("dynamodb",
                                    region_name=REGION,
                                    config=BOTO3_RETRY_CONFIG)
healthlake_client  = boto3.client("healthlake",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)
eventbridge_client = boto3.client("events",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)
cloudwatch_client  = boto3.client("cloudwatch",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)
sagemaker_runtime  = boto3.client("sagemaker-runtime",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)
lambda_client      = boto3.client("lambda",
                                  region_name=REGION,
                                  config=BOTO3_RETRY_CONFIG)

# --- Resource Names ---
# Fill these in with your actual resource names. The demo prints
# what it would write rather than failing if the resources do
# not exist; see run_demo() at the bottom.
HEALTHLAKE_DATASTORE_ID    = "trajectory-fhir-datastore"
COHORT_DATASET_BUCKET      = "trajectory-cohort-datasets"
HARMONIZED_BUCKET          = "trajectory-harmonized"
MODEL_ARTIFACT_BUCKET      = "trajectory-model-artifacts"
FORECAST_BUCKET            = "trajectory-forecasts"
COUNTERFACTUAL_BUCKET      = "trajectory-counterfactuals"
TRIAL_PRIOR_BUCKET         = "trajectory-trial-priors"
PATIENT_TRAJECTORIES_TABLE = "patient-trajectories"
TRAJECTORY_EVENT_BUS_NAME  = "trajectory-events-bus"
CLOUDWATCH_NAMESPACE       = "DiseaseProgressionTrajectory"

# --- Versioning ---
# Every surfaced trajectory and counterfactual carries the cohort
# definition version, the model version, and the trial-prior
# version active at the time of generation. This is how a future
# audit reconstructs which artifact produced which forecast on
# which day for which patient.
COHORT_DEFINITION_VERSION = "adpkd-cohort-v3"
MODEL_VERSION_TAG         = "adpkd-bayesian-hierarchical-v4"
TRIAL_PRIOR_VERSION       = "priors-2026q1"
PIPELINE_VERSION          = "trajectory-pipeline-v1.0"

# --- Forecast Horizon and Grid ---
# 60 months is a defensible upper bound for ADPKD trajectory
# forecasts. Beyond that, the long-horizon extrapolation is
# unreliable and the system positions itself as a "trajectory
# characterizer" rather than a future-state predictor. The
# 3-month grid step matches typical clinical visit cadence.
DEFAULT_FORECAST_HORIZON_MONTHS = 60
DEFAULT_FORECAST_GRID_STEP_MO   = 3

# --- Endpoint Definition ---
# eGFR < 15 mL/min/1.73 m^2 is the clinical threshold for
# considering renal replacement therapy. The endpoint is
# defined here as a configuration value so it can change per
# disease and per institution.
ENDPOINT_DEFINITION = {
    "name":                "renal_replacement_therapy_threshold",
    "loinc_code":          "48642-3",   # eGFR
    "threshold_value":     15.0,
    "threshold_direction": "below",      # endpoint hits when value drops below
    "display":             "eGFR < 15 (RRT consideration)",
}

# --- Disease Cohort Definition (ADPKD) ---
# In production this is a versioned, clinician-reviewed config
# stored in a separate config repo. Here it is inline so the demo
# is self-contained. Changing this definition retroactively
# reshapes every downstream forecast, which is why production
# treats it as a clinical artifact under change control.
ADPKD_COHORT_DEFINITION = {
    "name":                            "adpkd",
    "version":                         COHORT_DEFINITION_VERSION,
    "display":                         "Autosomal Dominant Polycystic Kidney Disease",
    "inclusion_icd10":                 ["Q61.2", "Q61.3"],
    "exclusion_icd10":                 ["Q61.4", "Q61.5", "Z94.0"],
    "minimum_observation_window_months": 24,
    "minimum_egfr_measurements":       6,
    "primary_outcome_loinc":           "48642-3",
    "time_zero_anchor":                "earliest_inclusion_icd_date",
    "trajectory_loincs": {
        "48642-3": "eGFR (mL/min/1.73 m^2)",
        "33914-3": "Total Kidney Volume (mL)",
        "8480-6":  "Systolic BP (mmHg)",
    },
    "relevant_drug_classes": [
        "tolvaptan",
        "acei_arb",
        "sglt2_inhibitor",
    ],
}

# --- LOINC Catalog ---
# Per-LOINC configuration for the trajectory pipeline. Production
# institutions maintain a curated catalog with thousands of entries.
LOINC_CATALOG = {
    "48642-3": {
        "display":          "eGFR",
        "canonical_unit":   "mL/min/1.73m2",
        "model_role":       "primary_outcome",
        "lower_clip":       1.0,
        "upper_clip":       150.0,
    },
    "33914-3": {
        "display":          "Total Kidney Volume",
        "canonical_unit":   "mL",
        "model_role":       "trajectory_feature",
        "lower_clip":       100.0,
        "upper_clip":       8000.0,
    },
    "8480-6": {
        "display":          "Systolic BP",
        "canonical_unit":   "mmHg",
        "model_role":       "covariate",
        "lower_clip":       60.0,
        "upper_clip":       240.0,
    },
}

# --- Acute Encounter Classes ---
# Encounter classes that get tagged "acute" and excluded from the
# chronic-trajectory training data.
ACUTE_ENCOUNTER_CLASSES = {"inpatient", "emergency", "observation"}

# --- Trial-Literature-Derived Effect Priors ---
# Calibrated prior beliefs about how each treatment bends the
# trajectory. Each entry encodes the slope-modifier point estimate
# and the credible-interval bounds derived from cited published
# trial evidence. Production maintains this registry as a versioned
# config with explicit clinical-advisor sign-off.
#
# IMPORTANT: The numeric values in this demo are illustrative
# approximations of trial-derived estimates. They are not exact
# transcriptions of the published trial output. A production
# system replaces these with the institution's curated prior
# library reviewed by a disease specialist.
TRIAL_DERIVED_EFFECT_PRIORS = {
    "tolvaptan": {
        "slope_modifier_mean":   0.30,
        "slope_modifier_sd":     0.06,
        "evidence_basis":        "TEMPO 3:4 (NCT00428948) and REPRISE (NCT02160145) trials",
        "applicable_diseases":   ["adpkd"],
        "version":               TRIAL_PRIOR_VERSION,
    },
    "sglt2_inhibitor": {
        "slope_modifier_mean":   0.25,
        "slope_modifier_sd":     0.08,
        "initial_dip_mean":      -3.0,
        "initial_dip_sd":        1.0,
        "evidence_basis":        "DAPA-CKD (NCT03036150) and EMPA-KIDNEY (NCT03594110) trials",
        "applicable_diseases":   ["adpkd", "ckd"],
        "version":               TRIAL_PRIOR_VERSION,
    },
    "acei_arb": {
        "slope_modifier_mean":   0.10,
        "slope_modifier_sd":     0.04,
        "evidence_basis":        "RENAAL (NCT00308347) and AASK trial evidence",
        "applicable_diseases":   ["adpkd", "ckd"],
        "version":               TRIAL_PRIOR_VERSION,
    },
}

# --- Bayesian Hierarchical Model Hyperparameters ---
# Population-level priors anchored to published CKD/ADPKD literature.
# The demo uses a simplified linear-mixed-effects formulation;
# production uses non-linear parametric forms (Gompertz for kidney
# volume) or fully Bayesian hierarchical specifications in PyMC.
BAYESIAN_MODEL_HYPERPARAMETERS = {
    "prior_population_slope_mean":      -0.25,
    "prior_population_slope_sd":         0.08,
    "prior_per_patient_slope_sd":        0.20,
    "prior_per_patient_intercept_sd":    8.0,
    "observation_noise_sd":              4.0,
    "num_posterior_samples":             200,
    "temporal_holdout_fraction":         0.20,
}

# --- Synthetic Data ---
SYNTHETIC_PATIENT_COUNT      = 14
SYNTHETIC_HISTORY_YEARS_MIN  = 3
SYNTHETIC_HISTORY_YEARS_MAX  = 9
SYNTHETIC_RANDOM_SEED        = 42

# Deploy-time guardrail. Any blank resource name is a deploy-time
# bug, not a runtime surprise.
for _name, _value in [
    ("HEALTHLAKE_DATASTORE_ID",    HEALTHLAKE_DATASTORE_ID),
    ("COHORT_DATASET_BUCKET",      COHORT_DATASET_BUCKET),
    ("HARMONIZED_BUCKET",          HARMONIZED_BUCKET),
    ("MODEL_ARTIFACT_BUCKET",      MODEL_ARTIFACT_BUCKET),
    ("FORECAST_BUCKET",            FORECAST_BUCKET),
    ("COUNTERFACTUAL_BUCKET",      COUNTERFACTUAL_BUCKET),
    ("TRIAL_PRIOR_BUCKET",         TRIAL_PRIOR_BUCKET),
    ("PATIENT_TRAJECTORIES_TABLE", PATIENT_TRAJECTORIES_TABLE),
    ("TRAJECTORY_EVENT_BUS_NAME",  TRAJECTORY_EVENT_BUS_NAME),
    ("CLOUDWATCH_NAMESPACE",       CLOUDWATCH_NAMESPACE),
]:
    assert _value, f"{_name} must be set before deploying."

def _to_decimal(value):
    """Convert numeric values to Decimal for DynamoDB-safe writes.

    DynamoDB rejects Python float at the SDK boundary. Pass
    everything numeric through this helper before any PutItem,
    BatchWriteItem, or UpdateItem call.
    """
    if value is None:
        return None
    if isinstance(value, Decimal):
        return value
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        return Decimal(str(round(float(value), 6)))
    if isinstance(value, str):
        return Decimal(value)
    raise TypeError(f"Cannot convert {type(value).__name__} to Decimal")

def _normal_cdf(z):
    """Standard normal CDF using the error function approximation."""
    return 0.5 * (1.0 + math.erf(z / math.sqrt(2.0)))

class MockHealthLake:
    """In-memory stand-in for an Amazon HealthLake FHIR datastore.

    Production uses boto3.client('healthlake') and the FHIR REST
    API to read Observation, Condition, and MedicationRequest
    resources keyed by patient. The mock stores resources as
    typed lists and provides search-by-patient-and-type semantics,
    which is what cohort qualification and harmonization need.
    """

    def __init__(self):
        self.observations  = []   # FHIR Observation dicts
        self.conditions    = []   # FHIR Condition dicts
        self.medications   = []   # FHIR MedicationRequest dicts

    def put_observation(self, observation):
        self.observations.append(dict(observation))

    def put_condition(self, condition):
        self.conditions.append(dict(condition))

    def put_medication(self, medication):
        self.medications.append(dict(medication))

    def search_observations(self, patient_id, loinc_code=None):
        out = []
        for obs in self.observations:
            if obs.get("subject_reference") != f"Patient/{patient_id}":
                continue
            if loinc_code and obs.get("code_loinc") != loinc_code:
                continue
            out.append(obs)
        out.sort(key=lambda o: o["effective_dt"])
        return out

    def search_conditions(self, patient_id):
        return [c for c in self.conditions
                if c.get("subject_reference") == f"Patient/{patient_id}"]

    def search_medications(self, patient_id):
        return [m for m in self.medications
                if m.get("subject_reference") == f"Patient/{patient_id}"]

    def list_all_patient_ids(self):
        ids = set()
        for c in self.conditions:
            ref = c.get("subject_reference", "")
            if ref.startswith("Patient/"):
                ids.add(ref[len("Patient/"):])
        return sorted(ids)

class MockS3:
    """In-memory stand-in for an S3 bucket.

    Production uses boto3.client('s3').get_object / put_object.
    The mock stores objects keyed by (bucket, key).
    """

    def __init__(self):
        self.objects = {}

    def put_object(self, Bucket, Key, Body, **kwargs):
        if isinstance(Body, str):
            Body = Body.encode("utf-8")
        self.objects[(Bucket, Key)] = Body
        return {"ETag": '"' + str(uuid.uuid4()) + '"'}

    def get_object(self, Bucket, Key, **kwargs):
        if (Bucket, Key) not in self.objects:
            raise KeyError(f"NoSuchKey: s3://{Bucket}/{Key}")
        body = self.objects[(Bucket, Key)]

        class _StreamingBody:
            def __init__(self, b):
                self._b = b
            def read(self):
                return self._b
        return {"Body": _StreamingBody(body)}

class MockTable:
    """In-memory stand-in for a DynamoDB table.

    Supports the operations the demo calls: batch_writer, put_item,
    query, get_item. Not a complete DynamoDB emulation; covers
    what this pipeline needs.
    """

    def __init__(self, name):
        self.name        = name
        self.items       = {}
        self.write_count = 0

    class _BatchWriter:
        def __init__(self, table):
            self.table = table
        def __enter__(self):
            return self
        def __exit__(self, *exc):
            return False
        def put_item(self, Item):
            pk = Item["patient_id"]
            sk = Item["disease_modelversion_generated_at"]
            self.table.items[(pk, sk)] = dict(Item)
            self.table.write_count += 1

    def batch_writer(self):
        return self._BatchWriter(self)

    def put_item(self, Item):
        pk = Item["patient_id"]
        sk = Item["disease_modelversion_generated_at"]
        self.items[(pk, sk)] = dict(Item)
        self.write_count += 1

class MockEventBus:
    """In-memory stand-in for EventBridge."""

    def __init__(self, name):
        self.name   = name
        self.events = []

    def put_events(self, Entries):
        self.events.extend(Entries)
        return {"FailedEntryCount": 0}

class MockCloudWatch:
    """In-memory stand-in for CloudWatch."""

    def __init__(self):
        self.metrics = defaultdict(list)

    def put_metric_data(self, Namespace, MetricData):
        for m in MetricData:
            self.metrics[f"{Namespace}/{m['MetricName']}"].append({
                "Value": m["Value"],
                "Unit":  m.get("Unit", "None"),
                "Time":  datetime.now(timezone.utc).isoformat(),
            })

def generate_synthetic_adpkd_cohort(
        patient_count=SYNTHETIC_PATIENT_COUNT,
        seed=SYNTHETIC_RANDOM_SEED):
    """Generate a synthetic ADPKD cohort with longitudinal trajectories.

    Production reads from an Amazon HealthLake FHIR datastore that
    has been ingesting from the institutional EHR. The demo
    synthesizes a small ADPKD cohort with varied disease severity,
    varied observation density, and varied treatment exposures so
    you can see the qualification, harmonization, training, and
    forecasting steps work end-to-end.

    Each patient has:
      - A diagnosis Condition with an inclusion ICD-10 code
      - A series of eGFR Observations spanning multiple years
      - A series of total kidney volume Observations (less frequent)
      - A series of systolic BP Observations
      - Optional medication history (ACEi/ARB, possibly SGLT2)
      - A few inpatient eGFR values that the chronic pipeline
        must exclude

    The synthetic eGFR trajectory is generated as:
        egfr(t) = baseline + slope_per_month * months_since_dx + noise
    where slope_per_month varies per patient (faster vs slower
    progressors) and is dampened if the patient is on a relevant
    medication. This lets the trained model recover both the
    population slope and the per-patient deviations.
    """
    rng = random.Random(seed)
    today = date.today()

    patients = []
    for i in range(patient_count):
        patient_id = f"patient-adpkd-{i+1:03d}"

        # Random history length and diagnosis date.
        history_years = rng.uniform(
            SYNTHETIC_HISTORY_YEARS_MIN, SYNTHETIC_HISTORY_YEARS_MAX)
        diagnosis_date = today - timedelta(days=int(history_years * 365.25))

        # Patient-specific disease severity. Slope is in mL/min/1.73m^2 per month.
        # Population mean is around -0.25; per-patient deviation is wide.
        true_slope_per_month = rng.gauss(-0.25, 0.20)
        # Clip to clinically plausible range.
        true_slope_per_month = max(min(true_slope_per_month, 0.0), -0.7)

        # Patient-specific baseline eGFR at diagnosis.
        true_baseline_egfr = rng.gauss(85.0, 12.0)
        true_baseline_egfr = max(min(true_baseline_egfr, 110.0), 50.0)

        # Treatment exposure flags. Most patients are on ACEi/ARB.
        on_acei_arb_from_month = rng.choice([0, 0, 0, 6, 12, None])
        on_sglt2_from_month    = rng.choice([None, None, None, 24, 36])
        on_tolvaptan_from_month = rng.choice([None, None, None, None, 18])

        patients.append({
            "patient_id":               patient_id,
            "diagnosis_date":           diagnosis_date,
            "history_years":            history_years,
            "true_slope_per_month":     true_slope_per_month,
            "true_baseline_egfr":       true_baseline_egfr,
            "on_acei_arb_from_month":   on_acei_arb_from_month,
            "on_sglt2_from_month":      on_sglt2_from_month,
            "on_tolvaptan_from_month":  on_tolvaptan_from_month,
        })

    return patients

def populate_synthetic_healthlake(patients, healthlake, seed=SYNTHETIC_RANDOM_SEED):
    """Populate the MockHealthLake with synthetic FHIR resources.

    Walks each synthetic patient and emits the Condition for the
    diagnosis, the Observations for eGFR / TKV / SBP across the
    history, and the MedicationRequest entries for the relevant
    drug classes. A few inpatient-context eGFR values are
    interspersed to exercise the chronic-vs-acute filtering.
    """
    rng = random.Random(seed + 1)
    today = date.today()

    for p in patients:
        # Diagnosis Condition (drives cohort qualification).
        healthlake.put_condition({
            "resourceType":      "Condition",
            "subject_reference": f"Patient/{p['patient_id']}",
            "code_icd10":        rng.choice(["Q61.2", "Q61.3"]),
            "onset_dt":          p["diagnosis_date"].isoformat(),
        })

        # eGFR observations (~ every 3 months in chronic care).
        cursor_months = 0.0
        history_months = p["history_years"] * 12.0
        while cursor_months <= history_months:
            obs_date = p["diagnosis_date"] + timedelta(days=int(cursor_months * 30.5))
            if obs_date > today:
                break
            # Apply treatment effects to the synthetic trajectory.
            effective_slope = p["true_slope_per_month"]
            if (p["on_acei_arb_from_month"] is not None
                    and cursor_months >= p["on_acei_arb_from_month"]):
                effective_slope *= (1.0 - 0.10)
            if (p["on_sglt2_from_month"] is not None
                    and cursor_months >= p["on_sglt2_from_month"]):
                effective_slope *= (1.0 - 0.25)
            if (p["on_tolvaptan_from_month"] is not None
                    and cursor_months >= p["on_tolvaptan_from_month"]):
                effective_slope *= (1.0 - 0.30)
            value = (p["true_baseline_egfr"]
                     + effective_slope * cursor_months
                     + rng.gauss(0.0, 3.0))
            value = max(value, 5.0)
            healthlake.put_observation({
                "resourceType":      "Observation",
                "subject_reference": f"Patient/{p['patient_id']}",
                "code_loinc":        "48642-3",
                "value_quantity":    round(value, 1),
                "value_unit":        "mL/min/1.73m2",
                "effective_dt":      obs_date.isoformat(),
                "encounter_class":   "ambulatory",
            })
            cursor_months += rng.uniform(2.5, 4.0)

        # Two acute-context eGFR values mid-history (must be excluded).
        if history_months > 30:
            for offset_months in (12.0, 12.2):
                d = p["diagnosis_date"] + timedelta(days=int(offset_months * 30.5))
                healthlake.put_observation({
                    "resourceType":      "Observation",
                    "subject_reference": f"Patient/{p['patient_id']}",
                    "code_loinc":        "48642-3",
                    "value_quantity":    round(p["true_baseline_egfr"] * 0.5
                                               + rng.gauss(0, 4.0), 1),
                    "value_unit":        "mL/min/1.73m2",
                    "effective_dt":      d.isoformat(),
                    "encounter_class":   "inpatient",
                })

        # SBP observations every 3 months.
        cursor_months = 0.0
        while cursor_months <= history_months:
            obs_date = p["diagnosis_date"] + timedelta(days=int(cursor_months * 30.5))
            if obs_date > today:
                break
            sbp = 130 + rng.gauss(0, 10)
            healthlake.put_observation({
                "resourceType":      "Observation",
                "subject_reference": f"Patient/{p['patient_id']}",
                "code_loinc":        "8480-6",
                "value_quantity":    round(sbp, 0),
                "value_unit":        "mmHg",
                "effective_dt":      obs_date.isoformat(),
                "encounter_class":   "ambulatory",
            })
            cursor_months += rng.uniform(2.5, 4.0)

        # MedicationRequest entries for any relevant exposures.
        for drug_class, start_month in [
                ("acei_arb", p["on_acei_arb_from_month"]),
                ("sglt2_inhibitor", p["on_sglt2_from_month"]),
                ("tolvaptan", p["on_tolvaptan_from_month"])]:
            if start_month is None:
                continue
            start_date = p["diagnosis_date"] + timedelta(days=int(start_month * 30.5))
            if start_date > today:
                continue
            healthlake.put_medication({
                "resourceType":         "MedicationRequest",
                "subject_reference":    f"Patient/{p['patient_id']}",
                "drug_class":           drug_class,
                "authored_on_dt":       start_date.isoformat(),
                "end_dt":               None,
            })

def _months_between(d1, d2):
    """Difference in months between two date objects."""
    if isinstance(d1, str):
        d1 = date.fromisoformat(d1[:10])
    if isinstance(d2, str):
        d2 = date.fromisoformat(d2[:10])
    return (d2.year - d1.year) * 12 + (d2.month - d1.month)

def define_disease_cohort(disease_definition, healthlake, s3, bucket):
    """Step 1: Identify patients who qualify for the disease cohort.

    See pseudocode Step 1 in the main recipe. Returns the qualified
    cohort as a list of dicts. Writes the cohort to S3 keyed by
    cohort definition name and version so downstream stages can
    pick it up by reference.

    Production splits this across a Glue job that runs the
    HealthLake search and a Lambda that applies the qualification
    rules. The demo runs both in process for clarity.
    """
    today = date.today()

    candidate_patients = healthlake.list_all_patient_ids()
    logger.info("Walking %d candidate patients for cohort %s/%s",
                len(candidate_patients),
                disease_definition["name"],
                disease_definition["version"])

    cohort = []
    for patient_id in candidate_patients:
        # Pull the patient's Conditions to check inclusion and
        # exclusion ICD-10 codes.
        conditions = healthlake.search_conditions(patient_id)
        icd_codes = [c.get("code_icd10") for c in conditions]

        # Exclusion check first: patients with any exclusion code
        # are dropped regardless of inclusion.
        if any(c in disease_definition["exclusion_icd10"] for c in icd_codes):
            continue

        # Inclusion check.
        inclusion_matches = [c for c in icd_codes
                             if c in disease_definition["inclusion_icd10"]]
        if not inclusion_matches:
            continue

        # Time-zero anchor: earliest date of an inclusion ICD-10.
        inclusion_dates = [
            date.fromisoformat(c["onset_dt"][:10])
            for c in conditions
            if c.get("code_icd10") in disease_definition["inclusion_icd10"]
            and c.get("onset_dt")
        ]
        if not inclusion_dates:
            continue
        time_zero_date = min(inclusion_dates)

        # Minimum observation window: time from diagnosis to today
        # must exceed the configured floor.
        observation_span_months = _months_between(time_zero_date, today)
        if observation_span_months < disease_definition["minimum_observation_window_months"]:
            continue

        # Minimum density of trajectory-relevant measurements.
        # Only chronic-context eGFR observations count.
        primary_loinc = disease_definition["primary_outcome_loinc"]
        all_egfr = healthlake.search_observations(patient_id, primary_loinc)
        chronic_egfr = [o for o in all_egfr
                        if o.get("encounter_class") not in ACUTE_ENCOUNTER_CLASSES]
        if len(chronic_egfr) < disease_definition["minimum_egfr_measurements"]:
            continue

        # Capture qualifying signals for auditability.
        cohort.append({
            "patient_id":                  patient_id,
            "qualified_by_icd":            sorted(set(inclusion_matches)),
            "time_zero_date":              time_zero_date.isoformat(),
            "observation_span_months":     observation_span_months,
            "chronic_egfr_count":          len(chronic_egfr),
            "cohort_definition_name":      disease_definition["name"],
            "cohort_definition_version":   disease_definition["version"],
            "qualified_at_ts":             datetime.now(timezone.utc).isoformat(),
        })

    # Persist the qualified cohort to S3 so the downstream
    # harmonization, training, and inference stages can consume
    # it without re-running the qualification logic.
    cohort_key = (f"cohorts/{disease_definition['name']}/"
                  f"{disease_definition['version']}/cohort.json")
    s3.put_object(
        Bucket=bucket, Key=cohort_key,
        Body=json.dumps(cohort, default=str))

    logger.info("Cohort %s/%s qualified %d of %d patients",
                disease_definition["name"],
                disease_definition["version"],
                len(cohort),
                len(candidate_patients))

    return cohort

def _convert_units(value, from_unit, to_unit, loinc_code):
    """Convert a numeric value between units for a specific LOINC code.

    The demo's ADPKD pipeline uses LOINCs whose canonical units
    match the source-system units, so the conversion is a passthrough.
    Production carries a per-LOINC conversion table covering
    creatinine (mg/dL <-> umol/L using the molecular-weight factor),
    glucose (mg/dL <-> mmol/L), hemoglobin (g/dL <-> g/L), and others.
    Quarantine the record rather than guess for unsupported conversions.
    """
    if value is None:
        return None
    if from_unit == to_unit:
        return float(value)
    return None

def harmonize_patient_trajectory(cohort_member, healthlake, s3,
                                  harmonized_bucket,
                                  disease_definition):
    """Step 2: Build the harmonized per-patient trajectory record.

    See pseudocode Step 2 in the main recipe. Returns the harmonized
    dict. Writes the dict to S3 keyed by patient under the cohort
    definition's prefix.
    """
    patient_id     = cohort_member["patient_id"]
    time_zero_date = date.fromisoformat(cohort_member["time_zero_date"])

    # Walk every trajectory-relevant LOINC for this disease.
    harmonized_observations = []
    for loinc_code, display in disease_definition["trajectory_loincs"].items():
        catalog = LOINC_CATALOG.get(loinc_code)
        if catalog is None:
            logger.warning("LOINC %s not in catalog, skipping", loinc_code)
            continue

        all_obs = healthlake.search_observations(patient_id, loinc_code)
        for obs in all_obs:
            # Convert to canonical UCUM unit (passthrough in this demo).
            canonical_value = _convert_units(
                obs["value_quantity"], obs["value_unit"],
                catalog["canonical_unit"], loinc_code)
            if canonical_value is None:
                logger.warning(
                    "unit conversion failed for %s %s: %s -> %s",
                    patient_id, loinc_code,
                    obs["value_unit"], catalog["canonical_unit"])
                continue

            # Clip to clinically plausible bounds. Out-of-range
            # values are typically transcription or unit-mismatch
            # errors. Production may quarantine instead of clip.
            clipped = max(catalog["lower_clip"],
                          min(catalog["upper_clip"], canonical_value))

            # Tag chronic-vs-acute context.
            encounter_class = obs.get("encounter_class", "ambulatory")
            context_tag = ("acute"
                           if encounter_class in ACUTE_ENCOUNTER_CLASSES
                           else "chronic")

            # Anchor to time-since-diagnosis in months.
            obs_date = date.fromisoformat(obs["effective_dt"][:10])
            months_from_zero = (obs_date - time_zero_date).days / 30.44

            harmonized_observations.append({
                "loinc_code":        loinc_code,
                "loinc_display":     catalog["display"],
                "value":             round(clipped, 3),
                "unit":              catalog["canonical_unit"],
                "collection_ts":     obs["effective_dt"],
                "months_from_zero":  round(months_from_zero, 3),
                "context_tag":       context_tag,
                "encounter_class":   encounter_class,
            })

    harmonized_observations.sort(key=lambda o: o["months_from_zero"])

    # Treatment timeline: align medication starts and stops to
    # the same time frame.
    harmonized_treatments = []
    for med in healthlake.search_medications(patient_id):
        drug_class = med.get("drug_class")
        if drug_class not in disease_definition["relevant_drug_classes"]:
            continue
        start_date = date.fromisoformat(med["authored_on_dt"][:10])
        start_month = (start_date - time_zero_date).days / 30.44
        end_month = None
        if med.get("end_dt"):
            end_date = date.fromisoformat(med["end_dt"][:10])
            end_month = (end_date - time_zero_date).days / 30.44
        harmonized_treatments.append({
            "drug_class":             drug_class,
            "start_months_from_zero": round(start_month, 2),
            "end_months_from_zero":   round(end_month, 2) if end_month else None,
        })
    harmonized_treatments.sort(key=lambda t: t["start_months_from_zero"])

    # Compose the final harmonized record.
    today = date.today()
    last_observation_months = max(
        (o["months_from_zero"] for o in harmonized_observations
         if o["context_tag"] == "chronic"),
        default=cohort_member["observation_span_months"])

    harmonized = {
        "patient_id":                cohort_member["patient_id"],
        "time_zero_date":            cohort_member["time_zero_date"],
        "observation_span_months":   cohort_member["observation_span_months"],
        "last_observation_months":   round(last_observation_months, 2),
        "observations":              harmonized_observations,
        "treatments":                harmonized_treatments,
        "cohort_definition_name":    cohort_member["cohort_definition_name"],
        "cohort_definition_version": cohort_member["cohort_definition_version"],
        "harmonized_at_ts":          datetime.now(timezone.utc).isoformat(),
    }

    # Persist to S3.
    key = (f"cohorts/{harmonized['cohort_definition_name']}/"
           f"{harmonized['cohort_definition_version']}/"
           f"harmonized/{patient_id}.json")
    s3.put_object(Bucket=harmonized_bucket, Key=key,
                  Body=json.dumps(harmonized, default=str))

    return harmonized

def harmonize_cohort(cohort, healthlake, s3, harmonized_bucket,
                     disease_definition):
    """Run harmonization across every cohort member."""
    out = []
    for member in cohort:
        h = harmonize_patient_trajectory(
            member, healthlake, s3, harmonized_bucket, disease_definition)
        out.append(h)
    logger.info("Harmonized %d cohort members for %s/%s",
                len(out),
                disease_definition["name"],
                disease_definition["version"])
    return out

class BayesianHierarchicalMixedEffects:
    """Pedagogical Bayesian hierarchical linear mixed-effects model.

    The model is:

        egfr_observed[i, t] ~ Normal(
            intercept[i] + slope[i] * months_from_zero[i, t]
              + treatment_effect[i, t],
            observation_noise_sd
        )

        intercept[i]  ~ Normal(prior_intercept_mean, prior_intercept_sd)
        slope[i]      ~ Normal(population_slope, prior_per_patient_slope_sd)
        population_slope ~ Normal(prior_population_slope_mean,
                                  prior_population_slope_sd)

    The treatment effect is a multiplicative slope modifier driven
    by trial-derived priors. For example, when tolvaptan is active,
    the effective slope becomes slope[i] * (1 - tolvaptan_modifier).

    Production uses PyMC or Stan to fit the joint posterior with
    proper NUTS sampling and rich diagnostic output. The demo uses
    closed-form normal-conjugate updates per patient (which is what
    the linear-Gaussian case allows analytically) plus an empirical
    estimate of the population slope. This is deliberately simpler
    than what a real model fit produces, but it captures the same
    structure: population shape, per-patient deviation, treatment
    modifier, and posterior uncertainty.
    """

    def __init__(self, hyperparameters, trial_priors):
        self.hyper       = dict(hyperparameters)
        self.trial_priors = dict(trial_priors)
        self.fitted      = False
        self.population_slope_mean = None
        self.population_slope_sd   = None
        self.per_patient_params    = {}    # patient_id -> dict
        self.training_summary      = {}

    def _treatment_modifier(self, treatments, months_from_zero):
        """Compute the active treatment slope modifier at a time point.

        Returns the multiplicative factor on the slope. A patient
        with no active treatments returns 1.0; a patient on
        tolvaptan returns approximately (1 - tolvaptan_mean).
        Multiple concurrent treatments compound multiplicatively
        as a simplification; production uses an additive
        log-slope formulation or a proper joint model.
        """
        modifier = 1.0
        for tx in treatments:
            if tx["start_months_from_zero"] > months_from_zero:
                continue
            if (tx["end_months_from_zero"] is not None
                    and tx["end_months_from_zero"] < months_from_zero):
                continue
            prior = self.trial_priors.get(tx["drug_class"])
            if not prior:
                continue
            modifier *= (1.0 - prior["slope_modifier_mean"])
        return modifier

    def fit(self, harmonized_cohort, primary_outcome_loinc):
        """Fit the model on the harmonized cohort.

        Returns the population-level summary plus per-patient
        posterior parameters. Production fits this with NUTS in
        PyMC/Stan and returns full posterior samples.
        """
        # Step 3a: extract per-patient (months_from_zero, value)
        # series for the primary outcome, chronic-context only.
        per_patient_series = {}
        for patient in harmonized_cohort:
            obs = [(o["months_from_zero"], o["value"], o["context_tag"])
                   for o in patient["observations"]
                   if o["loinc_code"] == primary_outcome_loinc
                   and o["context_tag"] == "chronic"]
            if len(obs) < 3:
                continue
            obs.sort(key=lambda x: x[0])
            per_patient_series[patient["patient_id"]] = {
                "series":     obs,
                "treatments": patient["treatments"],
            }

        if not per_patient_series:
            raise ValueError("No qualifying per-patient series for fit")

        # Step 3b: estimate population slope as the inverse-variance-
        # weighted mean of per-patient OLS slopes (with intercept).
        # The slopes are computed against the treatment-adjusted
        # months so that treatment effects do not bias the population
        # estimate.
        per_patient_ols = {}
        for pid, pdata in per_patient_series.items():
            xs = [m for (m, v, c) in pdata["series"]]
            ys = [v for (m, v, c) in pdata["series"]]
            n = len(xs)
            mean_x = sum(xs) / n
            mean_y = sum(ys) / n
            num = sum((xs[i] - mean_x) * (ys[i] - mean_y) for i in range(n))
            den = sum((xs[i] - mean_x) ** 2 for i in range(n))
            if den <= 0:
                continue
            slope = num / den
            intercept = mean_y - slope * mean_x
            # Residual variance for the inverse-variance weight.
            residuals = [ys[i] - (intercept + slope * xs[i]) for i in range(n)]
            resid_var = sum(r ** 2 for r in residuals) / max(n - 2, 1)
            slope_var = resid_var / max(den, 1e-9)
            per_patient_ols[pid] = {
                "slope":      slope,
                "intercept":  intercept,
                "slope_var":  max(slope_var, 1e-6),
                "n_obs":      n,
            }

        # Inverse-variance-weighted population slope.
        total_weight = sum(1.0 / o["slope_var"] for o in per_patient_ols.values())
        weighted_sum = sum(o["slope"] / o["slope_var"]
                           for o in per_patient_ols.values())
        empirical_pop_slope = weighted_sum / max(total_weight, 1e-9)
        empirical_pop_sd    = math.sqrt(1.0 / max(total_weight, 1e-9))

        # Step 3c: combine the empirical estimate with the literature
        # prior using a normal-conjugate update.
        prior_mean = self.hyper["prior_population_slope_mean"]
        prior_sd   = self.hyper["prior_population_slope_sd"]
        prior_var  = prior_sd ** 2
        emp_var    = empirical_pop_sd ** 2

        post_var  = 1.0 / (1.0 / prior_var + 1.0 / emp_var)
        post_mean = post_var * (prior_mean / prior_var
                                + empirical_pop_slope / emp_var)

        self.population_slope_mean = post_mean
        self.population_slope_sd   = math.sqrt(post_var)

        # Step 3d: per-patient posterior. Shrink each patient's
        # OLS slope toward the population posterior using the
        # configured per-patient slope SD as the prior.
        per_pat_prior_sd = self.hyper["prior_per_patient_slope_sd"]
        per_pat_prior_var = per_pat_prior_sd ** 2
        for pid, ols in per_patient_ols.items():
            patient_prior_var = post_var + per_pat_prior_var
            obs_var = ols["slope_var"]
            shrunk_var = 1.0 / (1.0 / patient_prior_var + 1.0 / obs_var)
            shrunk_mean = shrunk_var * (post_mean / patient_prior_var
                                        + ols["slope"] / obs_var)
            self.per_patient_params[pid] = {
                "slope_mean":     shrunk_mean,
                "slope_sd":       math.sqrt(shrunk_var),
                "intercept_mean": ols["intercept"],
                "intercept_sd":   self.hyper["prior_per_patient_intercept_sd"],
                "n_obs":          ols["n_obs"],
                "treatments":     per_patient_series[pid]["treatments"],
            }

        # Step 3e: temporal-holdout calibration check. For each
        # patient, hold out the most recent fraction of observations,
        # refit (here, just predict at held-out times using the
        # already-fit per-patient posterior), and compute coverage.
        calibration = self._compute_calibration(
            per_patient_series,
            self.hyper["temporal_holdout_fraction"])

        self.training_summary = {
            "cohort_size":              len(per_patient_series),
            "population_slope_mean":    round(self.population_slope_mean, 4),
            "population_slope_sd":      round(self.population_slope_sd, 4),
            "empirical_pop_slope":      round(empirical_pop_slope, 4),
            "literature_prior_slope":   prior_mean,
            "calibration":              calibration,
            "model_version":            MODEL_VERSION_TAG,
            "trial_prior_version":      TRIAL_PRIOR_VERSION,
            "trained_at_ts":            datetime.now(timezone.utc).isoformat(),
        }
        self.fitted = True
        return self.training_summary

    def _compute_calibration(self, per_patient_series, holdout_fraction):
        """Coverage of credible intervals on a temporal holdout.

        For each patient, hold out the last `holdout_fraction` of
        observations, predict at those times using the per-patient
        posterior, and compute the empirical coverage of the
        50/80/90/95% credible intervals.
        """
        coverage_50 = 0
        coverage_80 = 0
        coverage_90 = 0
        coverage_95 = 0
        total = 0
        for pid, pdata in per_patient_series.items():
            params = self.per_patient_params.get(pid)
            if not params:
                continue
            n = len(pdata["series"])
            n_hold = max(1, int(round(n * holdout_fraction)))
            holdout = pdata["series"][-n_hold:]
            for (m, v, c) in holdout:
                modifier = self._treatment_modifier(params["treatments"], m)
                pred_mean = (params["intercept_mean"]
                             + params["slope_mean"] * modifier * m)
                # Predictive SD combines slope uncertainty
                # propagated through `m`, intercept uncertainty,
                # and observation noise.
                pred_sd = math.sqrt(
                    (params["slope_sd"] * m) ** 2
                    + params["intercept_sd"] ** 2
                    + self.hyper["observation_noise_sd"] ** 2)
                # 50/80/90/95 z thresholds for normal CI.
                for z, target in [(0.674, 50), (1.282, 80),
                                  (1.645, 90), (1.960, 95)]:
                    lo = pred_mean - z * pred_sd
                    hi = pred_mean + z * pred_sd
                    if lo <= v <= hi:
                        if target == 50: coverage_50 += 1
                        elif target == 80: coverage_80 += 1
                        elif target == 90: coverage_90 += 1
                        else: coverage_95 += 1
                total += 1
        if total == 0:
            return {"coverage_50": None, "coverage_80": None,
                    "coverage_90": None, "coverage_95": None,
                    "holdout_count": 0}
        return {
            "coverage_50":   round(coverage_50 / total, 3),
            "coverage_80":   round(coverage_80 / total, 3),
            "coverage_90":   round(coverage_90 / total, 3),
            "coverage_95":   round(coverage_95 / total, 3),
            "holdout_count": total,
        }

def train_trajectory_model(harmonized_cohort, disease_definition,
                            s3, model_artifact_bucket):
    """Step 3: Train the per-disease trajectory model.

    See pseudocode Step 3 in the main recipe. Returns the trained
    model object plus a summary. Persists the artifact to S3.
    """
    model = BayesianHierarchicalMixedEffects(
        hyperparameters=BAYESIAN_MODEL_HYPERPARAMETERS,
        trial_priors=TRIAL_DERIVED_EFFECT_PRIORS)
    summary = model.fit(
        harmonized_cohort,
        primary_outcome_loinc=disease_definition["primary_outcome_loinc"])

    # Persist a serialized summary plus the per-patient parameters.
    # In production the artifact also contains the full posterior
    # samples and the calibration backtest history.
    artifact = {
        "summary":             summary,
        "per_patient_params":  {
            pid: {k: (round(v, 4) if isinstance(v, (int, float)) else v)
                  for k, v in params.items()
                  if k != "treatments"}
            for pid, params in model.per_patient_params.items()
        },
        "cohort_definition_name":    disease_definition["name"],
        "cohort_definition_version": disease_definition["version"],
        "model_version":             MODEL_VERSION_TAG,
        "trial_prior_version":       TRIAL_PRIOR_VERSION,
    }
    artifact_key = (f"models/{disease_definition['name']}/"
                    f"{MODEL_VERSION_TAG}/{TRIAL_PRIOR_VERSION}/artifact.json")
    s3.put_object(Bucket=model_artifact_bucket, Key=artifact_key,
                  Body=json.dumps(artifact, default=str))

    logger.info(
        "Trained model for %s: cohort=%d, pop_slope=%.4f, calibration_90=%s",
        disease_definition["name"],
        summary["cohort_size"],
        summary["population_slope_mean"],
        summary["calibration"]["coverage_90"])

    return model, summary

def _build_forecast_grid(from_months, horizon_months, step_months):
    """Return a list of months_from_zero for the forecast horizon."""
    grid = []
    cursor = from_months
    end = from_months + horizon_months
    while cursor <= end + 1e-6:
        grid.append(round(cursor, 2))
        cursor += step_months
    return grid

def _forecast_at_time(model, params, treatments, m_target,
                      treatment_override=None):
    """Predictive mean and SD for a single patient at a given time.

    Combines slope uncertainty propagated through m_target,
    intercept uncertainty, observation noise, and (when applicable)
    treatment-effect modifier uncertainty derived from the trial
    priors. Returns (mean, sd).
    """
    treatments_to_use = treatment_override or treatments
    modifier = model._treatment_modifier(treatments_to_use, m_target)
    pred_mean = (params["intercept_mean"]
                 + params["slope_mean"] * modifier * m_target)
    # Slope-uncertainty contribution (scaled by the lever arm m).
    slope_var_contrib = (params["slope_sd"] * modifier * m_target) ** 2
    # Treatment-effect uncertainty contribution. For each active
    # treatment, the trial-derived SD propagates through the slope
    # and the lever arm.
    tx_var_contrib = 0.0
    for tx in treatments_to_use:
        if tx["start_months_from_zero"] > m_target:
            continue
        if (tx["end_months_from_zero"] is not None
                and tx["end_months_from_zero"] < m_target):
            continue
        prior = model.trial_priors.get(tx["drug_class"])
        if not prior:
            continue
        tx_var_contrib += (prior["slope_modifier_sd"]
                           * params["slope_mean"] * m_target) ** 2
    pred_sd = math.sqrt(slope_var_contrib
                        + params["intercept_sd"] ** 2
                        + model.hyper["observation_noise_sd"] ** 2
                        + tx_var_contrib)
    return pred_mean, pred_sd

def infer_patient_trajectory(patient, model, endpoint_definition,
                              forecast_horizon_months=DEFAULT_FORECAST_HORIZON_MONTHS,
                              grid_step=DEFAULT_FORECAST_GRID_STEP_MO,
                              num_endpoint_samples=400,
                              seed=12345):
    """Step 4: Produce the per-patient trajectory and forecast.

    Returns a dict with: fitted_trajectory through the observed
    history, forecast at the future grid (median + credible
    intervals), and time-to-endpoint distribution.
    """
    rng = random.Random(seed + hash(patient["patient_id"]) % 10000)
    params = model.per_patient_params.get(patient["patient_id"])
    if params is None:
        return {
            "patient_id":  patient["patient_id"],
            "status":      "no_per_patient_params",
        }

    primary_loinc = endpoint_definition["loinc_code"]
    chronic_obs = [(o["months_from_zero"], o["value"])
                   for o in patient["observations"]
                   if o["loinc_code"] == primary_loinc
                   and o["context_tag"] == "chronic"]

    # 4a. Fitted trajectory through observed times.
    fitted = []
    for (m, v) in chronic_obs:
        pred_mean, pred_sd = _forecast_at_time(
            model, params, patient["treatments"], m)
        fitted.append({
            "months_from_zero": m,
            "observed":         round(v, 2),
            "fitted_mean":      round(pred_mean, 2),
            "fitted_p10":       round(pred_mean - 1.282 * pred_sd, 2),
            "fitted_p90":       round(pred_mean + 1.282 * pred_sd, 2),
        })

    # 4b. Forward forecast under current treatment continued.
    last_obs_months = patient["last_observation_months"]
    grid = _build_forecast_grid(
        last_obs_months, forecast_horizon_months, grid_step)
    forecast = []
    for m in grid:
        pred_mean, pred_sd = _forecast_at_time(
            model, params, patient["treatments"], m)
        forecast.append({
            "months_from_zero": m,
            "p10":              round(pred_mean - 1.282 * pred_sd, 2),
            "p50":              round(pred_mean, 2),
            "p90":              round(pred_mean + 1.282 * pred_sd, 2),
            "predictive_sd":    round(pred_sd, 2),
        })

    # 4c. Time-to-endpoint by Monte Carlo. Sample slope and
    # intercept from the per-patient posterior, simulate the
    # trajectory forward, find the first crossing of the threshold.
    threshold       = endpoint_definition["threshold_value"]
    direction_below = (endpoint_definition["threshold_direction"] == "below")
    crossings = []
    for _ in range(num_endpoint_samples):
        sampled_slope     = rng.gauss(params["slope_mean"], params["slope_sd"])
        sampled_intercept = rng.gauss(
            params["intercept_mean"], params["intercept_sd"])
        # Sample treatment-effect modifiers per-treatment per-draw.
        sampled_modifiers = {}
        for tx in patient["treatments"]:
            prior = model.trial_priors.get(tx["drug_class"])
            if prior:
                sampled_modifiers[tx["drug_class"]] = rng.gauss(
                    prior["slope_modifier_mean"], prior["slope_modifier_sd"])
            else:
                sampled_modifiers[tx["drug_class"]] = 0.0

        # Walk the forecast grid for this draw.
        crossing_m = None
        for m in grid:
            modifier = 1.0
            for tx in patient["treatments"]:
                if tx["start_months_from_zero"] > m:
                    continue
                if (tx["end_months_from_zero"] is not None
                        and tx["end_months_from_zero"] < m):
                    continue
                modifier *= (1.0 - sampled_modifiers.get(tx["drug_class"], 0.0))
            v = sampled_intercept + sampled_slope * modifier * m
            v += rng.gauss(0.0, model.hyper["observation_noise_sd"])
            if direction_below and v <= threshold:
                crossing_m = m
                break
            if not direction_below and v >= threshold:
                crossing_m = m
                break
        crossings.append(crossing_m)

    finite_crossings = [c for c in crossings if c is not None]
    fraction_crossing = len(finite_crossings) / len(crossings)
    sorted_crossings = sorted(finite_crossings)

    def _percentile(data, pct):
        if not data:
            return None
        idx = max(0, min(len(data) - 1, int(round(pct / 100.0 * len(data)))))
        return data[idx]

    time_to_endpoint = {
        "endpoint_definition":  endpoint_definition,
        "fraction_reaching_endpoint_in_horizon": round(fraction_crossing, 3),
        "p10_months":           _percentile(sorted_crossings, 10),
        "p50_months":           _percentile(sorted_crossings, 50),
        "p90_months":           _percentile(sorted_crossings, 90),
        "samples":              len(crossings),
        "horizon_months":       forecast_horizon_months,
    }

    return {
        "patient_id":              patient["patient_id"],
        "status":                  "ok",
        "fitted_trajectory":       fitted,
        "forecast":                forecast,
        "time_to_endpoint":        time_to_endpoint,
        "model_version":           MODEL_VERSION_TAG,
        "trial_prior_version":     TRIAL_PRIOR_VERSION,
        "cohort_definition_version": patient["cohort_definition_version"],
        "inferred_at_ts":          datetime.now(timezone.utc).isoformat(),
    }

def infer_all_trajectories(harmonized_cohort, model,
                            endpoint_definition, s3,
                            forecast_bucket):
    """Run inference for every cohort patient. Persist to S3."""
    results = []
    for patient in harmonized_cohort:
        r = infer_patient_trajectory(patient, model, endpoint_definition)
        if r.get("status") != "ok":
            continue
        key = (f"forecasts/{patient['cohort_definition_name']}/"
               f"{patient['patient_id']}/{r['inferred_at_ts'][:10]}.json")
        s3.put_object(Bucket=forecast_bucket, Key=key,
                      Body=json.dumps(r, default=str))
        results.append(r)
    logger.info("Inferred %d patient trajectories", len(results))
    return results

def _apply_treatment_change(base_timeline, change, anchor_months):
    """Apply a treatment change spec to a base timeline.

    Returns (new_timeline, change_applied). Supported changes:
      - {"add": {"drug_class": "tolvaptan", "start_offset_months": 0}}
      - {"stop": {"drug_class": "acei_arb"}}
      - None (no change; "current treatment continued")

    IMPORTANT: If the patient is already on the requested drug class
    (an active entry with no end date or an end date in the future),
    the "add" is a no-op and change_applied returns False. Without
    this reconciliation, duplicate entries compound multiplicatively
    in _treatment_modifier: (1 - 0.30) * (1 - 0.30) = 0.49 instead
    of the correct 0.70, producing a 25% over-attribution of
    treatment benefit.
    """
    new_timeline = [dict(t) for t in base_timeline]
    if change is None:
        return new_timeline, True
    if "add" in change:
        spec = change["add"]
        drug_class = spec["drug_class"]
        # Check pre-existing exposure: if the patient is already on
        # this drug class (active entry with no end or end in future),
        # return unchanged timeline with change_applied=False.
        already_active = any(
            t["drug_class"] == drug_class
            and (t["end_months_from_zero"] is None
                 or t["end_months_from_zero"] > anchor_months)
            for t in new_timeline
        )
        if already_active:
            return new_timeline, False
        new_timeline.append({
            "drug_class":             drug_class,
            "start_months_from_zero": (anchor_months
                                       + spec.get("start_offset_months", 0)),
            "end_months_from_zero":   None,
        })
    if "stop" in change:
        target = change["stop"]["drug_class"]
        for tx in new_timeline:
            if tx["drug_class"] == target and tx["end_months_from_zero"] is None:
                tx["end_months_from_zero"] = anchor_months
    new_timeline.sort(key=lambda t: t["start_months_from_zero"])
    return new_timeline, True

def _compose_assumption_disclosure(scenario, model):
    """Plain-language disclosure of what the forecast assumes."""
    base = (
        "Forecasts assume the specified treatment scenario continues "
        "through the forecast horizon, no acute clinical events "
        "disrupt the trajectory, and the patient's disease behavior "
        "remains comparable to the cohort. Forecasts are statistical "
        "projections; individual outcomes vary substantially.")
    if scenario.get("change") is None:
        return base + " This scenario assumes current treatment continued."
    if "add" in scenario["change"]:
        drug = scenario["change"]["add"]["drug_class"]
        prior = model.trial_priors.get(drug, {})
        evidence = prior.get("evidence_basis", "published trial evidence")
        return (
            f"{base} The treatment-effect modifier for {drug} is "
            f"derived from {evidence}. Effect size carries the "
            f"trial-derived credible interval; the forecast propagates "
            f"that uncertainty.")
    return base

def _compose_explanation_text(patient_id, forecast, time_to_endpoint,
                                scenario_name):
    """Plain-language narrative for the surfaced trajectory.

    The narrative is the product. A clinician should be able to
    read this in fifteen seconds and have the magnitude, the
    horizon, the uncertainty, and the assumption surfaced.
    """
    horizon_months = time_to_endpoint["horizon_months"]
    p10 = time_to_endpoint.get("p10_months")
    p50 = time_to_endpoint.get("p50_months")
    p90 = time_to_endpoint.get("p90_months")
    fraction = time_to_endpoint["fraction_reaching_endpoint_in_horizon"]
    five_year = next(
        (f for f in forecast if abs(f["months_from_zero"]
                                     - (forecast[0]["months_from_zero"] + 60))
            < 4),
        forecast[-1])

    parts = []
    parts.append(
        f"Under the {scenario_name} scenario, the patient's eGFR is "
        f"projected to reach approximately {five_year['p50']:.0f} "
        f"mL/min/1.73 m^2 at the {five_year['months_from_zero']:.0f}-month "
        f"forecast point, with a 90% credible interval of "
        f"{five_year['p10']:.0f} to {five_year['p90']:.0f}.")
    if p50 is not None:
        parts.append(
            f"Median time to the eGFR-under-15 (RRT consideration) "
            f"endpoint is approximately {p50:.0f} months "
            f"(P10 = {p10:.0f}, P90 = {p90:.0f}); "
            f"{fraction*100:.0f}% of posterior draws cross the threshold "
            f"within the {horizon_months}-month horizon.")
    else:
        parts.append(
            f"In the {horizon_months}-month horizon, fewer than 10% of "
            f"posterior draws reach the eGFR-under-15 threshold; the "
            f"median time to that endpoint is beyond the modeled horizon.")
    return " ".join(parts)

def evaluate_counterfactual_scenarios(patient, baseline_inference, model,
                                        scenarios, endpoint_definition,
                                        forecast_horizon_months=DEFAULT_FORECAST_HORIZON_MONTHS,
                                        grid_step=DEFAULT_FORECAST_GRID_STEP_MO,
                                        num_endpoint_samples=400):
    """Step 5: Compose counterfactual scenarios for one patient.

    Returns a payload containing the baseline scenario, every
    requested counterfactual scenario, and the assumption
    disclosure for each.
    """
    anchor = patient["last_observation_months"]
    grid = _build_forecast_grid(anchor, forecast_horizon_months, grid_step)
    params = model.per_patient_params.get(patient["patient_id"])
    if params is None:
        return None

    rng_seed = 23456 + hash(patient["patient_id"]) % 10000
    counterfactual_results = []

    for scenario in scenarios:
        modified_timeline, change_applied = _apply_treatment_change(
            patient["treatments"], scenario.get("change"), anchor)

        # Forward forecast under the modified timeline.
        scenario_forecast = []
        for m in grid:
            pred_mean, pred_sd = _forecast_at_time(
                model, params, modified_timeline, m,
                treatment_override=modified_timeline)
            scenario_forecast.append({
                "months_from_zero": m,
                "p10":              round(pred_mean - 1.282 * pred_sd, 2),
                "p50":              round(pred_mean, 2),
                "p90":              round(pred_mean + 1.282 * pred_sd, 2),
            })

        # Time-to-endpoint Monte Carlo under the modified timeline.
        rng = random.Random(rng_seed + hash(scenario["name"]) % 10000)
        threshold = endpoint_definition["threshold_value"]
        direction_below = (endpoint_definition["threshold_direction"] == "below")
        crossings = []
        for _ in range(num_endpoint_samples):
            sampled_slope = rng.gauss(params["slope_mean"], params["slope_sd"])
            sampled_intercept = rng.gauss(
                params["intercept_mean"], params["intercept_sd"])
            sampled_modifiers = {}
            for tx in modified_timeline:
                prior = model.trial_priors.get(tx["drug_class"])
                if prior:
                    sampled_modifiers[tx["drug_class"]] = rng.gauss(
                        prior["slope_modifier_mean"],
                        prior["slope_modifier_sd"])
                else:
                    sampled_modifiers[tx["drug_class"]] = 0.0
            crossing_m = None
            for m in grid:
                modifier = 1.0
                for tx in modified_timeline:
                    if tx["start_months_from_zero"] > m:
                        continue
                    if (tx["end_months_from_zero"] is not None
                            and tx["end_months_from_zero"] < m):
                        continue
                    modifier *= (1.0 - sampled_modifiers.get(tx["drug_class"], 0.0))
                v = (sampled_intercept + sampled_slope * modifier * m
                     + rng.gauss(0.0, model.hyper["observation_noise_sd"]))
                if direction_below and v <= threshold:
                    crossing_m = m; break
                if not direction_below and v >= threshold:
                    crossing_m = m; break
            crossings.append(crossing_m)
        finite = sorted([c for c in crossings if c is not None])
        fraction = len(finite) / len(crossings)

        def _pct(data, p):
            if not data:
                return None
            i = max(0, min(len(data) - 1, int(round(p / 100.0 * len(data)))))
            return data[i]

        time_to_endpoint = {
            "endpoint_definition":  endpoint_definition,
            "fraction_reaching_endpoint_in_horizon": round(fraction, 3),
            "p10_months":           _pct(finite, 10),
            "p50_months":           _pct(finite, 50),
            "p90_months":           _pct(finite, 90),
            "samples":              len(crossings),
            "horizon_months":       forecast_horizon_months,
        }

        explanation = _compose_explanation_text(
            patient["patient_id"],
            scenario_forecast,
            time_to_endpoint,
            scenario["name"])

        counterfactual_results.append({
            "scenario_name":          scenario["name"],
            "scenario_description":   scenario.get("description", ""),
            "forecast":               scenario_forecast,
            "time_to_endpoint":       time_to_endpoint,
            "assumption_disclosure":  _compose_assumption_disclosure(scenario, model),
            "explanation_text":       explanation,
        })

    return {
        "patient_id":                patient["patient_id"],
        "scenarios":                 counterfactual_results,
        "baseline_scenario_name":    "current_continued",
        "model_version":             MODEL_VERSION_TAG,
        "trial_prior_version":       TRIAL_PRIOR_VERSION,
        "cohort_definition_version": patient["cohort_definition_version"],
        "generated_at_ts":           datetime.now(timezone.utc).isoformat(),
    }

def deliver_trajectory_payloads(payloads, table, event_bus, cloudwatch,
                                  s3, counterfactual_bucket, run_id):
    """Persist per-patient counterfactual payloads.

    Surfaced records go to DynamoDB with the disease + model version
    + generated_at composite sort key. Full payloads (including
    every scenario forecast point) go to S3 as the analytic record.
    Pipeline-completion events go to EventBridge; metrics go to
    CloudWatch.
    """
    written = 0
    for payload in payloads:
        # 5a. Full payload to S3 (the analytic record).
        s3_key = (f"counterfactuals/adpkd/{payload['patient_id']}/"
                  f"{payload['generated_at_ts'][:19].replace(':', '-')}.json")
        s3.put_object(Bucket=counterfactual_bucket, Key=s3_key,
                      Body=json.dumps(payload, default=str))

        # 5b. Summarized record per scenario to DynamoDB. The sort
        # key combines disease, model version, and generated_at so
        # the latest record per patient is a Query with prefix and
        # Limit=1 in descending order.
        with table.batch_writer() as bw:
            for scen in payload["scenarios"]:
                tte = scen["time_to_endpoint"]
                forecast = scen["forecast"]
                # Use the forecast point closest to 60 months
                # past anchor as the surfaced "5-year" reference.
                anchor_m = forecast[0]["months_from_zero"]
                target_m = anchor_m + 60.0
                ref_point = min(forecast,
                                key=lambda f: abs(f["months_from_zero"] - target_m))
                sk = (f"adpkd#{MODEL_VERSION_TAG}#"
                      f"{payload['generated_at_ts']}#{scen['scenario_name']}")
                item = {
                    "patient_id":               payload["patient_id"],
                    "disease_modelversion_generated_at": sk,
                    "scenario_name":            scen["scenario_name"],
                    "model_version":            MODEL_VERSION_TAG,
                    "trial_prior_version":      TRIAL_PRIOR_VERSION,
                    "cohort_definition_version": payload["cohort_definition_version"],
                    "ref_point_month":          _to_decimal(ref_point["months_from_zero"]),
                    "ref_point_p10":            _to_decimal(ref_point["p10"]),
                    "ref_point_p50":            _to_decimal(ref_point["p50"]),
                    "ref_point_p90":            _to_decimal(ref_point["p90"]),
                    "fraction_to_endpoint":     _to_decimal(
                        tte["fraction_reaching_endpoint_in_horizon"]),
                    "p10_months_to_endpoint":   _to_decimal(tte.get("p10_months")),
                    "p50_months_to_endpoint":   _to_decimal(tte.get("p50_months")),
                    "p90_months_to_endpoint":   _to_decimal(tte.get("p90_months")),
                    "endpoint_display":         tte["endpoint_definition"]["display"],
                    "explanation_text":         scen["explanation_text"],
                    "assumption_disclosure":    scen["assumption_disclosure"],
                    "generated_at_ts":          payload["generated_at_ts"],
                    "run_id":                   run_id,
                    "pipeline_version":         PIPELINE_VERSION,
                }
                bw.put_item(Item=item)
                written += 1

    # 5c. EventBridge completion event (no PHI in the payload).
    event_bus.put_events(Entries=[{
        "Source":       "trajectory.adpkd",
        "DetailType":   "TrajectoryBatchCompleted",
        "EventBusName": TRAJECTORY_EVENT_BUS_NAME,
        "Time":         datetime.now(timezone.utc),
        "Detail":       json.dumps({
            "run_id":              run_id,
            "patients_processed":  len(payloads),
            "scenarios_written":   written,
            "model_version":       MODEL_VERSION_TAG,
            "trial_prior_version": TRIAL_PRIOR_VERSION,
            "pipeline_version":    PIPELINE_VERSION,
        }),
    }])

    # 5d. CloudWatch metrics. Patient and scenario counts plus
    # the calibration-coverage metric so the operations team can
    # track drift over time. Calibration drift is the single most
    # important operational signal.
    cloudwatch.put_metric_data(
        Namespace=CLOUDWATCH_NAMESPACE,
        MetricData=[
            {"MetricName": "PatientsProcessed",
             "Value":      float(len(payloads)),
             "Unit":       "Count"},
            {"MetricName": "ScenariosWritten",
             "Value":      float(written),
             "Unit":       "Count"},
        ])

    logger.info(
        "Delivered: %d patients, %d scenarios written to DynamoDB",
        len(payloads), written)
    return {
        "patients_delivered":  len(payloads),
        "scenarios_written":   written,
    }

def run_trajectory_pipeline(table, event_bus, cloudwatch, healthlake, s3,
                             disease_definition=ADPKD_COHORT_DEFINITION,
                             endpoint_definition=ENDPOINT_DEFINITION):
    """End-to-end pipeline orchestration.

    The demo wires up synthetic data; production starts with a
    HealthLake search for the disease cohort plus the harmonized
    S3 prefix populated by the streaming ingest job.
    """
    run_id = str(uuid.uuid4())
    print(f"\n=== Disease Progression Trajectory Pipeline run_id={run_id} ===\n")

    # --- Generate synthetic input data (production reads from HealthLake) ---
    patients = generate_synthetic_adpkd_cohort()
    populate_synthetic_healthlake(patients, healthlake)
    print(f"[input] {len(patients)} synthetic ADPKD patients in HealthLake")

    # --- Step 1: Define the cohort ---
    print("\n[step 1] define_disease_cohort")
    cohort = define_disease_cohort(
        disease_definition, healthlake, s3, COHORT_DATASET_BUCKET)
    print(f"  -> {len(cohort)} qualified cohort members")

    # --- Step 2: Harmonize ---
    print("\n[step 2] harmonize_cohort")
    harmonized = harmonize_cohort(
        cohort, healthlake, s3, HARMONIZED_BUCKET, disease_definition)
    chronic_obs_total = sum(
        len([o for o in p["observations"] if o["context_tag"] == "chronic"])
        for p in harmonized)
    acute_obs_total = sum(
        len([o for o in p["observations"] if o["context_tag"] == "acute"])
        for p in harmonized)
    print(f"  -> {len(harmonized)} harmonized records")
    print(f"  -> {chronic_obs_total} chronic observations, "
          f"{acute_obs_total} acute observations (excluded from training)")

    # --- Step 3: Train the trajectory model ---
    print("\n[step 3] train_trajectory_model")
    model, summary = train_trajectory_model(
        harmonized, disease_definition, s3, MODEL_ARTIFACT_BUCKET)
    print(f"  -> population slope: {summary['population_slope_mean']:.4f} +/- "
          f"{summary['population_slope_sd']:.4f} mL/min/1.73m^2 per month")
    print(f"  -> calibration coverage: 50%={summary['calibration']['coverage_50']}, "
          f"80%={summary['calibration']['coverage_80']}, "
          f"90%={summary['calibration']['coverage_90']}, "
          f"95%={summary['calibration']['coverage_95']}")

    # --- Step 4: Per-patient inference ---
    print("\n[step 4] infer_all_trajectories")
    inferences = infer_all_trajectories(
        harmonized, model, endpoint_definition, s3, FORECAST_BUCKET)
    crossing_patients = sum(
        1 for inf in inferences
        if inf["time_to_endpoint"]["fraction_reaching_endpoint_in_horizon"] >= 0.5)
    print(f"  -> {len(inferences)} per-patient forecasts produced")
    print(f"  -> {crossing_patients} patients with > 50% probability of "
          f"reaching eGFR < 15 within {DEFAULT_FORECAST_HORIZON_MONTHS} months")

    # --- Step 5: Counterfactual scenarios ---
    print("\n[step 5] evaluate_counterfactual_scenarios")
    scenarios_to_evaluate = [
        {"name":        "current_continued",
         "description": "Continue current treatment regimen.",
         "change":      None},
        {"name":        "start_tolvaptan_now",
         "description": "Add tolvaptan starting at the next visit.",
         "change":      {"add": {"drug_class":           "tolvaptan",
                                  "start_offset_months": 0}}},
        {"name":        "start_sglt2_now",
         "description": "Add SGLT2 inhibitor starting at the next visit.",
         "change":      {"add": {"drug_class":           "sglt2_inhibitor",
                                  "start_offset_months": 0}}},
    ]
    counterfactual_payloads = []
    for patient in harmonized:
        # Skip patients already on tolvaptan in the baseline; the
        # "start tolvaptan" scenario is meaningless for them.
        already_on_tolvaptan = any(
            t["drug_class"] == "tolvaptan" for t in patient["treatments"])
        scenarios = scenarios_to_evaluate
        if already_on_tolvaptan:
            scenarios = [s for s in scenarios_to_evaluate
                         if s["name"] != "start_tolvaptan_now"]

        baseline_inf = next(
            (i for i in inferences if i["patient_id"] == patient["patient_id"]),
            None)
        if baseline_inf is None:
            continue
        payload = evaluate_counterfactual_scenarios(
            patient, baseline_inf, model, scenarios, endpoint_definition)
        if payload is not None:
            counterfactual_payloads.append(payload)
    print(f"  -> {len(counterfactual_payloads)} per-patient counterfactual payloads")

    # --- Step 5b: Deliver to DynamoDB / EventBridge / CloudWatch ---
    print("\n[step 5b] deliver_trajectory_payloads")
    delivery = deliver_trajectory_payloads(
        counterfactual_payloads, table, event_bus, cloudwatch,
        s3, COUNTERFACTUAL_BUCKET, run_id)
    print(f"  -> wrote {delivery['scenarios_written']} scenario records to DynamoDB")
    print(f"  -> emitted {len(event_bus.events)} EventBridge events")
    print(f"  -> emitted CloudWatch metrics: "
          f"{sorted(cloudwatch.metrics.keys())}")

    return {
        "run_id":                   run_id,
        "cohort_size":              len(cohort),
        "model_summary":            summary,
        "counterfactual_payloads":  counterfactual_payloads,
    }

def run_demo():
    """Run the pipeline end-to-end against the in-memory mocks.

    No AWS resources are touched; every external dependency is a
    mock. Useful for sanity-checking the trajectory math and the
    counterfactual scenarios before wiring to real services.
    """
    table       = MockTable(PATIENT_TRAJECTORIES_TABLE)
    event_bus   = MockEventBus(TRAJECTORY_EVENT_BUS_NAME)
    cloudwatch  = MockCloudWatch()
    healthlake  = MockHealthLake()
    s3          = MockS3()

    result = run_trajectory_pipeline(
        table, event_bus, cloudwatch, healthlake, s3)

    print("\n=== Sample patient counterfactual summary ===")
    if result["counterfactual_payloads"]:
        sample = result["counterfactual_payloads"][0]
        print(f"  patient_id: {sample['patient_id']}")
        for scen in sample["scenarios"]:
            tte = scen["time_to_endpoint"]
            ref_point = scen["forecast"][min(20, len(scen["forecast"]) - 1)]
            print(f"  scenario: {scen['scenario_name']}")
            print(f"    eGFR forecast at month {ref_point['months_from_zero']}: "
                  f"P10={ref_point['p10']}, P50={ref_point['p50']}, "
                  f"P90={ref_point['p90']}")
            print(f"    fraction to endpoint in horizon: "
                  f"{tte['fraction_reaching_endpoint_in_horizon']*100:.0f}%")
            print(f"    median months to endpoint: {tte['p50_months']}")
            print(f"    explanation: {scen['explanation_text'][:120]}...")

    print("\n=== Sample DynamoDB record ===")
    if table.items:
        first_key = next(iter(table.items))
        sample_item = table.items[first_key]

        def _decimalify(o):
            if isinstance(o, Decimal):
                return str(o)
            if isinstance(o, datetime):
                return o.isoformat()
            return o
        print(json.dumps(sample_item, default=_decimalify, indent=2))

    return result

if __name__ == "__main__":
    run_demo()