"""
Generate a Gaboro/UHC-style stress-test CSV with all 8 flag categories.

Visit dates are included so RENEWAL_CRITICAL / ELEVATED / SOON / VISIT_REQUIRED
all fire correctly. Run once to regenerate (shipment dates are absolute so
the distribution stays stable over time).

Usage:
    python3 test-data/generate_gaboro_stress.py [--rows N]  # default 5000
"""

import csv
import random
import sys
from datetime import date, timedelta
from pathlib import Path

random.seed(42)
TODAY = date(2026, 6, 18)   # frozen so distribution is reproducible
N = int(sys.argv[2]) if len(sys.argv) > 2 and sys.argv[1] == "--rows" else 5000
OUTPUT = Path(__file__).parent / "gaboro-stress-test.csv"

DEVICES = [
    ("dexcom_g7",         "sensor",  0.45),
    ("freestyle_libre_3", "sensor",  0.25),
    ("freestyle_libre_2", "sensor",  0.15),
    ("dexcom_g6",         "sensor",  0.10),
    ("omnipod_5",         "pod",     0.05),
]
PAYERS = [
    ("UnitedHealth",    0.40),
    ("Medicare Part B", 0.30),
    ("Aetna",           0.12),
    ("BCBS - FL",       0.08),
    ("Medicaid - GA",   0.06),
    ("Cigna",           0.04),
]

# How each scenario maps to columns
# Scenario weights: all 8 flag categories + balanced overall
SCENARIOS = [
    # (name, weight, shipment_offset_range, visit_offset, qty, swo, pecos, pa, dx)
    # shipment_offset: days before TODAY; positive = past
    # visit_offset: days before TODAY for visit_date (None = omit)

    # SUPPLY_LAPSED — coverage ended long ago
    ("supply_lapsed",      0.12, (200, 400),  None,   3, None, None, None, None),
    # VISIT_REQUIRED — visit overdue (visit was >180d ago)
    ("visit_required",     0.08, (60, 90),    200,    3, None, None, None, None),
    # RENEWAL_CRITICAL — visit due in <=45 days (visit was 135-175d ago)
    ("renewal_critical",   0.10, (5, 15),     150,    3, None, None, None, None),
    # RENEWAL_ELEVATED — visit due in 46-60 days (visit was 120-134d ago)
    ("renewal_elevated",   0.10, (5, 15),     125,    3, None, None, None, None),
    # RENEWAL_SOON — visit due in 61-90 days (visit was 90-119d ago)
    ("renewal_soon",       0.10, (5, 15),     100,    3, None, None, None, None),
    # RESUPPLY_READY — coverage ending within 30d, docs all clean
    ("resupply_ready",     0.15, (8, 22),     None,   3, "On File", "Yes", "Not Required", "Yes"),
    # RESUPPLY_READY with doc gaps → DOCS_REQUIRED (amber)
    ("docs_required",      0.10, (8, 22),     None,   3, "Pending", None, None, None),
    # ACTIVE — recently shipped, plenty of coverage left, no visit urgency
    ("active",             0.15, (5, 15),     None,   9, None, None, None, None),
    # TRANSFER_PENDING — transfer flag
    ("transfer_pending",   0.05, (5, 15),     None,   3, None, None, None, "TRANSFER"),
    # NO_RECENT_SHIPMENT — shipment > 365 days ago
    ("no_recent_shipment", 0.05, (400, 500),  None,   3, None, None, None, None),
]

scenario_names    = [s[0] for s in SCENARIOS]
scenario_weights  = [s[1] for s in SCENARIOS]
scenario_data     = {s[0]: s[2:] for s in SCENARIOS}

dev_names   = [d[0] for d in DEVICES]
dev_weights = [d[2] for d in DEVICES]
dev_comp    = {d[0]: d[1] for d in DEVICES}
pay_names   = [p[0] for p in PAYERS]
pay_weights = [p[1] for p in PAYERS]


def rand_offset(lo, hi):
    return random.randint(lo, hi)


rows = []
for i in range(1, N + 1):
    pid    = f"GAB-{i:06d}"
    device = random.choices(dev_names, weights=dev_weights)[0]
    comp   = dev_comp[device]
    payer  = random.choices(pay_names, weights=pay_weights)[0]
    scen   = random.choices(scenario_names, weights=scenario_weights)[0]
    ship_range, visit_offset_days, qty, swo, pecos, pa, dx = scenario_data[scen]

    ship_lo, ship_hi = ship_range
    ship = TODAY - timedelta(days=rand_offset(ship_lo, ship_hi))
    visit = (TODAY - timedelta(days=visit_offset_days)) if visit_offset_days else None

    is_transfer = (dx == "TRANSFER")
    row = {
        "patient_id":         pid,
        "device_type":        device,
        "shipment_date":      ship.isoformat(),
        "quantity":           qty,
        "payer":              payer,
        "component":          comp,
        "visit_date":         visit.isoformat() if visit else "",
        "swo_status":         swo or "",
        "pecos_verified":     pecos or "",
        "pa_status":          pa or "",
        "diagnosis_on_file":  dx if dx and dx != "TRANSFER" else ("Yes" if dx == "TRANSFER" else ""),
        "transfer_from":      "PRIOR-SUPPLIER-001" if is_transfer else "",
    }
    rows.append(row)

FIELDNAMES = [
    "patient_id", "device_type", "shipment_date", "quantity", "payer", "component",
    "visit_date", "swo_status", "pecos_verified", "pa_status", "diagnosis_on_file",
    "transfer_from",
]

with open(OUTPUT, "w", newline="") as f:
    writer = csv.DictWriter(f, fieldnames=FIELDNAMES)
    writer.writeheader()
    writer.writerows(rows)

print(f"Wrote {OUTPUT} — {N:,} rows")
print("Scenario distribution (approximate):")
from collections import Counter
dist = Counter(random.choices(scenario_names, weights=scenario_weights, k=N))
for name, count in sorted(dist.items(), key=lambda x: -x[1]):
    print(f"  {name:<22} ~{count}")