"""Render PTB-XL records to PNGs matching phase1-cases.json case_ids.

Usage:
    cd apps/ekg-tutor/interpret
    .venv/bin/python scripts/render_ptbxl.py \
        --manifest ../api/scripts/phase1-cases.json \
        --ptbxl-root /path/to/ptb-xl-1.0.3 \
        --out /path/to/rendered/images

For each entry where ptb_xl_record is set, reads the WFDB record from
{ptbxl-root}/{ptb_xl_record}, renders a 12-lead + rhythm strip image with
the standard 25mm/s 10mm/mV calibration on a parchment background, saves
as {case_id}.png in the output directory.

This script is decoupled from the api so it can be run offline against a
local copy of PTB-XL on whichever machine has the data.

Dependencies: wfdb, matplotlib (already in pyproject.toml's optional
group `render` — install with `uv pip install -e .[render]` if you need to
run this script).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

try:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np
    import wfdb
except ImportError as e:
    print(
        f"render_ptbxl.py needs `wfdb`, `numpy`, and `matplotlib`. "
        f"Install with: uv pip install wfdb matplotlib numpy. ({e})",
        file=sys.stderr,
    )
    sys.exit(2)


LEAD_ORDER = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]


def render(record_path: Path, out_path: Path) -> None:
    """Render a 12-lead EKG to a PNG.

    Layout: 3 rows x 4 columns of 2.5-second strips for the 12 leads, with
    a 10-second rhythm strip (lead II) along the bottom. 25 mm/s sweep,
    10 mm/mV gain. Parchment background, salmon gridlines — the classic
    look most learners encounter in their training.
    """
    record = wfdb.rdrecord(str(record_path))
    fs = record.fs
    signal = record.p_signal  # shape (n_samples, n_leads)
    sig_names = [n.upper() for n in record.sig_name]

    # Map record leads onto the standard order; missing leads render blank.
    lead_index: dict[str, int] = {}
    for name in LEAD_ORDER:
        for i, s in enumerate(sig_names):
            if s == name or s == name.replace("AVR", "AVR"):
                lead_index[name] = i
                break

    seconds_per_panel = 2.5
    samples_per_panel = int(seconds_per_panel * fs)
    rhythm_lead = "II"

    fig = plt.figure(figsize=(12.5, 8.5), dpi=200)
    fig.patch.set_facecolor("#f8f1e0")

    # 3x4 grid for the 12 leads + 1 row for the rhythm strip
    gs = fig.add_gridspec(4, 4, height_ratios=[1, 1, 1, 1.4], hspace=0.25, wspace=0.05)

    for row in range(3):
        for col in range(4):
            lead = LEAD_ORDER[row * 4 + col]
            ax = fig.add_subplot(gs[row, col])
            _setup_axes(ax, seconds_per_panel)
            if lead in lead_index:
                y = signal[:samples_per_panel, lead_index[lead]]
                x = np.linspace(0, seconds_per_panel, len(y))
                ax.plot(x, y, color="#1a1a1a", linewidth=0.8)
            ax.text(0.02, 0.95, lead, transform=ax.transAxes,
                    fontsize=9, fontweight="bold", color="#1a1a1a", va="top")

    # Rhythm strip (lead II, full 10s)
    ax = fig.add_subplot(gs[3, :])
    rhythm_samples = min(int(10 * fs), signal.shape[0])
    _setup_axes(ax, rhythm_samples / fs)
    if rhythm_lead in lead_index:
        y = signal[:rhythm_samples, lead_index[rhythm_lead]]
        x = np.linspace(0, rhythm_samples / fs, len(y))
        ax.plot(x, y, color="#1a1a1a", linewidth=0.8)
    ax.text(0.005, 0.92, f"{rhythm_lead}  (rhythm)", transform=ax.transAxes,
            fontsize=9, fontweight="bold", color="#1a1a1a", va="top")

    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, dpi=200, facecolor="#f8f1e0", bbox_inches="tight")
    plt.close(fig)


def _setup_axes(ax, x_seconds: float) -> None:
    # 25 mm/s sweep, 10 mm/mV gain. The big-box grid is 5mm = 0.2s = 0.5mV.
    ax.set_facecolor("#f8f1e0")
    ax.set_xlim(0, x_seconds)
    ax.set_ylim(-1.5, 1.5)
    ax.set_xticks([])
    ax.set_yticks([])
    # Salmon major grid every 0.2s and 0.5mV; faint minor every 0.04s, 0.1mV.
    for i in np.arange(0, x_seconds + 0.001, 0.04):
        ax.axvline(i, color="#e8b8b0", linewidth=0.3, zorder=0)
    for i in np.arange(0, x_seconds + 0.001, 0.2):
        ax.axvline(i, color="#d4827a", linewidth=0.5, zorder=0)
    for i in np.arange(-1.5, 1.51, 0.1):
        ax.axhline(i, color="#e8b8b0", linewidth=0.3, zorder=0)
    for i in np.arange(-1.5, 1.51, 0.5):
        ax.axhline(i, color="#d4827a", linewidth=0.5, zorder=0)
    for side in ("top", "bottom", "left", "right"):
        ax.spines[side].set_color("#d4827a")
        ax.spines[side].set_linewidth(0.6)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True, help="phase1-cases.json")
    parser.add_argument("--ptbxl-root", required=True, help="PTB-XL extracted root directory")
    parser.add_argument("--out", required=True, help="output directory for rendered PNGs")
    args = parser.parse_args()

    manifest = json.loads(Path(args.manifest).read_text(encoding="utf-8"))
    ptbxl_root = Path(args.ptbxl_root)
    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    ok = 0
    skipped = 0
    failed = 0
    for entry in manifest["cases"]:
        rec = entry.get("ptb_xl_record")
        case_id = entry["id"]
        if not rec:
            print(f"skip {case_id} — no ptb_xl_record set")
            skipped += 1
            continue
        record_path = ptbxl_root / rec
        out_path = out_dir / f"{case_id}.png"
        try:
            render(record_path, out_path)
            print(f"ok   {case_id} → {out_path}")
            ok += 1
        except Exception as e:
            print(f"fail {case_id} — {e}", file=sys.stderr)
            failed += 1

    print(f"\nrender done: {ok} ok, {skipped} skipped, {failed} failed")
    return 1 if failed > 0 else 0


if __name__ == "__main__":
    raise SystemExit(main())