scripts/generate_data.py

"""Synthetic data generator for the RHTP evaluation dashboard.

The schemas in this script mirror the column structure and units of the
real data sources the RHTP evaluation will draw from:

  * Hospital ops data (CMS Hospital Cost Reports + state-supplied facility
    quarterly metrics)
  * Claims data (Montana Medicaid + duals enrollment files)
  * Workforce data (HRSA Area Health Resource File + Montana DPHHS licensure)
  * Prevention data (BRFSS, county vital statistics, Montana YRBS)
  * HIE / EHR data (Big Sky Care Connect participation roster + ONC certified
    EHR registry)
  * Geography (US Census TIGER county FIPS, MT DPHHS facility roster)

To use real data instead, replace the files in /data/* with files of the
same name and column schema -- everything else just works.
"""
from __future__ import annotations

import math
import os
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
RNG = np.random.default_rng(20260425)

# --------------------------------------------------------------------------
# 1. Geography: Montana counties
# --------------------------------------------------------------------------
# Approximate centroids (lat/lon) used for map markers.
# Population estimates align with 2020 census order of magnitude.
COUNTIES = [
    # (fips, name, lat, lon, pop, classification, region)
    ("30001", "Beaverhead",       45.10, -113.00,   9540, "Rural",  "Southwest"),
    ("30003", "Big Horn",         45.42, -107.51,  13062, "Tribal", "South Central"),
    ("30005", "Blaine",           48.55, -108.90,   6643, "Tribal", "North Central"),
    ("30007", "Broadwater",       46.32, -111.51,   7587, "Rural",  "Southwest"),
    ("30009", "Carbon",           45.22, -109.05,  10725, "Rural",  "South Central"),
    ("30011", "Carter",           45.59, -104.55,   1411, "Rural",  "Southeast"),
    ("30013", "Cascade",          47.40, -111.30,  84414, "Urban",  "North Central"),
    ("30015", "Chouteau",         47.90, -110.40,   5635, "Rural",  "North Central"),
    ("30017", "Custer",           46.30, -105.61,  11790, "Rural",  "Southeast"),
    ("30019", "Daniels",          48.80, -105.60,   1722, "Rural",  "Northeast"),
    ("30021", "Dawson",           47.27, -104.90,   8940, "Rural",  "Northeast"),
    ("30023", "Deer Lodge",       46.07, -112.90,   9421, "Rural",  "Southwest"),
    ("30025", "Fallon",           46.30, -104.40,   3076, "Rural",  "Southeast"),
    ("30027", "Fergus",           47.05, -109.20,  11446, "Rural",  "Central"),
    ("30029", "Flathead",         48.30, -114.30, 104357, "Urban",  "Northwest"),
    ("30031", "Gallatin",         45.55, -111.20, 118960, "Urban",  "Southwest"),
    ("30033", "Garfield",         47.20, -106.90,   1207, "Rural",  "Northeast"),
    ("30035", "Glacier",          48.70, -113.00,  13873, "Tribal", "Northwest"),
    ("30037", "Golden Valley",    46.40, -109.20,    813, "Rural",  "Central"),
    ("30039", "Granite",          46.40, -113.40,   3379, "Rural",  "Southwest"),
    ("30041", "Hill",             48.60, -110.10,  16481, "Tribal", "North Central"),
    ("30043", "Jefferson",        46.10, -112.00,  12222, "Rural",  "Southwest"),
    ("30045", "Judith Basin",     47.00, -110.30,   1939, "Rural",  "Central"),
    ("30047", "Lake",             47.60, -114.10,  30458, "Tribal", "Northwest"),
    ("30049", "Lewis and Clark",  47.00, -112.40,  70973, "Urban",  "West Central"),
    ("30051", "Liberty",          48.60, -111.00,   2347, "Rural",  "North Central"),
    ("30053", "Lincoln",          48.60, -115.40,  20230, "Rural",  "Northwest"),
    ("30055", "McCone",           47.60, -105.80,   1664, "Rural",  "Northeast"),
    ("30057", "Madison",          45.30, -111.90,   8857, "Rural",  "Southwest"),
    ("30059", "Meagher",          46.60, -110.90,   1958, "Rural",  "Central"),
    ("30061", "Mineral",          47.10, -115.00,   4563, "Rural",  "Northwest"),
    ("30063", "Missoula",         47.00, -113.90, 117922, "Urban",  "West Central"),
    ("30065", "Musselshell",      46.50, -108.40,   4633, "Rural",  "Central"),
    ("30067", "Park",             45.50, -110.50,  17191, "Rural",  "Southwest"),
    ("30069", "Petroleum",        47.10, -108.30,    524, "Rural",  "Central"),
    ("30071", "Phillips",         48.30, -107.90,   3939, "Tribal", "North Central"),
    ("30073", "Pondera",          48.20, -112.20,   6029, "Rural",  "North Central"),
    ("30075", "Powder River",     45.40, -105.60,   1682, "Rural",  "Southeast"),
    ("30077", "Powell",           46.90, -113.00,   7109, "Rural",  "West Central"),
    ("30079", "Prairie",          46.90, -105.40,   1067, "Rural",  "Southeast"),
    ("30081", "Ravalli",          46.10, -114.10,  44174, "Rural",  "West Central"),
    ("30083", "Richland",         47.80, -104.60,  10803, "Rural",  "Northeast"),
    ("30085", "Roosevelt",        48.30, -105.00,  10425, "Tribal", "Northeast"),
    ("30087", "Rosebud",          46.20, -106.70,   8295, "Tribal", "Southeast"),
    ("30089", "Sanders",          47.70, -115.10,  12253, "Rural",  "Northwest"),
    ("30091", "Sheridan",         48.70, -104.50,   3539, "Rural",  "Northeast"),
    ("30093", "Silver Bow",       45.90, -112.60,  35133, "Rural",  "Southwest"),
    ("30095", "Stillwater",       45.60, -109.40,   9642, "Rural",  "South Central"),
    ("30097", "Sweet Grass",      45.80, -109.90,   3737, "Rural",  "South Central"),
    ("30099", "Teton",            47.80, -112.20,   6204, "Rural",  "North Central"),
    ("30101", "Toole",            48.70, -111.70,   4974, "Rural",  "North Central"),
    ("30103", "Treasure",         46.20, -107.00,    718, "Rural",  "Southeast"),
    ("30105", "Valley",           48.40, -106.70,   7396, "Rural",  "Northeast"),
    ("30107", "Wheatland",        46.50, -109.80,   2081, "Rural",  "Central"),
    ("30109", "Wibaux",           46.90, -104.20,    969, "Rural",  "Southeast"),
    ("30111", "Yellowstone",      45.90, -108.30, 164731, "Urban",  "South Central"),
]

# Years of panel data: 5 baseline + 7 RHTP years
YEARS = list(range(2019, 2032))
QUARTERS = [(y, q) for y in YEARS for q in (1, 2, 3, 4)]


def write(df: pd.DataFrame, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(path, index=False)
    print(f"  wrote {path.relative_to(ROOT)}  ({len(df):,} rows)")


# --------------------------------------------------------------------------
# 2. County master and population panel
# --------------------------------------------------------------------------
def make_counties() -> pd.DataFrame:
    df = pd.DataFrame(
        COUNTIES,
        columns=["fips", "county_name", "lat", "lon", "population_2020",
                 "rurality", "region"],
    )
    df["state"] = "MT"
    df["state_fips"] = "30"
    return df


def make_population_panel(counties: pd.DataFrame) -> pd.DataFrame:
    """Annual population panel with mild realistic drift."""
    rows = []
    for _, c in counties.iterrows():
        # urban counties grow ~1.2% / yr, rural counties drift -0.3% to +0.3%
        if c["rurality"] == "Urban":
            trend = RNG.uniform(0.008, 0.018)
        else:
            trend = RNG.uniform(-0.005, 0.004)
        for y in YEARS:
            yrs_from_2020 = y - 2020
            pop = c["population_2020"] * math.exp(trend * yrs_from_2020)
            pop *= RNG.normal(1.0, 0.005)  # tiny annual noise
            rows.append({
                "fips": c["fips"],
                "year": y,
                "population": int(round(pop)),
                "median_age": round(RNG.normal(42 if c["rurality"] == "Rural" else 38, 1.5), 1),
                "unemployment_pct": round(RNG.normal(4.2 if c["rurality"] == "Urban" else 5.0, 0.6), 2),
                "medicaid_share": round(RNG.normal(
                    0.34 if c["rurality"] == "Tribal" else
                    0.27 if c["rurality"] == "Rural" else 0.20, 0.03), 3),
                "uninsured_share": round(RNG.normal(
                    0.10 if c["rurality"] == "Tribal" else
                    0.08 if c["rurality"] == "Rural" else 0.06, 0.015), 3),
            })
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------
# 3. Hospital roster
# --------------------------------------------------------------------------
HOSPITAL_TYPES = ["CAH", "PPS", "Tertiary", "Tribal/IHS", "Sole Community"]

# Curated representative roster (clearly synthetic, but plausible for MT).
HOSPITALS = [
    # (facility_id, name, county_fips, type, beds, lat_off, lon_off, ownership)
    ("MT001", "Yellowstone Regional Medical Center",   "30111", "Tertiary",       304, 0.02, -0.03, "Nonprofit"),
    ("MT002", "St. Vincent Heart & Health Campus",     "30111", "Tertiary",       286, -0.04, 0.02, "Nonprofit"),
    ("MT003", "Bozeman Health Deaconess",              "30031", "Tertiary",       190, 0.01, -0.02, "Nonprofit"),
    ("MT004", "Benefis Cascade Health",                "30013", "Tertiary",       265, 0.0, 0.02, "Nonprofit"),
    ("MT005", "Community Medical Center Missoula",     "30063", "PPS",            150, 0.03, -0.03, "Nonprofit"),
    ("MT006", "Providence St. Patrick Hospital",       "30063", "PPS",            213, -0.02, 0.04, "Nonprofit"),
    ("MT007", "Logan Health Flathead Valley",          "30029", "PPS",            180, -0.01, 0.03, "Nonprofit"),
    ("MT008", "St. Peter's Health",                    "30049", "PPS",            123, 0.0, 0.0, "Nonprofit"),
    ("MT009", "St. James Healthcare",                  "30093", "PPS",             71, 0.02, -0.01, "Nonprofit"),
    # Critical Access Hospitals (CAH) — most rural counties have one
    ("MT010", "Beaverhead Community Hospital",         "30001", "CAH",             14, 0.01, 0.0, "Nonprofit"),
    ("MT011", "Big Horn County Memorial",              "30003", "CAH",             20, 0.02, -0.02, "District"),
    ("MT012", "Northern Montana Health (Hill)",        "30041", "CAH",             49, 0.0, 0.01, "Nonprofit"),
    ("MT013", "Sweet Grass Family Medicine",           "30097", "CAH",             14, 0.0, 0.02, "Nonprofit"),
    ("MT014", "Stillwater Billings Clinic",            "30095", "CAH",             25, -0.02, 0.0, "Nonprofit"),
    ("MT015", "Carbon County Health & Hospital",       "30009", "CAH",             14, 0.0, 0.01, "Nonprofit"),
    ("MT016", "Park Memorial Livingston",              "30067", "CAH",             25, 0.01, 0.0, "Nonprofit"),
    ("MT017", "Madison Valley Medical Center",         "30057", "CAH",             14, 0.0, -0.01, "District"),
    ("MT018", "Granite County Medical Center",         "30039", "CAH",             14, 0.02, 0.0, "District"),
    ("MT019", "Powell County Memorial",                "30077", "CAH",             20, 0.0, 0.02, "District"),
    ("MT020", "Mineral Community Hospital",            "30061", "CAH",             14, 0.01, 0.0, "Nonprofit"),
    ("MT021", "Sanders Regional Medical",              "30089", "CAH",             14, 0.0, 0.01, "Nonprofit"),
    ("MT022", "Cabinet Peaks Medical Center",          "30053", "CAH",             20, 0.02, 0.0, "Nonprofit"),
    ("MT023", "Ravalli Marcus Daly Memorial",          "30081", "CAH",             25, -0.01, 0.0, "Nonprofit"),
    ("MT024", "Anaconda Community Hospital",           "30023", "CAH",             14, 0.0, 0.0, "Nonprofit"),
    ("MT025", "Pondera Medical Center",                "30073", "CAH",             14, 0.0, 0.02, "District"),
    ("MT026", "Teton Medical Center",                  "30099", "CAH",             14, 0.02, 0.0, "District"),
    ("MT027", "Marias Medical Center",                 "30101", "CAH",             14, 0.0, -0.02, "District"),
    ("MT028", "Liberty County Hospital",               "30051", "CAH",             14, 0.0, 0.01, "District"),
    ("MT029", "Big Sandy Medical Center",              "30015", "CAH",             14, 0.02, 0.0, "District"),
    ("MT030", "Central Montana Medical (Fergus)",      "30027", "PPS",             47, 0.01, 0.0, "Nonprofit"),
    ("MT031", "Wheatland Memorial Healthcare",         "30107", "CAH",             14, 0.0, 0.01, "District"),
    ("MT032", "Meagher County Memorial",               "30059", "CAH",             14, 0.01, 0.0, "District"),
    ("MT033", "Broadwater Health Center",              "30007", "CAH",             14, 0.0, 0.02, "Nonprofit"),
    ("MT034", "Roundup Memorial Healthcare",           "30065", "CAH",             14, 0.0, 0.0, "District"),
    ("MT035", "Mountainview Medical Center",           "30095", "CAH",             14, 0.02, 0.0, "Nonprofit"),
    ("MT036", "Phillips County Hospital",              "30071", "CAH",             14, 0.0, 0.01, "District"),
    ("MT037", "Frances Mahon Deaconess Hospital",      "30105", "CAH",             25, 0.0, -0.02, "Nonprofit"),
    ("MT038", "Daniels Memorial Healthcare",           "30019", "CAH",             14, 0.01, 0.0, "Nonprofit"),
    ("MT039", "Sheridan Memorial Hospital",            "30091", "CAH",             14, 0.0, 0.0, "Nonprofit"),
    ("MT040", "Roosevelt Medical Center",              "30085", "CAH",             14, 0.0, 0.02, "District"),
    ("MT041", "Northeast Montana Health (Richland)",   "30083", "CAH",             20, 0.02, 0.0, "Nonprofit"),
    ("MT042", "Glendive Medical Center",               "30021", "PPS",             47, 0.0, 0.0, "Nonprofit"),
    ("MT043", "Prairie Community Hospital",            "30079", "CAH",             14, 0.0, 0.02, "District"),
    ("MT044", "Fallon Medical Complex",                "30025", "CAH",             14, 0.02, 0.0, "Nonprofit"),
    ("MT045", "Holy Rosary Healthcare (Custer)",       "30017", "CAH",             20, 0.0, 0.0, "Nonprofit"),
    ("MT046", "Powder River Medical",                  "30075", "CAH",             14, 0.01, 0.01, "District"),
    ("MT047", "Carter County Health",                  "30011", "CAH",             14, 0.0, 0.0, "District"),
    ("MT048", "McCone County Health Center",           "30055", "CAH",             14, 0.0, 0.02, "District"),
    ("MT049", "Garfield County Health Center",         "30033", "CAH",             14, 0.02, 0.0, "District"),
    ("MT050", "Petroleum County Memorial",             "30069", "CAH",             14, 0.0, 0.0, "District"),
    ("MT051", "Treasure County Clinic",                "30103", "CAH",             14, 0.0, 0.01, "District"),
    ("MT052", "Wibaux County Health",                  "30109", "CAH",             14, 0.02, 0.0, "District"),
    ("MT053", "Golden Valley Health",                  "30037", "CAH",             14, 0.0, 0.0, "District"),
    ("MT054", "Judith Basin Medical",                  "30045", "CAH",             14, 0.0, 0.02, "District"),
    # Tribal / IHS facilities
    ("MT055", "Blackfeet Community Hospital (IHS)",    "30035", "Tribal/IHS",      24, 0.0, 0.0, "IHS"),
    ("MT056", "Crow / Northern Cheyenne IHS Hospital", "30003", "Tribal/IHS",      24, -0.02, 0.04, "IHS"),
    ("MT057", "Fort Belknap Service Unit",             "30005", "Tribal/IHS",      18, 0.02, 0.0, "IHS"),
    ("MT058", "Fort Peck Service Unit (Poplar)",       "30085", "Tribal/IHS",      18, 0.04, -0.04, "IHS"),
    ("MT059", "Rocky Boy Health Center",               "30041", "Tribal/IHS",      18, -0.04, 0.04, "Tribal"),
    ("MT060", "Northern Cheyenne Service Unit",        "30087", "Tribal/IHS",      18, 0.0, 0.0, "IHS"),
    ("MT061", "St. Luke Community (Lake)",             "30047", "Sole Community",  25, 0.04, 0.0, "Nonprofit"),
    ("MT062", "Clark Fork Valley Hospital",            "30089", "Sole Community",  20, 0.0, -0.04, "Nonprofit"),
    ("MT063", "Bitterroot Health",                     "30081", "Sole Community",  35, 0.04, 0.0, "Nonprofit"),
    ("MT064", "Barrett Hospital & Healthcare",         "30001", "Sole Community",  25, -0.02, 0.02, "Nonprofit"),
    ("MT065", "Missouri River Medical Center",         "30015", "Sole Community",  20, -0.02, 0.0, "Nonprofit"),
]


def make_hospitals(counties: pd.DataFrame) -> pd.DataFrame:
    cdf = counties.set_index("fips")
    rows = []
    for fid, name, fips, htype, beds, lat_off, lon_off, owner in HOSPITALS:
        c = cdf.loc[fips]
        rows.append({
            "facility_id": fid,
            "facility_name": name,
            "county_fips": fips,
            "county_name": c["county_name"],
            "facility_type": htype,
            "ownership": owner,
            "staffed_beds": beds,
            "lat": round(c["lat"] + lat_off, 4),
            "lon": round(c["lon"] + lon_off, 4),
            "rurality": c["rurality"],
            "region": c["region"],
            # CMS Provider Identifier (synthetic but in correct format)
            "ccn": f"27{1000 + int(fid[2:]):04d}",
        })
    return pd.DataFrame(rows)


# --------------------------------------------------------------------------
# 4. CoE / treatment assignment
# --------------------------------------------------------------------------
def assign_coe(hospitals: pd.DataFrame) -> pd.DataFrame:
    """Staggered CoE adoption.

    Cohort 1 (early adopters):  treated 2025-2026
    Cohort 2 (mid):             treated 2027-2028
    Cohort 3 (late):            treated 2029-2030
    Never-treated control:      ~25% of facilities never enroll
    """
    rows = []
    for _, h in hospitals.iterrows():
        # PPS / tertiary much more likely to be early adopters
        if h["facility_type"] in ("Tertiary", "PPS"):
            cohort = RNG.choice(["early", "mid", "late", "never"], p=[0.45, 0.30, 0.15, 0.10])
        elif h["facility_type"] == "Sole Community":
            cohort = RNG.choice(["early", "mid", "late", "never"], p=[0.30, 0.40, 0.20, 0.10])
        elif h["facility_type"] == "CAH":
            cohort = RNG.choice(["early", "mid", "late", "never"], p=[0.10, 0.30, 0.30, 0.30])
        else:  # tribal/IHS — different funding pathway, slower
            cohort = RNG.choice(["early", "mid", "late", "never"], p=[0.05, 0.20, 0.35, 0.40])

        if cohort == "early":
            enroll_q = (2025, RNG.integers(1, 5))
            recs_q = (2025, 4) if enroll_q[1] <= 2 else (2026, 1)
            milestones_total = 8
        elif cohort == "mid":
            enroll_q = (2026, RNG.integers(1, 5))
            recs_q = (2027, RNG.integers(1, 3))
            milestones_total = 8
        elif cohort == "late":
            enroll_q = (2028, RNG.integers(1, 5))
            recs_q = (2029, RNG.integers(1, 3))
            milestones_total = 8
        else:
            enroll_q = (None, None)
            recs_q = (None, None)
            milestones_total = 0

        rows.append({
            "facility_id": h["facility_id"],
            "coe_cohort": cohort,
            "enroll_year": enroll_q[0],
            "enroll_quarter": enroll_q[1],
            "recommendations_year": recs_q[0],
            "recommendations_quarter": recs_q[1],
            "milestones_total": milestones_total,
        })
    return pd.DataFrame(rows)


def quarter_index(y: int, q: int) -> int:
    return y * 4 + (q - 1)


# --------------------------------------------------------------------------
# 5. Facility-quarter operations panel
# --------------------------------------------------------------------------
def make_facility_quarter(
    hospitals: pd.DataFrame, coe: pd.DataFrame, pop: pd.DataFrame
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    coe_idx = coe.set_index("facility_id")
    pop_idx = pop.set_index(["fips", "year"])

    rows = []
    coe_status_rows = []
    tele_rows = []
    shared_rows = []

    # Treatment effect parameters (true, but unknown to model fitters)
    true_coe_margin_lift = 2.8     # pp on operating margin, gradual
    true_coe_ed_los = -16.0        # minutes ED LOS reduction
    true_coe_telehealth = 6.5      # pp telehealth share lift
    true_telespec_ed_los = -8.0
    true_sharedsvc_margin = 1.1

    for _, h in hospitals.iterrows():
        c = coe_idx.loc[h["facility_id"]]
        # Facility-specific random effects
        margin_re = RNG.normal(0, 1.5)
        ed_re = RNG.normal(0, 30)
        tele_re = RNG.normal(0, 2.0)

        # Baseline values by facility type
        if h["facility_type"] in ("Tertiary", "PPS"):
            margin_base = -3.0
            ed_los_base = 215
            tele_base = 6.0
            inpatient_factor = 0.62
        elif h["facility_type"] == "Sole Community":
            margin_base = -8.0
            ed_los_base = 245
            tele_base = 4.0
            inpatient_factor = 0.45
        elif h["facility_type"] == "CAH":
            margin_base = -16.5
            ed_los_base = 295
            tele_base = 3.5
            inpatient_factor = 0.30
        else:  # Tribal/IHS
            margin_base = -18.0
            ed_los_base = 320
            tele_base = 3.0
            inpatient_factor = 0.28

        # Telespec / shared services rollout (independent of CoE timing,
        # correlated with adoption cohort)
        tele_start_q = (RNG.integers(2024, 2030), RNG.integers(1, 5)) if RNG.random() < (
            0.85 if c["coe_cohort"] != "never" else 0.30) else (None, None)
        shared_start_q = (RNG.integers(2024, 2030), RNG.integers(1, 5)) if RNG.random() < (
            0.75 if c["coe_cohort"] != "never" else 0.20) else (None, None)
        ehr_modern_q = (RNG.integers(2025, 2031), RNG.integers(1, 5)) if RNG.random() < 0.80 else (None, None)

        if tele_start_q[0] is not None:
            tele_rows.append({
                "facility_id": h["facility_id"],
                "telestroke_year": tele_start_q[0],
                "telestroke_quarter": tele_start_q[1],
                "teleed_year": tele_start_q[0],
                "teleed_quarter": tele_start_q[1],
                "telecard_year": tele_start_q[0] + (1 if tele_start_q[1] >= 3 else 0),
                "telecard_quarter": tele_start_q[1],
            })
        if shared_start_q[0] is not None:
            shared_rows.append({
                "facility_id": h["facility_id"],
                "lab_share_year": shared_start_q[0],
                "lab_share_quarter": shared_start_q[1],
                "hr_share_year": shared_start_q[0] + 1,
                "hr_share_quarter": 1,
                "supply_share_year": shared_start_q[0],
                "supply_share_quarter": shared_start_q[1],
                "finance_share_year": shared_start_q[0] + (1 if shared_start_q[1] >= 3 else 0),
                "finance_share_quarter": shared_start_q[1],
            })

        # Build event-time treatment ramp for CoE
        if c["coe_cohort"] != "never" and pd.notna(c["recommendations_year"]):
            coe_event_q = quarter_index(int(c["recommendations_year"]),
                                        int(c["recommendations_quarter"]))
        else:
            coe_event_q = None

        for (y, q) in QUARTERS:
            qi = quarter_index(y, q)

            # CoE treatment indicator + intensity (milestones completed)
            if coe_event_q is not None and qi >= coe_event_q:
                qtrs_post = qi - coe_event_q
                # Sigmoidal ramp to full effect over ~6 quarters
                coe_intensity = 1 / (1 + math.exp(-(qtrs_post - 3) / 1.5))
                milestones_done = min(8, int(round(8 * coe_intensity + RNG.normal(0, 0.4))))
            else:
                coe_intensity = 0.0
                milestones_done = 0

            telespec_active = int(
                tele_start_q[0] is not None and qi >= quarter_index(*tele_start_q))
            sharedsvc_active = int(
                shared_start_q[0] is not None and qi >= quarter_index(*shared_start_q))
            ehr_modern = int(
                ehr_modern_q[0] is not None and qi >= quarter_index(*ehr_modern_q))

            # Outcomes
            yr_offset = y - 2019  # baseline is 2019
            margin_trend = -0.15 * yr_offset  # general downward pressure on rural margin
            margin = (margin_base + margin_re + margin_trend
                      + true_coe_margin_lift * coe_intensity
                      + true_sharedsvc_margin * sharedsvc_active
                      + RNG.normal(0, 1.6))

            ed_los_trend = -0.5 * yr_offset
            ed_los = (ed_los_base + ed_re + ed_los_trend
                      + true_coe_ed_los * coe_intensity
                      + true_telespec_ed_los * telespec_active
                      + RNG.normal(0, 18))

            telehealth = (tele_base + tele_re + 0.6 * yr_offset
                          + true_coe_telehealth * coe_intensity
                          + 4.5 * telespec_active
                          + RNG.normal(0, 1.4))
            telehealth = max(0.0, min(40.0, telehealth))

            # Volumes
            try:
                p = pop_idx.loc[(h["county_fips"], y), "population"]
            except KeyError:
                p = 5000

            ed_visits = max(80, int(0.25 * p * (1 if h["facility_type"] != "CAH" else 0.4)
                                    * inpatient_factor / 4 + RNG.normal(0, 60)))
            inpatient_days = max(20, int(h["staffed_beds"] * 90 * 0.55 + RNG.normal(0, 60)))
            total_visits = max(200, int(ed_visits * 4 + h["staffed_beds"] * 25 + RNG.normal(0, 200)))
            fte_count = max(15, int(h["staffed_beds"] * 4.2 + RNG.normal(0, 8)))

            payer_medicaid = max(0.05, min(0.65, RNG.normal(
                0.40 if h["rurality"] == "Tribal" else
                0.32 if h["rurality"] == "Rural" else 0.22, 0.04)))
            payer_medicare = max(0.10, min(0.65, RNG.normal(
                0.42 if h["rurality"] == "Rural" else 0.30, 0.04)))
            payer_uninsured = max(0.0, min(0.20, RNG.normal(
                0.09 if h["rurality"] == "Tribal" else 0.06, 0.02)))
            case_mix = round(RNG.normal(
                1.25 if h["facility_type"] == "Tertiary" else
                1.05 if h["facility_type"] == "PPS" else
                0.85, 0.07), 3)

            rows.append({
                "facility_id": h["facility_id"],
                "year": y,
                "quarter": q,
                "operating_margin_pct": round(margin, 2),
                "ed_los_min": round(ed_los, 1),
                "telehealth_share_pct": round(telehealth, 2),
                "inpatient_days": inpatient_days,
                "staffed_beds": h["staffed_beds"],
                "ed_visits": ed_visits,
                "total_visits": total_visits,
                "fte_count": fte_count,
                "payer_mix_medicaid": round(payer_medicaid, 3),
                "payer_mix_medicare": round(payer_medicare, 3),
                "payer_mix_uninsured": round(payer_uninsured, 3),
                "case_mix_index": case_mix,
            })

            coe_status_rows.append({
                "facility_id": h["facility_id"],
                "year": y,
                "quarter": q,
                "coe_enrolled": int(c["coe_cohort"] != "never"
                                    and pd.notna(c["enroll_year"])
                                    and qi >= quarter_index(int(c["enroll_year"]),
                                                            int(c["enroll_quarter"]))),
                "coe_recs_received": int(coe_event_q is not None and qi >= coe_event_q),
                "coe_milestones_completed": milestones_done,
                "coe_milestones_total": int(c["milestones_total"]),
                "telespec_active": telespec_active,
                "sharedsvc_active": sharedsvc_active,
                "ehr_modern_active": ehr_modern,
            })

    return (pd.DataFrame(rows), pd.DataFrame(coe_status_rows),
            pd.DataFrame(tele_rows), pd.DataFrame(shared_rows))


# --------------------------------------------------------------------------
# 6. Workforce panel
# --------------------------------------------------------------------------
def make_workforce(counties: pd.DataFrame, pop: pd.DataFrame) -> tuple[pd.DataFrame, ...]:
    rows = []
    awards_rows = []
    resid_rows = []
    train_rows = []
    support_rows = []

    pop_idx = pop.set_index(["fips", "year"])

    # Treatment exposure: define which counties are "high exposure" to
    # workforce interventions (rural counties + tribal counties more likely)
    treated_counties = []
    for _, c in counties.iterrows():
        if c["rurality"] == "Tribal":
            p_treat = 0.85
        elif c["rurality"] == "Rural":
            p_treat = 0.55
        else:
            p_treat = 0.10
        treated_counties.append(int(RNG.random() < p_treat))
    counties = counties.assign(workforce_treated=treated_counties)

    # True effects (with lag)
    np_per_award = 0.55          # NPs per 100k per cumulative service-commitment award (lag 1y)
    md_per_resid_slot = 0.32     # MDs per 100k per cumulative resid slot (lag 2y)
    rn_per_train = 0.06          # RNs per 100k per training participant (lag 1y)
    turnover_per_support = -0.018  # decline in turnover rate per support intensity (current year)

    for _, c in counties.iterrows():
        # Stable county-specific random effects
        np_re = RNG.normal(0, 6)
        md_re = RNG.normal(0, 8)
        rn_re = RNG.normal(0, 30)

        # Baselines tuned to MT realities (rural counties lower)
        if c["rurality"] == "Urban":
            np_base, md_base, rn_base = 95.0, 285.0, 980.0
            dent_base, emt_base, pa_base = 64.0, 95.0, 35.0
            turnover_base = 0.16
            mh_base = 6.6
        elif c["rurality"] == "Tribal":
            np_base, md_base, rn_base = 48.0, 95.0, 510.0
            dent_base, emt_base, pa_base = 28.0, 65.0, 14.0
            turnover_base = 0.27
            mh_base = 5.4
        else:  # Rural
            np_base, md_base, rn_base = 62.0, 145.0, 695.0
            dent_base, emt_base, pa_base = 41.0, 78.0, 22.0
            turnover_base = 0.21
            mh_base = 5.9

        # Generate per-year intervention counts (only for treated counties, after 2024)
        # Intervention timing: ramp from 2025 onward
        cumulative_awards = 0
        cumulative_resid = 0
        for y in YEARS:
            if c["workforce_treated"] and y >= 2025:
                yrs_in = y - 2024
                lam = 1.5 + 0.4 * yrs_in
                awards_yr = int(RNG.poisson(lam))
                resid_yr = int(RNG.poisson(0.6 + 0.18 * yrs_in))
                train_yr = int(RNG.poisson(8.0 + 1.4 * yrs_in))
                support_int = round(min(0.95, 0.10 + 0.12 * yrs_in + RNG.normal(0, 0.05)), 3)
            elif y >= 2025:
                awards_yr = int(RNG.poisson(0.4))
                resid_yr = int(RNG.poisson(0.1))
                train_yr = int(RNG.poisson(2.0))
                support_int = round(max(0.0, 0.05 + RNG.normal(0, 0.03)), 3)
            else:
                awards_yr = 0
                resid_yr = 0
                train_yr = int(RNG.poisson(1.0))
                support_int = round(max(0.0, RNG.normal(0.04, 0.02)), 3)

            cumulative_awards += awards_yr
            cumulative_resid += resid_yr

            # Compute outcomes with appropriate lags
            # Look up lagged interventions
            awards_lag1 = sum(
                a["awards_count"] for a in awards_rows
                if a["fips"] == c["fips"] and a["year"] == y - 1
            )
            resid_lag2 = sum(
                a["residency_slots"] for a in resid_rows
                if a["fips"] == c["fips"] and a["year"] == y - 2
            )
            train_lag1 = sum(
                a["training_participants"] for a in train_rows
                if a["fips"] == c["fips"] and a["year"] == y - 1
            )

            yr_offset = y - 2019
            np_density = (np_base + np_re + 0.4 * yr_offset
                          + np_per_award * awards_lag1
                          + RNG.normal(0, 4))
            md_density = (md_base + md_re + 0.2 * yr_offset
                          + md_per_resid_slot * resid_lag2 * 5  # convert to per-100k scale
                          + RNG.normal(0, 6))
            rn_density = (rn_base + rn_re + 1.5 * yr_offset
                          + rn_per_train * train_lag1
                          + RNG.normal(0, 22))

            dent_density = max(8.0, dent_base + RNG.normal(0, 5) + 0.05 * yr_offset)
            emt_density = max(20.0, emt_base + RNG.normal(0, 6) + 0.1 * yr_offset)
            pa_density = max(4.0, pa_base + RNG.normal(0, 4) + 0.15 * yr_offset
                             + 0.2 * awards_lag1)

            turnover = max(0.04, turnover_base + RNG.normal(0, 0.018)
                           - 0.003 * yr_offset
                           + turnover_per_support * support_int * 8)

            mh_score = round(mh_base + RNG.normal(0, 0.4) + 0.05 * yr_offset
                             + 0.3 * support_int, 2)

            try:
                p = pop_idx.loc[(c["fips"], y), "population"]
                u = pop_idx.loc[(c["fips"], y), "unemployment_pct"]
                m = pop_idx.loc[(c["fips"], y), "medicaid_share"]
            except KeyError:
                p, u, m = 5000, 4.5, 0.25

            rows.append({
                "fips": c["fips"],
                "year": y,
                "np_per_100k": round(np_density, 2),
                "md_per_100k": round(md_density, 2),
                "rn_per_100k": round(rn_density, 2),
                "dental_hyg_per_100k": round(dent_density, 2),
                "emt_per_100k": round(emt_density, 2),
                "pa_per_100k": round(pa_density, 2),
                "turnover_rate": round(turnover, 4),
                "provider_mh_score": mh_score,
                "population": int(p),
                "unemployment_pct": round(u, 2),
                "medicaid_share": round(m, 3),
                "workforce_treated": int(c["workforce_treated"]),
            })

            awards_rows.append({
                "fips": c["fips"],
                "year": y,
                "awards_count": awards_yr,
                "primary_specialty": RNG.choice(
                    ["Family Medicine", "Internal Medicine", "Pediatrics",
                     "Behavioral Health", "OB/GYN", "Dental", "Pharmacy"]),
            })
            resid_rows.append({
                "fips": c["fips"],
                "year": y,
                "residency_slots": resid_yr,
                "track_type": RNG.choice(
                    ["Rural Family Medicine", "Internal Medicine",
                     "Psychiatry", "General Surgery", "Pediatrics"]),
            })
            train_rows.append({
                "fips": c["fips"],
                "year": y,
                "training_participants": train_yr,
                "training_type": RNG.choice(
                    ["Resilience", "Telehealth Skills", "Substance Use",
                     "Cultural Competence", "Quality Improvement"]),
            })
            support_rows.append({
                "fips": c["fips"],
                "year": y,
                "support_intensity": support_int,
                "support_recipients": int(round(support_int * max(20, p / 1500))),
            })

    return (pd.DataFrame(rows),
            pd.DataFrame(awards_rows),
            pd.DataFrame(resid_rows),
            pd.DataFrame(train_rows),
            pd.DataFrame(support_rows))


# --------------------------------------------------------------------------
# 7. Claims (member-month + county-quarter aggregates)
# --------------------------------------------------------------------------
def make_claims(counties: pd.DataFrame, hospitals: pd.DataFrame, coe: pd.DataFrame
                ) -> tuple[pd.DataFrame, ...]:
    """Generate synthetic Medicaid duals member-month claims plus
    county-level aggregate claims metrics."""
    coe_idx = coe.set_index("facility_id")
    rural_counties = counties[counties["rurality"].isin(["Rural", "Tribal"])]

    # Member-month panel: smaller sample (5000 members observed across 2022-2031)
    n_members = 5000
    # Sample members weighted by county population
    weights = rural_counties["population_2020"] / rural_counties["population_2020"].sum()
    member_county = RNG.choice(rural_counties["fips"].values, size=n_members, p=weights)
    member_age = RNG.choice([55, 65, 72, 78, 85], size=n_members,
                            p=[0.15, 0.25, 0.25, 0.20, 0.15])
    member_sex = RNG.choice(["F", "M"], size=n_members, p=[0.55, 0.45])
    member_dual = RNG.choice(["Full-Benefit", "Partial-Benefit", "QMB-Only"],
                             size=n_members, p=[0.65, 0.20, 0.15])
    member_hcc = RNG.normal(1.6, 0.6, size=n_members).clip(0.4, 4.0).round(2)

    member_rows = []
    months = [(y, m) for y in range(2022, 2032) for m in range(1, 13)]

    # County-level VBC active flag rolls out 2026+
    counties_vbc_year = {f: 2026 + RNG.integers(0, 4) for f in counties["fips"]}

    true_vbc_pmpm = -42.0  # dollars, true effect
    for i in range(n_members):
        f = member_county[i]
        county_vbc_y = counties_vbc_year[f]
        for (y, m) in months:
            # Skip first months for late-enrolled members (synthetic enrollment)
            if y < 2022:
                continue
            yr_offset = y - 2022
            base = 1450.0  # full-benefit dual avg PMPM ($)
            if member_dual[i] == "Partial-Benefit":
                base = 980.0
            elif member_dual[i] == "QMB-Only":
                base = 760.0

            inflation = 1.0 + 0.025 * yr_offset
            risk_mult = (member_hcc[i] / 1.5) ** 0.85
            vbc_active = int(y >= county_vbc_y)
            pmpm = base * inflation * risk_mult + RNG.normal(0, 110)
            pmpm += vbc_active * (true_vbc_pmpm + RNG.normal(0, 30))
            pmpm = max(50, pmpm)

            inpat = max(0.0, pmpm * RNG.uniform(0.30, 0.45) + RNG.normal(0, 30))
            pharma = max(0.0, pmpm * RNG.uniform(0.13, 0.22) + RNG.normal(0, 15))
            bh = max(0.0, pmpm * RNG.uniform(0.06, 0.14) + RNG.normal(0, 12))

            member_rows.append({
                "member_id": f"M{1_000_000 + i:07d}",
                "year_month": f"{y}-{m:02d}",
                "year": y,
                "month": m,
                "county_fips": f,
                "age_band": "55-64" if member_age[i] < 65 else "65-74" if member_age[i] < 75 else "75+",
                "sex": member_sex[i],
                "dual_subtype": member_dual[i],
                "hcc_score": member_hcc[i],
                "vbc_attributed": vbc_active,
                "pmpm_total": round(pmpm, 2),
                "pmpm_inpatient": round(inpat, 2),
                "pmpm_pharmacy": round(pharma, 2),
                "pmpm_behavioral": round(bh, 2),
            })

    member_df = pd.DataFrame(member_rows)

    # County-quarter aggregates (TNT, ED high utilizers, pharmacy, outpatient)
    cq_rows = []
    ed_rows = []
    tnt_rows = []
    pharm_rows = []
    outpat_rows = []
    ems_rows = []
    vbc_rows = []

    for _, c in counties.iterrows():
        # County-specific EMS / Pharmacy / VBC rollout
        ems_modern_year = 2026 + RNG.integers(0, 4) if RNG.random() < 0.7 else None
        pharm_part_year = 2026 + RNG.integers(0, 4) if RNG.random() < 0.6 else None
        outpat_shift_start = 2026

        for y in range(2022, 2032):
            for q in (1, 2, 3, 4):
                yr_offset = y - 2022
                ems_active = int(ems_modern_year is not None and y >= ems_modern_year)
                pharm_active = int(pharm_part_year is not None and y >= pharm_part_year)
                vbc_active = int(y >= counties_vbc_year[c["fips"]])

                # ED high utilizer rate (per 1000 Medicaid members)
                base_ed_high = 38.0 if c["rurality"] == "Tribal" else 30.0 if c["rurality"] == "Rural" else 24.0
                ed_high = (base_ed_high - 0.5 * yr_offset
                           - 5.5 * ems_active - 3.0 * pharm_active - 2.5 * vbc_active
                           + RNG.normal(0, 2.2))

                # TNT (Treat-No-Transport) CPT use
                ems_runs = int(max(20, c["population_2020"] / 250 + RNG.normal(0, 25)))
                tnt_count = int(max(0, (0.04 + 0.10 * ems_active) * ems_runs + RNG.normal(0, 4)))

                # Pharmacy diagnostic / prescribing for Medicaid
                pharm_rx = int(max(0, c["population_2020"] / 1100 * (1.0 + 0.6 * pharm_active)
                                   + RNG.normal(0, 12)))
                total_rx = int(max(50, pharm_rx * 28 + RNG.normal(0, 80)))

                # Outpatient share
                outp_base = 0.62 if c["rurality"] == "Urban" else 0.55
                outp = max(0.30, min(0.85, outp_base + 0.012 * (y - outpat_shift_start)
                                     + 0.025 * vbc_active + RNG.normal(0, 0.018)))

                ed_rows.append({
                    "fips": c["fips"], "year": y, "quarter": q,
                    "ed_high_util_per_1k": round(max(2.0, ed_high), 2),
                })
                tnt_rows.append({
                    "fips": c["fips"], "year": y, "quarter": q,
                    "tnt_count": tnt_count, "total_ems_runs": ems_runs,
                    "tnt_rate": round(tnt_count / max(1, ems_runs), 4),
                })
                pharm_rows.append({
                    "fips": c["fips"], "year": y, "quarter": q,
                    "pharm_rx_count": pharm_rx, "total_rx_count": total_rx,
                    "pharm_share": round(pharm_rx / max(1, total_rx), 4),
                })
                outpat_rows.append({
                    "fips": c["fips"], "year": y, "quarter": q,
                    "outpatient_share": round(outp, 4),
                })

                cq_rows.append({
                    "fips": c["fips"], "year": y, "quarter": q,
                    "ems_modern_active": ems_active,
                    "pharmacist_program_active": pharm_active,
                    "vbc_active": vbc_active,
                    "outpatient_shift": round(outp, 4),
                })

        ems_rows.append({
            "fips": c["fips"],
            "modernization_year": ems_modern_year,
            "agency_count": int(max(1, c["population_2020"] / 4500)),
        })
        vbc_rows.append({
            "fips": c["fips"],
            "vbc_start_year": counties_vbc_year[c["fips"]],
            "attributed_lives_2031": int(c["population_2020"] * 0.06 + RNG.normal(0, 50)),
            "active_vbc_pcps_2031": int(max(0, c["population_2020"] / 6000 + RNG.normal(0, 2))),
        })

    return (member_df, pd.DataFrame(cq_rows), pd.DataFrame(ed_rows),
            pd.DataFrame(tnt_rows), pd.DataFrame(pharm_rows),
            pd.DataFrame(outpat_rows), pd.DataFrame(ems_rows),
            pd.DataFrame(vbc_rows))


# --------------------------------------------------------------------------
# 8. Prevention panel (county-year)
# --------------------------------------------------------------------------
def make_prevention(counties: pd.DataFrame, pop: pd.DataFrame) -> tuple[pd.DataFrame, ...]:
    pop_idx = pop.set_index(["fips", "year"])
    prev_rows = []
    school_rows = []
    mobile_rows = []
    chap_rows = []
    crisis_rows = []

    # Treatment effect parameters
    a1c_per_school = 0.45      # pp improvement in A1c control per school site
    bp_per_chap = 0.35         # pp BP control per CHAP provider per 100k
    bh_ed_per_crisis = -0.55   # reduction in BH ED admissions per crisis space
    sui_per_crisis = -0.06     # per 100k per crisis space (small, slow)

    for _, c in counties.iterrows():
        # Baseline values
        if c["rurality"] == "Tribal":
            a1c_base, bp_base, bmi_base = 38.0, 41.0, 24.0
            w30_base = 71.0
            bh_ed_base = 28.0
            sui_base = 30.0
            student_mh_base = 32.0
        elif c["rurality"] == "Urban":
            a1c_base, bp_base, bmi_base = 56.0, 60.0, 38.0
            w30_base = 84.0
            bh_ed_base = 14.0
            sui_base = 21.0
            student_mh_base = 22.0
        else:
            a1c_base, bp_base, bmi_base = 47.0, 51.0, 30.0
            w30_base = 76.0
            bh_ed_base = 19.0
            sui_base = 26.0
            student_mh_base = 27.0

        # County-specific RE
        a1c_re = RNG.normal(0, 1.8)
        bp_re = RNG.normal(0, 2.0)
        bh_re = RNG.normal(0, 2.5)
        sui_re = RNG.normal(0, 4.0)

        # Rollout: school sites scale up over 2025-2031, mobile rollout 2025+
        treated = int(c["rurality"] != "Urban" and RNG.random() < 0.65)
        if treated:
            chap_active_year = 2026 + int(RNG.integers(0, 4))
        else:
            chap_active_year = None

        for y in YEARS:
            yr_offset = y - 2019
            if treated and y >= 2025:
                rollout = (y - 2024)
                school = int(max(0, RNG.poisson(0.7 + 0.55 * rollout)))
                mobile_runs = int(max(0, RNG.poisson(8 + 6 * rollout)))
                mobile_visitors = int(mobile_runs * RNG.uniform(2.4, 3.2))
                crisis_spaces = int(max(0, RNG.poisson(0.4 + 0.4 * rollout)))
                mobile_crisis_teams = int(max(0, RNG.poisson(0.2 + 0.3 * rollout)))
                chap_per_100k = round(0.0 if chap_active_year is None or y < chap_active_year
                                      else 1.5 + 0.6 * (y - chap_active_year) + RNG.normal(0, 0.3), 2)
            else:
                school = int(RNG.poisson(0.05))
                mobile_runs = int(RNG.poisson(2 + 0.5 * yr_offset))
                mobile_visitors = int(mobile_runs * 2.5)
                crisis_spaces = int(RNG.poisson(0.05))
                mobile_crisis_teams = 0
                chap_per_100k = 0.0

            # Outcomes
            a1c = (a1c_base + a1c_re + 0.3 * yr_offset
                   + a1c_per_school * school + RNG.normal(0, 1.4))
            bp = (bp_base + bp_re + 0.4 * yr_offset
                  + bp_per_chap * chap_per_100k + RNG.normal(0, 1.8))
            bmi = max(15.0, min(60.0, bmi_base + RNG.normal(0, 1.5) + 0.2 * yr_offset))
            w30 = max(40.0, min(98.0, w30_base + RNG.normal(0, 1.6) + 0.3 * yr_offset
                                + 0.4 * mobile_runs / 10))
            bh_ed = max(2.0, bh_ed_base + bh_re - 0.1 * yr_offset
                        + bh_ed_per_crisis * crisis_spaces + RNG.normal(0, 1.6))
            sui = max(1.0, sui_base + sui_re - 0.05 * yr_offset
                      + sui_per_crisis * crisis_spaces + RNG.normal(0, 3.5))
            student_mh = max(8.0, student_mh_base + RNG.normal(0, 2.0) - 0.4 * yr_offset
                             - 0.5 * school)

            try:
                p = pop_idx.loc[(c["fips"], y), "population"]
            except KeyError:
                p = 5000

            prev_rows.append({
                "fips": c["fips"], "year": y,
                "a1c_control_pct": round(a1c, 2),
                "bp_control_pct": round(bp, 2),
                "bmi_control_pct": round(bmi, 2),
                "w30_ch_pct": round(w30, 2),
                "bh_ed_per_100k": round(bh_ed, 2),
                "suicide_per_100k_3yr": round(sui, 2),
                "student_mh_risk_pct": round(student_mh, 2),
                "chap_providers_per_100k": chap_per_100k,
                "crisis_safe_spaces": crisis_spaces,
                "school_sites_count": school,
                "mobile_runs_count": mobile_runs,
                "population": int(p),
                "prevention_treated": treated,
            })

            school_rows.append({"fips": c["fips"], "year": y, "school_health_sites": school})
            mobile_rows.append({"fips": c["fips"], "year": y,
                                "mobile_runs": mobile_runs,
                                "mobile_unique_visitors": mobile_visitors})
            crisis_rows.append({"fips": c["fips"], "year": y,
                                "crisis_safe_spaces": crisis_spaces,
                                "mobile_crisis_teams": mobile_crisis_teams})

        chap_rows.append({
            "fips": c["fips"],
            "chap_active_year": chap_active_year,
            "chap_providers_at_2031": (round(0 if chap_active_year is None else
                                       1.5 + 0.6 * (2031 - chap_active_year), 2)),
        })

    return (pd.DataFrame(prev_rows),
            pd.DataFrame(school_rows),
            pd.DataFrame(mobile_rows),
            pd.DataFrame(chap_rows),
            pd.DataFrame(crisis_rows))


# --------------------------------------------------------------------------
# 9. Technology / data panel
# --------------------------------------------------------------------------
def make_technology(hospitals: pd.DataFrame, coe: pd.DataFrame
                    ) -> tuple[pd.DataFrame, ...]:
    coe_idx = coe.set_index("facility_id")
    hie_rows = []
    ehr_rows = []
    bed_rows = []
    bh_wait_rows = []
    dash_rows = []

    for _, h in hospitals.iterrows():
        # HIE (Big Sky Care Connect): adoption staggered 2024-2030
        hie_year = 2024 + int(RNG.integers(0, 7))
        hie_quarter = int(RNG.integers(1, 5))
        hie_active_q = quarter_index(hie_year, hie_quarter)

        # EHR modernization: 2025-2031
        ehr_year = 2025 + int(RNG.integers(0, 7)) if RNG.random() < 0.85 else None
        ehr_quarter = int(RNG.integers(1, 5)) if ehr_year is not None else None
        ehr_active_q = quarter_index(ehr_year, ehr_quarter) if ehr_year is not None else None

        # Bed registry: 2025-2030
        registry_year = 2025 + int(RNG.integers(0, 6)) if RNG.random() < 0.75 else None
        registry_quarter = int(RNG.integers(1, 5)) if registry_year is not None else None
        registry_active_q = quarter_index(registry_year, registry_quarter) if registry_year is not None else None

        ehr_rows.append({
            "facility_id": h["facility_id"],
            "ehr_modernization_year": ehr_year,
            "ehr_modernization_quarter": ehr_quarter,
            "ehr_system": RNG.choice(["Epic", "Oracle (Cerner)", "MEDITECH",
                                      "athenahealth", "Athena (Cloud)", "Athena Practice"]),
            "hitech_certified": int(ehr_year is not None),
        })

        for (y, q) in QUARTERS:
            qi = quarter_index(y, q)
            hie_part = int(qi >= hie_active_q)
            ehr_modern = int(ehr_active_q is not None and qi >= ehr_active_q)
            registry = int(registry_active_q is not None and qi >= registry_active_q)

            # Records exchanged: ramps after participation
            if hie_part:
                qpost = qi - hie_active_q
                records = int(max(50, 200 + 60 * qpost + 30 * h["staffed_beds"] / 14
                                  + RNG.normal(0, 100)))
            else:
                records = 0

            # BH bed wait (lower with registry; treat psychiatric/transfer cases)
            base_wait = 18.0 if h["facility_type"] == "Tertiary" else 30.0
            wait_avg = max(2.0, base_wait - 4.5 * registry - 1.0 * (y - 2024)
                           + RNG.normal(0, 3))
            wait_p90 = wait_avg * RNG.uniform(2.0, 2.6)
            transfers_attempted = int(max(2, RNG.poisson(8 + h["staffed_beds"] / 14)))
            transfers_completed = int(transfers_attempted * (0.55 + 0.20 * registry +
                                                             RNG.normal(0, 0.05)))
            transfers_completed = max(0, min(transfers_attempted, transfers_completed))

            # Bed postings (only when registry connected)
            postings = int(max(0, h["staffed_beds"] * 1.5 * registry + RNG.normal(0, 5)))

            hie_rows.append({
                "facility_id": h["facility_id"], "year": y, "quarter": q,
                "hie_participating": hie_part,
                "hie_records_exchanged": records,
            })
            bed_rows.append({
                "facility_id": h["facility_id"], "year": y, "quarter": q,
                "registry_connected": registry,
                "bed_postings": postings,
            })
            bh_wait_rows.append({
                "facility_id": h["facility_id"], "year": y, "quarter": q,
                "bh_wait_hours_avg": round(wait_avg, 2),
                "bh_wait_hours_p90": round(wait_p90, 2),
                "transfers_attempted": transfers_attempted,
                "transfers_completed": transfers_completed,
            })

    # Statewide dashboard usage
    for (y, q) in QUARTERS:
        if y < 2026:
            users = int(max(0, RNG.poisson(2)))
            facilities = int(max(0, RNG.poisson(2)))
        else:
            qpost = (y - 2026) * 4 + (q - 1)
            users = int(max(5, 60 + 12 * qpost + RNG.normal(0, 8)))
            facilities = int(min(len(hospitals), 8 + 4 * qpost + RNG.normal(0, 2)))
        dash_rows.append({
            "year": y, "quarter": q,
            "active_users": users,
            "facilities_using": facilities,
            "total_facilities": len(hospitals),
        })

    return (pd.DataFrame(hie_rows), pd.DataFrame(ehr_rows),
            pd.DataFrame(bed_rows), pd.DataFrame(bh_wait_rows),
            pd.DataFrame(dash_rows))


# --------------------------------------------------------------------------
# 10. Run everything
# --------------------------------------------------------------------------
def main() -> None:
    print("RHTP synthetic data generator")
    print(f"   target dir: {DATA}")

    print("\n1. Geography ...")
    counties = make_counties()
    write(counties, DATA / "geography" / "montana_counties.csv")
    pop = make_population_panel(counties)
    write(pop, DATA / "geography" / "county_population_panel.csv")
    hospitals = make_hospitals(counties)
    write(hospitals, DATA / "geography" / "montana_hospitals.csv")

    print("\n2. CoE / facility intervention timing ...")
    coe = assign_coe(hospitals)
    write(coe, DATA / "facilities" / "coe_implementation.csv")

    print("\n3. Facility-quarter operations ...")
    fq, coe_status, tele, shared = make_facility_quarter(hospitals, coe, pop)
    write(fq, DATA / "facilities" / "facility_quarter_ops.csv")
    write(coe_status, DATA / "facilities" / "facility_quarter_treatment_status.csv")
    write(tele, DATA / "facilities" / "telehealth_activation.csv")
    write(shared, DATA / "facilities" / "shared_services.csv")

    print("\n4. Workforce panel ...")
    wf, awards, resid, train, support = make_workforce(counties, pop)
    write(wf, DATA / "workforce" / "workforce_county_year.csv")
    write(awards, DATA / "workforce" / "service_commitment_awards.csv")
    write(resid, DATA / "workforce" / "residency_slots.csv")
    write(train, DATA / "workforce" / "training_participants.csv")
    write(support, DATA / "workforce" / "workforce_supports.csv")

    print("\n5. Claims (member-month + county-quarter) ...")
    member, cq, ed, tnt, pharm, outpat, ems, vbc = make_claims(counties, hospitals, coe)
    write(member, DATA / "claims" / "member_month_pmpm.csv")
    # Also write parquet for fast loading in the dashboard. CSV is the
    # canonical schema (used by replacements with real data); parquet is
    # an internal cache the loader prefers when present.
    parquet_path = DATA / "claims" / "member_month_pmpm.parquet"
    member.to_parquet(parquet_path, index=False)
    print(f"  wrote {parquet_path.relative_to(ROOT)}  ({len(member):,} rows)")
    write(cq, DATA / "claims" / "county_quarter_treatment_status.csv")
    write(ed, DATA / "claims" / "ed_high_utilizers.csv")
    write(tnt, DATA / "claims" / "tnt_cpt_usage.csv")
    write(pharm, DATA / "claims" / "pharmacy_prescribing.csv")
    write(outpat, DATA / "claims" / "outpatient_share.csv")
    write(ems, DATA / "claims" / "ems_modernization.csv")
    write(vbc, DATA / "claims" / "vbc_attribution.csv")

    print("\n6. Prevention panel ...")
    prev, school, mobile, chap, crisis = make_prevention(counties, pop)
    write(prev, DATA / "prevention" / "prevention_county_year.csv")
    write(school, DATA / "prevention" / "school_sites.csv")
    write(mobile, DATA / "prevention" / "mobile_care_runs.csv")
    write(chap, DATA / "prevention" / "chap_rollout.csv")
    write(crisis, DATA / "prevention" / "crisis_infrastructure.csv")

    print("\n7. Technology / data panel ...")
    hie, ehr, bed, bh_wait, dash = make_technology(hospitals, coe)
    write(hie, DATA / "technology" / "hie_participation.csv")
    write(ehr, DATA / "technology" / "ehr_modernization.csv")
    write(bed, DATA / "technology" / "bed_registry.csv")
    write(bh_wait, DATA / "technology" / "bh_bed_wait_time.csv")
    write(dash, DATA / "technology" / "dashboard_usage.csv")

    print("\nDone.")


if __name__ == "__main__":
    main()