app/rules/ecg_rules.py

from typing import Any, Dict, List


def apply_rules(
    patient_context: Dict[str, Any],
    model_output: Dict[str, Any],
) -> Dict[str, Any]:
    """
    Simple rule set over ECG model outputs.
    """
    label = model_output.get("label")
    score = float(model_output.get("score", 0.0))
    hr = model_output.get("hr")
    hr_int = int(hr) if hr is not None else None

    explanations: List[str] = []
    alert_level = "none"
    age = patient_context.get("age")
    has_prior_stroke = bool(patient_context.get("has_prior_stroke"))

    # Arrhythmia label with high confidence
    arrhythmia_labels = {"arrhythmia", "suspected_afib", "afib"}
    if label in arrhythmia_labels and score >= 0.85:
        alert_level = "escalate"
        explanations.append("High-confidence arrhythmia detected.")
    elif label in arrhythmia_labels and score >= 0.6:
        alert_level = "notify"
        explanations.append("Arrhythmia suspected; monitor closely.")

    # Heart rate based thresholds
    if hr_int is not None:
        if hr_int >= 140:
            alert_level = "escalate"
            explanations.append("Severe tachycardia (hr >= 140).")
        elif hr_int >= 120 and alert_level == "none":
            alert_level = "notify"
            explanations.append("Tachycardia (hr >= 120).")

    # Patient risk factors escalate one level
    if has_prior_stroke and alert_level != "none":
        alert_level = "escalate"
        explanations.append("Prior stroke history: escalate.")
    elif has_prior_stroke and alert_level == "none":
        alert_level = "notify"
        explanations.append("Prior stroke history: monitor closely.")

    if isinstance(age, int) and age >= 75 and alert_level == "notify":
        alert_level = "escalate"
        explanations.append("Age >= 75 with concerning signal: escalate.")

    # If no triggers, add baseline explanation
    if not explanations:
        explanations.append("No rule-based alerts triggered.")

    return {"alert_level": alert_level, "explanations": explanations}