"""Write example_referrals.csv: a synthetic outpatient referrals extract with planted problems.

Planted: one exact duplicate row, two reused referral IDs, inconsistent priority spellings,
missing priorities, three impossible ages, three received dates in the future, and waits
missing for referrals not yet seen. Run with: python make_example_extract.py
"""

from pathlib import Path

import numpy as np
import pandas as pd

SEED = 20260828
ROWS = 2_000
OUT = Path(__file__).with_name("example_referrals.csv")
SPECIALTIES = ["Cardiology", "Dermatology", "ENT", "Gastroenterology", "Ophthalmology", "Urology"]
SOURCES = ["GP", "Consultant", "Emergency department", "Other"]


def base_extract(rng: np.random.Generator) -> pd.DataFrame:
    received = pd.Timestamp("2025-01-01") + pd.to_timedelta(rng.integers(0, 365, ROWS), unit="D")
    waits = rng.gamma(shape=2.0, scale=30.0, size=ROWS).round()
    return pd.DataFrame(
        {
            "referral_id": [f"R{n:06d}" for n in rng.choice(900_000, ROWS, replace=False)],
            "received_date": received.strftime("%Y-%m-%d"),
            "specialty": rng.choice(SPECIALTIES, ROWS),
            "source": rng.choice(SOURCES, ROWS, p=[0.7, 0.15, 0.1, 0.05]),
            "priority": rng.choice(["Routine", "Urgent"], ROWS, p=[0.8, 0.2]),
            "age_at_referral": rng.integers(0, 96, ROWS).astype(float),
            "wait_days": np.where(rng.random(ROWS) < 0.15, np.nan, waits),
        }
    )


def plant_problems(df: pd.DataFrame, rng: np.random.Generator) -> pd.DataFrame:
    rows = rng.choice(ROWS, 40, replace=False)
    df.loc[rows[:6], "priority"] = ["urgent", "URGENT", "Routine ", "routine", "Urgent ", "urgent"]
    df.loc[rows[6:30], "priority"] = np.nan
    df.loc[rows[30:33], "age_at_referral"] = [-1.0, 142.0, -3.0]
    df.loc[rows[33:36], "received_date"] = ["2027-02-14", "2026-11-30", "2027-01-05"]
    df.loc[rows[36:38], "referral_id"] = df.loc[rows[38:40], "referral_id"].to_numpy()
    return pd.concat([df, df.iloc[[int(rows[0])]]], ignore_index=True)


def main() -> None:
    rng = np.random.default_rng(SEED)
    extract = plant_problems(base_extract(rng), rng)
    extract.to_csv(OUT, index=False, float_format="%.0f")
    print(f"Wrote {len(extract):,} rows to {OUT.name}")


if __name__ == "__main__":
    main()