Core data engineering, one real project.

# Databricks notebook — Week 14 starter scaffold (structure only; fill the TODOs)
# Bronze contract: land raw, add source + ingest timestamp, NO transforms.

import requests, json
from datetime import datetime
from pyspark.sql import functions as F

CATALOG = "workspace"           # Free Edition default
SCHEMA  = "pricing_integrity"    # TODO: create if needed
BRONZE  = "/Volumes/workspace/pricing_integrity/bronze"  # TODO: adjust to your volume/path
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {CATALOG}.{SCHEMA}")

# --- 1. Land one CMS hospital pricing file (raw) --------------------------
# TODO: from the CMS hospital index, paste ONE hospital's machine-readable file URL.
CMS_FILE_URL = "<PASTE_A_HOSPITAL_MRF_URL>"
raw = requests.get(CMS_FILE_URL, timeout=60)
landing = f"{BRONZE}/cms_mrf/ingest_date={datetime.utcnow():%Y-%m-%d}"
dbutils.fs.put(f"{landing}/raw.txt", raw.text, overwrite=True)
# TODO: read it back as text/CSV/JSON depending on the file's real format and
#       add columns: source = 'cms_mrf', ingest_ts = current_timestamp(). Keep it RAW.

# --- 2. Pull NPPES providers (raw) ---------------------------------------
# TODO: set city/state to match your CMS hospital so records overlap.
nppes = requests.get(
    "https://npiregistry.cms.hhs.gov/api/",
    params={"version": "2.1", "city": "<CITY>", "state": "<ST>", "limit": 50},
    timeout=60,
).json()
nppes_rows = nppes.get("results", [])
nppes_df = (spark.createDataFrame([json.dumps(r) for r in nppes_rows], "string")
            .withColumnRenamed("value", "raw_json")
            .withColumn("source", F.lit("nppes_api"))
            .withColumn("ingest_ts", F.current_timestamp()))
nppes_df.write.mode("overwrite").format("delta").save(f"{BRONZE}/nppes_raw")

# --- 3. Profile the mess (fill in) ---------------------------------------
# TODO: for BOTH sources print column names, types, null-rate, and 5 sample values.
#       Write down: nested pipe columns, bad prices, free-text payers, NPI dupes, drift.

# --- 4/5. Dimensional skeleton (empty, typed Delta) ----------------------
# TODO: replace the '...' column lists with the columns YOU chose during design.
#       Keep surrogate key + business key + audit columns on each dimension.
spark.sql(f"""CREATE TABLE IF NOT EXISTS {CATALOG}.{SCHEMA}.dim_provider (
  provider_sk BIGINT,          -- surrogate key (generate later)
  npi STRING,                  -- business/natural key
  -- TODO: provider_name, taxonomy, ... (from your NPPES profiling)
  source STRING, ingest_ts TIMESTAMP
) USING DELTA""")
# TODO: dim_procedure (business key = CPT/HCPCS/DRG code), dim_facility (CCN/EIN),
#       dim_payer (payer_name + plan_name) — same pattern as dim_provider.

# TODO: fact_charge — one row per charge event. Include the 4 dimension surrogate
#       keys, your MEASURE columns (mark additive vs non-additive in a comment),
#       a degenerate charge/line id, and audit columns. Create empty & typed only.
# NOTE: do NOT load the fact this week — population comes as silver is built.

# --- 6. Idempotency check ------------------------------------------------
# TODO: re-run the CREATE cell and confirm no error and row counts stay 0.
for t in ["dim_provider"]:  # TODO: add your other 4 tables
    print(t, spark.table(f"{CATALOG}.{SCHEMA}.{t}").count())

# Databricks Free Edition — PySpark/Delta — Week 15 STARTER SCAFFOLD
# Provider Pricing Integrity: lock the fact GRAIN + build CONFORMED dims.
# Fill every TODO. This scaffold intentionally does NOT solve the lab.

from pyspark.sql import functions as F

SCHEMA = "ppi_silver"
spark.sql(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA}")

# --- 0. Grain declared IN CODE as the single source of truth --------------
# TODO: write your grain as ONE sentence in the README, then list the columns
#       that make one fact row unique. Decide and DEFEND this in the charter.
GRAIN_KEYS = [
    # "hospital_ccn", "billing_code", "code_system", "payer_plan", "charge_type",
]  # TODO fill

# --- 1. Load W14 silver pricing lines -------------------------------------
pricing = spark.table(f"{SCHEMA}.silver_pricing_lines")

# --- 2. Compose the grain key ---------------------------------------------
# TODO: coalesce each key col to a sentinel, then sha2(concat_ws('||',...),256)
pricing_keyed = pricing  # TODO .withColumn("grain_key", ...)

# --- 3. GRAIN-UNIQUENESS TEST (must reach 0 after quarantine) --------------
dupes = None  # TODO groupBy("grain_key").count().filter(F.col("count") > 1)

# --- 4. Quarantine grain violators (build for failure, be LOUD) -----------
REJECT = f"{SCHEMA}.rejected_pricing"
# TODO: append violators to REJECT with reason='GRAIN_VIOLATION',
#       rejected_at=current_timestamp(), source_file; then print count+sample.

# --- 5. Conformed dimensions (shared by pricing AND future claims) --------
# NOTE: Type-1 / overwrite only. SCD2 history is next week (W16). Do NOT add it.
providers = spark.table(f"{SCHEMA}.silver_providers")
# TODO dim_provider: select conformed attrs (npi, provider_name, org_name,
#      city, state, taxonomy); add provider_sk; write Delta (overwrite).
# TODO dim_service : conform on billing_code+code_system; keep the most
#      frequent description per code; add service_sk; write Delta (overwrite).

# --- 6. Idempotency: 10 reruns => identical dim/fact counts ----------------
# TODO: rewrite the fact with provider_sk/service_sk + degenerate keys,
#       mode('overwrite'); print counts; assert they don't move across runs.

print("SCAFFOLD ONLY — steps above are not implemented yet")

# W16_silver_scd2_scaffold.py
# SCAFFOLD ONLY - sets up structure + the two-step SCD2 pattern.
# You implement every TODO. Do NOT paste a full solution here; the point is
# that YOU can defend each MERGE clause in the live hot seat.
#
# Runs on Databricks Free Edition (Delta + PySpark). Assumes W15 landed a
# grain-correct silver feed. Here we historize it into an SCD2 dimension.

from pyspark.sql import functions as F
from delta.tables import DeltaTable

# --- 0. Config -------------------------------------------------------------
CATALOG_PATH   = "workspace.default"            # adjust to your Free Edition schema
DIM_TABLE      = f"{CATALOG_PATH}.silver_provider_dim"   # SCD2 target
# For the PRICE dimension, re-run this file with these swapped:
#   DIM_TABLE = f"{CATALOG_PATH}.silver_price_dim"
BUSINESS_KEY   = ["npi"]                          # natural/business key that IDENTIFIES the entity
# TODO: for price_dim the business key is a COMPOSITE (e.g. npi + code_type + billing_code + payer)
#       - decide the key at the grain W15 fixed, and justify it in your why-note.

# Columns whose CHANGE should open a new version (Type 2 / tracked):
TRACKED_COLS   = ["provider_name", "practice_city", "practice_state", "taxonomy_code"]
# TODO: for price_dim, TRACKED_COLS is the price fields (gross_charge, discounted_cash,
#       negotiated_rate...). Decide per column: Type 2 (track) vs Type 1 (overwrite,
#       e.g. a corrected typo you never want in history). Write that decision down.

# --- 1. Load the incoming snapshot ----------------------------------------
# `snapshot_ts` is the CMS republish date for THIS run - the effective date of any change.
def load_snapshot(snapshot_ts: str):
    src = (spark.table(f"{CATALOG_PATH}.silver_provider_clean")  # W15 output
                .withColumn("snapshot_ts", F.lit(snapshot_ts).cast("timestamp")))
    # TODO: drop/repair rows that violate the business key (null NPI) BEFORE hashing -
    #       route them to a *_rejected table with reason + timestamp (build-for-failure).
    return src

# --- 2. Change-detection hash ---------------------------------------------
# A stable hash over ONLY the tracked columns lets us detect a real change
# without comparing every field, and prevents spurious new versions when an
# untracked column wiggles.
def add_row_hash(df):
    # TODO: sha2(concat_ws over TRACKED_COLS, coalescing nulls to a sentinel like '<NULL>')
    #       Why coalesce? concat_ws silently drops nulls -> two different rows could hash equal.
    return df  # TODO: return df.withColumn("row_hash", F.sha2(...))

# --- 3. First (seed) load --------------------------------------------------
# If the dimension doesn't exist yet, every incoming row is a brand-new
# current version. Open the interval, leave the end open, mark is_current.
def seed_dimension(df):
    seeded = (df.withColumn("effective_start_ts", F.col("snapshot_ts"))
                .withColumn("effective_end_ts",   F.lit(None).cast("timestamp"))  # open interval
                .withColumn("is_current",         F.lit(True))
                .withColumn("surrogate_key",      F.expr("uuid()")))
    # TODO: write `seeded` as a Delta table at DIM_TABLE. Guard it: only seed if
    #       the DeltaTable does not already exist. This branch runs ONCE.
    pass

# --- 4. Incremental SCD2 MERGE (the two-step pattern) ---------------------
# You CANNOT both close an old row and insert its replacement in a single
# MERGE, because MERGE matches one target row per source key. So:
#   Step A: MERGE - for keys whose row_hash CHANGED, set end_ts + is_current=False.
#   Step B: INSERT - the changed rows AND the truly-new keys, as new current versions.
def apply_scd2(df):
    dim = DeltaTable.forName(spark, DIM_TABLE)
    current = spark.table(DIM_TABLE).filter("is_current = true")
    # TODO: join df to `current` on BUSINESS_KEY; classify each incoming row as
    #       'new_key' | 'changed' | 'unchanged' by comparing row_hash.
    #       (unchanged -> do NOTHING; that is what makes reruns idempotent.)

    # ----- Step A: close out the superseded current rows -----
    # TODO: dim.alias("t").merge(changed.alias("s"), key match AND t.is_current = true)
    #        .whenMatchedUpdate(set = {"effective_end_ts": "s.snapshot_ts",
    #                                   "is_current": "false"}).execute()
    #       Prove your interval [start, end) is half-open: no gaps, no overlaps.

    # ----- Step B: insert new current versions (changed + new keys) -----
    # TODO: to_insert = changed U new_key, with effective_start_ts = snapshot_ts,
    #       effective_end_ts = NULL, is_current = true, surrogate_key = uuid();
    #       then append to DIM_TABLE.
    pass

# --- 5. Point-in-time (as-of) validation ----------------------------------
def as_of(as_of_ts: str):
    return spark.sql(f"""
        SELECT * FROM {DIM_TABLE}
        WHERE effective_start_ts <= TIMESTAMP('{as_of_ts}')
          AND (effective_end_ts  >  TIMESTAMP('{as_of_ts}') OR effective_end_ts IS NULL)
    """)
    # TODO: assert exactly ONE row per business key for any as_of_ts. Two rows ->
    #       your intervals overlap -> bug is in the Step A/B boundary handling.

# --- 6. Orchestration ------------------------------------------------------
if __name__ == "__main__":
    # TODO: if DIM_TABLE missing -> seed_dimension(add_row_hash(load_snapshot(SNAP1)))
    #        else               -> apply_scd2(add_row_hash(load_snapshot(SNAP2)))
    # TODO: append an audit-log row: run_ts, snapshot_ts, rows_in, versions_closed,
    #        versions_opened, rows_rejected  (append-only _audit table).
    # TODO: IDEMPOTENCY TEST - apply_scd2 on the SAME snapshot 10x and assert
    #        versions_opened == 0 after the first application.
    pass

# Databricks Free Edition | PySpark + Spark SQL | Delta
# Week 17 - GOLD: price-by-facility-procedure mart + completeness KPI + semantic view
# CONTRACT: gold = business-ready, SERVE-ONLY. No cleaning here. If silver is dirty, fix silver.

from pyspark.sql import functions as F

CATALOG   = "main"          # TODO: set your schema names
SILVER    = f"{CATALOG}.silver"
GOLD      = f"{CATALOG}.gold"
RUN_ID    = dbutils.widgets.get("run_id") if dbutils.widgets.getAll() else "manual_run"  # TODO

# --- 1. CURRENT-view dimensions from SCD2 silver -----------------------------
# Silver providers/facilities are SCD2 (is_current flag from W16). Gold uses CURRENT only.
silver_fac = spark.table(f"{SILVER}.dim_facility_scd2").where("is_current = true")
# TODO: assign a surrogate key + keep business key (ccn). Build dim_provider, dim_procedure the same way.
# TODO: PROVE (a count) why carrying ALL scd2 versions would fan out the fact. Write it in your note.

# --- 2. Fact: declare grain, then build it -----------------------------------
# GRAIN: one row = facility x billing_code x payer x plan x charge_setting  (write this in the note)
silver_price = spark.table(f"{SILVER}.pricing_clean")
# TODO: join silver_price to the 3 dims on BUSINESS keys; keep ONLY surrogate keys + measures
#       + degenerate attrs (payer, plan, setting) + run_id + ingest_ts.
#       Measures as decimal(12,2). No null-fixing here.
fact = None  # TODO build me

# --- 3. Enforce the grain (fail loudly) --------------------------------------
# TODO: find rows whose grain key is duplicated; route them to GOLD.gold_rejected
#       with reason='grain_violation', the key, and rejected_ts. Print an alert + count.

# --- 4. Idempotent write ------------------------------------------------------
# TODO: write `fact` with .option("replaceWhere", f"run_id = '{RUN_ID}'") (or MERGE on grain key)
# fact.write.format("delta").mode("overwrite").option("replaceWhere", ...).saveAsTable(f"{GOLD}.fact_price_by_facility_procedure")

# --- 5. Canonical completeness KPI (define ONCE) -----------------------------
# The ONLY place this metric is defined. The view and any report must read this table.
kpi_sql = f"""
  -- TODO: per facility, completeness_pass = has non-null negotiated_price for EVERY required code.
  -- Aggregate to: run_id, run_ts, facilities_total, facilities_passing, pct_passing
  SELECT NULL AS run_id  -- replace me
"""
# spark.sql(kpi_sql).write.format("delta").mode("overwrite")\
#      .option("replaceWhere", f"run_id = '{RUN_ID}'").saveAsTable(f"{GOLD}.kpi_pricing_completeness")

# --- 6. Semantic view (business-named, serve-only) ---------------------------
spark.sql(f"""
CREATE OR REPLACE VIEW {GOLD}.vw_price_transparency AS
SELECT
    -- TODO: expose facility_name, procedure_description, payer, plan,
    --       negotiated_price, pricing_complete_flag. HIDE surrogate keys + silver internals.
    NULL AS facility_name
FROM {GOLD}.fact_price_by_facility_procedure f
-- TODO: join dims for the human-readable names; join kpi for the complete flag
""")

# --- 7. Reconcile gold -> silver (audit) -------------------------------------
# TODO: compare distinct facility count (and distinct facility+code pairs) gold vs silver.
#       Log deltas to GOLD.gold_reconciliation(run_id, metric, silver_value, gold_value, delta)
#       with a stated tolerance. A non-zero unexplained delta is a bug, not a rounding note.

# --- 8. Idempotency proof (run this AFTER the build works) --------------------
# TODO: loop the whole build 10x; assert identical fact.count() AND identical pct_passing each time.
#       Capture the 10 identical numbers in your why-note.

# Week 18 - MDM I: reference data, crosswalks & match keys (STARTER SCAFFOLD)
# BUSINESS FRAMING: duplicate provider records make Maria's cost + compliance
# rollups double-count. This week we build governed reference data and the
# candidate pairs W19 will score. We do NOT score or merge here.

from pyspark.sql import functions as F

SILVER_NPPES   = "silver.nppes_provider"        # landed+typed from earlier weeks
SILVER_PRICING = "silver.pricing_facility"      # pricing facility rows
DIM_REF        = "silver.dim_provider_ref"
XWALK          = "silver.xwalk_provider"
PAIRS          = "silver.candidate_pairs"
REJECTED       = "silver.rejected_records"
AUDIT          = "silver.audit_log"

# ---------------------------------------------------------------
# 1) STANDARDIZE match attributes into the reference dimension
# ---------------------------------------------------------------
src = spark.table(SILVER_NPPES)

# TODO: casefold + trim + strip punctuation on name/address columns.
#       Hint: F.upper(F.trim(F.regexp_replace(col, r"[^A-Za-z0-9 ]", "")))
# TODO: normalize USPS street suffixes (ST->STREET, AVE->AVENUE, ...).
#       Build a small mapping and apply it; do NOT hardcode one example.
# TODO: derive zip5 (validate/zero-pad) and keep state.
# TODO: split entity_type_code: 1 = individual (last/first), 2 = organization
#       (org name) -> a single std_name column chosen by entity type.
ref = (src
    .withColumn("npi", F.trim(F.col("npi")))
    # .withColumn("std_name", ...)   # TODO
    # .withColumn("std_addr", ...)   # TODO
    # .withColumn("zip5",     ...)   # TODO
)

# ---------------------------------------------------------------
# 2) IDEMPOTENT upsert of the reference dimension (MERGE on NPI)
# ---------------------------------------------------------------
# TODO: write `ref` to DIM_REF the first run (Delta, partitionBy state),
#       then convert to a MERGE keyed on NPI so reruns update in place.
#       WHY note: NPI is a stable natural key -> MERGE, not append.

# ---------------------------------------------------------------
# 3) CROSSWALK skeleton (W19 fills match_status)
# ---------------------------------------------------------------
# TODO: build xwalk_provider = one row per pricing facility source record:
#       (source_system, source_record_id, npi_or_null, facility_key,
#        match_status='unmatched'). Most pricing rows have NO clean NPI.

# ---------------------------------------------------------------
# 4) BLOCKING KEYS (>=2 deterministic passes)
# ---------------------------------------------------------------
# TODO: block_key_1 = soundex(std_name) concat state
# TODO: block_key_2 = zip5 concat substring(std_name,1,3)
# blocked = ref.withColumn("block_key_1", ...).withColumn("block_key_2", ...)

# ---------------------------------------------------------------
# 5) QUARANTINE un-blockable records (loud, not silent)
# ---------------------------------------------------------------
# TODO: rows with no name AND no zip -> REJECTED with
#       reason='no_usable_block_key', ingest_ts, source. Count + audit.

# ---------------------------------------------------------------
# 6) CANDIDATE PAIRS within blocks (NO cartesian) + skew guard
# ---------------------------------------------------------------
def pairs_for(df, key):
    # TODO: compute block sizes; flag/cap blocks over a threshold
    #       (e.g. >500) as 'oversized_block' before joining.
    a = df.select(F.col("npi").alias("id_a"), F.col(key).alias("bk"))
    b = df.select(F.col("npi").alias("id_b"), F.col(key).alias("bk"))
    return (a.join(b, "bk")
              .where(F.col("id_a") < F.col("id_b"))   # kill symmetric dups
              .withColumn("block_key", F.lit(key)))
# TODO: union pairs_for(blocked,"block_key_1") and block_key_2, dedup, write PAIRS.

# ---------------------------------------------------------------
# 7) AUDIT + reduction ratio, then rerun 10x to prove idempotency
# ---------------------------------------------------------------
# TODO: log records_standardized, records_quarantined, n_blocks,
#       oversized_blocks, candidate_pairs, reduction_ratio to AUDIT.
# TODO: run this notebook 10x; assert DIM_REF count and PAIRS count are
#       identical every run. Paste the proof into your why-note.

# silver_to_golden_entity_resolution.py
# Databricks Free Edition + Delta + PySpark. Run as a notebook.
# W18 produced SILVER provider records with a standardized name/address + a blocking_key.
# This week: pairs -> similarity features -> confidence -> cluster -> golden + review queue.

from pyspark.sql import functions as F, Window

SILVER = "workspace.default.silver_providers"        # standardized + blocking_key from W18
GOLDEN = "workspace.default.golden_providers"
LINEAGE = "workspace.default.golden_lineage"
REVIEW = "workspace.default.provider_review_queue"

silver = spark.table(SILVER)  # cols: src_key, npi, name_std, addr_std, state, postal, ingest_ts, source

# --- 1) Candidate PAIRS within each block (NO full cross join) ---
l = silver.alias("l"); r = silver.alias("r")
pairs = (l.join(r, (F.col("l.blocking_key") == F.col("r.blocking_key")) &
                  (F.col("l.src_key") < F.col("r.src_key")))   # dedupe + no self-pair
         .select(F.col("l.src_key").alias("key_left"), F.col("r.src_key").alias("key_right"),
                 F.col("l.name_std").alias("name_l"), F.col("r.name_std").alias("name_r"),
                 F.col("l.addr_std").alias("addr_l"), F.col("r.addr_std").alias("addr_r"),
                 F.col("l.state").alias("state_l"), F.col("r.state").alias("state_r"),
                 F.col("l.postal").alias("postal_l"), F.col("r.postal").alias("postal_r")))
# TODO: log pairs.count() and the max block size; if one block explodes, tighten the W18 blocking_key.

# --- 2) Similarity FEATURES (Spark built-ins only) ---
def norm_lev(a, b):
    # 1 - (edit_distance / max_len)  -> in [0,1]; higher = more similar
    maxlen = F.greatest(F.length(a), F.length(b), F.lit(1))
    return F.lit(1.0) - (F.levenshtein(a, b) / maxlen)

feat = (pairs
        .withColumn("f_name", norm_lev(F.col("name_l"), F.col("name_r")))
        .withColumn("f_addr", norm_lev(F.col("addr_l"), F.col("addr_r")))
        .withColumn("f_soundex", (F.soundex(F.col("name_l")) == F.soundex(F.col("name_r"))).cast("double"))
        .withColumn("f_state", (F.col("state_l") == F.col("state_r")).cast("double")))
# TODO: add a postal-prefix feature and, if you like, a token-overlap feature for org names.

# --- 3) Combine into ONE confidence (defend your weights) ---
# TODO: replace this naive equal weighting after you hand-label ~30 pairs.
conf = feat.withColumn("match_confidence",
        (F.col("f_name")*0.4 + F.col("f_addr")*0.3 + F.col("f_soundex")*0.2 + F.col("f_state")*0.1))

# --- 4) DECISION bands ---
AUTO, REVIEW_LO = 0.90, 0.65   # TODO: justify these from your labeled sample (precision vs recall).
accepted = conf.filter(F.col("match_confidence") >= AUTO)
to_review = conf.filter((F.col("match_confidence") >= REVIEW_LO) & (F.col("match_confidence") < AUTO))

# --- 5) CLUSTER accepted pairs into entity groups (transitive closure) ---
# TODO: implement iterative label propagation: each node's label = min(own, neighbors');
#       iterate until labels stop changing. Then golden_provider_id = a DETERMINISTIC hash of
#       the sorted src_keys in the cluster (F.sha2), NOT monotonically_increasing_id().

# --- 6) SURVIVORSHIP -> one golden row per cluster ---
# TODO: per golden_provider_id, choose fields by rule:
#       address/phone = most-recent ingest_ts; name = most-complete (longest non-null);
#       npi = source-priority. Write GOLDEN, and write LINEAGE (golden_id -> src_key -> field_won).

# --- 7) REVIEW QUEUE (governance artifact) ---
review_rows = (to_review
    .withColumn("reason", F.lit("confidence_in_review_band"))
    .withColumn("status", F.lit("pending"))
    .withColumn("queued_ts", F.current_timestamp()))
# TODO: MERGE review_rows into REVIEW on (key_left, key_right) so reruns don't duplicate the queue.

# --- 8) IDEMPOTENCY ---
# TODO: MERGE golden + lineage on their deterministic keys; then run this notebook 10x and assert
#       golden count, golden_provider_id set, and REVIEW count are identical every run.

# --- KPI line for the Charter ---
# TODO: print { collapsed = input_rows - golden_rows, golden_rows, review_rows, pct_change_vs_last }

# Databricks Free Edition notebook — Week 20 STARTER SCAFFOLD
# Purpose: ground your data-flow map in REAL tables + row counts from W15-W19,
#          not a fictional diagram. This only INVENTORIES; you write the doc.
#
# It does NOT decide ELT/incremental for you and does NOT draw the map.
# Fill every TODO. The prose (ADRs, failure plan, idempotency plan) is yours.

from pyspark.sql import functions as F

# TODO 1: point these at YOUR schema/catalog from earlier weeks.
#         On Free Edition this is likely the default 'workspace' / hive_metastore.
SCHEMA = "provider_pricing"          # TODO: your schema name
LAYER_PREFIXES = {                    # TODO: adjust to how you actually named tables
    "bronze": "bronze_",
    "silver": "silver_",
    "gold":   "gold_",
    "reject": "rejected_",           # quarantine tables from W16-W19
}

# --- Inventory every table and tag it with its intended layer ------------------
rows = []
tables = spark.catalog.listTables(SCHEMA)
for t in tables:
    fqn = f"{SCHEMA}.{t.name}"
    layer = next((L for L, p in LAYER_PREFIXES.items() if t.name.startswith(p)), "UNCLASSIFIED")
    try:
        cnt = spark.table(fqn).count()
    except Exception as e:
        cnt = -1  # e.g., view or permission issue — note it in the doc
    # TODO 2: also capture partition columns so your idempotency plan can
    #         justify 'replace-partition' where you use it.
    part_cols = None  # TODO: read via DESCRIBE DETAIL <fqn> -> partitionColumns
    rows.append((layer, t.name, cnt, part_cols))

inv = spark.createDataFrame(rows, ["layer", "table", "row_count", "partition_cols"]) \
           .orderBy("layer", "table")

print("=== TABLE INVENTORY (paste into ARCHITECTURE.md data-flow map) ===")
inv.show(truncate=False)

# TODO 3: For EACH layer, write the CONTRACT you will defend in the doc:
#   bronze = raw landed, append-only, source + ingest_ts, NO transforms
#   silver = validated + typed + conformed + SCD2 historized + quarantine
#   gold   = business-ready, serve-only (KPIs for Maria), NO history
# Then write the transform BOUNDARY crossed between each pair.

# TODO 4: Flag any UNCLASSIFIED table — an unowned table is a contract gap
#         and a great thing to name in your 'why' note.
unclassified = inv.filter(F.col("layer") == "UNCLASSIFIED")
print("UNCLASSIFIED tables (resolve or justify in the doc):")
unclassified.show(truncate=False)

# NOTE: This scaffold produces EVIDENCE. The four architecture-doc parts
#       (data-flow map, tool-selection log, failure plan, idempotency plan)
#       are written by YOU in ARCHITECTURE.md — that's the deliverable.

# Week 21 - Bronze ingestion scaffold (Databricks Free Edition + Delta)
# GOAL: land NPPES API + one CMS MRF file into Bronze, append-only, with
# source + ingest metadata, LOUD failures, and idempotent-per-run writes.
# CONTRACT: Bronze = raw landed, NO transforms. Keep the ugly payload.

import time, json, uuid, random, requests
from datetime import datetime, timezone
from pyspark.sql import Row

# ---- run identity (lineage) shared by every write this run ----
BATCH_ID = str(uuid.uuid4())
INGEST_DATE = datetime.now(timezone.utc).strftime("%Y-%m-%d")
CATALOG = "workspace"  # Free Edition default; adjust to your workspace
BRONZE_PROVIDERS = f"{CATALOG}.default.bronze_nppes_providers"
BRONZE_MRF       = f"{CATALOG}.default.bronze_cms_mrf"
AUDIT            = f"{CATALOG}.default.ingest_audit_log"
NPPES_URL = "https://npiregistry.cms.hhs.gov/api/"

def now_utc():
    return datetime.now(timezone.utc).isoformat()

def log_audit(source, status, rows, reason=""):
    """One queryable row per source per run outcome. This is your alerting."""
    row = [Row(batch_id=BATCH_ID, source=source, status=status,
               row_count=int(rows), reason=reason[:500],
               ingest_ts=now_utc(), ingest_date=INGEST_DATE)]
    (spark.createDataFrame(row)
          .write.format("delta").mode("append").saveAsTable(AUDIT))

# -----------------------------------------------------------------
# 1) NPPES API pull  --  handle the 429 (G6)
# -----------------------------------------------------------------
def fetch_nppes_page(params, max_retries=5):
    """Return the raw 'results' list for ONE page.
    TODO(G6): loop with an attempt counter:
      - resp = requests.get(NPPES_URL, params=params, timeout=30)
      - if 200: return resp.json().get('results', [])
      - if 429: sleep for int(Retry-After) if header present,
                else backoff = (2 ** attempt) + random.random(); then retry
      - if 5xx: retryable -> backoff + retry
      - if other 4xx: PERMANENT -> raise (do not swallow)
      - if attempts exhausted: raise (caller logs FAILED and the job goes red)
    """
    raise NotImplementedError("Implement 429/backoff + return results list")

def ingest_nppes(states=("CA",), pages=5, limit=200):
    try:
        records = []
        for st in states:
            for i in range(pages):
                params = {"version": "2.1", "state": st,
                          "limit": limit, "skip": i * limit}
                page = fetch_nppes_page(params)
                if not page:
                    break
                # TODO: tag each raw record with its source state, keep it RAW
                records.extend(page)
        if not records:
            log_audit("nppes_api", "EMPTY", 0, "no results returned")
            return
        # TODO(Bronze contract): build rows with ONE raw payload column +
        #   metadata. Example shape (fill in):
        #   rows = [Row(_payload=json.dumps(r), _source="nppes_api",
        #               _source_file=None, _ingest_ts=now_utc(),
        #               _batch_id=BATCH_ID, _ingest_date=INGEST_DATE)
        #           for r in records]
        #   df = spark.createDataFrame(rows)
        df = None  # TODO build df
        # TODO(idempotency): append across days, overwrite TODAY's partition:
        #   (df.write.format("delta").mode("overwrite")
        #      .option("replaceWhere", f"_ingest_date = '{INGEST_DATE}'")
        #      .partitionBy("_ingest_date").saveAsTable(BRONZE_PROVIDERS))
        log_audit("nppes_api", "SUCCESS", df.count() if df else 0)
    except Exception as e:
        log_audit("nppes_api", "FAILED", 0, reason=str(e))
        raise  # FAIL LOUD

# -----------------------------------------------------------------
# 2) CMS MRF file  --  detect the file that never showed up (G6)
# -----------------------------------------------------------------
def ingest_mrf(landing_path):
    """Land one hospital MRF file already uploaded to a Volume/DBFS path.
    TODO(silent-skip): BEFORE reading, verify the file exists (dbutils.fs.ls).
       If absent:
         log_audit('cms_mrf','MISSING_SOURCE',0,f'expected {landing_path}')
         raise FileNotFoundError(landing_path)   # 'no data' != 'clean run'
    TODO(Bronze): read RAW (text/csv/json per the file) with NO cleaning;
       add _source='cms_mrf', _source_file=landing_path, _ingest_ts,
       _batch_id, _ingest_date; write with the SAME replaceWhere pattern;
       then log SUCCESS with the row count.
    """
    raise NotImplementedError

# ---- entrypoint (uncomment as you implement) ----
# ingest_nppes(states=("CA","NY"), pages=5)
# ingest_mrf("/Volumes/workspace/default/landing/hospital_mrf.json")
# display(spark.table(AUDIT).orderBy("ingest_ts"))

# W22 · Provider Pricing Integrity · Medallion layer-contract enforcement
# Databricks Free Edition | Delta Lake | PySpark. This is a STARTER scaffold.
# Fill every TODO. Do not paste a full solution here.

from pyspark.sql import functions as F, types as T

# --- CONFIG ---------------------------------------------------------------
BRONZE  = "workspace.pricing.bronze_hospital_charges"   # landed in W21 (raw, append-only)
SILVER  = "workspace.pricing.silver_hospital_charges"
REJECT  = "workspace.pricing.silver_rejects"
RUN_TS  = F.current_timestamp()
REQUIRED_BRONZE_META = ["_source_file", "_ingest_ts", "_source_system"]

bronze = spark.read.table(BRONZE)

# === CONTRACT 1: BRONZE IS RAW + APPEND-ONLY =============================
# TODO 1: assert every name in REQUIRED_BRONZE_META is in bronze.columns;
#         raise AssertionError naming the missing column (LOUD failure).
# TODO 2: assert payload/business columns are still StringType (no casting
#         leaked into bronze). Iterate bronze.schema and check .dataType.
# TODO 3: assert append-only: today's count >= the count you logged last run
#         (store last count in a small Delta audit table and compare).

# === CONTRACT 2: SILVER IS CLEAN + CONFORMED (typing half; historize=W23+) =
silver_typed = (
    bronze
    # TODO 4: cast standard_charge -> DECIMAL(12,2); strip '$', handle 'N/A'/ranges
    # TODO 5: trim + upper() provider name; state -> 2 letters; ZIP -> lpad 5
    # TODO 6: keep the ORIGINAL raw columns too, so audit can trace back to bronze
)

# TODO 7: DECLARE THE SILVER PRIMARY-KEY GRAIN in one line, e.g.
# PK = ["npi", "payer", "plan", "billing_code", "charge_type"]
# Document this same key in LAYER_CONTRACTS.md.
PK = []  # <-- fill in

# --- QUALITY GATE + QUARANTINE (no silent drops) -------------------------
gate = (
    F.col("npi").isNotNull()
    & F.col("standard_charge").isNotNull()
    & (F.col("standard_charge") > 0)
    # TODO 8: add NOT-NULL checks for each PK component + any range rules
)

passed = silver_typed.filter(gate)
rejected = (
    silver_typed.filter(~gate)
    .withColumn("_reject_reason", F.lit("TODO: name the SPECIFIC failing rule"))
    .withColumn("_reject_ts", RUN_TS)
)

# TODO 9: dedupe `passed` within batch on PK, keeping latest _ingest_ts.
# TODO 10: write `rejected` (append) to REJECT.
# TODO 11: write `passed` to SILVER with a DELIBERATE idempotency strategy
#          (MERGE on PK  OR  overwrite partition). Add a one-line comment
#          justifying the choice. Rerun 10x -> identical SILVER counts.
# TODO 12: reconcile -> assert passed.count() + rejected.count() == input count.

# === CONTRACT 3: GOLD IS SERVE-ONLY =====================================
# TODO 13: create GOLD as a VIEW over SILVER ONLY (no bronze ref, no cast/
#          trim/coalesce). If you need to clean it, it belongs in silver.
# TODO 14: write a contract test that FAILS if the gold definition text
#          references BRONZE or contains a cleaning function.

# Week 23 - CDC & Incremental I: watermark-based bronze->silver load
# Databricks Free Edition + Delta + PySpark. Run cells top to bottom.
from pyspark.sql import functions as F

CATALOG_PATH = "main.default"          # adjust to your schema
BRONZE = f"{CATALOG_PATH}.bronze_cms_pricing"
SILVER = f"{CATALOG_PATH}.silver_cms_pricing"
WATERMARK = f"{CATALOG_PATH}.etl_watermark"
SOURCE, LAYER = "cms_pricing", "silver"

# --- 1. Control / state table (durable high-water mark) -------------------
spark.sql(f"""
CREATE TABLE IF NOT EXISTS {WATERMARK} (
  source STRING, layer STRING, watermark_column STRING,
  watermark_value TIMESTAMP, run_ts TIMESTAMP, rows_processed LONG
) USING DELTA
""")
# Bootstrap one row so the FIRST run processes everything (epoch watermark).
if spark.table(WATERMARK).where(f"source='{SOURCE}' AND layer='{LAYER}'").count() == 0:
    spark.sql(f"""INSERT INTO {WATERMARK} VALUES
      ('{SOURCE}','{LAYER}','ingest_timestamp', TIMESTAMP'1970-01-01', NULL, 0)""")

def get_watermark():
    # TODO: return the current watermark_value for (SOURCE, LAYER) as a Python datetime.
    #       Read from WATERMARK, filter, .collect(). This is your resume point.
    raise NotImplementedError

def set_watermark(new_value, rows):
    # TODO: UPDATE the WATERMARK row for (SOURCE, LAYER):
    #       watermark_value=new_value, run_ts=current_timestamp(), rows_processed=rows.
    #       Call this ONLY on the success path (see step below).
    raise NotImplementedError

# --- 2. Land a CMS pricing 'period' to bronze (append-only) ---------------
def land_period(df_period, period_id, source_file):
    # Bronze contract: raw + provenance ONLY. No dedup, no typing here.
    (df_period
       .withColumn("source_file", F.lit(source_file))
       .withColumn("period_id", F.lit(period_id))
       .withColumn("ingest_timestamp", F.current_timestamp())
       .write.mode("append").saveAsTable(BRONZE))
# TODO: land 3 periods (2025-07/08/09). For 08 add a few new facilities and
#       change prices on a few existing rows; for 09 change a few more. This
#       simulates CMS republishing without waiting a month.

# --- 3. Incremental read: only rows newer than the watermark --------------
wm = get_watermark()                       # TODO returns datetime
bronze = spark.table(BRONZE)
incremental = bronze.where(F.col("ingest_timestamp") > F.lit(wm))
full_count, inc_count = bronze.count(), incremental.count()
print(f"full={full_count}  incremental={inc_count}  (saving={full_count-inc_count})")

# --- 4. Process ONLY the slice into silver, then advance on SUCCESS -------
try:
    if inc_count == 0:
        print("No new data - watermark stays put.")           # do NOT advance
    else:
        # TODO: apply your W22 silver conform/type logic to `incremental` only.
        silver_df = incremental  # <-- replace with real transform
        # NOTE: append double-counts on rerun. Intentional - W24 adds idempotent MERGE.
        silver_df.write.mode("append").saveAsTable(SILVER)
        # Advance watermark to the max ingest_timestamp we actually processed.
        new_wm = incremental.agg(F.max("ingest_timestamp")).first()[0]
        set_watermark(new_wm, inc_count)                       # success path ONLY
        print(f"Advanced watermark -> {new_wm} ({inc_count} rows)")
except Exception as e:
    # Failure path: DO NOT advance the watermark. Next run re-attempts the slice.
    print(f"Run failed, watermark unchanged: {e}")
    raise

# --- 5. TODO: prove it --------------------------------------------------
# (a) rerun with no new data -> assert inc_count == 0, watermark unchanged.
# (b) land period 2025-10 -> rerun -> assert only its rows flow, watermark advances.
# (c) land 1 late record (event date 2025-06, ingest now) -> show it IS processed;
#     comment why an event-time watermark on last_updated would have skipped it.
# Then read `spark.table(WATERMARK)` history as evidence for your 'why' note.

# Week 24 - Idempotent silver load via Delta MERGE + 10x proof harness
# STARTER SCAFFOLD: complete the TODOs. Do NOT ship this as-is.
from pyspark.sql import functions as F, Window
from delta.tables import DeltaTable

NATURAL_KEY  = "npi"                 # NPPES National Provider Identifier = the true natural key
BRONZE_TABLE = "bronze.nppes_providers"
SILVER_TABLE = "silver.dim_provider"
# A BACKFILL re-runs an OLD batch id. An incremental run uses the newest one.
BATCH_ID     = "2026-08"             # TODO: set to the CMS/NPPES month you are (re)processing
TRACKED_COLS = ["provider_name", "address_line1", "state", "taxonomy_code"]  # attrs that define a 'change'

def read_bronze_batch(spark, batch_id):
    df = spark.table(BRONZE_TABLE).where(F.col("source_batch") == batch_id)
    # MERGE requires <= 1 source row per key. Bronze can hold the same NPI twice in one file.
    # TODO: keep only the latest row per NATURAL_KEY using a window over ingest_ts (row_number == 1).
    return df  # TODO

def conform(df):
    # Must be DETERMINISTIC: same input -> same output every run, or the 10x test will flip.
    # TODO: type-cast, trim, upper-case state, standardize phone.
    # TODO: keep audit timestamps in METADATA columns only, and exclude them from TRACKED_COLS.
    return df  # TODO

def upsert_silver(spark, src):
    if not spark.catalog.tableExists(SILVER_TABLE):
        src.write.format("delta").saveAsTable(SILVER_TABLE)   # first load
        return
    tgt = DeltaTable.forName(spark, SILVER_TABLE)
    change_pred = " OR ".join([f"t.{c} <> s.{c}" for c in TRACKED_COLS])
    (tgt.alias("t")
        .merge(src.alias("s"), f"t.{NATURAL_KEY} = s.{NATURAL_KEY} AND t.is_current = true")
        # TODO: whenMatchedUpdate(condition=change_pred, ...) to CLOSE the old current row (SCD2 step 1)
        # TODO: whenNotMatchedInsertAll() to add brand-new providers
        # NOTE: the 'insert new version of a changed provider' step (SCD2 step 2) is a second write -
        #       decide where it lives and why (write it in your why-note).
        .execute())

def run_load(spark, batch_id):
    upsert_silver(spark, conform(read_bronze_batch(spark, batch_id)))

def prove_idempotent(spark, batch_id, runs=10):
    seen = []
    for i in range(runs):
        run_load(spark, batch_id)
        m = (spark.table(SILVER_TABLE)
               .agg(F.count("*").alias("rows"),
                    F.countDistinct(NATURAL_KEY).alias("keys"),
                    F.sum(F.col("is_current").cast("int")).alias("current_rows"))
               .collect()[0])
        seen.append((m["rows"], m["keys"], m["current_rows"]))
        print(f"run {i+1}: rows={m['rows']} keys={m['keys']} current={m['current_rows']}")
    # TODO: assert every tuple in `seen` is identical to seen[0] -> FAIL LOUD with the run number.
    # TODO: assert current_rows == distinct keys (exactly one current row per provider = no dup collapse).
    return seen

# silver_transform_spark.py  (Databricks Free Edition notebook)
# STARTER SCAFFOLD — lands structure + proves lazy eval. You fill the TODOs.
# Goal: re-express the Week 19-24 silver transforms in the PySpark DataFrame API.

from pyspark.sql import functions as F
from datetime import datetime

BRONZE_PROVIDERS = "workspace.bronze.nppes_providers"   # <- adjust to your catalog/schema
SILVER_PROVIDERS = "workspace.silver.providers"
REJECTED_TABLE   = "workspace.silver.providers_rejected"

# 1) READ bronze -> DataFrame (this is a transformation-ish setup; no job yet)
bronze = spark.read.format("delta").table(BRONZE_PROVIDERS)
bronze.printSchema()   # printSchema reads metadata only

# 2) BUILD the transform chain — ALL transformations, still lazy.
#    Map your prior SQL/pandas silver logic onto DataFrame calls.
conformed = (
    bronze
    # TODO: select + withColumnRenamed to the silver CONTRACT columns only
    # .select(...)
    # TODO: cast types — NPI to long/string(10), price to double, dates to date
    # .withColumn("npi", F.col("npi").cast("string"))
    # TODO: conform enum-like field with when/otherwise
    # .withColumn("entity_type",
    #     F.when(F.col("entity_type_code") == 1, "Individual")
    #      .when(F.col("entity_type_code") == 2, "Organization")
    #      .otherwise("Unknown"))
    # TODO: null-handling for missing/blank prices (na.fill or F.coalesce)
)

# 3) PROVE laziness: nothing above ran. Time the first ACTION.
t0 = datetime.now();  print("chain built, no job yet:", t0)
n_in = conformed.count()          # <-- ACTION #1 (triggers the scan)
t1 = datetime.now();  print("after count():", t1, "rows:", n_in, "elapsed:", t1 - t0)
# TODO: open the Spark UI -> Jobs. How many jobs fired so far? Match your prediction.

# 4) QUALITY GATE — define the predicate ONCE, split into valid/rejected.
# TODO: replace with your real rule set (10-digit NPI, price non-null & >= 0, ...)
is_valid = (F.length("npi") == 10) & F.col("price").isNotNull() & (F.col("price") >= 0)

df_valid = conformed.filter(is_valid)
df_rejected = (
    conformed.filter(~is_valid)
    .withColumn("rejection_reason", F.lit("TODO: derive specific reason per rule"))
    .withColumn("validation_ts", F.current_timestamp())
)

# TODO: assert conservation — valid + rejected must equal input (delta == 0)
# assert df_valid.count() + df_rejected.count() == n_in

# 5) READ THE PLAN (do NOT collect() the data to the driver!)
df_valid.explain(True)   # save this output to /architecture/spark_plans_w25.txt

# 6) WRITE — apply YOUR W24 idempotency decision (overwrite partition / MERGE).
# TODO: write df_valid to SILVER_PROVIDERS (idempotent mode)
# TODO: append df_rejected to REJECTED_TABLE

# 7) KPIs for Maria
# TODO: pct_complete = df_valid.count() / n_in; print quarantined count by reason

from pyspark.sql import functions as F
from pyspark.sql.functions import broadcast

# --- config: point these at YOUR silver tables from W21-W25 ---
SILVER_PRICING   = "workspace.silver.cms_pricing"
SILVER_PROVIDERS = "workspace.silver.nppes_providers"
GOLD_KPI         = "workspace.gold.pricing_kpis"
REJECTED         = "workspace.silver.rejected_pricing"

pricing   = spark.table(SILVER_PRICING)
providers = spark.table(SILVER_PROVIDERS)

# 1) PROFILE THE JOIN PROBLEM (evidence for your why-note)
# TODO: print pricing.count(), providers.count(),
#       null rate of the NPI key in pricing, and # distinct facility names with no NPI.

# 2) MATCHING UDF the MDM step (W29) will REUSE.
#    NOTE: this runs row-at-a-time on the EXECUTORS and is the deliberate
#    exception to 'prefer a built-in'. Keep it pure + deterministic.
from pyspark.sql.types import StringType

def _normalize_key(name, address):
    # TODO: lowercase; strip suffixes (inc, llc, pc); collapse whitespace;
    #       standardize address abbreviations (st->street, ste->suite).
    #       Return a stable blocking/match key string.
    return None

normalize_key = F.udf(_normalize_key, StringType())
# TODO: also register for SQL reuse: spark.udf.register("normalize_key", _normalize_key)

# 3) EXACT-KEY JOIN + QUARANTINE
# TODO: choose the join type deliberately (inner vs left-outer) and join on NPI.
# matched = pricing.join(providers, on="npi", how="???")
# TODO: use a LEFT-ANTI join to capture unmatched pricing rows ->
#       add reason="no_npi_match", source_file, quarantined_at=current_timestamp(),
#       then append to REJECTED. (Do NOT .collect() to check membership.)

# 4) BROADCAST DECISION
# TODO: identify which side is small enough to broadcast (a reference/lookup
#       or a filtered provider slice) and wrap it in broadcast(...).
#       In one comment, state what ships to executors vs stays on driver.

# 5) GOLD KPIs (groupBy/agg + one window function)
# TODO: pricing-completeness % per facility; quarantined count by reason;
#       freshness minutes from max ingest ts; row_number() over provider
#       partitions ordered by completeness (MDM preview).

# 6) IDEMPOTENT WRITE + PROOF
# TODO: write GOLD_KPI deliberately (overwrite snapshot OR MERGE on kpi key),
#       then run this build 10x and assert identical counts + KPI values.
#       Append a row to your audit log (run ts, input counts, match rate, quarantined).

# W27 - Spark skew: diagnose in the UI, fix with salt/broadcast, prove it stayed correct.
# STARTER SCAFFOLD. Fill every TODO. Do NOT paste in a full solution - the point is to MEASURE.

from pyspark.sql import functions as F

SILVER = "workspace.default.silver_providers"  # TODO: point at YOUR W25/W26 silver provider table
providers = spark.read.table(SILVER)

# --- 1. Rebuild the SAME blocking key you used in W26 (keep it identical so counts match) --------
def with_block_key(df):
    # TODO: recreate your real W26 blocking key (e.g. org-name token + state). Use YOUR columns.
    return df.withColumn(
        "block_key",
        F.concat_ws("|", F.upper(F.col("org_name_token")), F.col("state"))
    )

blocked = with_block_key(providers)

# --- 2. MEASURE FIRST: how skewed is the key? Name the health systems that dominate ------------
# TODO: show top 20 block_key by count; screenshot it; write down which orgs (HCA/Kaiser/...) are hot.
# blocked.groupBy("block_key").count().orderBy(F.desc("count")).show(20, truncate=False)

# --- 3. BASELINE the candidate-pair self-join with AQE OFF so the skew is visible in the UI -----
# spark.conf.set("spark.sql.adaptive.enabled", False)           # TODO: enable to expose raw skew
# spark.conf.set("spark.sql.shuffle.partitions", 200)           # TODO: record whatever value you use
l = blocked.alias("l")
r = blocked.alias("r")
candidate_pairs = (l.join(r, on="block_key")
                    .where(F.col("l.npi") < F.col("r.npi")))    # keep each unordered pair once
# TODO: force + time it (candidate_pairs.count()). In the Spark UI find the join stage and record:
#       max task time, median task time, largest-task shuffle read. Screenshot the event timeline.
baseline_count = None  # TODO: capture this - you will assert against it later.

# --- 4. FIX A: SALT the hot key --------------------------------------------------------------
SALT_N = 8  # TODO: choose N deliberately; justify with max-task time. Try a few, keep the evidence.
# TODO: left side gets a random salt 0..N-1; right side is REPLICATED across all N salts so salted
#       keys still meet. Join on (block_key, salt). Re-run, re-time, re-read the UI vs baseline.

# --- 5. FIX B: BROADCAST the small side ------------------------------------------------------
# TODO: pick a genuinely small side (CMS hospital list, or a filtered subset under the broadcast
#       threshold), wrap it F.broadcast(small_df), re-run, re-time. Note: this KILLS the shuffle.

# --- 6. AQE ON: does Databricks' automatic skew-join splitting alone fix it? ------------------
# spark.conf.set("spark.sql.adaptive.enabled", True)
# TODO: re-run the untouched baseline with AQE on; record result. Decide AQE-only vs salt vs
#       broadcast and WRITE DOWN why in your one-pager.

# --- 7. CORRECTNESS GATE --------------------------------------------------------------------
# TODO: assert your fixed join produces the SAME count as baseline_count.
#       assert fixed_count == baseline_count, f"fix changed output: {fixed_count} != {baseline_count}"

# silver_dq_gate.py — Silver Data Quality gate (Databricks Free Edition + Delta + PySpark)
# Run per CMS republish. Flag -> split -> quarantine -> summarize -> alert. Idempotent by run_id.

from pyspark.sql import functions as F

# ---- 0. Run context ---------------------------------------------------------
# run_id makes writes idempotent: a rerun REPLACES this run's rejects/summary,
# it does not append. TODO: derive run_id from the CMS file version / ingest date.
RUN_ID = "TODO_run_id"            # e.g. '2026-09-cms-v1'
SOURCE_FILE = "TODO_source_path"

# silver_conformed = your typed/conformed silver DF from prior weeks (+ your bad-row deck UNIONed in)
silver_conformed = spark.read.table("silver_provider_pricing_conformed")  # TODO: your table
rows_in = silver_conformed.count()

# ---- 1. Expectations catalog ------------------------------------------------
# Each rule: id, class, severity ('warn'|'fail'), a PASS expression (True = good),
# a per-rule quarantine BUDGET (fraction), and a one-line business reason.
# TODO: write 8+ SPECIFIC rules. Two examples are stubbed — replace/extend them.
EXPECTATIONS = [
    {"id": "npi_len10", "cls": "schema", "sev": "fail", "budget": 0.02,
     "expr": F.col("npi").rlike("^[0-9]{10}$"),
     "why": "A non-10-digit NPI cannot join to NPPES -> provider KPI is wrong."},
    {"id": "price_nonneg_numeric", "cls": "domain", "sev": "fail", "budget": 0.05,
     "expr": (F.col("negotiated_price").cast("double").isNotNull()) & (F.col("negotiated_price").cast("double") >= 0),
     "why": "Negative / non-numeric prices corrupt the compliance cost report."},
    # TODO: npi_not_null (completeness), (npi,cpt) uniqueness (grain — needs a window/groupBy),
    #       cpt_in_accepted_set (domain), payer_name_not_blank (completeness), gross_charge rules...
]
ACCEPTED_CPT = ["TODO", "fill", "your", "accepted", "code", "set"]  # TODO domain list

# ---- 2. Flag pass (no filtering — every input row kept) ----------------------
df = silver_conformed
for rule in EXPECTATIONS:
    df = df.withColumn(f"dq__{rule['id']}", rule["expr"])   # True = passed

# Collect failed rule ids into an array + validity flag.
# TODO: build dq_failed_rules as an array of the ids where dq__<id> is False.
#   hint: F.array(*[F.when(~F.col(f"dq__{r['id']}"), F.lit(r['id'])) for r in EXPECTATIONS])
#          then F.array_except(..., F.array(F.lit(None))) to drop nulls.
df = df.withColumn("dq_failed_rules", F.array())            # TODO replace
df = df.withColumn("dq_is_valid", F.size("dq_failed_rules") == 0)
assert df.count() == rows_in, "flag pass must not drop rows"

# ---- 3. Split ---------------------------------------------------------------
clean = df.filter("dq_is_valid")
rejected = df.filter("NOT dq_is_valid")
assert clean.count() + rejected.count() == rows_in, "split must conserve rows"

# ---- 4. Quarantine payload + write (idempotent on run_id) -------------------
rejected_out = (rejected
    .withColumn("dq_failure_reason", F.col("dq_failed_rules"))
    .withColumn("rejected_at", F.current_timestamp())
    .withColumn("source_file", F.lit(SOURCE_FILE))
    .withColumn("run_id", F.lit(RUN_ID)))
# TODO: write to Delta silver_rejected so a RERUN of this run_id REPLACES, not appends.
#   option A: rejected_out.write.format('delta').mode('overwrite')
#             .option('replaceWhere', f"run_id = '{RUN_ID}'").saveAsTable('silver_rejected')
#   option B: MERGE on a natural key + run_id. Pick one; justify in the why-note.

# TODO: write `clean` forward to your silver serving table for this run.

# ---- 5. DQ run summary (feeds KPI: % passing pricing-completeness) -----------
# TODO: build one summary row per rule: run_id, rule_id, rows_in, rows_failed, pass_rate.
#       write to Delta dq_run_summary keyed on (run_id, rule_id), idempotent.

# ---- 6. Threshold alert -----------------------------------------------------
# TODO: for each rule where sev=='fail', compute quarantine_rate = rows_failed/rows_in.
#       if quarantine_rate > rule['budget']: append a dq_alerts row AND print a LOUD message.
#       Decide: does a breach halt the run (raise) or continue with alert logged? Write it down.

# audit_reconcile_scaffold.py
# Week 29 - Reconciliation & Audit scaffold for the Provider Pricing Integrity pipeline.
# Databricks Free Edition + Delta + PySpark. This is a SCAFFOLD: it sets up the run_id,
# the audit_log table, and the reconciliation SHAPE. You fill the TODOs with your own
# counts, control totals, tolerance calls, and failure handling. Do NOT expect this to run
# a full pipeline as-is.

import uuid, time
from datetime import datetime
from pyspark.sql import functions as F

AUDIT_TABLE = "pricing.audit_log"   # TODO: point at your schema

# --- 1. One run_id for the whole run, threaded through every layer -------------
RUN_ID = datetime.utcnow().strftime("%Y%m%dT%H%M%S") + "-" + uuid.uuid4().hex[:6]
print(f"RUN_ID = {RUN_ID}")

# --- 2. Create the ledger once (idempotent DDL). Grain: one row per (run_id, layer) ---
spark.sql(f"""
CREATE TABLE IF NOT EXISTS {AUDIT_TABLE} (
  run_id STRING, layer STRING, source STRING,
  rows_in LONG, rows_out LONG, rows_rejected LONG,
  watermark_value STRING, control_total_num DECIMAL(18,2), distinct_key_count LONG,
  duration_seconds DOUBLE, status STRING, error_msg STRING, run_ts TIMESTAMP
) USING DELTA
""")

# --- 3. Idempotent audit writer: MERGE on (run_id, layer), never append blindly ---
def write_audit_row(row: dict):
    src = spark.createDataFrame([row])
    src.createOrReplaceTempView("_audit_src")
    # TODO: complete the MERGE so re-running the same (run_id, layer) UPDATES in place.
    spark.sql(f"""
    MERGE INTO {AUDIT_TABLE} t
    USING _audit_src s
    ON t.run_id = s.run_id AND t.layer = s.layer
    WHEN MATCHED THEN UPDATE SET *
    WHEN NOT MATCHED THEN INSERT *
    """)  # <- confirm this leaves exactly ONE row per (run_id, layer) after 2 runs

# --- 4. Instrument one layer. Call this pattern for bronze, silver, gold ----------
def instrument_layer(layer, source, df_in, df_out, df_rejected, watermark):
    t0 = time.time()
    rows_in       = df_in.count()
    rows_out      = df_out.count()
    rows_rejected = df_rejected.count() if df_rejected is not None else 0
    # TODO: control total(s). Clean currency BEFORE summing or bronze will never balance.
    control_total = df_out.agg(F.sum("standard_charge")).first()[0]  # TODO: cast/clean
    distinct_keys = df_out.agg(F.countDistinct("npi")).first()[0]
    write_audit_row({
        "run_id": RUN_ID, "layer": layer, "source": source,
        "rows_in": rows_in, "rows_out": rows_out, "rows_rejected": rows_rejected,
        "watermark_value": str(watermark),
        "control_total_num": control_total, "distinct_key_count": distinct_keys,
        "duration_seconds": round(time.time() - t0, 2),
        "status": "PENDING", "error_msg": None, "run_ts": datetime.utcnow(),
    })
    return rows_in, rows_out, rows_rejected

# --- 5. The completeness identity check. Counts MUST balance to 0 -----------------
def assert_completeness(rows_in, rows_out, rows_rejected, rows_dropped_by_dedup):
    # TODO: the identity. Any nonzero residual is a FAILED run, not a rounding note.
    residual = rows_in - (rows_out + rows_rejected + rows_dropped_by_dedup)
    if residual != 0:
        # TODO: write a FAIL audit row + raise loudly so gold is NOT published.
        raise ValueError(f"RECONCILE FAIL: {residual} rows unaccounted for in run {RUN_ID}")
    return True

# --- 6. Control-total tolerance. Counts tol=0; money tol is YOUR justified number ---
def check_control_total(bronze_total, silver_total, money_tolerance):
    # TODO: variance = bronze_total - silver_total; the only allowed gap is charges you
    #       quarantined for being null/uncleanable. Quantify it, compare, PASS/FAIL.
    pass

# TODO (lab): wire bronze->silver->gold through instrument_layer(), run assert_completeness()
#   and check_control_total(), flip run status, then run the whole thing 3-10x and prove
#   the audit_log holds one row per (run_id, layer) with identical numbers every time.

# test_pipeline.py — four-tier test suite for Provider Pricing Integrity (Project 1)
# Assumes a conftest.py with:
#   @pytest.fixture(scope="session")
#   def spark():  # local[2] + Delta configs on your machine; the notebook `spark` on Databricks
#       ...
# and importable code: from transforms import parse_price, normalize_npi, tag_quarantine
# plus your pipeline entrypoints: run_bronze_to_silver(spark, src), run_full_pipeline(spark, src)

import pytest

# ----------------------------------------------------------------------
# TIER 1 — UNIT: pure functions, no Spark table, run in milliseconds
# ----------------------------------------------------------------------
@pytest.mark.parametrize("raw,expected", [
    ("$1,234.00", 1234.00),
    ("1234",      1234.00),
    # TODO: add a null/'' case and assert it is flagged (not silently 0.0)
])
def test_parse_price(raw, expected):
    from transforms import parse_price
    assert parse_price(raw) == expected

def test_normalize_npi_strips_and_cases():
    from transforms import normalize_npi
    # TODO: assert '  1234567890 ' -> '1234567890' (whitespace/dtype normalized)
    ...

# ----------------------------------------------------------------------
# TIER 2 — INTEGRATION: bronze->silver on the frozen fixture (real Spark)
# ----------------------------------------------------------------------
def test_bad_row_is_quarantined_with_reason(spark):
    src = "tests/fixtures/bronze_sample"          # the ~50-row deterministic mess
    silver, rejected = None, None                 # TODO: run_bronze_to_silver(spark, src)
    # TODO: assert the null-price row is ABSENT from silver
    # TODO: assert it is PRESENT in rejected with reason=='MISSING_PRICE' and a non-null quarantine_ts
    ...

def test_scd2_closes_old_and_opens_new(spark):
    # TODO: ingest a provider, then ingest an updated version of the same NPI
    # TODO: assert exactly one row has is_current=true, and the prior row has valid_to set
    ...

# ----------------------------------------------------------------------
# TIER 3 — SMOKE: does it run, and is the OUTPUT CONTRACT intact?
# ----------------------------------------------------------------------
def test_gold_contract_and_pk(spark):
    gold = None                                   # TODO: run_full_pipeline(spark, "tests/fixtures/bronze_sample").gold
    # TODO: assert set(gold.columns) == EXPECTED_GOLD_COLUMNS (your written layer contract)
    # TODO: assert gold PK golden_provider_id is unique and non-null; assert gold is non-empty
    ...

# ----------------------------------------------------------------------
# TIER 4 — E2E: KPIs compute + IDEMPOTENCY (run twice, identical)
# ----------------------------------------------------------------------
def test_pipeline_is_idempotent(spark):
    src = "tests/fixtures/bronze_sample"
    first  = None                                 # TODO: run_full_pipeline(spark, src); capture counts
    second = None                                 # TODO: run_full_pipeline(spark, src) AGAIN (same src)
    # TODO: assert first.silver_count == second.silver_count
    # TODO: assert first.golden_record_count == second.golden_record_count  # dedupe didn't drift
    ...

def test_kpis_are_computed_and_sane(spark):
    kpis = None                                   # TODO: run_full_pipeline(spark, "tests/fixtures/bronze_sample").kpis
    # TODO: assert 0 <= kpis['pct_completeness_pass'] <= 100
    # TODO: assert kpis['duplicates_collapsed'] >= 1  # the fixture's two-spelling provider must merge
    ...