#!/usr/bin/env python3
"""Red-flag rule engine + scanner for patient questions.

Rules classify an utterance into alert findings (mode=alert) vs grounded.
Produces red_flags[] entries for the router, and lets us measure coverage
over the real patient-question dataset.
"""
import json
import re

# ---------- A. duration / temporal flags ----------
_DUR_RE = [
    # 持续/反复 + 数字时间 → 慢性化红旗 (>=2周 or 反复次数>=3)
    (r"(?:持续|一直|老是|总是|反复|很多|几个月|半年|一年|长期)",
     "chronicity"),
    (r"\d+开头的数字度[dD天周月年数]", "chronicity"),  # placeholder fallback
    # 急性新发
    (r"(?:突然|突发|骤然|一下子|猛地|几分钟内|几小时内)", "acuteness"),
    # 进行性加重
    (r"(?:越来越(?:严重|疼|痛|差)|进行性|加重|恶化|变差|恶化)", "progression"),
]

# numeric-duration → chronicity
#   - >=2 weeks OR >=1 month/year → chronic (例外: 孕周/年龄/发烧"X天"短时长)
#   - Chinese + Arabic numerals both supported
_CN_NUM = {"一":1,"两":2,"二":2,"三":3,"四":4,"五":5,"六":6,"七":7,"八":8,"九":9,"十":10,"半":0.5}
_DUR_UNIT = r"(?:个)?(周|星期|天|日|个月|月|年)"
_DUR_RE = re.compile(r"([0-9]{1,3}(?:\.[0-9])?|[一二两三四五六七八九十半]+?)\s*" + _DUR_UNIT)


def _parse_cn_num(s):
    if s.replace(".", "", 1).isdigit():
        return float(s)
    if s == "半":
        return 0.5
    total = 0
    if "十" in s:
        a, _, b = s.partition("十")
        total = (_CN_NUM.get(a, 0) if a else 1) * 10 + _CN_NUM.get(b, 0)
    else:
        total = sum(_CN_NUM.get(ch, 0) for ch in s)
    return total


def _duration_flags(text):
    flags = []
    lo = text.lower()
    seen = set()
    for m in _DUR_RE.finditer(text):
        n = _parse_cn_num(m.group(1))
        unit = m.group(2)
        seg_start = max(0, m.start() - 14)
        ctx = text[seg_start:m.start()]
        # skip pregnancy/gestation duration
        if re.search(r"(?:怀孕|孕|妊娠|孕周|胎)", ctx):
            continue
        # skip age ("2岁" handled separately; "半岁/月龄" may appear but not chronic)
        if re.search(r"(?:岁|出生|生后|月龄)", ctx) and unit in ("周", "天", "个月"):
            continue
        is_chronic = False
        if unit in ("年",):
            is_chronic = True
        elif unit in ("个月", "月"):
            is_chronic = n >= 1
        elif unit in ("周", "星期"):
            is_chronic = n >= 2
        elif unit in ("天", "日"):
            # 发热/咳嗽短时长不算慢性，但>=2周(>=14天)或伴随"持续"算
            if n >= 14:
                is_chronic = True
            elif re.search(r"(?:持续|一直|老是|总是|反复)", ctx) and n >= 7:
                is_chronic = True
        if is_chronic and "chronicity" not in seen:
            flags.append(({"type": "chronicity", "severity": "high"}, "duration:" + m.group(0)))
            seen.add("chronicity")
    # recurrence (Chinese/Arabic)
    rec = re.compile(r"([0-9]{1,2}|[一二两三四五六七八九十]+)\s*次")
    for m in rec.finditer(text):
        n = _parse_cn_num(m.group(1))
        if n >= 3 and re.search(r"反复|复发|老是|常|每次|每个月|经常", text):
            flags.append(({"type": "chronicity", "severity": "high"}, "recurrent:" + m.group(0)))
            break
    if re.search(r"(?:突然|突发|骤然|一下子|猛地|爆裂)", lo):
        flags.append(({"type": "acuteness", "severity": "high"}, "acute-onset"))
    if re.search(r"(?:越来越(?:严重|疼|痛|差)|进行性加重|加重|恶化)", lo):
        flags.append(({"type": "progression", "severity": "high"}, "worsening"))
    return flags


# ---------- B. dangerous-signal entity flags ----------
_ENTITY_RULES = [
    (r"(?:剧烈|此生最痛|爆裂样|刀劈样).{0,6}(头痛|头疼)", "redflag_pain", "worst-ever-headache"),
    (r"(?:突?发)?(?:言语不清|说话.*?含糊|口齿不清|嘴巴歪|口角歪斜|面瘫|单侧肢体无力|一侧手脚麻木|吐字不清|视物重影|构音障碍)", "neuro", "acute-neuro-defecit"),
    (r"(?:胸痛|胸闷).{0,10}(?:压榨|憋闷|放射|冷汗|濒死|到|放射)", "cardiac", "chest-pain"),
    (r"呼吸困难|气促|喘不上气|憋气", "respiratory", "dyspnea"),
    (r"(?:呕血|吐血|咯血|黑便|柏油样便)", "gi_bleed", "gi-bleeding"),
    (r"(?:高热|高烧|发烧|发热)\D{0,6}(?:39(?:\.\d)?|[4-9]\d)\s*(?:度|℃|°c|摄氏度)|发烧?(?:39(?:\.\d)?|[4-9]\d)\s*度", "fever", "high-fever"),
    (r"脖子僵硬|颈项强直|抽搐|惊厥|意识(?:模糊|不清|丧失)|昏(?:迷|睡)", "neuro_alert", "meningitis-seizure-altered"),
    (r"(?:血尿|无尿|少尿)", "renal", "hematuria"),
    (r"(?:莫名|不明原因|突然|止不住).{0,6}(?:出血|流血)", "bleed", "spontaneous-bleeding"),
    (r"(?:突然)?心慌|心悸|心跳快.{0,6}(?:闷|痛|晕|歇)|(?:心慌|心悸).{0,6}(?:冒汗|冷汗|晕)", "cardiac", "palpitations"),
    (r"血压高\D{0,8}(?:一直|总是|老是|持续)|(?:一直|总是|老是|持续)\D{0,6}血压(?:高|150|\d{3}/)", "chronicity", "sustained-hypertension"),
    (r"(?:头晕|头昏).{0,8}(?:站不稳|晕倒|摔倒|天旋地转|走不了)", "neuro", "dizziness-unsteady"),
    (r"(?:吃什么都吐|一直吐|不停吐|反复呕吐|吐个不停|吃什么吐什么)", "gi", "persistent-vomiting"),
    (r"(?:背痛|腰痛)\D{0,10}(?:发热|发烧)|(?:发热|发烧)\D{0,4}(?:背痛|腰痛)", "infect", "back-pain-with-fever"),
    (r"(?:皮疹|荨麻疹).{0,8}(?:唇舌(?:肿|发麻)|喉头|呼吸困难)|(?:速发过敏|过敏性休克)", "anaphylaxis", "anaphylaxis"),
]


def _entity_flags(text):
    flags = []
    for pat, cat, tag in _ENTITY_RULES:
        if re.search(pat, text, re.IGNORECASE):
            flags.append(({"type": cat, "severity": "high"}, tag))
    return flags


# ---------- C. population / status flags ----------
_POP_RULES = [
    (r"怀孕|孕妇|妊娠", "pregnancy"),
    (r"(?:婴?儿|新生(?:儿|宝宝))|不足?\d*个?月?宝宝", "infant"),
    (r"\d+\s*(?:岁|周岁)", "elderly_anyage"),
    (r"\d\s*(?:岁|周岁|个月大|月龄)", "child_anyage"),
    (r"化疗|放疗|移植|免疫抑制|免疫力差|hiv|艾滋病", "immunocompromised"),
    (r"抗凝|抗血小板|华法林|阿司匹林|氯吡格雷|利伐沙班|达比加群", "anticoagulant"),
    (r"刚(?:手术|做完手术|出院)", "recent_surgery"),
]


def _pop_flags(text):
    flags = []
    child_age = None
    for pat, tag in _POP_RULES:
        m = re.search(pat, text, re.IGNORECASE)
        if m:
            if tag == "elderly_anyage":
                for am in re.finditer(r"(\d+)\s*(?:岁|周岁)", text):
                    if int(am.group(1)) >= 75:
                        flags.append(({"type": "population", "severity": "medium"}, "elderly:" + am.group(0)))
                        break
            elif tag == "child_anyage":
                for am in re.finditer(r"(\d)\s*(?:岁|周岁)", text):
                    child_age = int(am.group(1))
                    break
            elif tag == "infant":
                flags.append(({"type": "population", "severity": "medium"}, "infant"))
            else:
                flags.append(({"type": "population", "severity": "medium"}, tag))
    if child_age is not None and child_age <= 12:
        flags.append(({"type": "population", "severity": "medium"}, f"child:{child_age}岁"))
    return flags


# ---------- intent + category (minimal, for coverage stats) ----------
# Keys MUST match the 30 per-category FAISS index dirs in data/rag/ so the
# router output aligns with multi-RAG lookups.
_CATEGORY_KEYS = {
    "brain-and-nerves": [r"headache|头痛|头疼|偏头痛|脑|神经|癫痫|中风|stroke|头晕|失眠"],
    "heart-and-blood-vessel-disease": [r"心|heart|cardiac|血压|胸痛|blood pressure|心慌|心悸"],
    "diabetes": [r"糖尿|diabetes|血糖|insulin|胰岛素"],
    "pregnancy-and-childbirth": [r"怀孕|孕|pregnan|分娩|生产|baby|宝宝"],
    "childrens-health": [r"小儿|儿童|child|婴儿|幼童|新生儿"],
    "cancer": [r"癌|cancer|tumor|肿瘤|chemotherapy|化疗"],
    "allergies-and-asthma": [r"哮喘|asthma|过敏|allerg|鼻炎|花粉|喷嚏"],
    "mental-health": [r"焦虑|抑郁|depress|anxiet|情绪|精神|恐慌|惊恐"],
    "bones-joints-and-muscles": [r"关节炎|关节|骨|osteo|骨刺|腰背痛|膝盖|肩|肌肉"],
    "gastrointestinal-system": [r"胃|肠|消化|腹痛|反酸|胃食|diarrhea|便秘|nausea|腹胀"],
    "lung-disease": [r"肺|呼吸|咳嗽|cough|lung|慢性阻塞|copd|哮喘"],
    "skin-hair-and-nails": [r"皮肤|皮疹|skin|湿疹|荨麻疹|脱发|指甲"],
    "kidneys-and-urinary-system": [r"肾|尿|kidney|膀胱|泌尿|结石"],
    "blood-disorders": [r"贫血|anemia|血小|凝血|出血|血友病"],
    "infections-and-vaccines": [r"感染|传染|vaccin|疫苗|流感|flu|发烧|发热"],
    "womens-health-issues": [r"月经|卵巢|子宫|乳房|breast|menopaus|妇科|绝经|痛经"],
    "mens-health-issues": [r"前列腺|阳痿|睾丸|prostat|男性"],
    "eyes-and-vision": [r"眼|eye|视力|青光眼|白内障|近视"],
    "ear-nose-and-throat": [r"耳|鼻|喉|耳痛|耳鸣|咽|扁桃体"],
    "hormones": [r"甲状腺|thyroid|激素|内分泌|睾酮|雌激素"],
    "sleep": [r"失眠|睡眠|打鼾|sleep|insomnia"],
    "senior-health": [r"老年|老|跌|痴呆|失禁|骨质疏松|髋|老人"],
    "travel-health": [r"旅行|出差|travel|疫苗"],
    "diet-and-weight": [r"饮食|减肥|体重|diet|营养|肥胖"],
    "liver-disease": [r"肝|liver|肝炎|肝硬化|转氨酶"],
    "arthritis": [r"风湿|arthritis|类风湿|痛风"],
    "autoimmune-disease": [r"自身免疫|autoimmun|狼疮|lupus"],
    "surgery": [r"手术|surgery|术后|切口|麻醉"],
    "hiv-and-aids": [r"\bhiv|aids|艾滋病"],
    "general-health": [r".*"],
}
_ORDER = [c for c in _CATEGORY_KEYS if c != "general-health"] + ["general-health"]

_INTENT_RULES = [
    ("self_care", r"(?:自己|在家|非药|免吃药|缓解|舒缓|不做手术|自行|anything.*on my own|how to feel better|do to feel better)"),
    ("when_seek_care", r"(?:就医|看医生|急诊|去医院|call.*doctor|see.*doctor|何时.*(?:就医|看)|emergency|什么时候.*医院)"),
    ("symptoms", r"(?:症状|symptom)"),
    ("cause", r"(?:原因|导致|造成|cause|why.*(?:得|有|需要))"),
    ("diagnosis", r"(?:检查|诊断|测试|确诊|test|screen|检测)"),
    ("treatment", r"(?:治疗|treat|用药|吃药|方案|procedure|手术|therapy)"),
    ("prevention", r"(?:预防|prevent|避免|防范|keep from)"),
    ("prognosis", r"(?:恢复|预后|康复|会(?:不会)?好|能否.*好|after|手术后|恢复期)"),
    ("safety_risk", r"(?:安全|副作用|风险|危险|side effect|complicat)"),
    ("diet", r"(?:饮食|吃|食物|diet|food|营养)"),
    ("lifestyle", r"(?:运动|锻炼|exercise|活动量|生活方式)"),
    ("definition", r""),
]


def scan(text):
    """Return (red_flags, sub_intents) for a patient utterance."""
    red = []
    for data, tag in _duration_flags(text):
        data["trigger"] = tag
        red.append(data)
    for data, tag in _entity_flags(text):
        data["trigger"] = tag
        red.append(data)
    for data, tag in _pop_flags(text):
        data["trigger"] = tag
        red.append(data)
    # severity cap: population flags alone don't force alert
    has_alert = any(f["type"] in ("chronicity", "acuteness", "progression") or
                   f["type"] not in ("population",) for f in red)
    mode = "alert" if has_alert else "grounded"

    # category
    lo = text.lower()
    cat = "general-health"
    for c in _ORDER:
        if any(re.search(p, lo) for p in _CATEGORY_KEYS[c]):
            cat = c
            break
    # intent
    intent = "definition"
    for it, pat in _INTENT_RULES:
        if pat and re.search(pat, lo):
            intent = it
            break
    subs = [{"intent": intent, "entity": "", "category": cat,
             "level": "basics", "tmpl": intent, "risk": "high" if red else "low"}]
    return red, subs, mode


def main():
    rows = [json.loads(l) for l in open("data/patient_questions_clean.jsonl", encoding="utf-8")]
    n_alert = n_ground = 0
    alert_by_type = {}
    print(f"total questions: {len(rows)}")
    for r in rows:
        red, _, mode = scan(r["question"])
        if mode == "alert":
            n_alert += 1
            for f in red:
                alert_by_type[f["type"]] = alert_by_type.get(f["type"], 0) + 1
        else:
            n_ground += 1
    print(f"mode=alert : {n_alert} ({100*n_alert/len(rows):.1f}%)")
    print(f"mode=ground: {n_ground} ({100*n_ground/len(rows):.1f}%)")
    print("alert by type:")
    for k, v in sorted(alert_by_type.items(), key=lambda x: -x[1]):
        print(f"  {v:5d}  {k}")


if __name__ == "__main__":
    main()