Synthetic ECG reading

The prompt

Synthetic ECG reading.

The environment simulates a 12-lead ECG as time series. The model describes the strip.

It can read the raw samples, and it may run raw Python on them. Limit that Python to the standard library, math and statistics. Do not let it call analysis libraries.

Generate the signal with NeuroKit2, nk.ecg_simulate(), so the traces are realistic.

A patient state looks like:

rhythm: sinus | AF | flutter (2:1, 3:1) | 1st/2nd/3rd-degree AV block | SVT | ...
rate: 48 bpm, HRV: moderate
intervals: PR 0.21 s, QRS 0.09 s, QT 0.40 s
axis: +70°
conduction: normal | LBBB | RBBB
ectopy: PVC burden, positions
ischemia: none | ST elevation in leads II, III, aVF (+2 mm)
other: peaked T, U waves, low voltage, ...

"""Synthetic 12-lead ECG reading environment.

The environment simulates a realistic 12-lead ECG time series (generated with
NeuroKit2's ``nk.ecg_simulate``) for a randomly drawn *patient state*. The model
inspects the raw samples and describes the strip.

The model can:

* read the raw samples directly from the observation (``samples``), and
* run Python on them by submitting a move prefixed with ``run:``. That Python is
  executed in a restricted sandbox that only allows the standard library
  (``math``, ``statistics``, and a handful of safe pure-stdlib modules); analysis
  libraries such as ``numpy``/``scipy``/``neurokit2`` are refused.

Any other move is treated as the model's interpretation/report of the ECG. It is
compared against the ground-truth findings and produces the episode ``reward``
and ``done``.
"""

from __future__ import annotations

import builtins
import contextlib
import io
import json
import math
import random
import re
import statistics
from typing import Any, Dict, List, Optional, Tuple

import neurokit2 as nk
import numpy as np

FS = 250
DURATION = 10.0
N_SAMPLES = int(FS * DURATION)
LEADS = ["I", "II", "III", "aVR", "aVL", "aVF", "V1", "V2", "V3", "V4", "V5", "V6"]
LEAD_ORDER = LEADS

# --------------------------------------------------------------------------- #
# Restricted Python sandbox (what the model is allowed to run)
# --------------------------------------------------------------------------- #

_SAFE_BUILTINS = {
    "abs": abs, "all": all, "any": any, "bool": bool, "dict": dict,
    "divmod": divmod, "enumerate": enumerate, "filter": filter, "float": float,
    "int": int, "len": len, "list": list, "map": map, "max": max, "min": min,
    "pow": pow, "range": range, "reversed": reversed, "round": round,
    "set": set, "slice": slice, "sorted": sorted, "str": str, "sum": sum,
    "tuple": tuple, "zip": zip, "print": print, "repr": repr,
    "isinstance": isinstance, "issubclass": issubclass, "hasattr": hasattr,
    "getattr": getattr, "type": type, "format": format, "chr": chr, "ord": ord,
    "hex": hex, "oct": oct, "bin": bin, "hash": hash, "iter": iter,
    "next": next, "object": object, "property": property, "staticmethod": staticmethod,
    "classmethod": classmethod, "vars": vars, "callable": callable, "dir": dir,
    "id": id, "complex": complex, "frozenset": frozenset, "bytes": bytes,
    "bytearray": bytearray, "memoryview": memoryview,
}

_ALLOWED_MODULES = {
    "math", "statistics", "random", "itertools", "functools", "operator",
    "collections", "decimal", "fractions", "re", "string", "json", "datetime",
    "time", "heapq", "bisect", "cmath", "enum", "typing", "dataclasses",
    "copy", "numbers", "array", "struct", "textwrap",
}


def _restricted_import(name: str, globals=None, locals=None, fromlist=(), level=0):
    top = name.split(".")[0]
    if top not in _ALLOWED_MODULES:
        raise ImportError(
            f"module {top!r} is not available; only the standard library "
            "(math, statistics, ...) is allowed"
        )
    return builtins.__import__(name, globals, locals, fromlist, level)


def _sandbox_builtins() -> Dict[str, Any]:
    env = dict(_SAFE_BUILTINS)
    env["__import__"] = _restricted_import
    env["__name__"] = "__main__"
    env["__build_class__"] = builtins.__build_class__
    return env


def run_python(
    code: str,
    samples: Dict[str, List[float]],
    fs: int,
    duration: float,
    leads: List[str],
) -> Tuple[str, Optional[str], Optional[str]]:
    """Execute restricted Python, returning (stdout, result_repr, error)."""
    globals_ = {
        "__builtins__": _sandbox_builtins(),
        "ecg": samples,
        "samples": samples,
        "leads": leads,
        "fs": fs,
        "duration": duration,
        "math": math,
        "statistics": statistics,
    }
    out = io.StringIO()
    result_repr: Optional[str] = None
    error: Optional[str] = None
    stripped = code.strip()
    if not stripped:
        return "", None, None
    try:
        try:
            node = compile(stripped, "<ecg>", "eval")
            with contextlib.redirect_stdout(out):
                value = eval(node, globals_)
            result_repr = repr(value)
        except SyntaxError:
            with contextlib.redirect_stdout(out):
                exec(compile(stripped, "<ecg>", "exec"), globals_)
    except Exception as exc:  # noqa: BLE001 - surface any error to the model
        error = f"{type(exc).__name__}: {exc}"
    stdout = out.getvalue()
    if len(stdout) > 4000:
        stdout = stdout[:4000] + "\n... (truncated)"
    if result_repr and len(result_repr) > 4000:
        result_repr = result_repr[:4000] + "... (truncated)"
    return stdout, result_repr, error


# --------------------------------------------------------------------------- #
# ECG synthesis
# --------------------------------------------------------------------------- #


def _boxcar(x: np.ndarray, n: int) -> np.ndarray:
    if n <= 1:
        return x.copy()
    return np.convolve(x, np.ones(n) / n, mode="same")


def _extract_template(rng: np.random.Generator) -> Tuple[np.ndarray, int, int]:
    """Return a clean single cardiac-cycle beat (lead-II morphology)."""
    pre = int(0.30 * FS)
    post = int(0.55 * FS)
    seed = int(rng.integers(0, 2**31 - 1))
    sig = nk.ecg_simulate(
        duration=12,
        sampling_rate=FS,
        heart_rate=60,
        noise=0.003,
        method="ecgsyn",
        random_state=seed,
    )
    arr = np.asarray(sig, dtype=float)
    _, info = nk.ecg_peaks(arr, sampling_rate=FS)
    rpeaks = info["ECG_R_Peaks"].astype(int)
    cycles = []
    for i in range(1, len(rpeaks) - 1):
        a = int(rpeaks[i]) - pre
        b = int(rpeaks[i]) + post
        if a < 0 or b > len(arr):
            continue
        cycles.append(arr[a:b])
    if not cycles:
        cycles = [arr[0:pre + post]]
    beat = np.mean(cycles, axis=0)
    beat = beat - float(np.median(beat[: int(0.06 * FS)]))
    return beat, pre, post


def _r_peak(beat: np.ndarray) -> int:
    return int(np.argmax(beat))


def _no_p(beat: np.ndarray, rp: int) -> np.ndarray:
    out = beat.copy()
    out[: rp - int(0.04 * FS)] = 0.0
    return out


def _p_only(beat: np.ndarray, rp: int) -> np.ndarray:
    out = beat.copy()
    out[rp - int(0.04 * FS):] = 0.0
    return out


def _widen_qrs(beat: np.ndarray, rp: int, width: Optional[int] = None) -> np.ndarray:
    out = beat.copy()
    a = rp - int(0.03 * FS)
    b = rp + int(0.06 * FS)
    width = width or max(3, int(0.02 * FS))
    out[a:b] = _boxcar(out[a:b], width)
    return out


def _lengthen_pr(beat: np.ndarray, rp: int) -> Tuple[np.ndarray, int]:
    out = beat.copy()
    qrs_on = rp - int(0.04 * FS)
    shift = int(0.06 * FS)
    segment = out[qrs_on:].copy()
    out[qrs_on + shift:] = segment[: len(out) - (qrs_on + shift)]
    out[qrs_on:qrs_on + shift] = 0.0
    return out, rp + shift


def _pvc_template(rng: np.random.Generator, pre: int, post: int, rp: int) -> np.ndarray:
    L = pre + post
    t = np.arange(L)
    sign = 1.0 if rng.random() < 0.5 else -1.0
    qrs = sign * 1.25 * np.exp(-0.5 * ((t - rp) / int(0.05 * FS)) ** 2)
    twave = -sign * 0.5 * np.exp(-0.5 * ((t - (rp + int(0.22 * FS))) / int(0.06 * FS)) ** 2)
    return qrs + twave


def _st_window(rp: int) -> Tuple[int, int]:
    return rp + int(0.05 * FS), rp + int(0.16 * FS)


def _peaked_t(beat: np.ndarray, rp: int) -> np.ndarray:
    out = beat.copy()
    out[rp + int(0.12 * FS):rp + int(0.22 * FS)] *= 1.6
    return out


def _inverted_t(beat: np.ndarray, rp: int) -> np.ndarray:
    out = beat.copy()
    out[rp + int(0.12 * FS):rp + int(0.24 * FS)] *= -1.0
    return out


def _u_wave(beat: np.ndarray, rp: int) -> np.ndarray:
    out = beat.copy()
    t = np.arange(len(beat))
    center = rp + int(0.28 * FS)
    out += 0.06 * np.exp(-0.5 * ((t - center) / int(0.025 * FS)) ** 2)
    return out


def _low_voltage(beat: np.ndarray) -> np.ndarray:
    return beat * 0.45


def _lbbb(beat: np.ndarray, rp: int) -> np.ndarray:
    return _widen_qrs(beat, rp)


def _rbbb(beat: np.ndarray, rp: int) -> np.ndarray:
    out = _widen_qrs(beat, rp)
    t = np.arange(len(beat))
    rprime = rp + int(0.11 * FS)
    out += 0.5 * np.exp(-0.5 * ((t - rprime) / int(0.015 * FS)) ** 2)
    return out


# --------------------------------------------------------------------------- #
# Lead derivation (one clean beat -> 12 coherent leads)
# --------------------------------------------------------------------------- #


def _derive_12lead(s: np.ndarray, axis_deg: float) -> np.ndarray:
    theta = math.radians(axis_deg)

    def frontal(lead_angle_deg: float) -> np.ndarray:
        return s * math.cos(theta - math.radians(lead_angle_deg))

    out = np.zeros((len(LEADS), len(s)))
    out[LEAD_ORDER.index("I")] = frontal(0.0)
    out[LEAD_ORDER.index("II")] = frontal(60.0)
    out[LEAD_ORDER.index("III")] = frontal(120.0)
    out[LEAD_ORDER.index("aVR")] = frontal(-150.0)
    out[LEAD_ORDER.index("aVL")] = frontal(-30.0)
    out[LEAD_ORDER.index("aVF")] = frontal(90.0)

    precordial = {
        "V1": -0.32, "V2": -0.15, "V3": 0.08, "V4": 0.40, "V5": 0.62, "V6": 0.70,
    }
    for name, gain in precordial.items():
        out[LEAD_ORDER.index(name)] = s * gain
    return out


def _apply_st_elevation(sig: np.ndarray, rpeaks: List[int], leads: List[str], mm: float) -> None:
    amplitude = mm / 10.0
    for rp in rpeaks:
        a, b = _st_window(rp)
        if b > N_SAMPLES:
            b = N_SAMPLES
        for name in leads:
            if name in LEAD_ORDER:
                sig[LEAD_ORDER.index(name), a:b] += amplitude


def _flutter_waves(sig: np.ndarray, ventricular_rate: int, ratio: int) -> None:
    atrial_rate = ventricular_rate * ratio
    t = np.arange(N_SAMPLES)
    saw = 2.0 * ((t * atrial_rate / 60.0 / FS) % 1.0) - 1.0
    for name in ["II", "III", "aVF"]:
        sig[LEAD_ORDER.index(name)] += 0.10 * saw


def _fibrillatory_baseline(sig: np.ndarray, rng: random.Random) -> None:
    t = np.arange(N_SAMPLES)
    noise = 0.035 * np.sin(2 * math.pi * 5.0 * t / FS + float(rng.random()))
    noise += 0.025 * np.sin(2 * math.pi * 8.0 * t / FS + float(rng.random()))
    for i in range(len(LEADS)):
        sig[i] += noise


# --------------------------------------------------------------------------- #
# Patient state
# --------------------------------------------------------------------------- #

_RHYTHMS = [
    "sinus", "sinus bradycardia", "sinus tachycardia", "atrial fibrillation",
    "atrial flutter (2:1)", "atrial flutter (3:1)", "first-degree AV block",
    "second-degree AV block (Mobitz I)", "third-degree AV block", "SVT",
]


def _pick(rng: random.Random, seq: List[Any], weights: Optional[List[float]] = None) -> Any:
    if weights is None:
        return rng.choice(seq)
    return rng.choices(seq, weights=weights, k=1)[0]


def sample_state(task_id: str) -> Tuple[Dict[str, Any], int]:
    """Return a patient state dict and the seed used."""
    seed: int
    override: Optional[Dict[str, Any]] = None
    task_id = (task_id or "").strip()
    if task_id == "":
        seed = random.SystemRandom().randint(0, 2**31 - 1)
    elif task_id.isdigit():
        seed = int(task_id) % (2**31)
    else:
        try:
            parsed = json.loads(task_id)
            if isinstance(parsed, dict) and "rhythm" in parsed:
                override = parsed
                seed = random.SystemRandom().randint(0, 2**31 - 1)
            else:
                seed = (sum(ord(c) for c in task_id) * 104729) % (2**31)
        except json.JSONDecodeError:
            seed = (sum(ord(c) for c in task_id) * 104729) % (2**31)

    rng = random.Random(seed)

    rhythm = _pick(
        rng,
        _RHYTHMS,
        weights=[18, 8, 8, 12, 8, 5, 6, 6, 4, 6],
    )

    def rate_for(rhythm: str) -> int:
        if rhythm == "sinus bradycardia":
            return rng.randint(42, 54)
        if rhythm == "sinus tachycardia":
            return rng.randint(101, 140)
        if rhythm == "atrial fibrillation":
            return rng.randint(70, 150)
        if rhythm == "atrial flutter (2:1)":
            return 150
        if rhythm == "atrial flutter (3:1)":
            return 100
        if rhythm == "first-degree AV block":
            return rng.randint(50, 80)
        if rhythm == "second-degree AV block (Mobitz I)":
            return rng.randint(45, 70)
        if rhythm == "third-degree AV block":
            return rng.randint(30, 45)
        if rhythm == "SVT":
            return rng.randint(150, 200)
        return rng.randint(55, 100)

    rate = rate_for(rhythm)

    hrv = "high" if rhythm == "atrial fibrillation" else _pick(rng, ["low", "moderate"])

    conduction = _pick(rng, ["normal", "LBBB", "RBBB"], weights=[80, 10, 10])
    if rhythm in ("SVT", "third-degree AV block"):
        conduction = "normal"

    ectopy = _pick(
        rng,
        ["none", "occasional PVC", "frequent PVCs", "PVC bigeminy", "PVC couplet"],
        weights=[62, 14, 10, 8, 6],
    )

    ischemia = _pick(
        rng,
        ["none", "none", "none", "inferior", "anterior", "lateral", "anterolateral"],
    )
    ischemia_leads = {
        "inferior": ["II", "III", "aVF"],
        "anterior": ["V1", "V2", "V3", "V4"],
        "lateral": ["I", "aVL", "V5", "V6"],
        "anterolateral": ["V2", "V3", "V4", "V5", "V6"],
    }
    st_leads: List[str] = ischemia_leads.get(ischemia, [])
    st_mm = rng.choice([1, 2, 2, 3]) if st_leads else 0

    axis = _pick(rng, [-15, 0, 15, 30, 45, 60, 75, 90, 105])

    other: List[str] = []
    roll = rng.random()
    if roll < 0.12:
        other.append("peaked T waves")
    elif roll < 0.22:
        other.append("T wave inversion")
    if rng.random() < 0.08:
        other.append("U waves")
    if rng.random() < 0.08:
        other.append("low voltage")

    state: Dict[str, Any] = {
        "age": rng.randint(42, 84),
        "sex": _pick(rng, ["male", "female"]),
        "rhythm": rhythm,
        "rate": rate,
        "hrv": hrv,
        "conduction": conduction,
        "ectopy": ectopy,
        "ischemia": ischemia,
        "st_leads": st_leads,
        "st_mm": st_mm,
        "axis_deg": axis,
        "other": other,
    }

    if override:
        for key, value in override.items():
            if key in state and value is not None:
                state[key] = value

    return state, seed


# --------------------------------------------------------------------------- #
# Signal assembly
# --------------------------------------------------------------------------- #

_NO_P_RHYTHMS = {
    "atrial fibrillation", "atrial flutter (2:1)", "atrial flutter (3:1)", "SVT",
}


def _build_beat_sequence(state: Dict[str, Any], rng: random.Random):
    rhythm = state["rhythm"]
    rate = state["rate"]
    rr = FS * 60.0 / rate
    start = int(0.25 * FS)
    end = N_SAMPLES - int(0.25 * FS)
    qrs_times: List[int] = []
    qrs_kinds: List[str] = []
    p_times: List[int] = []

    if rhythm == "third-degree AV block":
        atrial_rate = rng.randint(75, 100)
        a_rr = FS * 60.0 / atrial_rate
        t = start
        while t < end:
            p_times.append(t)
            t += int(a_rr)
        t = start
        while t < end:
            qrs_times.append(t)
            qrs_kinds.append("escape")
            t += int(rr)
        return qrs_times, qrs_kinds, p_times

    default_kind = "nop" if rhythm in _NO_P_RHYTHMS else "normal"
    t = start
    beat_no = 0
    while t < end:
        qrs_times.append(t)
        beat_no += 1
        if rhythm == "second-degree AV block (Mobitz I)" and beat_no % 4 == 0:
            qrs_kinds.append("dropped")
        else:
            qrs_kinds.append(default_kind)
        if rhythm == "atrial fibrillation":
            step = int(rr * rng.uniform(0.7, 1.4))
        else:
            step = int(rr)
        t += max(int(0.3 * FS), step)

    n = len(qrs_times)
    ectopy = state["ectopy"]
    if ectopy == "occasional PVC" and n >= 4:
        qrs_kinds[rng.randrange(1, n - 1)] = "pvc"
    elif ectopy == "frequent PVCs" and n >= 6:
        for p in rng.sample(range(1, n - 1), k=3):
            qrs_kinds[p] = "pvc"
    elif ectopy == "PVC bigeminy" and n >= 2:
        for i in range(1, n, 2):
            qrs_kinds[i] = "pvc"
    elif ectopy == "PVC couplet" and n >= 4:
        pos = rng.randrange(1, n - 2)
        qrs_kinds[pos] = "pvc"
        qrs_kinds[pos + 1] = "pvc"

    return qrs_times, qrs_kinds, p_times


def _stamp_1d(sig: np.ndarray, tpl: np.ndarray, center: int, pre: int) -> None:
    L = len(tpl)
    start = center - pre
    if start >= N_SAMPLES or start + L <= 0:
        return
    s0 = max(0, start)
    e0 = min(N_SAMPLES, start + L)
    ts = s0 - start
    te = ts + (e0 - s0)
    if te <= ts:
        return
    sig[s0:e0] = tpl[ts:te]


def _measure_intervals(beat: np.ndarray) -> Tuple[float, float, float]:
    """Approximate PR, QRS, QT (seconds) from a clean beat (lead II)."""
    fs = FS
    rp = int(np.argmax(beat))
    base = float(np.median(beat))

    q_idx = int(np.argmin(beat[max(0, rp - int(0.08 * fs)): rp])) + max(0, rp - int(0.08 * fs))
    s_idx = int(np.argmin(beat[rp: rp + int(0.08 * fs)])) + rp

    p_lo = max(0, rp - int(0.28 * fs))
    p_hi = max(p_lo + 1, rp - int(0.04 * fs))
    p_peak = int(np.argmax(beat[p_lo:p_hi])) + p_lo

    t_lo = s_idx
    t_hi = min(len(beat), rp + int(0.36 * fs))
    t_end = t_lo
    for i in range(t_lo, t_hi):
        if beat[i] > base + 0.02:
            t_end = i

    pr = (rp - p_peak) / fs
    qrs = (s_idx - q_idx) / fs
    qt = (t_end - q_idx) / fs
    return pr, qrs, qt


def _describe_axis(deg: float) -> str:
    if deg < -15:
        return f"{deg:+d}° (left axis deviation)"
    if deg > 90:
        return f"{deg:+d}° (right axis deviation)"
    return f"{deg:+d}° (normal axis)"


def _ischemia_label(state: Dict[str, Any]) -> str:
    if not state["st_mm"]:
        return "none"
    return f"ST elevation in {', '.join(state['st_leads'])} (+{state['st_mm']} mm)"


def _build_answer(state: Dict[str, Any], pr: float, qrs: float, qt: float) -> Dict[str, Any]:
    return {
        "rhythm": state["rhythm"],
        "rate": state["rate"],
        "hrv": state["hrv"],
        "intervals": {"PR": round(pr, 2), "QRS": round(qrs, 2), "QT": round(qt, 2)},
        "axis": _describe_axis(state["axis_deg"]),
        "conduction": state["conduction"],
        "ectopy": state["ectopy"],
        "ischemia": _ischemia_label(state),
        "other": state["other"],
    }


# --------------------------------------------------------------------------- #
# Scoring
# --------------------------------------------------------------------------- #

_RHYTHM_KEYWORD = {
    "sinus": "sinus",
    "sinus bradycardia": "sinus",
    "sinus tachycardia": "sinus",
    "atrial fibrillation": "atrial fibrillation",
    "atrial flutter (2:1)": "flutter",
    "atrial flutter (3:1)": "flutter",
    "first-degree AV block": "av block",
    "second-degree AV block (Mobitz I)": "av block",
    "third-degree AV block": "av block",
    "SVT": "svt",
}


def _findings(state: Dict[str, Any]) -> List[str]:
    findings = [_RHYTHM_KEYWORD[state["rhythm"]]]
    rate = state["rate"]
    if rate < 60:
        findings.append("bradycardia")
    elif rate > 100:
        findings.append("tachycardia")
    if state["conduction"] != "normal":
        findings.append(state["conduction"].lower())
    if state["ectopy"] != "none":
        findings.append("pvc")
    if state["st_mm"]:
        findings.append("st elevation")
        findings.append(state["ischemia"])
    for other in state["other"]:
        findings.append(other)
    return findings


def _score(text: str, findings: List[str], state: Dict[str, Any]) -> Tuple[float, List[str], List[str]]:
    lower = text.lower()
    matched: List[str] = []
    missing: List[str] = []
    for f in findings:
        if f in ("bradycardia", "tachycardia"):
            m = re.search(r"(\d{2,3})\s*bpm", lower)
            if m:
                reported = int(m.group(1))
                if abs(reported - state["rate"]) <= 15:
                    matched.append(f)
                else:
                    missing.append(f)
            else:
                missing.append(f)
        elif f in lower:
            matched.append(f)
        else:
            missing.append(f)
    reward = round(len(matched) / max(1, len(findings)), 3)
    return reward, matched, missing


# --------------------------------------------------------------------------- #
# Game
# --------------------------------------------------------------------------- #


class ECGGame:
    """Synthetic 12-lead ECG reading environment."""

    def __init__(self) -> None:
        self.samples: Dict[str, List[float]] = {}
        self.leads: List[str] = list(LEADS)
        self.fs = FS
        self.duration = DURATION
        self.state: Dict[str, Any] = {}
        self.answer: Dict[str, Any] = {}
        self.task_id = ""
        self.last_move: Optional[str] = None

    def reset(self, task_id: str = "") -> Dict[str, Any]:
        self.task_id = task_id
        state, seed = sample_state(task_id)
        self.state = state

        rng_np = np.random.default_rng(seed)
        rng = random.Random(seed)

        clean, pre, post = _extract_template(rng_np)
        rp = _r_peak(clean)

        # --- build the "normal" (ventricular) beat template -----------------
        base = clean
        if state["conduction"] == "LBBB":
            base = _lbbb(clean, rp)
        elif state["conduction"] == "RBBB":
            base = _rbbb(clean, rp)

        rp_now = _r_peak(base)

        if state["rhythm"] == "first-degree AV block":
            base, rp_now = _lengthen_pr(base, rp_now)

        for feature in state["other"]:
            if feature == "peaked T waves":
                base = _peaked_t(base, rp_now)
            elif feature == "T wave inversion":
                base = _inverted_t(base, rp_now)
            elif feature == "U waves":
                base = _u_wave(base, rp_now)
            elif feature == "low voltage":
                base = _low_voltage(base)

        templates: Dict[str, np.ndarray] = {
            "normal": base,
            "nop": _no_p(base, rp_now),
            "p": _p_only(clean, rp),
            "escape": _no_p(_widen_qrs(clean, rp), rp),
            "pvc": _pvc_template(rng_np, pre, post, rp),
        }

        qrs_times, qrs_kinds, p_times = _build_beat_sequence(state, rng)

        s = np.zeros(N_SAMPLES)
        for pt in p_times:
            _stamp_1d(s, templates["p"], pt, pre)
        for t, kind in zip(qrs_times, qrs_kinds):
            template = templates.get(kind)
            if template is not None:
                _stamp_1d(s, template, t, pre)

        sig = _derive_12lead(s, state["axis_deg"])

        if state["st_mm"]:
            _apply_st_elevation(sig, qrs_times, state["st_leads"], state["st_mm"])

        if state["rhythm"] == "atrial flutter (2:1)":
            _flutter_waves(sig, state["rate"], 2)
        elif state["rhythm"] == "atrial flutter (3:1)":
            _flutter_waves(sig, state["rate"], 3)
        elif state["rhythm"] == "atrial fibrillation":
            _fibrillatory_baseline(sig, rng)

        sig = sig + rng_np.normal(0.0, 0.006, size=sig.shape)

        self.samples = {
            name: [round(float(v), 4) for v in sig[LEAD_ORDER.index(name)]]
            for name in LEADS
        }

        pr, qrs, qt = _measure_intervals(base)
        if state["conduction"] in ("LBBB", "RBBB"):
            qrs = max(qrs, 0.12)
        if state["rhythm"] == "first-degree AV block":
            pr = pr + 0.06

        self.answer = _build_answer(state, pr, qrs, qt)
        self.last_move = None

        return self._observation(note=self._prompt())

    def step(self, action: Any) -> Tuple[Dict[str, Any], bool, Optional[str]]:
        move = getattr(action, "move", action)
        if not isinstance(move, str):
            move = str(move)
        self.last_move = move
        stripped = move.strip()

        if stripped.lower() == "help":
            obs = self._observation(note=self._prompt())
            obs["accepted"] = True
            return obs, True, None

        if stripped.lower().startswith(("run:", "code:")):
            code = stripped.split(":", 1)[1]
            stdout, result, error = run_python(
                code, self.samples, self.fs, self.duration, self.leads
            )
            obs = self._observation(note="ran code")
            obs["stdout"] = stdout
            obs["result"] = result
            obs["run_error"] = error
            obs["accepted"] = True
            return obs, True, error

        findings = _findings(self.state)
        reward, matched, missing = _score(stripped, findings, self.state)
        obs = self._observation(note="interpretation recorded")
        obs["reward"] = reward
        obs["done"] = True
        obs["answer"] = self.answer
        obs["matched"] = matched
        obs["missing"] = missing
        return obs, True, None

    def _observation(self, note: str) -> Dict[str, Any]:
        return {
            "env_id": "synthetic-ecg",
            "task_id": self.task_id,
            "reward": 0.0,
            "done": False,
            "legal_moves": self._legal_moves(),
            "last_move": self.last_move,
            "note": note,
            "question": "Describe this 12-lead ECG: rhythm, rate, intervals, axis, "
                        "conduction, ectopy, and any ischemic or other findings.",
            "age": self.state.get("age"),
            "sex": self.state.get("sex"),
            "fs": self.fs,
            "duration": self.duration,
            "leads": self.leads,
            "samples": self.samples,
        }

    def _prompt(self) -> str:
        return (
            "You are shown a synthetic 12-lead ECG (250 Hz, 10 s). Read the raw "
            "samples in `samples` and/or run Python on them with `run:<code>` "
            "(stdlib, math, statistics only). Submit your interpretation as a "
            "plain-text move to end the episode."
        )

    def _legal_moves(self) -> List[str]:
        return [
            "run:statistics.mean(ecg['II'])",
            "run:[round(min(v),2) for v in ecg.values()]",
            "run:sorted(leads)",
            "help",
        ]


GAME_CLASS = "ECGGame"