"""
    Support for dependency analysis of metamodel IRs.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Iterable, Optional, cast
from itertools import count
from more_itertools import peekable
from relationalai.semantics.metamodel import builtins, metamodel as mm
from relationalai.semantics.metamodel.pprint import format_with_id
from relationalai.util.structures import OrderedSet
from relationalai.shims.helpers import VarFinder

from .visitor import Visitor

#--------------------------------------------------
# Useful node categories
#--------------------------------------------------

BINDERS = (mm.Lookup, mm.Construct, mm.Aggregate, mm.Exists, mm.Data, mm.Not)
COMPOSITES = (mm.Logical, mm.Sequence, mm.Union, mm.Match, mm.Until, mm.Wait)

#--------------------------------------------------
# Public API
#--------------------------------------------------

@dataclass
class DependencyInfo():
    """
    Represents the result of performing binding and dependency analysis on a metamodel tree.

    All dicts are keyed by task ids.
    """
    # input vars for a task
    input_bindings: dict[int, OrderedSet[mm.Var]] = field(default_factory=dict)
    # output vars for a task
    output_bindings: dict[int, OrderedSet[mm.Var]] = field(default_factory=dict)
    # clusters of dependency that each task participates on
    dependency_clusters: dict[int, Cluster] = field(default_factory=dict)
    # keep track of the parents of each task
    parent: dict[int, mm.Task] = field(default_factory=dict)
    # keep track of replacements that were made during a rewrite
    replacements: dict[int, mm.Task] = field(default_factory=dict)
    # keep track of which logicals are effectful
    effectful: set[int] = field(default_factory=set)

    def task_inputs(self, node: mm.Task) -> Optional[OrderedSet[mm.Var]]:
        """ The input variables for this task, if any. """
        if node.id in self.input_bindings:
            return self.input_bindings[node.id]
        return None

    def has_inputs(self, node: mm.Task):
        return node.id in self.input_bindings

    def task_outputs(self, node: mm.Task) -> Optional[OrderedSet[mm.Var]]:
        """ The output variables for this task, if any. """
        if node.id in self.output_bindings:
            return self.output_bindings[node.id]
        return None

    def has_outputs(self, node: mm.Task):
        return node.id in self.output_bindings

    def task_dependencies(self, task: mm.Task) -> OrderedSet[mm.Task]:
        """
        All dependencies for this task, if any. This includes tasks in outer contexts,
        not only siblings.
        """

        deps = OrderedSet()
        cluster = self.dependency_clusters.get(task.id, None)
        parent = self.parent.get(task.id)

        if cluster:
            self._collect_deps(cluster, deps)
        while parent:
            cluster = self.dependency_clusters.get(parent.id, None)
            if cluster:
                self._collect_deps(cluster, deps)
                parent = self.parent[cluster.content[0].id]
            else:
                parent = None
        return self._with_replacements(deps)

    def local_dependencies(self, task: mm.Task) -> Optional[OrderedSet[mm.Task]]:
        """ Similar to task_dependencies but returns only dependencies that are siblings. """

        if task.id in self.dependency_clusters:
            deps = OrderedSet()
            self._collect_deps(self.dependency_clusters[task.id], deps)
            return self._with_replacements(deps)
        return None

    def replaced(self, original: mm.Task, replacement: mm.Task):
        """
        Inform that, during some pass, this original task was replaced with this replacement
        task. This affects the info of .task_dependencies(...) since it will answer with
        the replacement when the original is in the set of dependencies.
        """
        self.replacements[original.id] = replacement

    #
    # Implementation details
    #
    def _collect_deps(self, cluster: Cluster, deps: OrderedSet[mm.Task]):
            queue = []
            # start with the cluster dependencies, because cluster represents the task we
            # care about
            queue.extend(cluster.dependencies)
            seen = set()
            while queue:
                cluster = queue.pop()
                if cluster.id in seen:
                    continue
                seen.add(cluster.id)
                deps.update(cluster.content)
                queue.extend(cluster.dependencies)

    def _with_replacements(self, deps: OrderedSet[mm.Task]) -> OrderedSet[mm.Task]:
        # Return deps with all tasks that need replacements (because they are in the
        # replacements dict) replaced.
        if any([dep.id in self.replacements for dep in deps]):
            # Only allocate and compute replacements if there's a dep that was replaced
            info = OrderedSet()
            for dep in deps:
                info.add(self.replacements.get(dep.id, dep))
            return info
        return deps


def analyze_bindings(task: mm.Task) -> DependencyInfo:
    """
    Perform just the binding analysis, skipping the dependency analysis. This is useful
    for passes that only require knowing the input/output sets for nodes, but do not need
    the more expensive dependency information.

    The returned DependencyInfo object will only have input and output bindings filled.
    """
    binding = BindingAnalysis()
    binding(task)
    return binding.info

def analyze(task: mm.Task) -> DependencyInfo:
    """
    Perform binding and dependency analysis on this task. Note that negated tasks do not
    bind variables outside the negation, unlike when using `analyze_constraints`.

    The returned DependencyInfo object will have all dictionaries filled.
    """
    # perform binding analysis to get inputs and outputs
    info = analyze_bindings(task)

    # TODO - if the toplevel task needs inputs, that's a groundness error

    # now perform the dependency analysis
    dependency = DependencyAnalysis(info)
    dependency(task)
    return info

def analyze_constraints(task: mm.Task) -> DependencyInfo:
    """
    Perform binding and dependency analysis on this task, but tracking only the variables
    that are relevant for constraints (i.e. variables that are inputs to tasks, but not
    outputs of any task, since those are the ones that need to be grounded by the logicals
    above them).

    The returned DependencyInfo object will have all dictionaries filled, but input_bindings
    will include constraint-relevant variables.
    """
    binding = BindingAnalysis(track_constraints=True)
    binding(task)
    info = binding.info

    # now perform the dependency analysis
    dependency = DependencyAnalysis(info)
    dependency(task)
    return info

#--------------------------------------------------
# Dependency Analysis
#--------------------------------------------------

# id generator for Clusters
_global_id = peekable(count(0))
def next_id():
    return next(_global_id)

class Cluster():
    def __init__(self, info: DependencyInfo, task: mm.Task):
        """ Create a cluster starting with only this task. """
        self.id = next_id()
        self.info = info
        # exists and lookups for nullary relations are always required (i.e. everything
        # should depend on them)
        self.required = isinstance(task, mm.Exists) or (isinstance(task, mm.Lookup) and not task.args)
        # this is a binders cluster, which is a candidate to being merged
        self.mergeable = not self.required and isinstance(task, BINDERS)
        # this is a cluster that will only hold an effect
        self.effectful = isinstance(task, mm.Update) or task.id in info.effectful
        # this is a cluster that will only hold a composite
        self.composite = isinstance(task, COMPOSITES)
        # content is either a single task or a set of tasks
        self.content: OrderedSet[mm.Task] = OrderedSet([task])
        # combined inputs and outputs for all tasks in the cluster
        self.inputs: OrderedSet[mm.Var] = OrderedSet(info.input_bindings.get(task.id))
        self.outputs: OrderedSet[mm.Var] = OrderedSet(info.output_bindings.get(task.id))
        # eventually we will compute dependencies between clusters
        self.dependencies: OrderedSet[Cluster] = OrderedSet()

    def __str__(self) -> str:
        if isinstance(self.content, mm.Task):
            return str(self.content.id)
        else:
            return ', '.join(str(node.id) for node in self.content)

    def __eq__(self, other):
        return isinstance(other, Cluster) and other.id == self.id

    def __hash__(self):
        return hash(self.id)

    def depends_on(self, other: Cluster):
        """ Assert that this cluster depends on the other cluster. """
        if self in other.dependencies:
            # prevent cycles caused by bugs, like the union hoisting pets in the
            # relationship7 test
            print("Warning: there is a cycle in the dependency graph. This is likely a bug.")
            # k = ','.join(str(x.id) for x in self.content)
            # v = ','.join(str(x.id) for x in other.content)
            # print(f"{k} --> {v}")
        else:
            self.dependencies.add(other)

    def shares_variable(self, other: Cluster):
        """ Returns True iff this cluster and the other cluster have at least one var in common. """
        return (
            any(i in other.inputs or i in other.outputs for i in self.inputs) or
            any(o in other.inputs or o in other.outputs for o in self.outputs)
        )

    def try_merge(self, other: Cluster):
        """
        Verify that this cluster and the other cluster can be merged. If so, merge the
        other cluster into this cluster and return True. Otherwise, return False and the
        clusters are left unmodified.
        """

        # 1. only mergeable
        if not (self.mergeable and other.mergeable):
            return False

        # 2. share some variable
        if not self.shares_variable(other):
            return False

        # 3. all inputs are covered by outputs within the cluster
        if not all(v in self.outputs or v in other.outputs for v in self.inputs):
            return False
        if not all(v in self.outputs or v in other.outputs for v in other.inputs):
            return False

        # ok, can merge
        self.merge(other)
        return True

    def try_merge_group(self, others: list[Cluster]):
        """
        Verify that this cluster and the other clusters can all be merged together. If so,
        merge the other clusters into this cluster and return True. Otherwise, return False
        and all clusters are left unmodified.
        """
        assert(len(others) > 0)

        # 1. only mergeable
        if not (self.mergeable and all(o.mergeable for o in others)):
            return False

        # 2. share some variable
        if not self.shares_variable(others[0]):
            return False

        # 3. all inputs are covered by outputs within the newly formed cluster
        if len(others) == 1:
            others_outputs = others[0].outputs
        else:
            others_outputs = OrderedSet()
            [others_outputs.update(o.outputs) for o in others]
        if not all(v in self.outputs or v in others_outputs for v in self.inputs):
            return False
        for other in others:
            if not all(v in self.outputs or v in others_outputs for v in other.inputs):
                return False

        # ok, can merge
        self.merge(others)
        return True

    def merge(self, other: Cluster|Iterable[Cluster]):
        """
        Merge the other cluster(s) into this one. This assumes that the merge makes sense.
        In general, we should use try_merge or try_merge_group instead.
        """
        # merge the other cluster's content, inputs and outputs
        if isinstance(other, Cluster):
            self.content.update(other.content)
            self.inputs.update(other.inputs)
            self.outputs.update(other.outputs)
            # update dependencies, ensuring we remove the other from the self
            self.dependencies.update(other.dependencies)
            if other in self.dependencies:
                self.dependencies.remove(other)
        else:
            for o in other:
                self.content.update(o.content)
                self.inputs.update(o.inputs)
                self.outputs.update(o.outputs)
                self.dependencies.update(o.dependencies)
                if o in self.dependencies:
                    self.dependencies.remove(o)

class DependencyAnalysis(Visitor):
    """
    A visitor to perform dependency analysis in logicals and store the result in a
    DependencyInfo object.

    The dependency analysis is performed for Logical nodes in 3 steps:

    1. form clusters of children that are mutually dependent, because they are binders (like
    lookups and aggregates) that have variables in common.

    2. compute dependencies between the clusters based on input and outputs.

    3. attempt to merge clusters that could form larger clusters.

    Consider the following example, where we are computing dependencies for the numbered nodes:

        Logical
    |1|    Edges(edges)
    |2|    i(edges, i)
    |3|    Logical ⇑[v]
              Logical ⇑[j=None]
                j(edges, j)
              count([edges, j], [i], [v])
    |4|    i < 10
    |5|    i < v
    |6|    Logical ⇑[v_2=None]
              max([i], [], [v, v_2])
    |7|    → output[i](v, v_2 as 'v2')

    In step 1., we merge (1,2,4) because they have `edges` and `i` in common, but nothing
    else. In particular, (5) is not merged as it also depends on v.

    In step 2. we compute dependencies:
      (1,2,4)
      (3) -> (1,2,4)
      (5) -> (1,2,4), (3)
      (6) -> (1,2,4), (3), (5)
      (7) -> (1,2,4), (3), (5), (6)

    For this example, step 3. does not change anything.

    This result can be interpreted as, in order to compute task 6, tasks (1,2,4), (3) and (5)
    must hold. This means that a compiler pass that extracts the logical in task 6 into its
    own top-level logical must bring these dependencies together.

    To illustrate the need for step 3, consider this example:

        Logical
    |1|    Edge(edge)
    |2|    i(edge, i)
    |3|    Edge(edge_2)
    |4|    j(edge_2, j)
    |5|    i = j
    |6|    Logical ⇑[res=None]
            .....
    |7|    → output[edge, j, i](i, res)

    The result of step 2 is the following:
    (1,2)
    (3,4)
    (5) -> (1,2), (3,4)
    (6) -> (1,2), (3,4), (5)
    (7) -> (1,2), (3,4), (5), (6)

    Step 3 observes that task (5) is a bridge between (1,2) and (3,4), so that merging those
    3 clusters together yields a consistent cluster. So the result is:

    (1,2,3,4,5)
    (6) -> (1,2,3,4,5)
    (7) -> (1,2,3,4,5), (6)

    """
    def __init__(self, info: DependencyInfo):
        self.info = info
        super().__init__()

    def _visit_node(self, node: mm.Node, parent: Optional[mm.Node]=None):
        # keep track of parents of all nodes
        if parent and isinstance(parent, mm.Task):
            self.info.parent[node.id] = parent
        return super()._visit_node(node, parent)


    def enter_logical(self, node: mm.Logical, parent: Optional[mm.Node]=None):
        # quick check to see if it's worth it computing clusters at all
        some_child_has_bindings = False
        for child in node.body:
            if child.id in self.info.input_bindings or child.id in self.info.output_bindings:
                some_child_has_bindings = True
                break
            if not VarFinder().find_vars(child):
                # If the child has no variables, then we still need to analyze dependencies,
                # because it will act as a filter and other tasks will depend on it.
                some_child_has_bindings = True
                break

        if some_child_has_bindings:
            # print(mm.node_to_string(node, print_ids=True))
            # compute clusters for the nodes based on inputs/outputs and shared variables
            clusters = self.compute_clusters(node)
            # compute the dependencies between those clusters
            self.compute_dependencies(clusters)
            # attempt to further merge clusters
            self.merge_clusters(clusters)
            # index the clusters by tasks participating in those clusters, and record it
            self.index(clusters)
            # self._print_debug_info(node, clusters)

    def enter_not(self, node: mm.Not, parent: Optional[mm.Node]=None):
        # For Not nodes, we need to record that the child task is in the same cluster as the
        # Not itself.
        current_cluster = self.info.dependency_clusters.get(node.id, None)
        if current_cluster:
            self.info.dependency_clusters[node.task.id] = current_cluster

    def compute_clusters(self, task: mm.Logical) -> list[Cluster]:
        """
        Cluster the children of the logical, storing together children that are mutually
        dependent.
        """
        # create initial clusters
        clusters:list[Cluster] = [Cluster(self.info, child) for child in task.body]

        # iterate clustering until nothing changes
        merging = True
        while merging:
            merging = False
            cs = list(clusters)
            while cs:
                # last c1 merged some c2s, so cs was modified, restart
                if merging:
                    break
                c1 = cs.pop()
                for c2 in cs:
                    if c1 is c2:
                        continue
                    if c1.try_merge(c2):
                        clusters.remove(c2)
                        merging = True
        return clusters


    def compute_dependencies(self, clusters: list[Cluster]):
        """
        Traverse the clusters finding dependencies between them, based on input and output
        variables used by tasks within the clusters.
        """
        def has_dependency(c1: Cluster, c2: Cluster):
            # c2 is a required cluster, everything depends no it
            if c2.required:
                return True
            # if c1 has an effect and c2 is mergeable (basically it contains only binders)
            # then c2 behaves like a filter, so c1 must depend on it, even if it does not
            # have variables in common (this may bring other dependencies).
            if c1.effectful and c2.mergeable:
                return True

            # If c1 is mergeable and c2 is effectful, then c1 does not depend on c2. The
            # dependency is the other way around.
            if c1.mergeable and c2.effectful:
                return False

            # if c1 has an effect and c2 is a composite without hoisted variables or with a
            # hoisted variable, then c2 behaves like a filter and c1 depends on it.
            if c1.effectful and c2.composite and not c2.effectful:
                task = c2.content.first()
                assert(isinstance(task, COMPOSITES))

                # If the c2 composite defines a new scope, or if it defines non-optional
                # variables, then it acts as a filter.
                if task.scope or not task.optional:
                    return True

            # if c1 is a composite and c2 binds its hoisted vars, c1 can't depend on c2
            # (dependency is the other way around)
            if c1.composite and c1.outputs:
                for v in c1.outputs:
                    if c2.outputs and v in c2.outputs:
                        return False
                    if c2.inputs and v in c2.inputs:
                        return False

            # c1 does not depend on c2 if one of its output vars is an input to c2
            # (dependency is the other way around)
            if (c1.outputs and c2.inputs):
                # optimization for any([v in c2.inputs for v in c1_outputs])):
                for v in c1.outputs:
                    if v in c2.inputs:
                        return False

            # c1 depends on c2 if one of its input vars is an output of c2
            if (c1.inputs and c2.outputs):
                # optimization for any([v in c2.outputs for v in c1_inputs])):
                for v in c1.inputs:
                    if v in c2.outputs:
                        return True

            # c1 is a composite with hoisted variables; it depends on c2 if c2 is a
            # composite that does not have hoisted vars, hence behaving like a filter.
            if c1.composite and c2.composite:
                c1task = c1.content.first()
                assert(isinstance(c1task, COMPOSITES))
                if not c1task.scope and not c2.effectful:
                    c2task = c2.content.first()
                    assert(isinstance(c2task, COMPOSITES))
                    if c2task.scope:
                        return True

            # c1 is a composite with hoisted variables; it depends on c2 if c2 is mergeable
            # and they share variables, hence behaving like a filter
            if c1.composite and c2.mergeable and c1.shares_variable(c2):
                c1task = c1.content.first()
                assert(isinstance(c1task, COMPOSITES))
                # Unless c1 introduces a new scope, in which case its variables are local
                if not c1task.scope and not c1task.optional:
                    return True

            return False

        cs = list(clusters)
        while cs:
            c1 = cs.pop()
            for c2 in cs:
                if c1 is c2:
                    continue

                if has_dependency(c1, c2):
                    c1.depends_on(c2)
                if has_dependency(c2, c1):
                    c2.depends_on(c1)


    def merge_clusters(self, clusters: list[Cluster]):
        """
        Traverse clusters trying to merge multiple clusters if they together form a larger
        cluster.
        """
        # iterate clustering until nothing changes
        merging = True
        while merging:
            merging = False
            cs = list(clusters)
            while cs:
                if merging:
                    break
                c = cs.pop()
                if c.dependencies and c.mergeable:
                    deps = list(c.dependencies)
                    if c.try_merge_group(deps):
                        # remove the deps from clusters
                        for d in deps:
                            clusters.remove(d)
                        # rewire other clusters to the new node
                        # this is not very efficient but should not happen often
                        for c2 in clusters:
                            if c2 is not c:
                                for d in deps:
                                    if d in c2.dependencies:
                                        c2.dependencies.add(c)
                                        c2.dependencies.remove(d)
                        merging = True


    def index(self, clusters: list[Cluster]):
        """
        Index clusters by task, and record in the info
        """
        for dep_node in clusters:
            for n in dep_node.content:
                self.info.dependency_clusters[n.id] = dep_node


    def _print_debug_info(self, node, clusters: list[Cluster]):
        print(format_with_id(node))

        print("dependencies")
        for dep_node in clusters:
            k = ','.join(str(x.id) for x in dep_node.content)
            print(f"({k}) ->")
            for v in dep_node.dependencies:
                v = ','.join(str(x.id) for x in v.content)
                print(f"    ({v})")
        print()
        print("clusters")
        for c in clusters:
            print(f"{c}")
            if c.inputs:
                print(f"    inputs: {','.join(str(v.name) for v in c.inputs)}")
            if c.outputs:
                print(f"    outputs: {','.join(str(v.name) for v in c.outputs)}")
        print()

class BindingAnalysis(Visitor):
    """
    Visitor to perform binding analysis, i.e. figure out for each task in the tree, which
    variables it binds as input and output.
    """
    def __init__(self, track_constraints: bool = False):
        self.info = DependencyInfo()
        # a stack of variables grounded by the last logical being visited
        self._grounded: list[OrderedSet[mm.Var]] = []
        self._track_constraints = track_constraints
        self._negated = False
        super().__init__()

    def input(self, key: mm.Task, val: Optional[mm.Value|Iterable[mm.Value]]):
        """ Assert that this task binds this variable(s) as input. If we're only interested
        in bindings, and the current scope is negated, we ignore the input.
        """
        if self._track_constraints or not self._negated:
            self._register(self.info.input_bindings, key, val)


    def output(self, key: mm.Task, val: Optional[mm.Value|Iterable[mm.Value]]):
        """ Assert that this task binds this variable(s) as output. """
        self._register(self.info.output_bindings, key, val)


    def _register(self, map: dict[int, OrderedSet[mm.Var]], key: mm.Task, val: Optional[mm.Value|Iterable[mm.Value]]):
        """ Register key.id -> val in this map, assuming the map holds ordered sets of vals. """
        # Check prerequisites for registration - if there are no `Var` objects, then we
        # don't register anything
        if val is None:
            return
        if isinstance(val, Iterable):
            if not any(isinstance(v, mm.Var) for v in val):
                return
        else:
            if not isinstance(val, mm.Var):
                return

        if key.id not in map:
            map[key.id] = OrderedSet()
        if isinstance(val, Iterable):
            for v in val:
                if isinstance(v, mm.Var):
                    map[key.id].add(v)
        else:
            assert isinstance(val, mm.Var)
            map[key.id].add(val)

    def exit(self, node: mm.Node, parent: Optional[mm.Node]=None):
        if parent and node.id in self.info.effectful:
            self.info.effectful.add(parent.id)
        return node

    #
    # Composite tasks
    #
    def enter_logical(self, node: mm.Logical, parent: Optional[mm.Node]):
        """Set up the grounded set before visiting children."""
        # compute variables grounded by children of this logical
        grounds = OrderedSet()
        grounded_by_ancestors = None
        if self._grounded:
            # grounded variables inherited from ancestors or siblings
            grounded_by_ancestors = self._grounded[-1]
            grounds.update(grounded_by_ancestors)

        potentially_grounded = OrderedSet()
        for child in node.body:
            # leaf constructs that ground variables
            if isinstance(child, mm.Lookup):
                # special case eq because it can be input or output
                # TODO: this is similar to what's done below in visit_lookup, modularize
                if child.relation == builtins.Builtins.core.eq:
                    x, y = child.args[0], child.args[1]
                    # Compute input/output vars of the equality
                    if isinstance(x, mm.Var) and not isinstance(y, mm.Var):
                        # Variable x is potentially grounded by other expressions at
                        # level in the Logical. If it is, then we should mark it as
                        # input (which is done later).
                        potentially_grounded.add((child, x, x))
                    elif not isinstance(x, mm.Var) and isinstance(y, mm.Var):
                        potentially_grounded.add((child, y, y))
                    elif isinstance(x, mm.Var) and isinstance(y, mm.Var):
                        # mark as potentially grounded, if any is grounded in other atoms then we later ground both
                        potentially_grounded.add((child, x, y))
                else:
                    # grounds only outputs
                    for idx, f in enumerate(child.relation.fields):
                        arg = child.args[idx]
                        if not f.input and isinstance(arg, mm.Var):
                            grounds.add(arg)
            elif isinstance(child, mm.Aggregate):
                # grounds output args
                for idx, f in enumerate(child.aggregation.fields):
                    arg = child.args[idx]
                    if not f.input and isinstance(arg, mm.Var):
                        grounds.add(arg)
            elif isinstance(child, mm.Construct):
                # grounds the output var
                grounds.add(child.id_var)

        # equalities where both sides are already grounded mean that both sides are input
        for child, x, y in potentially_grounded:
            if x in grounds and y in grounds:
                self.input(child, x)
                self.input(child, y)

        # deal with potentially grounded vars up to a fixpoint
        changed = True
        while changed:
            changed = False
            for child, x, y in potentially_grounded:
                if x in grounds and y not in grounds:
                    self.input(child, x)
                    self.output(child, y)
                    grounds.add(y)
                    changed = True
                elif y in grounds and x not in grounds:
                    self.input(child, y)
                    self.output(child, x)
                    grounds.add(x)
                    changed = True

        # Push grounds onto stack for children to use
        self._grounded.append(grounds)

    def exit_logical(self, node: mm.Logical, parent: Optional[mm.Node]):
        """Clean up after visiting children and set inputs/outputs."""
        # Pop from grounded stack since we've now finished the Logical at this level
        self._grounded.pop()

        # Get the variables grounded in the level above, which represents being grounded by
        # ancestors
        grounded_by_ancestors = self._grounded[-1] if self._grounded else None

        # If this Logical introduces a new scope, then it should not input/output anything
        if node.scope:
            return

        # The inputs to this Logical are the variables grounded by an ancestor while being
        # used inside the logical
        inputs = OrderedSet()
        if grounded_by_ancestors:
            vars = VarFinder().find_vars(node)
            inputs = grounded_by_ancestors & vars

        # The outputs from this logical are the variables that are output by tasks inside
        # the Logical
        outputs = OrderedSet()
        for child in node.body:
            child_outputs = self.info.task_outputs(child)
            if child_outputs:
                outputs.update(child_outputs)

        self.input(node, inputs)
        self.output(node, outputs - inputs)


    def exit_union(self, node: mm.Union, parent: Optional[mm.Node]):
        # inputs taken from all children
        for child in node.tasks:
            self.input(node, self.info.task_inputs(child))
        # outputs are vars declared as hoisted
        self.output(node, node.outputs)


    def exit_match(self, node: mm.Match, parent: Optional[mm.Node]):
        # inputs taken from all children
        for child in node.tasks:
            self.input(node, self.info.task_inputs(child))
        # outputs are vars declared as hoisted
        self.output(node, node.outputs)


    def exit_require(self, node: mm.Require, parent: Optional[mm.Node]):
        # inputs taken from the domain and all check tasks
        self.input(node, self.info.task_inputs(node.domain))
        self.input(node, self.info.task_inputs(node.check))

    #
    # Logical tasks
    #
    def enter_not(self, node: mm.Not, parent: Optional[mm.Node]):
        # input_bindings do not propagate into a Not.
        self._negated = not self._negated

    def exit_not(self, node: mm.Not, parent: Optional[mm.Node]):
        # restore negation state
        self._negated = not self._negated

        # Not gets inputs from its child task, but does not have outputs
        self.input(node, self.info.task_inputs(node.task))


    def exit_exists(self, node: mm.Exists, parent: Optional[mm.Node]):
        # exists variables are local, so they are ignored
        self.input(node, self.info.task_inputs(node.task))


    #
    # Leaf tasks
    #

    def visit_update(self, node: mm.Update, parent: Optional[mm.Node]):
        assert parent is not None
        self.info.effectful.add(parent.id)
        # register variables being used as arguments to the update, it's always considered an input
        self.input(node, node.args)


    def visit_lookup(self, node: mm.Lookup, parent: Optional[mm.Node]):
        def register(node, field, arg):
            if isinstance(arg, mm.Var):
                if field.input:
                    self.input(node, arg)
                else:
                    self.output(node, arg)

        if node.relation == builtins.Builtins.core.eq:
            # Most cases are covered already at the parent level if the equality is part of
            # a Logical. The remaining cases are when the equality is a child of a
            # non-Logical, or if its variables are not ground elsewhere in the Logical.
            if self.info.task_inputs(node) or self.info.task_outputs(node):
                # already covered
                pass
            else:
                x, y = node.args[0], node.args[1]
                grounds = self._grounded[-1] if self._grounded else OrderedSet()
                if isinstance(x, mm.Var):
                    if x in grounds:
                        self.input(node, x)
                    else:
                        self.output(node, x)
                        if y in grounds or not isinstance(y, mm.Var):
                            grounds.add(x)
                if isinstance(y, mm.Var):
                    if y in grounds:
                        self.input(node, y)
                    else:
                        self.output(node, y)
                        if x in grounds or not isinstance(x, mm.Var):
                            grounds.add(y)
        else:
            # register variables depending on the input flag of the relation bound to the lookup
            for idx, f in enumerate(node.relation.fields):
                register(node, f, node.args[idx])

    def visit_construct(self, node: mm.Construct, parent: Optional[mm.Node]):
        # values are inputs, id_var is an output
        vars = cast(list[mm.Var], list(filter(lambda v: isinstance(v, mm.Var), node.values)))
        for v in vars:
            self.input(node, v)
        self.output(node, node.id_var)


    def visit_aggregate(self, node: mm.Aggregate, parent: Optional[mm.Node]):
        # register projection and group as inputs
        for v in node.projection:
            self.input(node, v)
        for v in node.group:
            self.input(node, v)

        # register variables depending on the input flag of the aggregation relation
        for idx, f in enumerate(node.aggregation.fields):
            arg = node.args[idx]
            if isinstance(arg, mm.Var):
                if f.input:
                    self.input(node, arg)
                else:
                    self.output(node, arg)
