🏥 Tutorial 30: Tribal Healthcare¶

# Workspace configuration checklist (configure via Fabric Admin Portal)
hipaa_workspace_config = {
    "workspace_name": "ws_tribal_healthcare_phi",
    "description": "HIPAA-compliant workspace for IHS tribal health analytics",
    "license_mode": "Premium",  # Required for sensitivity labels + RLS
    "settings": {
        # Access control
        "workspace_access": "Restricted",        # No automatic access
        "guest_access": "Disabled",               # No external users
        "export_data": "Disabled",                # Prevent PHI export
        "download_reports": "Disabled",           # Prevent PHI download
        "print_reports": "Disabled",              # Prevent PHI printing
        "copy_paste_visuals": "Disabled",         # Prevent clipboard leakage

        # Data protection
        "sensitivity_label": "Highly Confidential - PHI",
        "encryption_at_rest": "AES-256",          # Fabric default
        "encryption_in_transit": "TLS 1.3",       # Fabric default

        # Governance
        "purview_integration": "Enabled",
        "audit_logging": "Enabled",
        "data_lineage_tracking": "Enabled",
    },
    "roles": {
        "Admin": ["tribal_health_data_admin@contoso.com"],
        "Member": ["tribal_health_analyst@contoso.com"],
        "Contributor": ["tribal_health_engineer@contoso.com"],
        "Viewer": ["ihs_area_director@contoso.com"],
    }
}

print("Workspace Configuration:")
for key, value in hipaa_workspace_config["settings"].items():
    print(f"  {key}: {value}")

# Column masking configuration for the Silver/Gold semantic model
column_masking_rules = {
    # Full masking: Replace with fixed value
    "patient_id_hash": {
        "mask_type": "full",
        "replacement": "****MASKED****",
        "visible_to": ["tribal_health_data_admin"],
    },
    # Partial masking: Show first 3 chars of facility
    "facility_name_std": {
        "mask_type": "partial",
        "prefix_chars": 3,
        "replacement_char": "X",
        "visible_to": ["tribal_health_data_admin", "tribal_health_analyst"],
    },
    # Null masking: Return NULL for medication details
    "medication_prescribed": {
        "mask_type": "null",
        "visible_to": ["tribal_health_data_admin"],
    },
    # Age group only (never expose exact date of birth)
    "encounter_date": {
        "mask_type": "year_only",
        "visible_to": ["tribal_health_data_admin", "tribal_health_analyst"],
    },
}

print("Column Masking Rules:")
for col_name, rule in column_masking_rules.items():
    print(f"  {col_name}: {rule['mask_type']} (visible to: {rule['visible_to']})")

# Application-level audit log schema
# Written by every notebook that touches PHI data
hipaa_audit_schema = {
    "audit_id": "UUID",
    "audit_timestamp": "ISO 8601",
    "action": "BRONZE_INGESTION | SILVER_TRANSFORM | GOLD_AGGREGATE | REPORT_ACCESS",
    "table_name": "Target table accessed",
    "batch_id": "Pipeline batch identifier",
    "records_affected": "Count of records read/written",
    "records_quarantined": "Count of non-compliant records",
    "phi_violations_detected": "Count of PHI masking violations",
    "consent_violations_detected": "Count of missing consent flags",
    "hipaa_compliant": "Boolean - overall compliance status",
    "user_identity": "Service principal or user UPN",
    "source_notebook": "Notebook name that performed the action",
    "data_classification": "PHI - HIPAA Protected",
    "access_justification": "Business reason for data access",
}

print("HIPAA Audit Log Fields:")
for field, description in hipaa_audit_schema.items():
    print(f"  {field}: {description}")

# Tribal data sovereignty verification before any analytics query
def verify_tribal_data_access(tribal_affiliation: str, use_case: str) -> bool:
    """
    Check whether the specified use case is permitted for the given tribe's data.
    Consult the tribal_data_governance table before proceeding.
    """
    governance = spark.table("lh_bronze.tribal_data_governance")

    permitted = governance.filter(
        (col("tribal_affiliation") == tribal_affiliation) &
        (col("permitted_use_case") == use_case) &
        (col("consent_active") == True) &
        (col("consent_expiry_date") >= current_date())
    ).count()

    if permitted == 0:
        raise PermissionError(
            f"TRIBAL DATA SOVEREIGNTY BLOCK: Use case '{use_case}' "
            f"is not permitted for {tribal_affiliation} data. "
            f"Contact the tribal data governance office."
        )
    return True

# NPIRS reporting categories relevant to population health
npirs_data_categories = {
    "User Population": "Unduplicated patient counts by facility, age, gender",
    "Workload": "Visit counts by encounter type and provider type",
    "Diabetes Audit": "Annual diabetes care metrics (A1C, LDL, BP control)",
    "Immunization": "Childhood and adult immunization rates by service unit",
    "Behavioral Health": "Mental health and substance abuse encounter counts",
    "Dental": "Dental visit rates and sealant application counts",
    "Maternal/Child Health": "Prenatal care, birth outcomes, well-child visits",
}

for category, description in npirs_data_categories.items():
    print(f"  {category}: {description}")

# SHIR configuration for IHS RPMS connectivity
# Reference: Tutorial 23 - SHIR & Data Gateways

shir_config = {
    "runtime_name": "SHIR-IHS-RPMS-01",
    "host": "ihs-etl-server.ihs.gov",
    "os": "Windows Server 2022",
    "network_zone": "IHS DMZ",

    # RPMS Cache database connection
    "jdbc_driver": "com.intersystems.jdbc.IRISDriver",
    "jdbc_url": "jdbc:IRIS://rpms-db.ihs.gov:1972/RPMS",
    "authentication": "LDAP (IHS Active Directory)",
    "tls_version": "1.3",
    "certificate": "IHS-issued PKI certificate",

    # Outbound connectivity
    "fabric_endpoint": "https://onelake.dfs.fabric.microsoft.com",
    "proxy": "http://proxy.ihs.gov:8080",
    "firewall_rules": [
        "Allow TCP 443 outbound to *.fabric.microsoft.com",
        "Allow TCP 443 outbound to *.dfs.fabric.microsoft.com",
        "Allow TCP 443 outbound to login.microsoftonline.com",
    ],
}

print("SHIR Configuration:")
for key, value in shir_config.items():
    if isinstance(value, list):
        print(f"  {key}:")
        for item in value:
            print(f"    - {item}")
    else:
        print(f"  {key}: {value}")

# JDBC connection properties for RPMS Cache database
jdbc_properties = {
    "driver": "com.intersystems.jdbc.IRISDriver",
    "url": "jdbc:IRIS://rpms-db.ihs.gov:1972/RPMS",
    "user": "${IHS_SERVICE_ACCOUNT}",        # From Azure Key Vault
    "password": "${IHS_SERVICE_PASSWORD}",    # From Azure Key Vault
    "fetchSize": "10000",
    "queryTimeout": "3600",
    "ssl": "true",
    "sslProtocols": "TLSv1.3",

    # RPMS-specific settings
    "readOnly": "true",                       # Never write back to RPMS
    "applicationName": "FabricTribalHealth",  # For RPMS audit trail
}

# Example: Read PCC encounter data via JDBC
# (In production, use Fabric Data Pipeline with SHIR connection)
rpms_query = """
    SELECT
        e.encounter_id,
        e.patient_dfn,
        e.encounter_date,
        e.encounter_type,
        e.facility_code,
        d.icd10_code,
        d.icd10_description,
        p.provider_npi,
        p.provider_type
    FROM PCC_ENCOUNTER e
    JOIN PCC_DIAGNOSIS d ON e.encounter_id = d.encounter_id
    JOIN PCC_PROVIDER p ON e.provider_ien = p.provider_ien
    WHERE e.encounter_date >= '2023-01-01'
    AND e.encounter_date < '2024-01-01'
"""

print("RPMS Query configured for encounter extraction")
print(f"Query targets: PCC_ENCOUNTER, PCC_DIAGNOSIS, PCC_PROVIDER")

# Schema definition for Bronze layer ingestion
# Reference: tribal_health_schema.json

from pyspark.sql.types import *

tribal_health_schema = StructType([
    StructField("encounter_id", StringType(), False),       # UUID, required
    StructField("patient_id", StringType(), False),         # PAT-XXXXXXXX, required
    StructField("encounter_date", DateType(), False),       # ISO date, required
    StructField("encounter_type", StringType(), True),      # outpatient, inpatient, etc.
    StructField("facility_id", StringType(), True),         # IHS facility code
    StructField("facility_name", StringType(), True),       # Facility display name
    StructField("service_unit", StringType(), True),        # IHS service unit
    StructField("area_office", StringType(), True),         # IHS area office (12 regions)
    StructField("provider_id", StringType(), True),         # NPI (10-digit)
    StructField("provider_type", StringType(), True),       # MD, NP, PA, etc.
    StructField("icd10_code", StringType(), True),          # ICD-10 diagnosis
    StructField("icd10_description", StringType(), True),   # Diagnosis description
    StructField("diagnosis_category", StringType(), True),  # High-level category
    StructField("procedure_code", StringType(), True),      # CPT code
    StructField("procedure_description", StringType(), True),
    StructField("medication_prescribed", StringType(), True),
    StructField("medication_ndc", StringType(), True),      # NDC 5-4-2 format
    StructField("insurance_type", StringType(), True),      # IHS_CONTRACT, MEDICAID, etc.
    StructField("insurance_id", StringType(), True),
    StructField("visit_duration_minutes", IntegerType(), True),
    StructField("referral_flag", BooleanType(), True),
    StructField("referral_destination", StringType(), True),
    StructField("follow_up_required", BooleanType(), True),
    StructField("follow_up_date", DateType(), True),
    StructField("phi_masked", BooleanType(), True),         # PHI masking flag
    StructField("hipaa_consent", BooleanType(), True),      # HIPAA consent flag
    StructField("tribal_affiliation", StringType(), True),  # Federally recognized tribe
    StructField("community_health_rep_id", StringType(), True),
    StructField("telehealth_flag", BooleanType(), True),
    StructField("emergency_flag", BooleanType(), True),
])

print(f"Schema fields: {len(tribal_health_schema.fields)}")
print(f"Required fields: {sum(1 for f in tribal_health_schema.fields if not f.nullable)}")

# HIPAA compliance verification -- runs before ANY data is written
# This code is from notebook 07_bronze_tribal_health.py

REQUIRE_PHI_MASKED = True
REQUIRE_HIPAA_CONSENT = True

# Read raw source data
df_raw = spark.read \
    .schema(tribal_health_schema) \
    .parquet(f"{SOURCE_PATH}*.parquet")

# Check phi_masked flag
phi_violations = df_raw.filter(
    (col("phi_masked").isNull()) | (col("phi_masked") == False)
).count()

# Check hipaa_consent flag
consent_violations = df_raw.filter(
    (col("hipaa_consent").isNull()) | (col("hipaa_consent") == False)
).count()

print("HIPAA Compliance Check:")
print(f"  PHI Masking Violations: {phi_violations}")
print(f"  Consent Violations: {consent_violations}")

# Separate compliant from non-compliant records
df_compliant = df_raw.filter(
    (col("phi_masked") == True) & (col("hipaa_consent") == True)
)

df_quarantined = df_raw.filter(
    (col("phi_masked") != True) | (col("hipaa_consent") != True) |
    col("phi_masked").isNull() | col("hipaa_consent").isNull()
)

# CRITICAL: Halt if zero compliant records
if df_compliant.count() == 0:
    raise Exception("HIPAA BLOCK: Zero compliant records. Ingestion halted.")

# Add ingestion metadata and write to Bronze
df_bronze = df_compliant \
    .withColumn("_ingested_at", current_timestamp()) \
    .withColumn("_source_file", input_file_name()) \
    .withColumn("_batch_id", lit(BATCH_ID)) \
    .withColumn("_run_id", lit(RUN_ID)) \
    .withColumn("_load_date", current_timestamp().cast("date"))

# Write to Delta table (append-only -- never overwrite Bronze)
df_bronze.write \
    .format("delta") \
    .mode("append") \
    .option("mergeSchema", "true") \
    .partitionBy("_load_date") \
    .saveAsTable("lh_bronze.bronze_tribal_health_encounters")

# Write HIPAA audit log entry
audit_entry = spark.createDataFrame([{
    "audit_id": RUN_ID,
    "audit_timestamp": datetime.utcnow().isoformat(),
    "action": "BRONZE_INGESTION",
    "table_name": "lh_bronze.bronze_tribal_health_encounters",
    "records_ingested": df_bronze.count(),
    "records_quarantined": df_quarantined.count(),
    "phi_violations_detected": phi_violations,
    "consent_violations_detected": consent_violations,
    "hipaa_compliant": True,
    "data_classification": "PHI - HIPAA Protected",
}])

audit_entry.write \
    .format("delta") \
    .mode("append") \
    .saveAsTable("lh_bronze.tribal_health_hipaa_audit_log")

print(f"Bronze: {df_bronze.count():,} records ingested")
print(f"Quarantine: {df_quarantined.count():,} records blocked")
print(f"Audit log: Entry written for run {RUN_ID}")

# Patient ID hashing -- second line of defense after Bronze consent check
# SHA-256 produces a 64-character hex string

df_phi = df_bronze \
    .withColumn(
        "patient_id_hash",
        when(
            length(col("patient_id")) == 64,
            col("patient_id")                   # Already hashed
        ).otherwise(
            sha2(col("patient_id"), 256)        # Hash raw patient_id
        )
    )

# Verify no raw SSN patterns exist in any string column
ssn_pattern = r"^\d{3}-\d{2}-\d{4}$"
string_columns = [f.name for f in df_bronze.schema.fields if isinstance(f.dataType, StringType)]

for col_name in string_columns:
    ssn_count = df_bronze.filter(col(col_name).rlike(ssn_pattern)).count()
    if ssn_count > 0:
        raise Exception(f"PHI VIOLATION: {col_name} contains {ssn_count} SSN-like values!")

# FHIR R4 Encounter class mapping
# Reference: http://terminology.hl7.org/CodeSystem/v3-ActCode

ENCOUNTER_TO_FHIR_CLASS = {
    "INPATIENT": "IMP",           # Inpatient encounter
    "OUTPATIENT": "AMB",          # Ambulatory
    "EMERGENCY": "EMER",          # Emergency
    "TELEHEALTH": "VR",           # Virtual
    "HOME_VISIT": "HH",           # Home health
    "DENTAL": "AMB",              # Ambulatory (dental)
    "BEHAVIORAL_HEALTH": "AMB",   # Ambulatory (behavioral)
    "SUBSTANCE_ABUSE": "AMB",     # Ambulatory (substance abuse)
    "IMMUNIZATION": "AMB",        # Ambulatory (immunization)
    "PHARMACY": "AMB",            # Ambulatory (pharmacy)
    "URGENT_CARE": "EMER",        # Emergency (urgent)
}

fhir_class_expr = create_map([lit(x) for pair in ENCOUNTER_TO_FHIR_CLASS.items() for x in pair])

df_fhir = df_phi \
    .withColumn("encounter_type_clean", upper(trim(col("encounter_type")))) \
    .withColumn(
        "fhir_encounter_class",
        coalesce(fhir_class_expr[upper(trim(col("encounter_type")))], lit("AMB"))
    ) \
    .withColumn("fhir_resource_type", lit("Encounter")) \
    .withColumn(
        "fhir_service_type",
        when(col("encounter_type_clean") == "DENTAL", lit("dental"))
        .when(col("encounter_type_clean").isin("BEHAVIORAL_HEALTH", "SUBSTANCE_ABUSE"), lit("mental-health"))
        .when(col("encounter_type_clean") == "IMMUNIZATION", lit("immunization"))
        .when(col("encounter_type_clean") == "PHARMACY", lit("pharmacy"))
        .otherwise(lit("general-practice"))
    )

# ICD-10 format validation: letter + 2 digits + optional decimal + up to 4 digits
ICD10_PATTERN = r"^[A-Z]\d{2}(\.\d{1,4})?$"

df_icd = df_fhir \
    .withColumn("icd10_code_clean", upper(trim(col("icd10_code")))) \
    .withColumn("icd10_valid", col("icd10_code_clean").rlike(ICD10_PATTERN)) \
    .withColumn(
        "icd10_chapter",
        when(col("icd10_code_clean").rlike("^[AB]"), lit("Infectious Diseases"))
        .when(col("icd10_code_clean").startswith("E"), lit("Endocrine/Metabolic"))
        .when(col("icd10_code_clean").startswith("F"), lit("Mental/Behavioral"))
        .when(col("icd10_code_clean").startswith("I"), lit("Circulatory System"))
        .when(col("icd10_code_clean").startswith("J"), lit("Respiratory System"))
        .when(col("icd10_code_clean").startswith("K"), lit("Digestive System"))
        .when(col("icd10_code_clean").startswith("M"), lit("Musculoskeletal"))
        .when(col("icd10_code_clean").startswith("Z"), lit("Health Status/Services"))
        .otherwise(lit("Other"))
    ) \
    .withColumn("is_diabetes_related", col("icd10_code_clean").rlike("^E1[0-4]")) \
    .withColumn("is_behavioral_health", col("icd10_code_clean").rlike("^F[0-9]")) \
    .withColumn("is_substance_use", col("icd10_code_clean").rlike("^F1[0-9]"))

# Validation report
valid_icd = df_icd.filter(col("icd10_valid") == True).count()
invalid_icd = df_icd.filter(col("icd10_valid") == False).count()

print(f"ICD-10 Validation: {valid_icd:,} valid, {invalid_icd:,} invalid")
print(f"Validation rate: {valid_icd / (valid_icd + invalid_icd) * 100:.1f}%")

# Deduplicate by composite key: patient + date + diagnosis
before_dedup = df_icd.count()

df_deduped = df_icd.dropDuplicates(["patient_id_hash", "encounter_date", "icd10_code_clean"])

after_dedup = df_deduped.count()
print(f"Deduplication: {before_dedup:,} -> {after_dedup:,} ({before_dedup - after_dedup:,} duplicates removed)")

# Write Silver layer with data quality scoring
df_silver = df_deduped \
    .withColumn("_dq_score",
        when(col("encounter_id").isNotNull(), lit(15)).otherwise(lit(0)) +
        when(col("patient_id_hash").isNotNull(), lit(15)).otherwise(lit(0)) +
        when(col("encounter_date").isNotNull(), lit(10)).otherwise(lit(0)) +
        when(col("icd10_valid") == True, lit(15)).otherwise(lit(0)) +
        when(col("facility_name").isNotNull(), lit(10)).otherwise(lit(0)) +
        when(col("provider_id").isNotNull(), lit(10)).otherwise(lit(0)) +
        when(col("insurance_type").isNotNull(), lit(10)).otherwise(lit(0)) +
        when(col("tribal_affiliation").isNotNull(), lit(5)).otherwise(lit(0)) +
        when(col("service_unit").isNotNull(), lit(5)).otherwise(lit(0)) +
        when(col("area_office").isNotNull(), lit(5)).otherwise(lit(0))
    ) \
    .withColumn("_silver_timestamp", current_timestamp())

df_silver.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \
    .partitionBy("encounter_date") \
    .saveAsTable("lh_silver.silver_tribal_health_encounters")

# Optimize for Direct Lake queries
spark.sql("OPTIMIZE lh_silver.silver_tribal_health_encounters ZORDER BY (patient_id_hash, area_office_std)")

# Build Patient 360 view: encounter history aggregation
df_encounter_summary = df_silver \
    .groupBy("patient_id_hash") \
    .agg(
        count("*").alias("total_encounters"),
        countDistinct("encounter_date").alias("distinct_visit_days"),
        min("encounter_date").alias("first_encounter_date"),
        max("encounter_date").alias("last_encounter_date"),
        countDistinct("facility_name_std").alias("facilities_visited"),
        countDistinct("provider_id").alias("distinct_providers"),

        # Encounter type breakdown
        sum(when(col("encounter_type_std") == "OUTPATIENT", 1).otherwise(0)).alias("outpatient_visits"),
        sum(when(col("encounter_type_std") == "INPATIENT", 1).otherwise(0)).alias("inpatient_visits"),
        sum(when(col("encounter_type_std") == "EMERGENCY", 1).otherwise(0)).alias("emergency_visits"),
        sum(when(col("encounter_type_std") == "TELEHEALTH", 1).otherwise(0)).alias("telehealth_visits"),
        sum(when(col("encounter_type_std") == "DENTAL", 1).otherwise(0)).alias("dental_visits"),
        sum(when(col("encounter_type_std") == "BEHAVIORAL_HEALTH", 1).otherwise(0)).alias("behavioral_health_visits"),
        sum(when(col("encounter_type_std") == "IMMUNIZATION", 1).otherwise(0)).alias("immunization_visits"),

        # Visit metrics
        avg("visit_duration_minutes").alias("avg_visit_duration_min"),
        sum(when(col("referral_flag") == True, 1).otherwise(0)).alias("total_referrals"),
        sum(when(col("follow_up_required") == True, 1).otherwise(0)).alias("follow_ups_required"),
    )

# Diagnosis summary per patient with chronic condition flags
df_diagnosis_summary = df_silver \
    .filter(col("icd10_code_clean").isNotNull()) \
    .groupBy("patient_id_hash") \
    .agg(
        countDistinct("icd10_code_clean").alias("distinct_diagnoses"),
        collect_set("icd10_chapter").alias("diagnosis_chapters"),

        # Chronic condition flags critical for tribal health
        max(col("is_diabetes_related").cast("int")).alias("has_diabetes"),
        max(col("is_behavioral_health").cast("int")).alias("has_behavioral_health"),
        max(col("is_substance_use").cast("int")).alias("has_substance_use"),

        # Top diagnoses
        collect_set("diagnosis_category").alias("diagnosis_categories"),
    ) \
    .withColumn("has_diabetes", col("has_diabetes").cast("boolean")) \
    .withColumn("has_behavioral_health", col("has_behavioral_health").cast("boolean")) \
    .withColumn("has_substance_use", col("has_substance_use").cast("boolean"))

# Medication list and adherence indicators
df_medication_summary = df_silver \
    .filter(col("medication_prescribed").isNotNull()) \
    .groupBy("patient_id_hash") \
    .agg(
        countDistinct("medication_prescribed").alias("distinct_medications"),
        collect_set("medication_prescribed").alias("medication_list"),
        countDistinct("medication_ndc").alias("distinct_ndc_codes"),
    )

# Assemble the Patient 360 view
df_patient_360 = df_encounter_summary \
    .join(df_diagnosis_summary, "patient_id_hash", "left") \
    .join(df_medication_summary, "patient_id_hash", "left") \
    .join(df_demographics, "patient_id_hash", "left")

# Calculate derived metrics
df_patient_360 = df_patient_360 \
    .withColumn("days_since_last_visit",
        datediff(current_date(), col("last_encounter_date"))
    ) \
    .withColumn("encounters_per_year",
        when(col("patient_tenure_days") > 0,
            round(col("total_encounters") * 365.0 / col("patient_tenure_days"), 2)
        ).otherwise(col("total_encounters").cast("double"))
    ) \
    .withColumn("ed_utilization_rate",
        when(col("total_encounters") > 0,
            round(col("emergency_visits") * 100.0 / col("total_encounters"), 2)
        ).otherwise(lit(0.0))
    ) \
    .withColumn("chronic_condition_count",
        coalesce(col("has_diabetes").cast("int"), lit(0)) +
        coalesce(col("has_behavioral_health").cast("int"), lit(0)) +
        coalesce(col("has_substance_use").cast("int"), lit(0))
    ) \
    .withColumn("care_complexity",
        when(col("chronic_condition_count") >= 3, lit("High"))
        .when(col("chronic_condition_count") >= 2, lit("Medium"))
        .when(col("chronic_condition_count") >= 1, lit("Low"))
        .otherwise(lit("Routine"))
    )

# Write Gold Patient 360
df_patient_360.write \
    .format("delta") \
    .mode("overwrite") \
    .option("overwriteSchema", "true") \
    .saveAsTable("lh_gold.gold_tribal_patient_360")

spark.sql("OPTIMIZE lh_gold.gold_tribal_patient_360 ZORDER BY (patient_id_hash, primary_area_office)")

# Diabetes prevalence calculation per service unit
df_diabetes = df_silver \
    .filter(col("is_diabetes_related") == True) \
    .groupBy("service_unit_std") \
    .agg(
        countDistinct("patient_id_hash").alias("diabetes_patients"),
        count("*").alias("diabetes_encounters"),
    )

df_service_unit_pop = df_silver \
    .groupBy("service_unit_std", "area_office_std") \
    .agg(
        countDistinct("patient_id_hash").alias("total_patients"),
        count("*").alias("total_encounters"),
    )

df_pop_health = df_service_unit_pop \
    .join(df_diabetes, "service_unit_std", "left") \
    .fillna(0, subset=["diabetes_patients", "diabetes_encounters"]) \
    .withColumn("diabetes_prevalence_rate",
        round(col("diabetes_patients") * 100.0 / col("total_patients"), 2)
    )

# Behavioral health utilization -- critical for tribal communities
df_behavioral = df_silver \
    .filter(col("is_behavioral_health") == True) \
    .groupBy("service_unit_std") \
    .agg(
        countDistinct("patient_id_hash").alias("behavioral_health_patients"),
        count("*").alias("behavioral_health_encounters"),
    )

# Substance use disorders -- often co-occurring
df_substance = df_silver \
    .filter(col("is_substance_use") == True) \
    .groupBy("service_unit_std") \
    .agg(
        countDistinct("patient_id_hash").alias("substance_use_patients"),
        count("*").alias("substance_use_encounters"),
    )

# Join to population
df_pop_health = df_pop_health \
    .join(df_behavioral, "service_unit_std", "left") \
    .join(df_substance, "service_unit_std", "left") \
    .fillna(0) \
    .withColumn("behavioral_health_utilization_rate",
        round(col("behavioral_health_patients") * 100.0 / col("total_patients"), 2)
    ) \
    .withColumn("substance_use_rate",
        round(col("substance_use_patients") * 100.0 / col("total_patients"), 2)
    )

# Area office-level community health KPIs
df_community_kpis = df_silver \
    .groupBy("area_office_std") \
    .agg(
        countDistinct("patient_id_hash").alias("total_patients"),
        count("*").alias("total_encounters"),
        countDistinct("facility_name_std").alias("total_facilities"),
        countDistinct("provider_id").alias("total_providers"),

        # Chronic disease prevalence
        round(
            countDistinct(when(col("is_diabetes_related"), col("patient_id_hash"))) * 100.0 /
            countDistinct("patient_id_hash"), 2
        ).alias("diabetes_prevalence_pct"),
        round(
            countDistinct(when(col("is_behavioral_health"), col("patient_id_hash"))) * 100.0 /
            countDistinct("patient_id_hash"), 2
        ).alias("behavioral_health_pct"),

        # Access indicators
        round(
            sum(when(col("encounter_type_std") == "EMERGENCY", 1).otherwise(0)) * 100.0 /
            count("*"), 2
        ).alias("ed_encounter_rate_pct"),
        round(
            sum(when(col("encounter_type_std") == "TELEHEALTH", 1).otherwise(0)) * 100.0 /
            count("*"), 2
        ).alias("telehealth_rate_pct"),

        # Insurance gap
        round(
            sum(when(col("insurance_type_std") == "UNINSURED", 1).otherwise(0)) * 100.0 /
            count("*"), 2
        ).alias("uninsured_rate_pct"),
    )

# Write community KPIs
df_community_kpis.write \
    .format("delta") \
    .mode("overwrite") \
    .saveAsTable("lh_gold.gold_tribal_community_kpis")

# Write population health table
df_pop_health.write \
    .format("delta") \
    .mode("overwrite") \
    .saveAsTable("lh_gold.gold_tribal_population_health")

spark.sql("OPTIMIZE lh_gold.gold_tribal_population_health ZORDER BY (service_unit_std, area_office_std)")

print("Gold layer tables written:")
print(f"  gold_tribal_patient_360: {spark.table('lh_gold.gold_tribal_patient_360').count():,} patients")
print(f"  gold_tribal_population_health: {spark.table('lh_gold.gold_tribal_population_health').count():,} service units")
print(f"  gold_tribal_community_kpis: {spark.table('lh_gold.gold_tribal_community_kpis').count():,} area offices")