agent/data/sample_reports/generate_samples.py

"""
Generates sample text-based blood test report PDFs for the Clinical RAG demo.

Each report is a fictional patient with a few markers flagged low/high that map onto
entries in agent/data/supplements_kb.json, so the agent's extraction + supplement
RAG flow has something realistic to work with.

Run:
    source venv/bin/activate
    python agent/data/sample_reports/generate_samples.py
"""

from pathlib import Path

from fpdf import FPDF

OUT_DIR = Path(__file__).parent

# Column widths (mm) for the results table
COL_WIDTHS = (70, 22, 20, 38, 20)
HEADERS = ("Test Name", "Result", "Units", "Reference Range", "Flag")


def make_report(filename, title, patient, rows, comments=""):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=15)
    pdf.add_page()

    pdf.set_font("Helvetica", "B", 14)
    pdf.cell(0, 10, title, new_x="LMARGIN", new_y="NEXT")

    pdf.set_font("Helvetica", "", 10)
    for line in patient:
        pdf.cell(0, 6, line, new_x="LMARGIN", new_y="NEXT")
    pdf.ln(4)

    # Table header
    pdf.set_font("Helvetica", "B", 9)
    for w, h in zip(COL_WIDTHS, HEADERS):
        pdf.cell(w, 7, h, border=1)
    pdf.ln()

    # Table rows
    pdf.set_font("Helvetica", "", 9)
    for test_name, value, unit, ref_range, flag in rows:
        pdf.cell(COL_WIDTHS[0], 7, test_name, border=1)
        pdf.cell(COL_WIDTHS[1], 7, value, border=1, align="R")
        pdf.cell(COL_WIDTHS[2], 7, unit, border=1)
        pdf.cell(COL_WIDTHS[3], 7, ref_range, border=1)
        pdf.cell(COL_WIDTHS[4], 7, flag, border=1, align="C")
        pdf.ln()

    if comments:
        pdf.ln(6)
        pdf.set_font("Helvetica", "B", 10)
        pdf.cell(0, 6, "Comments:", new_x="LMARGIN", new_y="NEXT")
        pdf.set_font("Helvetica", "", 10)
        pdf.multi_cell(0, 6, comments)

    out_path = OUT_DIR / filename
    pdf.output(str(out_path))
    print(f"Wrote {out_path}")


# ── Sample 1: Vitamin D + B12 deficiency ───────────────────────────────────────
make_report(
    "sample_report_1_vitamin_d_b12.pdf",
    "Comprehensive Vitamin & Metabolic Panel",
    [
        "Patient: John Doe       DOB: 1985-03-14    Sex: Male",
        "Collected: 2026-06-10   Reported: 2026-06-12",
        "Ordering Physician: Dr. A. Sharma",
    ],
    [
        ("Vitamin D, 25-Hydroxy", "14",   "ng/mL", "30-100",     "LOW"),
        ("Vitamin B12",          "180",  "pg/mL", "200-900",    "LOW"),
        ("Hemoglobin",           "14.2", "g/dL",  "13.5-17.5",  "NORMAL"),
        ("Calcium, Total",       "9.2",  "mg/dL", "8.5-10.5",   "NORMAL"),
        ("TSH",                  "2.1",  "mIU/L", "0.4-4.0",    "NORMAL"),
        ("Folate, Serum",        "6.5",  "ng/mL", "2.7-17.0",   "NORMAL"),
        ("Potassium",            "4.3",  "mEq/L", "3.5-5.0",    "NORMAL"),
    ],
    comments="Patient reports ongoing fatigue, low mood, and occasional muscle aches over the past two months.",
)

# ── Sample 2: Iron-deficiency anemia + low Folate ──────────────────────────────
make_report(
    "sample_report_2_iron_folate.pdf",
    "Complete Blood Count & Nutrient Panel",
    [
        "Patient: Jane Smith     DOB: 1990-11-02    Sex: Female",
        "Collected: 2026-06-08   Reported: 2026-06-10",
        "Ordering Physician: Dr. R. Mehta",
    ],
    [
        ("Hemoglobin",           "10.8", "g/dL",  "12.0-15.5",  "LOW"),
        ("Ferritin",             "8",    "ng/mL", "10-120",     "LOW"),
        ("Folate, Serum",        "2.1",  "ng/mL", "2.7-17.0",   "LOW"),
        ("Vitamin D, 25-Hydroxy","32",   "ng/mL", "30-100",     "NORMAL"),
        ("Potassium",            "4.1",  "mEq/L", "3.5-5.0",    "NORMAL"),
        ("Vitamin B12",          "410",  "pg/mL", "200-900",    "NORMAL"),
    ],
    comments="Patient reports persistent tiredness, occasional dizziness on standing, and brittle nails.",
)

# ── Sample 3: Magnesium, Zinc, lipid panel, B6, thyroid ────────────────────────
make_report(
    "sample_report_3_magnesium_lipids_thyroid.pdf",
    "Wellness Panel - Minerals, Lipids & Thyroid",
    [
        "Patient: Raj Patel      DOB: 1978-07-22    Sex: Male",
        "Collected: 2026-06-05   Reported: 2026-06-07",
        "Ordering Physician: Dr. S. Iyer",
    ],
    [
        ("Magnesium",            "1.4",  "mg/dL", "1.7-2.2",    "LOW"),
        ("Zinc",                 "55",   "mcg/dL","60-120",     "LOW"),
        ("Triglycerides",        "210",  "mg/dL", "below 150",  "HIGH"),
        ("HDL Cholesterol",      "38",   "mg/dL", "above 40",   "LOW"),
        ("Vitamin B6",           "3.2",  "ng/mL", "5-50",       "LOW"),
        ("TSH",                  "5.8",  "mIU/L", "0.4-4.0",    "HIGH"),
        ("Calcium, Total",       "9.4",  "mg/dL", "8.5-10.5",   "NORMAL"),
    ],
    comments="Patient reports muscle cramps, poor sleep, and recent weight gain.",
)