# ruff: noqa: ERA001

from __future__ import annotations

import contextlib
import logging
import warnings
from collections import deque
from dataclasses import field, dataclass
from functools import partial
from typing import TYPE_CHECKING, Any

from kele.control import HookMixin
from kele.grounder.grounded_rule_ds.grounded_ds_utils import unify_all_terms
from kele.syntax import Assertion, Formula
from ._nodes import _AssertionNode, _ConnectiveNode, _FlatCompoundTermNode, _OperatorNode, _QuestionRuleNode, _RootNode, \
    _RuleNode
from ._nodes._tupletable import NameToObject, _TupleTable
from .rule_check import RuleCheckGraph

if TYPE_CHECKING:
    from collections.abc import Mapping, Sequence

    from kele.config import Config
    from kele.control import InferencePath
    from kele.equality import Equivalence
    from kele.syntax import FACT_TYPE, GROUNDED_TYPE_FOR_UNIFICATION, CompoundTerm, Constant, Rule, SankuManagementSystem, Variable
    from kele.syntax.base_classes import _QuestionRule

logger = logging.getLogger(__name__)


@dataclass(slots=True)
class RuleRuntimeContext:
    """Per-run runtime context for one `GroundedRule.check_grounding()` execution.

    This object stores *round-local* state only (tables produced during this run).
    It is created at the beginning of `check_grounding()` and discarded afterwards.

    Why:
    - Keep `GroundedRule` as an orchestrator (static graph + cross-run state).
    - Make Nodes independent from `GroundedRule`'s private fields by passing explicit inputs.
    - Avoid implicit singletons/broadcasts for per-run tables.
    """

    # Step-3: merged/total table produced in this run (delta)
    non_action_grounding_total_table: _TupleTable | None = None
    action_grounding_total_table: _TupleTable | None = None
    total_table: _TupleTable | None = None

    # Used in step-3: whether tables were aligned for total-table computation in this run
    binding_consistency_for_total_table: bool = False

    # Action-assertion new facts produced in this run
    action_new_fact_list: list[Any] = field(default_factory=list)


class GroundedRule(HookMixin):
    """
    管理单条规则的实例化状态，负责 term-level unify、合并变量候选表与 check 阶段的执行。

    GroundedRule 不保存完整展开后的 grounded rules，而是维护变量候选表与执行状态，
    在 check 时再按需展开并生成新的事实。
    """

    def __init__(self, rule: Rule, equivalence: Equivalence, sk_system_handler: SankuManagementSystem,
                 args: Config, inference_path: InferencePath) -> None:
        super().__init__()
        if rule.unsafe_variables:
            raise TypeError(f"""Rule {rule!s} is unsafe because it contains unsafe variables {[str(u) for u in rule.unsafe_variables]}.\n
                            This error likely appears because the rule was not added to RuleBase or did not go through preprocessing.
                            """)
        if not self._is_conjunctive_body(rule):
            warnings.warn(f"""
                          Rule {rule!s} body must be a conjunction of positive and negative assertions; this rule does not meet the requirement.\n
                          This warning likely appears because the rule was not added to RuleBase or did not go through preprocessing.\n
                          For more information, see: #TODO
                          """, stacklevel=5)  # TODO: add engine tutorial URL
        if not self._is_conjunctive_head(rule):
            warnings.warn(
                f"Rule {rule!s} head contains connectives other than AND, which may prevent correct fact generation.",
                stacklevel=5,
            )

        self.args = args
        self.rule = rule
        self.rule_checker = RuleCheckGraph(self, self.args)  # risk: 如果fact里面里面是forall的很难处理，实操时候需要先实例化fact再推。
        self.is_concept_compatible_binding = partial(self.rule.is_concept_compatible_binding,
                                              fuzzy_match=self.args.grounder.conceptual_fuzzy_unification)

        self.inference_path: InferencePath = inference_path

        self.equivalence = equivalence
        self.sk_system_handler = sk_system_handler

        self._past_df_prefix_sum: list[_TupleTable] = []  # 前缀和会有较大的浪费，即第二个df包含第一个df，以此类推
        # FIXME: 这个应当逐步挪到conn去维护
        # 另外这个不进入ctx是希望它被持续地保存下来，而不是一个临时的contextvar

        columns = set()
        # 引擎只使用部分nodes进入grounding过程，部分node只做substitution。此外，substitution nodes中的action op相关nodes也会影响grounded rule的生成
        # 但其不增加已通过semi-naive策略生成的grounded rule的数量，也不需要进行unification，因此不记录与之相关的past_df
        for assertion_node in self.rule_checker.non_action_grounding_nodes + self.rule_checker.action_grounding_nodes:
            columns |= assertion_node.grounding_arguments
            self._past_df_prefix_sum.append(_TupleTable(column_name=tuple(columns)))

        self._back_up_total_table: _TupleTable  # 记录每一轮的total table，便于debug时输出grounded rules有哪些
        self._past_rule_table: _TupleTable = _TupleTable(self.rule.free_variables)  # 记录历史以来所有匹配成功的rule级别的table，
        # 以减少重复计算（但逐渐作用会减少）
        # TODO: 比如此时total已经在exec_check之后了，可能逐渐这个就删除了
        self._last_non_action_grounding_total_table: _TupleTable = _TupleTable(self._past_df_prefix_sum[-1].raw_column) \
            if self.rule_checker.action_grounding_nodes else _TupleTable.create_empty_table_with_emptyset()
        # HACK: 重构不充分，早期设计的的算法是在一轮内完成所有的事情，而现在是先后两轮（一轮的话，action term的匹配需要等mid table出来再unify）
        # 两轮的话，也就是说action的new，它实际上应该是用前一轮的non action total table，而不是当前轮的
        # 这个pr就不再进行一轮这种级别的重构了，已经蛮大了，因此ad-hoc地先留一个_last_non_action_grounding_total_table变量

    def unify(self, terms: list[CompoundTerm[Constant | CompoundTerm] | Constant]) -> None:
        """
        对传入的 term 进行 unify，仅生成实例化候选值，不检查事实的正确性。

        - FlatCompoundTermNode 会先将常量替换为等价类代表元；
        - unify 仅走 term-level，后续由 AssertionNode 进行 join 与 check。

        :params terms: 用于实例化的事实 term 列表
        """
        node_queue: deque[_RootNode | _OperatorNode | _FlatCompoundTermNode] = deque()
        root: _RootNode = self.rule_checker.graph_root
        node_queue.append(root)

        while node_queue:
            cur_node = node_queue.popleft()
            if isinstance(cur_node, _FlatCompoundTermNode):
                cur_node.process_equiv_represent_elem()
            elif isinstance(cur_node, (_OperatorNode, _RootNode)):
                node_queue.extend(cur_node.query_for_children())

        for t in terms:
            self._unify_single(t)
        self._start_passing_process()

    def check_grounding(self) -> list[FACT_TYPE]:
        """
        执行 check 阶段并返回新事实。

        AssertionNode 会执行 action assertion 的计算并生成额外事实，最终由 RuleNode 汇总。

        :returns list[FACT_TYPE]: check得到的新事实
        """
        self._validate_non_action_grounding_nodes()
        ctx = RuleRuntimeContext()

        new_facts: list[FACT_TYPE] = []
        action_start_index = len(self.rule_checker.non_action_grounding_nodes)

        non_action_grounding_delta_table, prefix_sum_for_action_nodes = self._join_nodes(
            nodes=self.rule_checker.non_action_grounding_nodes,
            start_index=0,
        )

        # 完成non action grounding的全流程后，记录此刻的table情况
        if self.rule_checker.non_action_grounding_nodes:
            ctx.non_action_grounding_total_table = non_action_grounding_delta_table
        else:  # 规则前提不含 free vars 时，total table 为空列。  TODO：考虑是否将这种特殊规则单独出来，不走unify的流程
            ctx.non_action_grounding_total_table = _TupleTable(column_name=())
            prefix_sum_for_action_nodes = _TupleTable(column_name=())

        # 将non action grounding的table广播给action grounding
        self._broadcast_total_table(self._last_non_action_grounding_total_table, self.rule_checker.action_grounding_nodes)
        self._last_non_action_grounding_total_table = ctx.non_action_grounding_total_table

        # ==============action grounding nodes====================
        # Action assertions compute action results during exec_check.

        # 下面的join nodes对应这三步
        # self.variable_binding_consistency_check()  # 这里要检查non action grounding算出来的表、和action grounding nodes
        # 中非action部分实例化的表，它们之间的semi join
        # self.check_assertion_truth(self.rule_checker.action_grounding_nodes)
        # self._preprocess_node_bindings(self.rule_checker.action_grounding_nodes,
        #                                start_index=len(self.rule_checker.non_action_grounding_nodes))

        action_grounding_delta_table, action_grounding_past_table = self._join_nodes(
            nodes=self.rule_checker.action_grounding_nodes,
            start_index=action_start_index,
            previous_new_table=ctx.non_action_grounding_total_table,
            previous_prefix_sum=prefix_sum_for_action_nodes,
        )

        if self.rule_checker.action_grounding_nodes:
            ctx.action_grounding_total_table = action_grounding_delta_table
        else:
            # keep correct columns for downstream substitution
            ctx.action_grounding_total_table = _TupleTable(column_name=())
        # notice: Combine none-action (A) and action (B) deltas should also consider semi-naive cross terms:
        # Δ(A∧B) = (ΔA ⋈ B_old) ∪ (ΔB ⋈ A_old) ∪ (ΔA ⋈ ΔB)
        merged_total_table = self._merge_total_tables(
            ctx.non_action_grounding_total_table,
            action_grounding_delta_table,
            non_action_old_table=prefix_sum_for_action_nodes,
            action_old_table=action_grounding_past_table,
        ) \
            if self.rule_checker.non_action_grounding_nodes else _TupleTable(column_name=())  # 有action肯定有non-action，所以这里不影响

        if self.rule_checker.non_action_grounding_nodes:  # 有action肯定有non-action，所以这里不影响
            if self.args.executor.anti_join_used_facts and hasattr(self, "_past_rule_table"):
                merged_total_table = merged_total_table.anti_join(self._past_rule_table)
                # 现在没有提前在俩table上anti，一定程度上削弱了anti的作用。但从保守的角度来说，这个是rule级别的table anti，那
                # 没有合并到rule级别前确实不能随意anti，容易漏接

            ctx.total_table = merged_total_table
        else:  # 规则前提不含free vars时，total table为空列。  TODO：考虑是否将这种特殊规则单独出来，不走unify的流程
            ctx.total_table = _TupleTable(column_name=())

        # 结束后将total table传给substitution nodes做校验
        self._broadcast_total_table(ctx.total_table, self.rule_checker.substitution_nodes)
        self._back_up_total_table = ctx.total_table

        if logger.isEnabledFor(logging.DEBUG):
            logger.debug(
                "Grounded rule final total table broadcast: rule=%s summary=%s",
                self.rule,
                ctx.total_table.debug_summary(),
            )

        # =============substitution nodes========================
        for node in self.rule_checker.substitution_nodes:
            node.exec_check()

        self._execute_queue(ctx, new_facts)
        self._append_action_new_facts(ctx, new_facts)

        return new_facts

    def _validate_non_action_grounding_nodes(self) -> None:
        for node in self.rule_checker.non_action_grounding_nodes:
            if (not node.action_assertion and
                not node.negated_assertion and
                not set(node.content.free_variables) >= set(node.grounding_arguments)):
                raise ValueError(
                    "Non-action grounding node has grounding vars not covered by content.free_variables."
                    ""
                )  # TODO: 要注意因为封装了_join_nodes，且variable_binding_consistency_check
            # 依赖于传入nodes的缘故，会导致一致性其实没办法检查non action grounding算出来的表、和action grounding nodes之间的semi join
            # 能支持的话多少有利于减少匹配数，代价是封装稍微不清晰一点，要有个单独的variable_binding_consistency_check调用

    def _execute_queue(self, ctx: RuleRuntimeContext, new_facts: list[FACT_TYPE]) -> None:
        # Step-4: connective / rule node evaluation
        execute_queue: deque[_ConnectiveNode | _RuleNode] = deque()

        for node in self.rule_checker.non_action_grounding_nodes:
            node.pass_tf_index(table=ctx.total_table)
            execute_queue.extend(node.query_for_children())

        for node in self.rule_checker.action_grounding_nodes:
            node.pass_tf_index()
            execute_queue.extend(node.query_for_children())
            ctx.action_new_fact_list.extend(node.get_action_result())

        for node in self.rule_checker.substitution_nodes:
            node.pass_tf_index()
            execute_queue.extend(node.query_for_children())

        while execute_queue:
            cur_node = execute_queue.popleft()
            if isinstance(cur_node, _ConnectiveNode) and cur_node.ready_to_execute:  # AssertionNode总是最先被执行，而是否可以执行ConnectiveNode，由
                # ConnectiveNode的ready_for_execute属性控制。注意可以执行即要求它的父节点都向它传递了TfIndexs。
                cur_node.exec_check()
                cur_node.pass_tf_index()
                execute_queue.extend(cur_node.query_for_children())
            elif isinstance(cur_node, _RuleNode):
                # 一般来说对于单个rule，RuleNode只有一个，所以直接赋值即可  # FIXME: 会多的以后，先留着
                new_facts.extend(cur_node.exec_check())

    def _append_action_new_facts(self, ctx: RuleRuntimeContext, new_facts: list[FACT_TYPE]) -> None:
        if ctx.action_new_fact_list:
            new_facts += ctx.action_new_fact_list

        if self.rule_checker.non_action_grounding_nodes:
            binding_table = ctx.non_action_grounding_total_table
            if binding_table is None:
                raise RuntimeError("non_action_grounding_total_table is required for action execution.")
            for node in self.rule_checker.action_grounding_nodes:
                node.exec_action(binding_table)
                new_facts.extend(node.get_action_result())

    def _join_nodes(
        self,
        *,
        nodes: Sequence[_AssertionNode],
        start_index: int,
        previous_new_table: _TupleTable | None = None,
        previous_prefix_sum: _TupleTable | None = None,
    ) -> tuple[_TupleTable, _TupleTable]:
        """对一系列给定的nodes，执行变量绑定一致性检查（通过semi join）、exec_check、再次semi join、合并table"""
        if not nodes:
            empty_table = _TupleTable(column_name=())
            return empty_table, previous_prefix_sum or empty_table

        # Step-1: variable binding consistency check -> check table
        self.variable_binding_consistency_check(nodes)  # HACK: 这部分暂时仍然需要合并整个total table的方式去干，以保证past是在check的
        # 范围内，而不会在下面check的时候丢失past，从而导致后续还要补一次check。而由于当我们加入证明的advanced算法之后，这一次check就不需要补了
        # 且现在补一次check有点破坏AssertionNode的功能，它就是个纯HACK，还是有点破坏性的HACK。所以我打算在这里仍然承担较高的内存开销，
        # 把这个开销留在advanced的代码进来之后，再补（如果后面实在内存爆了，我将单独写一个check函数在这里，然后exec_check的地方大不了重新check
        # 一遍；或者这里用semi join，但step3后面再接一轮完全的check，而不是只check past里被选出来的部分。换句话说，通过时间代价降低HACK对代码
        # 流程的影响。因为如果简写past与total取差集的话，前面还有找地方记录第一次check的true等等，回头这些又都用不上）。
        # 正常的流程，就是说termnode和assertion之间，有一次"汇总分发"，这里先后执行一致性检查（semi join）、exec_check。
        # assertion node和conn node之间，再执行一次一致性检查、然后拆分再pass tf table

        # self._broadcast_total_table_hack()  # 这个是一定会有的，哪怕是对df做原地semi join的变更。所以此HACK破坏性不大。现在是semi join +
        # check + semi_naive的分发；之后改成semi_join直接分发即可。（为了方便，这里直接写两个分发，回头这个分发删除即可
        # 我们此时先跳过一致性检查，这个hack暂时先不加。等我们试着对比用total table的hack时再看情况

        # 最终来说，step1需要分发当前轮的 freevar + unknown，加起来才是本轮的"freevar"。对它们加起来进行semi join后，
        # 发回去作为to_be_checked

        # Step-2: exec assertions to check T/F
        past_true_tables = [node.past_freevar_table.true for node in nodes]
        self.check_assertion_truth(nodes)
        # HACK: 这里面对past的处理不够完美，已在内部注释
        # HACK: 这里还引入了一个假设是：每个断言的table能够包含content的所有free vars。但特殊优化下这一点不一定可以保证。
        # 不过据我目前估计，如果出现这种情况，大概率实例化部分要改为backjumping的风格，而不会是现在这样。现在应该仅有action nodes存在这种

        # Step-3: compute total table for this run
        past_tables, new_tables = self._preprocess_node_bindings(
            nodes,
            start_index=start_index,
            past_true_tables=past_true_tables,
        )

        # self._semi_join()  # 这个函数单独封装出来，哪怕是semi naive时候，也不妨用semi join先降一波内存开销

        new_table, next_prefix_sum = self._semi_naive(
            past_tables=past_tables,
            new_tables=new_tables,
            start_index=start_index,
            initial_new_table=previous_new_table,
            initial_last_table=previous_prefix_sum,
        )
        return new_table, next_prefix_sum

    def _unify_single(self, term: CompoundTerm[Constant | CompoundTerm] | Constant) -> None:
        node_queue: deque[_RootNode | _OperatorNode | _FlatCompoundTermNode] = deque()
        root: _RootNode = self.rule_checker.graph_root
        node_queue.append(root)

        while node_queue:
            cur_node = node_queue.popleft()
            if isinstance(cur_node, _FlatCompoundTermNode) and not cur_node.only_substitution:  # FIXME: 带着一个term往下走，似乎要多判断很多次这个
                # 另外这个only的判断可能也有待优化
                cur_node.exec_unify(term, allow_unify_with_nested_term=self.args.grounder.allow_unify_with_nested_term)
                # 对_FlatCompoundTermNode进行unification操作
            elif isinstance(cur_node, (_OperatorNode, _RootNode)):
                node_queue.extend(cur_node.query_for_children(term))

    def receive_true_table(self, true_table: _TupleTable) -> None:
        """
        接收从AssertionNode传递来的true indexs
        """
        self._past_rule_table = self._past_rule_table.concat_table(true_table)

    def _start_passing_process(self) -> None:
        """
        启动传递过程，将 freevar_table 传递到 _AssertionNode。
        """
        node_queue: list[_RootNode | _OperatorNode | _FlatCompoundTermNode | _AssertionNode] = []
        root: _RootNode = self.rule_checker.graph_root
        node_queue.append(root)

        while node_queue:
            cur_node = node_queue.pop()
            if isinstance(cur_node, _FlatCompoundTermNode):
                cur_node.pass_freevar_to_child()
                node_queue.extend(cur_node.query_for_child())
            elif isinstance(cur_node, (_OperatorNode, _RootNode)):
                node_queue.extend(cur_node.query_for_children())

    def _preprocess_node_bindings(
        self,
        nodes: Sequence[_AssertionNode],
        *,
        start_index: int = 0,
        past_true_tables: Sequence[_TupleTable] | None = None,
        ) -> tuple[list[_TupleTable], list[_TupleTable]]:
        """
        预处理和计算 rule 级别变量绑定相关的操作，包括收集各节点的绑定与代表元替换。
        TODO: 当semi join实现后，当前操作可能提前或调整为简单合并某处的代码。而equiv操作可能会进一步前置。但总之此刻先合在一起吧，阅读上也方便
        """
        if past_true_tables is None:  # past_true_tables的目的是，当exec_check前置时，它check后的past就不再是真实的past了，所以
            # 需要靠传入一个past来覆盖。而这个函数本身是可以不考虑exec_check的，这时候只要自行获取true/false即可，取决于节点的情况
            # HACK: 这里面有许多地方取了true，是基于此时DNF下的情况的。但广泛意义上可能要优化一下
            past_true_tables = [node.past_freevar_table.true for node in nodes]

        past_tables: list[_TupleTable] = []
        new_tables: list[_TupleTable] = []

        for idx, assertion_node in enumerate(nodes):
            past_table = past_true_tables[idx]  # 注释在之前的代码里，为了方便来回切换，我先单独放一块
            new_table = assertion_node.tf_table.true

            if past_table.height > 0 and not past_table.raw_column:  # noqa: SIM108 和下面的意思一样
                new_table = _TupleTable(column_name=())
            else:
                new_table = new_table  # .anti_join(past_table)。TODO: 现在不能随便删，因为等价关系比较麻烦，anti的话  # noqa: PLW0127
                # 无法区分是不是当前轮又一次成功了的。所以这个还是要靠更高层的标记才能看。这里先暂时放着叭
                # TODO: 如果在高层处理了的话，那这个函数的past入参可能就不用要了
            past_tables.append(past_table)
            new_tables.append(new_table)

        # for assertion_node in nodes:
        #     past_tables.append(assertion_node.past_freevar_table.true)  # 这里取past true就够了，因为past true + unknown已经
        #     # 在之前被取出和check了。如果没被check的unknown说明也不需要在semi naive引入
        #     # TODO: 但最自然的形式的话，这里应该是unknown被check的true+上一轮的true。不过我在想，unknown的如果在新的时候才被发现是true
        #     # 似乎完全可以当做new去对待，因为作为past的意义是它不会被使用，但是如果它true了，那它就是新的就要被使用。
        #     new_tables.append(assertion_node.tf_table.true)  # 由于check过程前置，所以这里直接取check好的true，而不是
        #     # 像以前一样取所有的待check的（unknown） + 已知true（重新check一遍）
            # TODO: 理论上这里最好可以anti join掉exec_check前的true。并且上面的past应该是取
            # FIXME: 另外，如果表是一层一层的，那么有了下层后，上层应当予以移除。比方说这时候一致性后的表就没用了。

        prefix_sum_tables = self._past_df_prefix_sum[
            start_index:start_index + len(new_tables)
        ]
        new_tables, past_tables, prefix_sum_tables = self._equiv_represent_replaced(  # TODO: 这里如果移除new_tables可以避免
            # 从term nodes后进行了重复的replace。但有点风险，具体来说是：
            # freevar因为合并之前替换过等价类而可以不再替换；但以后那里面可能有unknown的，不替换会导致一定的不对齐
            new_tables,
            past_tables,
            prefix_sum_tables,
        )
        # HACK: 这里有点突兀，不过本身也是应该逐步挪到conn的
        self._past_df_prefix_sum[
            start_index:start_index + len(prefix_sum_tables)
        ] = prefix_sum_tables
        return past_tables, new_tables

    def _build_equiv_mapping(self, tables: list[_TupleTable]) -> tuple[dict[str, str], dict[int, list[str]]]:
        unique_values_by_table: dict[int, list[str]] = {}
        unique_values: set[str] = set()
        for table in tables:
            values = table.unique_values()
            unique_values_by_table[id(table)] = values
            unique_values.update(values)

        mapping: dict[str, str] = {}
        for old in unique_values:
            old_elem = NameToObject.get_item_by_name(old)
            rep_elem = self.equivalence.get_represent_elem(old_elem)
            if old_elem != rep_elem:
                mapping[old] = NameToObject.get_item_name(rep_elem)
        return mapping, unique_values_by_table

    @staticmethod
    def _subset_mapping(
        mapping: dict[str, str],
        unique_values_by_table: dict[int, list[str]],
        table: _TupleTable,
    ) -> dict[str, str]:
        values = unique_values_by_table.get(id(table), [])
        if not values or not mapping:
            return {}
        return {old: mapping[old] for old in values if old in mapping}

    @staticmethod
    def semi_naive_calculation(
            delta_left: _TupleTable,
            delta_right: _TupleTable,
            *,
            old_left: _TupleTable,
            old_right: _TupleTable,
            empty_column_name: tuple[Any, ...] | None = None,
    ) -> _TupleTable:
        """Compute semi-naive delta join: Δ(A∧B) from deltas and old tables."""
        terms: list[_TupleTable] = []
        if delta_left.height > 0 and old_right.height > 0:
            terms.append(delta_left.union_table(old_right))
        if old_left.height > 0 and delta_right.height > 0:
            terms.append(old_left.union_table(delta_right))
        if delta_left.height > 0 and delta_right.height > 0:
            terms.append(delta_left.union_table(delta_right))
        if not terms:
            if empty_column_name is None:
                empty_column_name = tuple(
                    set(delta_left.raw_column) | set(delta_right.raw_column)
                )
            return _TupleTable(column_name=empty_column_name)
        if len(terms) == 1:
            return terms[0]
        return terms[0].concat_table(*terms[1:])

    def _semi_naive(
        self,
        *,
        past_tables: Sequence[_TupleTable],
        new_tables: Sequence[_TupleTable],
        start_index: int = 0,
        initial_new_table: _TupleTable | None = None,
        initial_last_table: _TupleTable | None = None,
    ) -> tuple[_TupleTable, _TupleTable]:
        """TODO: 目前版本是(A1+A1')(A2+A2')...的，还没加入设计的其他公式"""
        if not new_tables:
            empty_table = _TupleTable(column_name=())
            return empty_table, initial_last_table or empty_table
        if len(past_tables) != len(new_tables):
            raise ValueError("Semi-naive expects past/new table lists with the same length.")

        t1_last = initial_last_table if initial_last_table is not None else self._past_df_prefix_sum[start_index]  # 对应前j-1个的总last，此时是第0个
        t1_new = initial_new_table if initial_new_table is not None else new_tables[0]  # 对应前j-1个总new，此时是第0个
        next_prefix_sum = t1_last

        base_new_tables = new_tables
        loop_start = 0
        if initial_new_table is None:
            self._past_df_prefix_sum[start_index] = self._past_df_prefix_sum[start_index].concat_table(t1_new)  # 更新为第T轮的prefix sum
            loop_start = 1

        for i in range(loop_start, len(base_new_tables)):
            # 对于任意两项，其可以拆解为A1A2 + A1A2' + A1'A2 + A1'A2'。t1_last对应A1，t1new对应A1'。t2同理
            t2_last = past_tables[i]
            t2_new = base_new_tables[i]

            new_table = self.semi_naive_calculation(
                t1_new,
                t2_new,
                old_left=t1_last,
                old_right=t2_last,
                empty_column_name=self._past_df_prefix_sum[start_index + i].raw_column,
            )

            nxt_prefix_sum = self._past_df_prefix_sum[start_index + i]

            if new_table.height > 0:
                self._past_df_prefix_sum[start_index + i] = self._past_df_prefix_sum[start_index + i].concat_table(new_table)
                # 更新prefix sum，用于下次

            t1_last = nxt_prefix_sum  # 继续为i+1项做准备而迭代
            t1_new = new_table
            next_prefix_sum = t1_last
        return t1_new, next_prefix_sum

    def _equiv_represent_replaced(
        self,
        *table_groups: Sequence[_TupleTable],
    ) -> tuple[list[_TupleTable], ...]:
        """对齐等价类代表元并按输入分组返回。"""
        groups: Sequence[Sequence[_TupleTable]] = table_groups
        aligned_groups = [list(group) for group in groups]
        tables_to_align = [
            table for group in aligned_groups for table in group
        ]
        if not tables_to_align:
            return tuple(aligned_groups)

        # HACK: 博洋觉得这里有点不好看
        mapping, unique_values_by_table = self._build_equiv_mapping(
            tables_to_align
        )
        updated_groups = [
            [
                table.update_equiv_element(
                    self._subset_mapping(mapping, unique_values_by_table, table)
                )
                for table in group
            ]
            for group in aligned_groups
        ]
        return tuple(updated_groups)

    @staticmethod
    def variable_binding_consistency_check(nodes: Sequence[_AssertionNode] | None = None) -> None:
        """
            Step-1: variable binding consistency check.
        """
        # TODO: 规则前提不含 grounding free vars 时，检查域是空列的单行表 {()}。
        with contextlib.suppress(MemoryError):
            # self._semi_join()  # HACK: 在有advanced的情况下，这里需要被传递一系列设计好的node/df，然后直接执行即可
            # 在没有的情况下，为了和以后更相似，这里选择只考虑new的。然后past的部分先通过仍然保留计算完整的total table作为代价来保留；
            # 虽然还有个替代方案是直接先跳过一致性检查，这样的话，考虑到后面的check还是需要在step2之前有一步用于分配待check的df，
            # 比方说可能是下面的_broadcast_df_to_be_checkd_exec_check_hack或者_broadcast_total_table_hack。这时候
            # 直接分配new+past_unknown即可，不过unknown此时为空可以忽略
            # 而如果past不用total table，且还没有纳入advanced时，判断哪些是past部分的将会比较奇怪。但由于内存在一些示例上遇到了瓶颈，
            # 考虑临时先允许跳过这个步骤。

            # self._broadcast_df_to_be_checkd()  # 暂且跳过，之后advanced时，会有n次concat，所以到时候concat做一个缓存
            # 此刻_broadcast_df_to_be_check的功能，在semi_join里面相当于已经实现掉了。而此刻没有broadcast时，相当于可以用freevar_table
            # 来正常执行

            # 存储结构换一下，成list[list]
            # ctx.binding_consistency_for_total_table = True
            pass

        # if self.rule_checker.grounding_nodes:
        #     total_table = self._calc_total_table(nodes)  # 先注释掉，思路是通过action的exec将其merge完整，后期可能action会单独处理。依赖的假设是
        #     # 我们没有表达2+3的需求（对于action op来说），任何时候都是想说5。真的非要表达2+3时，建议单独设置一个print_arithmetic_op这种算子
        #
        #     if self.args.executor.anti_join_used_facts and hasattr(self, "_past_rule_table"):
        #         # 在config中开启anti_join_used_facts、且已经有备份的情况下，执行anti join操作
        #         total_table = total_table.anti_join(self._past_rule_table)
        #
        #     self.total_table = total_table
        # else:  # 规则前提不含free vars时，total table为空列。  TODO：考虑是否将这种特殊规则单独出来，不走unify的流程
        #     self.total_table = _TupleTable(column_name=())
        #
        # self._broadcast_total_table(self.total_table)
        # self._back_up_total_table = self.__dict__.pop('total_table')

    @staticmethod
    def check_assertion_truth(nodes: Sequence[_AssertionNode]) -> None:
        """检查每个assertion级别的变量绑定的正确性"""
        for node in nodes:
            if isinstance(node, _AssertionNode) and node.ready_to_execute:
                node.exec_check()  # 现在exec_join被取消了，直接check即可。如果有semi join，那它会赋一个to be check的值直接覆盖掉
                # total table或freevar_table带来的small table，所以可以统一用exec_check

    @staticmethod
    def _is_conjunctive_body(rule: Rule) -> bool:
        """判断当前规则的body的连接词只出现and, not，且not只能作用在Assertion上（换言之就是没有OR的DNF格式）"""
        is_standard = True
        cur_term_queue: deque[Formula | Assertion | None] = deque([rule.body])
        while cur_term_queue:
            cur_term = cur_term_queue.popleft()
            if isinstance(cur_term, Formula):
                if cur_term.connective not in {'AND', 'NOT'}:
                    is_standard = False
                    break
                if cur_term.connective == 'NOT' and not isinstance(cur_term.formula_left, Assertion):
                    # NOT必须作用在Assertion上，不能作用在其他公式上
                    is_standard = False
                    break
                cur_term_queue.append(cur_term.formula_left)
                cur_term_queue.append(cur_term.formula_right)
        return is_standard

    @staticmethod
    def _is_conjunctive_head(rule: Rule) -> bool:
        """判断当前规则的head的连接词只出现and"""
        is_standard = True
        cur_term_queue: deque[Formula | Assertion | None] = deque([rule.head])
        while cur_term_queue:
            cur_term = cur_term_queue.popleft()
            if isinstance(cur_term, Formula):
                if cur_term.connective != 'AND':
                    is_standard = False
                    break
                cur_term_queue.append(cur_term.formula_left)
                cur_term_queue.append(cur_term.formula_right)
        return is_standard

    def _merge_total_tables(
        self,
        delta_non_action: _TupleTable,
        delta_action: _TupleTable,
        *,
        non_action_old_table: _TupleTable,
        action_old_table: _TupleTable,
    ) -> _TupleTable:
        """
        Merge none-action (A) and action (B) deltas for substitution stage.

        Δ(A∧B) = (ΔA ⋈ B_old) ∪ (ΔB ⋈ A_old) ∪ (ΔA ⋈ ΔB)
        """
        if not self.rule_checker.action_grounding_nodes:
            return delta_non_action

        return self.semi_naive_calculation(
            delta_non_action,
            delta_action,
            old_left=non_action_old_table,
            old_right=action_old_table,
        )

    @staticmethod
    def _broadcast_total_table(total_table: _TupleTable, nodes: list[_AssertionNode]) -> None:
        """
        将规则级别总表广播到 AssertionNode / RuleNode。

        :params: total_table (_TupleTable): 等待储存的总表
        """
        for cur_assertion in nodes:  # self.rule_checker.action_grounding_nodes + self.rule_checker.substitution_nodes:
            cur_assertion.broadcast_total_table(total_table)

    def print_all_grounded_rules(self) -> list[str]:
        """为了调试方便，打印所有实例化后的规则（注意控制内存开销）"""
        grounded_rule_text = []
        for combination in self._back_up_total_table.iter_rows():
            rule_text = self.rule.replace_variable(combination)
            grounded_rule_text.append(str(rule_text))

        return grounded_rule_text

    def total_table_unique_height(self) -> int:
        """
        用于日志/调试输出的去重规则行数。
        """
        return self._back_up_total_table.unique_height()

    def get_question_solutions(self) -> tuple[list[Mapping[Variable, Constant | CompoundTerm]], _QuestionRule | None]:
        """获取问题规则节点（如果存在）"""
        if isinstance(self.rule_checker.rule_node, _QuestionRuleNode):
            question_node = self.rule_checker.rule_node
            return question_node.solutions, question_node.question_rule
        return [], None

    def reset(self) -> None:  # fixme: 此函数实现仅起过渡作用，令当前commit与之前的代码表现一致。后续会逐渐用更多的reset替代这个reset
        """重置GroundedRule的状态，用于新一轮（iteration）的推理。_past_prefix_sum变量和_backup变量不应删除"""
        self.rule_checker.reset()


class GroundedRuleDS:
    """
    维护全局 GroundedRule 的生命周期与当前轮次的实例化输入。

    GroundedRuleDS 不追求存储“最终 grounded rule”，而是提供实例化、合并与复用的管理入口。
    """

    def __init__(self, equivalence: Equivalence, sk_system_handler: SankuManagementSystem, args: Config, inference_path: InferencePath) -> None:
        # FIXME：这里虽然因为对齐需要传入equivalence，但具体到底是否应该在init传入，以及是否应该换用其他方式调用都有待进一步讨论  # noqa: TD004
        self.grounded_rule_pool: dict[Rule, GroundedRule] = {}
        self.inference_path = inference_path
        self.equivalence = equivalence
        self.sk_system_handler = sk_system_handler
        self.args = args
        self.current_grounded_rule_terms: Sequence[tuple[GroundedRule, Sequence[GROUNDED_TYPE_FOR_UNIFICATION]]] | None = None

    def _add_rule(self, rule: Rule) -> GroundedRule:
        """
        将一条规则纳入 GroundedRuleDS 管理，并返回对应 GroundedRule。
        """
        if rule not in self.grounded_rule_pool:
            grounded_rule = GroundedRule(rule, self.equivalence, self.sk_system_handler, self.args, self.inference_path)
            self.grounded_rule_pool[rule] = grounded_rule

        return self.grounded_rule_pool[rule]
        # 这里转化为GroundedRule的时候，就已经将rule的图结构生成好了。暂时没有明确的处理范畴约束，先置空，仅
        # 转一下GroundedRule。但是以后可能会留存历史以来所有选中的Rule，并进行恰当的事实选取以进行额外的实例化

    def start(self, cur_rules_facts: Sequence[tuple[Rule, Sequence[GROUNDED_TYPE_FOR_UNIFICATION]]]) -> None:
        """
        设置本轮需要执行的cur_rules和facts，及一些可能需要的初始化操作

        :raises RuntimeError: grounding过程还未结束时，再次调用start
        """
        if self.current_grounded_rule_terms is not None:
            raise RuntimeError("Grounding process is not ended")

        self.current_grounded_rule_terms = []
        for rule, facts in cur_rules_facts:
            cur_rule = self._add_rule(rule)
            cur_rule.reset()
            self.current_grounded_rule_terms.append((cur_rule, facts))

    def end(self) -> None:
        """grounding过程结束，移除cur_rules和facts。"""
        self.current_grounded_rule_terms = None

    @staticmethod
    def _unify(cur_rule: GroundedRule, facts: Sequence[GROUNDED_TYPE_FOR_UNIFICATION]) -> None:
        """对单条规则使用对应事实进行实例化。"""
        useful_terms: list[CompoundTerm[Constant | CompoundTerm] | Constant] = []
        for single_fact in facts:
            useful_terms.extend(unify_all_terms(single_fact))

        cur_rule.unify(useful_terms)

    def _grounding_term_level(self) -> None:
        """遍历当前轮次规则并完成 term-level 实例化。"""
        if self.current_grounded_rule_terms is not None:
            for rule_tuple in self.current_grounded_rule_terms:
                single_rule, single_facts = rule_tuple
                self._unify(single_rule, single_facts)

    def exec_grounding(self) -> None:
        """TODO: 这个函数还可以优化。可能是对self.current_grounded_rule_terms做一些调整，如加入过往选择的rule？"""
        self._grounding_term_level()

    def get_corresponding_grounded_rules(self, abstract_rules: list[Rule]) -> list[GroundedRule]:
        """
        取出给定rule的GroundedRule，用于下一步executor的执行
        1. 如果grounding信息的存储方式是现在逐条的check graph的话，记录好二者的映射关系即可；
        2. 如果存储方式是一张大图的话，直接return规则末端的节点。
        此外，如果某条rule在本阶段没有得到实际的实例化，那是可以不return的。  risk: 这里有个权衡，是本阶段没有 or 至今没有
        TODO: executor.check时，要注意避免check已check的部分（这个和used还不完全一样，有可能没过全局的check所有没有use）。一个可能的策略
        是所有check过的都丢到下一个节点存储，当前节点移除。但还是会存在，比如op(x)=y中，x=1是used，不需要二次实例化。但x=1曾经被check过后，
        不代表日后不能被check（因为可能y变了）。但另一方面，如果等价类够快的话，检查是否已check可能开销check一下差不多。

        risk: 从上面的分析来看，这个函数的返回值约束不会出现太大的问题，因为无论是GroundedRule还是末端代表规则的节点，它们都能够索引到
        规则本身的信息 + 实例化信息，这样后续executor进行check时总是有办法快速适配的。但总之还是多加小心一些
        """
        return [self.grounded_rule_pool[r] for r in abstract_rules]

    def reset(self) -> None:
        """
        清空grounded_rule_pool，这是为了防止重复创建Inference_engine带来的错误
        """
        self.grounded_rule_pool.clear()


class GroundedProcess:
    """执行 grounding 的上下文管理器。"""
    def __init__(self, grounded_structure: GroundedRuleDS, cur_rules_terms: Sequence[tuple[Rule, Sequence[GROUNDED_TYPE_FOR_UNIFICATION]]]):
        self.grounded_structure = grounded_structure
        self.cur_terms_facts = cur_rules_terms

    def __enter__(self) -> GroundedRuleDS:
        self.grounded_structure.start(cur_rules_facts=self.cur_terms_facts)
        return self.grounded_structure

    def __exit__(self, exc_type, exc_value, traceback) -> bool | None:  # type: ignore[no-untyped-def]  # noqa: ANN001
        if exc_type:  # XXX: 暂时没有对异常做特殊处理，且强制返回了True。以后使用时根据实际出现的异常逐步优化
            raise exc_value

        self.grounded_structure.end()
        return True
