"""
=============================================================================
第 4 讲 · 工具模块：一个小型的约束求解器（三值逻辑、回溯求解、最小冲突集、关键条件）
=============================================================================

对应讲义：lectures/lecture04_accountable_ai.md 第 3、4 节

-----------------------------------------------------------------------------
【这个模块做什么？】
-----------------------------------------------------------------------------
VERDICT 和 SatIR 把临床试验的入组条件写成 SMT（可满足性模理论）公式，交给专业的求解器（例如 Z3）。
这里从零实现一个"玩具版"，只支持本讲需要的那一小部分，但每个概念都和真实的 SMT 一一对应：

    公式（formula）  ：由原子条件（变量 ≤ 常数、布尔变量为真……）用 与 / 或 / 非 组合而成
    事实（facts）    ：每个变量的取值；**取值未知**（病历里没写）用 None 表示
    三值逻辑         ：在部分未知的事实下求值，结果是 True / False / None（未知）
    可满足性（SAT）  ：能不能给未知的变量赋值，让公式为真？能 → 给出一个"模型"（赋值）
    最小冲突集       ：公式为假时，最少需要哪几条已知事实就足以推出"假"？—— 决策的**理由**
    关键条件         ：最少改变哪几条已知事实，决策就会翻转？—— 可行动的建议（MaxSAT 的思想）

【公式的写法】用嵌套的元组表示（简单、易打印）：
    ("bool", "hiv")                  布尔变量 hiv 为真
    ("le", "ecog", 1)                ecog ≤ 1
    ("ge", "age", 18)                age ≥ 18
    ("and", [f1, f2, ...])           所有子公式都为真
    ("or",  [f1, f2, ...])           至少一个子公式为真
    ("not", f)                       f 为假
每个公式还可以带一个"标签"，方便解释：("label", "ECOG 0-1", f)

【边界】真实的 SMT 求解器支持实数、整数线性算术、位向量、数组、量词……，用的是 CDCL 等高效算法，
能处理成千上万个变量。这里的求解器用简单的回溯，只适合变量很少（几个未知数）的情况 ——
一条入组条件、一个病人，正好如此。
"""

import itertools

TRUE, FALSE, UNKNOWN = True, False, None


# =============================================================================
# 第 1 部分：三值逻辑求值
# =============================================================================
def k_and(values):
    """
    三值逻辑（Kleene 逻辑）的"与"。

    【规则】只要有一个 False → False；全都是 True → True；否则（有 None、没有 False）→ None。
    【例子】k_and([True, None]) → None（还不知道）；k_and([False, None]) → False（已经确定是假）
    【为什么】"有一条入组条件不满足"就足以判定不合格，不管其他条件是否未知。
    """
    values = list(values)
    if any(v is False for v in values):
        return False
    if all(v is True for v in values):
        return True
    return None


def k_or(values):
    """三值逻辑的"或"：有一个 True → True；全都是 False → False；否则 → None。"""
    values = list(values)
    if any(v is True for v in values):
        return True
    if all(v is False for v in values):
        return False
    return None


def k_not(v):
    """三值逻辑的"非"：True ↔ False，None 仍是 None。"""
    return None if v is None else (not v)


def evaluate(f, facts):
    """
    在（可能部分未知的）事实下对公式求值。

    【参数】f：公式；facts：dict 变量名 → 值（None 或者缺这个键都表示未知）
    【返回】True / False / None
    【Python 小知识】f[0] 是公式的类型标签；用 if 分支逐类处理，递归地处理子公式。
    """
    kind = f[0]
    if kind == "label":
        return evaluate(f[2], facts)
    if kind == "const":
        return f[1]
    if kind in ("bool", "le", "ge", "eq"):
        v = facts.get(f[1])
        if v is None:
            return None
        if kind == "bool":
            return bool(v)
        if kind == "le":
            return v <= f[2]
        if kind == "ge":
            return v >= f[2]
        return v == f[2]
    if kind == "and":
        return k_and(evaluate(g, facts) for g in f[1])
    if kind == "or":
        return k_or(evaluate(g, facts) for g in f[1])
    if kind == "not":
        return k_not(evaluate(f[1], facts))
    raise ValueError(f"unknown formula kind {kind!r}")


def variables(f):
    """公式里出现的全部变量名（集合）。"""
    kind = f[0]
    if kind == "label":
        return variables(f[2])
    if kind == "const":
        return set()
    if kind in ("bool", "le", "ge", "eq"):
        return {f[1]}
    if kind in ("and", "or"):
        return set().union(*[variables(g) for g in f[1]]) if f[1] else set()
    if kind == "not":
        return variables(f[1])
    raise ValueError(kind)


def thresholds(f, var):
    """公式里和变量 var 比较过的所有常数（用来为整数变量挑选候选值）。"""
    kind = f[0]
    if kind == "label":
        return thresholds(f[2], var)
    if kind in ("le", "ge", "eq") and f[1] == var:
        return {f[2]}
    if kind in ("and", "or"):
        out = set()
        for g in f[1]:
            out |= thresholds(g, var)
        return out
    if kind == "not":
        return thresholds(f[1], var)
    return set()


def candidates(f, var, types):
    """
    为一个未知变量挑选候选取值。

    【思路】整数变量的取值有无穷多个，但公式只关心它和几个常数的大小关系：
        把每个常数 c 以及 c − 1、c + 1 作为候选，就足以覆盖所有"不同的情况"。
        （这就是 SMT 里"理论"的作用：整数算术的推理，可以归结到有限个关键点上。）
    【参数】types：dict 变量名 → "bool" 或 "int"
    """
    if types.get(var, "int") == "bool":
        return [False, True]
    cs = set()
    for c in thresholds(f, var):
        cs |= {c - 1, c, c + 1}
    return sorted(cs) or [0]


# =============================================================================
# 第 2 部分：可满足性（回溯求解）
# =============================================================================
def solve(f, facts, types, stats=None, prefer=None):
    """
    能不能给未知的变量赋值，使公式为真？

    【参数】
        f     : 公式
        facts : 已知的事实（未知的变量取 None 或缺省）
        types : 变量的类型
        stats : 传入 dict 时，记录搜索了多少个节点（stats["nodes"]）
        prefer: dict 变量名 → 偏好的值；候选值按"离偏好值多近"的顺序尝试（找关键条件时，优先找最小的改动）
    【返回】(是否可满足, 赋值)。可满足时，赋值是一个 dict：未知变量 → 选定的值（"假设"）。
    【思路】回溯搜索：
        1. 先在当前（部分）赋值下求值：已经是 True → 成功；已经是 False → 这条路走不通，回头；
        2. 否则挑一个还没赋值的未知变量，依次尝试它的候选值，递归。
        第 1 步的"提前判断"就是剪枝：三值逻辑让我们在变量还没全部赋值时就能发现矛盾。
    【边界】只适合未知变量很少的情况（最坏情况是指数级的）。
    """
    unknown = sorted(v for v in variables(f) if facts.get(v) is None)
    current = dict(facts)

    def bt(i):
        if stats is not None:
            stats["nodes"] = stats.get("nodes", 0) + 1
        val = evaluate(f, current)
        if val is True:
            return True
        if val is False or i == len(unknown):
            return False
        var = unknown[i]
        cands = candidates(f, var, types)
        if prefer and prefer.get(var) is not None and types.get(var, "int") == "int":
            cands = sorted(cands, key=lambda c: (abs(c - prefer[var]), c))
        for c in cands:
            current[var] = c
            if bt(i + 1):
                return True
        current[var] = None
        return False

    ok = bt(0)
    if not ok:
        return False, {}
    return True, {v: current[v] for v in unknown if current.get(v) is not None}


# =============================================================================
# 第 3 部分：解释 —— 最小冲突集与关键条件
# =============================================================================
def minimal_conflict(f, facts, types):
    """
    公式不可满足时，找出一个**最小冲突集**：一组已知事实，单凭它们（其余事实都当作未知）就足以让公式不可满足；
    而且去掉其中任何一条，就不再不可满足。

    【返回】list[变量名]（按字母序）；公式可满足时返回空列表
    【思路】"逐个删除"算法：从全部已知事实开始，依次尝试把每条事实改成未知；如果仍然不可满足，
        说明这条事实不是必需的，就真的删掉它。最后剩下的就是一个最小冲突集。
        这对应 SMT 求解器的 unsat core（不可满足核心），是"为什么不合格"的**可验证的理由**。
    """
    if solve(f, facts, types)[0]:
        return []
    keep = {k: v for k, v in facts.items() if v is not None and k in variables(f)}
    for var in sorted(keep):
        trial = dict(keep)
        trial[var] = None
        if not solve(f, trial, types)[0]:
            keep = trial
    return sorted(k for k, v in keep.items() if v is not None)


def pivotal_conditions(f, facts, types, want, max_size=3):
    """
    关键条件（MaxSAT 的思想）：最少需要改变哪几条已知事实，才能让决策变成 want？

    【参数】
        want     : True = 想让公式变为可满足（不合格 → 合格）；False = 想让公式的否定可满足（合格 → 不合格）
        max_size : 最多考虑同时改变几条事实
    【返回】list[(要改变的变量的元组, 改成什么值的 dict)]：所有**最小**的改变方案；找不到时返回空列表
    【思路】把每条已知事实看成一条"软约束"（权重为 1），把试验条件看成"硬约束"。
        MaxSAT 要找的是：满足硬约束的前提下，违反（改变）的软约束最少。
        这里用最朴素的做法：按大小 1、2、3 依次枚举要"放开"的事实组合，放开的变量交给求解器重新赋值；
        第一个找到解的大小就是最小的改变数，返回这一大小的所有方案。
    【边界】真实的 VERDICT 用 MaxSMT 求解器；这里的枚举只适合事实很少的情况。
    """
    target = f if want else ("not", f)
    known = sorted(k for k, v in facts.items() if v is not None and k in variables(f))
    for size in range(1, max_size + 1):
        found = []
        for combo in itertools.combinations(known, size):
            trial = dict(facts)
            for var in combo:
                trial[var] = None
            ok, model = solve(target, trial, types, prefer={v: facts[v] for v in combo})
            if ok and all(model.get(v) is not None and model[v] != facts[v] for v in combo):
                found.append((combo, {v: model[v] for v in combo}))
        if found:
            return found
    return []


def describe(f):
    """把公式写成人能读的字符串（标签优先）。"""
    kind = f[0]
    if kind == "label":
        return f[1]
    if kind == "const":
        return str(f[1])
    if kind == "bool":
        return f[1]
    if kind == "le":
        return f"{f[1]} ≤ {f[2]}"
    if kind == "ge":
        return f"{f[1]} ≥ {f[2]}"
    if kind == "eq":
        return f"{f[1]} = {f[2]}"
    if kind == "not":
        return f"NOT ({describe(f[1])})"
    sep = " AND " if kind == "and" else " OR "
    return "(" + sep.join(describe(g) for g in f[1]) + ")"


def failing_clauses(f, facts):
    """顶层是"与"的公式里，在已知事实下已经为假的那些子句（附标签），用来解释"哪条条件不满足"。"""
    if f[0] == "label":
        return failing_clauses(f[2], facts) or ([f[1]] if evaluate(f, facts) is False else [])
    if f[0] != "and":
        return [describe(f)] if evaluate(f, facts) is False else []
    return [describe(g) for g in f[1] if evaluate(g, facts) is False]


if __name__ == "__main__":
    # 一个小例子：入组条件 = 年龄 18~65 且 ECOG ≤ 1 且 不是 HIV 阳性
    f = ("and", [("label", "age 18-65", ("and", [("ge", "age", 18), ("le", "age", 65)])),
                 ("label", "ECOG 0-1", ("le", "ecog", 1)),
                 ("label", "no HIV", ("not", ("bool", "hiv")))])
    types = {"age": "int", "ecog": "int", "hiv": "bool"}
    for facts in [{"age": 54, "ecog": 1, "hiv": False}, {"age": 54, "ecog": 1, "hiv": None},
                  {"age": 54, "ecog": 2, "hiv": None}]:
        ok, model = solve(f, facts, types)
        print(f"事实 {facts}")
        print(f"  三值求值：{evaluate(f, facts)}；可满足：{ok}；需要的假设：{model}")
        if not ok:
            print(f"  最小冲突集：{minimal_conflict(f, facts, types)}；"
                  f"关键条件：{pivotal_conditions(f, facts, types, want=True)}")