from __future__ import annotations

from relationalai.semantics.metamodel import ir, helpers
from relationalai.semantics.metamodel.visitor import Rewriter, collect_by_type, collect, filtered_collect
from relationalai.semantics.metamodel.compiler import Pass, group_tasks
from relationalai.semantics.metamodel.util import OrderedSet, ordered_set
from relationalai.semantics.metamodel import dependency

# This rewrite pass handles aggregations and ranks to ensure that their dependencies are
# self-contained in the required format for emitting Rel/LQP. The expected format is that
# aggregations and ranks are each contained in their own Logical (which hoists the output
# vars), with all dependencies pulled into the same Logical.
#
# For example,
#
# Logical ^[v::Int128]
#     ... <dependencies> ...
#     sum([foo::Foo], [], [a::Int128, v::Int128])
#
# Firstly, the pass ensures that all dependencies
# required by an aggregation/rank are contained locally in the same Logical as the
# aggregation/rank. For example, in the following Logical:
#
# Logical ^[v::Int128]
#     Logical ^[a=None, foo=None]
#         Foo(foo::Foo)
#         a(foo::Foo, a::Int128)
#     sum([foo::Foo], [], [a::Int128, v::Int128])
# common(foo::Foo, a::Int128)
#
# The aggregation `sum` depends on the relation `common`. So, the lookup for `common` needs
# to be pulled into the same Logical as the aggregation.
#
# Logical ^[v::Int128]
#     Logical
#         Logical ^[a=None, foo=None]
#             Foo(foo::Foo)
#             a(foo::Foo, a::Int128)
#         common(foo::Foo, a::Int128)
#     sum([foo::Foo], [], [a::Int128, v::Int128])
# common(foo::Foo, a::Int128)
#
# Secondly, the pass separates Logicals containing more than one aggregation/rank into
# separate Logicals, each containing a single aggregation/rank.
#
# Thirdly, the pass removes dependencies which have been pulled into the body of the
# aggregation/rank and are not needed in the outer scope (see `DependencyRemover`). So, the
# `common` lookup is removed from the outer Logical.
#
# Logical ^[v::Int128]
#     Logical
#         Logical ^[a=None, foo=None]
#             Foo(foo::Foo)
#             a(foo::Foo, a::Int128)
#         common(foo::Foo, a::Int128)
#     sum([foo::Foo], [], [a::Int128, v::Int128])
#
# Lastly, the pass renames variables introduced inside aggregation/rank bodies to ensure
# they do not clash with variables outside.

class HandleAggregationsAndRanks(Pass):
    def __init__(self):
        super().__init__()

    #--------------------------------------------------
    # Public API
    #--------------------------------------------------
    def rewrite(self, model: ir.Model, options:dict={}) -> ir.Model:
        dep_info = dependency.analyze_constraints(model.root)

        # 1. Rewrite to pull dependencies into the same Logical as the aggregation/rank, and
        # to separate multiple aggregations/ranks into separate Logicals.
        r = AggregationsRanksRewriter(dep_info)
        result = r.walk(model)

        # 2. Rewrite to remove dependencies which have been pulled into the body of an
        # aggregation/rank Logical and which are no longer needed in the outer scope.
        remover = DependencyRemover(r.handled_dependencies)
        result = remover.walk(result)

        # 3. Rewrite to make deep copies of Logicals containing aggregations/ranks, to
        # ensure all tasks have unique ids for correctness in later passes.
        deep_copier = AggregationsRanksDeepCopy()
        result = deep_copier.walk(result)

        # 4. Rewrite to rename variables inside aggregation/rank bodies to ensure they do
        # not clash with variables outside.
        rn = AggregationsRanksVarRenameRewriter()
        result = rn.walk(result)

        return result

# The AggregationsRanksRewriter ensures that each aggregation and rank is contained
# in its own Logical, with all dependencies pulled into the same Logical.
class AggregationsRanksRewriter(Rewriter):
    def __init__(self, dep_info):
        super().__init__()
        self.info: dependency.DependencyInfo = dep_info
        self.handled_dependencies: dict[int, set[int]] = {}

    def handle_logical(self, node: ir.Logical, parent: ir.Node):
        groups = group_tasks(node.body, {
            "aggregates_and_ranks": (ir.Aggregate, ir.Rank),
        })

        aggregates_and_ranks = groups["aggregates_and_ranks"]

        # If there are no aggregates or ranks, then just recurse into the Logical body
        if not aggregates_and_ranks:
            return super().handle_logical(node, parent)

        agg_rank_logicals = []
        handled_deps = set()
        for agg_rank in aggregates_and_ranks:
            # Gather all dependencies of the logical containing the aggregate/rank
            agg_deps = self.info.task_dependencies(agg_rank)

            # Reconstruct the body of the Logical containing the aggregate/rank, starting
            # with the existing body
            body = ordered_set()

            # agg_body is the inner body containing the dependencies of the aggregate/rank
            agg_body = ordered_set()
            for t in node.body:
                if isinstance(t, (ir.Output, ir.Update)):
                    # Outputs and Updates need to be kept in the outer body, rather than
                    # nested inside the aggregate/rank body.
                    body.add(t)
                elif t not in aggregates_and_ranks:
                    agg_body.add(t)

            # Add all other dependencies
            for dep in agg_deps:
                # HACK: there are bugs in the dependency analysis that can cause cycles.
                # Avoid these cycles because otherwise they can cause infinite recursion.
                if agg_rank in self.info.task_dependencies(dep) or node in self.info.task_dependencies(dep):
                    continue

                # HACK: the dependency analysis does not handle well cases where multiple
                # identical lookups bind the same variables. If the lookup already exists
                # somewhere in agg_body, then don't add it.
                if isinstance(dep, ir.Lookup):
                    if collect(lambda t, p:
                        isinstance(t, ir.Lookup) and
                        isinstance(dep, ir.Lookup) and
                        t.relation == dep.relation and
                        t.args == dep.args,
                        *agg_body
                    ):
                        continue

                # HACK: a Construct node is never going to act as a filter, unless the
                # agg/rank depends on the output variable.
                if isinstance(dep, ir.Construct):
                    agg_rank_inputs = self.info.task_inputs(agg_rank) or set()
                    if dep.id_var not in agg_rank_inputs:
                        continue

                # Move the dependency into the aggregate/rank body, keeping track of its id
                handled_deps.add(dep.id)
                agg_body.add(dep)

            body.add(ir.Logical(node.engine, tuple(), tuple(agg_body)))

            # Add the actual aggregate/rank
            body.add(agg_rank)

            # Construct the final Logical holding the aggregate/rank contents.
            # Output variables need to be hoisted
            if isinstance(agg_rank, ir.Aggregate):
                output_vars = helpers.get_agg_outputs(agg_rank)
            else:
                assert isinstance(agg_rank, ir.Rank)
                output_vars = [a for a in agg_rank.args]
                output_vars.append(agg_rank.result)

            agg_logical = node.reconstruct(
                hoisted=tuple(output_vars),
                body=tuple(body)
            )
            agg_rank_logicals.append(agg_logical)

        if len(agg_rank_logicals) == 1:
            # If there's only one, no need to create a parent Logical
            result = agg_rank_logicals[0]
        else:
            # Otherwise, create a parent Logical, ensuring all body vars are hoisted
            hoisted = OrderedSet()
            for agg_rank_logical in agg_rank_logicals:
                hoisted.update(agg_rank_logical.hoisted)

            result = ir.Logical(
                engine=node.engine,
                hoisted=tuple(hoisted),
                body=tuple(agg_rank_logicals)
            )

        # Rewrite the children
        result = super().handle_logical(result, parent)

        # Keep track of the dependencies that are handled for the Logical at this level.
        # This is needed for the later DependencyRemover rewriter to discard dependencies
        # which are no longer needed at the outer level.
        self.handled_dependencies[result.id] = handled_deps
        return result

# The DependencyRemover removes the dependencies which have been pulled into the body of an
# aggregation/rank Logical by the AggregationsRanksRewriter, and are no longer needed in the
# outer scope. It relies on the handled_dependencies mapping computed by the
# AggregationsRanksRewriter to know which dependencies to consider for removal.
class DependencyRemover(Rewriter):
    def __init__(self, handled_dependencies: dict[int, set[int]]):
        super().__init__()
        self.handled_dependencies = handled_dependencies

        # Keep track of ancestor aggregate/ranks to know when the rewriter is inside a
        # nested aggregate/rank.
        self.ancestor_aggregates_ranks = set()

    def handle_logical(self, node: ir.Logical, parent: ir.Node):
        # Populate ancestor aggregates/ranks
        current_aggs_ranks = set(t for t in node.body if isinstance(t, (ir.Aggregate, ir.Rank)))
        self.ancestor_aggregates_ranks.update(current_aggs_ranks)

        # Rewrite the children
        logical = super().handle_logical(node, parent)

        # After processing children, remove the current agg/rank logicals from the ancestor set, since they won't be ancestors for nested levels.
        for current_agg_rank in current_aggs_ranks:
            if current_agg_rank in self.ancestor_aggregates_ranks:
                self.ancestor_aggregates_ranks.remove(current_agg_rank)

        # If we are in a nested aggregate/rank, then wait until the outermost level. It is
        # tricky to keep track of dependencies for nested aggs/ranks.
        # TODO: in the future, we may be able to handle nested aggs/ranks better.
        if self.ancestor_aggregates_ranks:
            return logical

        body_tasks: OrderedSet[ir.Node] = OrderedSet.from_iterable(logical.body)

        # Collect Logical nodes which represent aggregate/rank computations. Include nested
        # ones, since they may still have pulled dependencies from this level.
        #
        # However, don't include Logical nodes which are nested under a negation. These do
        # not bind any vars in the outer scope, and so should not be considered for
        # dependency cleanup.
        agg_rank_logicals = filtered_collect(
            lambda t, p: t.id in self.handled_dependencies,
            lambda t, p: isinstance(t, ir.Not),
            logical
        )

        if not agg_rank_logicals:
            # No aggregations/ranks with dependencies to remove at this level
            return logical

        # For each agg_rank_logical, if another agg_rank_logical strictly contains it, then
        # only keep the outermost one. For example, consider
        # Logical
        #     Logical
        #         sum(..., [x::Int128])
        #     max(..., [y::Int128])
        #     sum(..., [x::Int128])
        # In this case, the outer sum has been pulled into the dependencies of the max.
        # Thus, we want to consider the outer sum as a candidate for removal, not as a
        # target aggregate to keep.
        _nested_agg_rank_logicals = set()
        for a1 in agg_rank_logicals:
            for a2 in agg_rank_logicals:
                if a1 == a2:
                    continue
                if collect(lambda t, p: t.id == a1.id, a2):
                    _nested_agg_rank_logicals.add(a1.id)

        agg_rank_logicals = OrderedSet.from_iterable(a for a in agg_rank_logicals if a.id not in _nested_agg_rank_logicals)

        # Quick check to see if cleaning up dependencies is safe. If there are multiple
        # agg/rank nodes with different group-by variables, then it is unsafe.
        # TODO: we should be able to handle this better
        _previous_group_by_vars = None
        for agg_rank_logical in agg_rank_logicals:
            aggs_ranks = [t for t in agg_rank_logical.body if isinstance(t, (ir.Aggregate, ir.Rank))]
            for agg_or_rank in aggs_ranks:
                group_by_vars = set(agg_or_rank.group)
                if _previous_group_by_vars is not None and group_by_vars != _previous_group_by_vars:
                    # Different group-by variables, cannot safely remove dependencies
                    return logical
                _previous_group_by_vars = group_by_vars

        # The current removal candidates are the set of tasks which have been pulled into
        # the body of an aggregation/rank Logical.
        current_removal_candidates: set[int] = set()
        for a in agg_rank_logicals:
            current_removal_candidates.update(self.handled_dependencies[a.id])

        to_keep = OrderedSet.from_iterable(t for t in body_tasks if t.id not in current_removal_candidates)
        removal_candidate_tasks = body_tasks - to_keep

        # Out of the tasks to keep, find all variables that they contain, minus those
        # which are already bound inside of the aggregation/rank body. For all of these
        # variables, we need to keep the tasks which bind them.
        to_keep_vars = set(filtered_collect(
            (lambda t, p: isinstance(t, ir.Var)),
            (lambda t, p: t in agg_rank_logicals),
            *(to_keep - agg_rank_logicals))
        )

        # Variables which are output or grouped by in aggregates/ranks are already handled.
        # We do not need to keep tasks which bind these variables.
        body_vars_per_agg_rank = []

        # Collect all variables bounded by the aggregate or rank task
        for a in agg_rank_logicals:
            aggs_ranks = [t for t in a.body if isinstance(t, (ir.Aggregate, ir.Rank))]
            for agg_or_rank in aggs_ranks:
                bounded_vars = set(agg_or_rank.group)
                if isinstance(agg_or_rank, ir.Aggregate):
                    bounded_vars.update(helpers.get_agg_outputs(agg_or_rank))
                else:
                    assert isinstance(agg_or_rank, ir.Rank)
                    bounded_vars.update(helpers.vars(agg_or_rank.args))
                    bounded_vars.add(agg_or_rank.result)
                bounded_vars.difference_update(set(agg_or_rank.projection))
                body_vars_per_agg_rank.append(bounded_vars)

        agg_rank_body_vars = set.intersection(*body_vars_per_agg_rank) if body_vars_per_agg_rank else set()

        # Out of the current removal candidates, we need to keep those which bind needed vars
        # Do this in a fixpoint loop because keeping a task can introduce new needed vars,
        # which can in turn require keeping more tasks.
        needed_vars = to_keep_vars - agg_rank_body_vars
        changed = True
        while changed:
            changed = False
            for task in removal_candidate_tasks:
                task_vars = set(collect_by_type(ir.Var, task))
                if task_vars & needed_vars and task not in to_keep:
                    to_keep.add(task)
                    needed_vars.update(task_vars - agg_rank_body_vars)
                    changed = True

        if len(to_keep) == len(body_tasks):
            # All tasks are needed, nothing to remove
            return logical

        # Reconstruct the new body from the existing body tasks, to preserve order.
        new_body = []
        for t in body_tasks:
            if t in to_keep:
                new_body.append(t)

        # Reconstruct the logical with only the needed tasks
        return ir.Logical(
            engine=logical.engine,
            hoisted=logical.hoisted,
            body=tuple(new_body)
        )

# The AggregationsRanksVarRenameRewriter renames variables inside aggregation/rank
# bodies to ensure they do not clash with variables outside. It is careful to keep certain
# variables unrenamed because they need to interact with the outside, namely the group-by
# variables and output variables.
class AggregationsRanksVarRenameRewriter(Rewriter):
    class RenameRewriter(Rewriter):
        def __init__(self, to_keep: set[ir.Var], suffix: str):
            super().__init__()
            self.to_keep = to_keep
            self.suffix = suffix
            self.renamed_vars: dict[ir.Var, ir.Var] = {}

        # TODO: the Rewriter should handle Default nodes already, but for some reason it is
        # not behaving correctly. As a workaround, handle Default nodes specifically.
        def handle_default(self, node: ir.Default, parent):
            if node.var in self.to_keep:
                return node

            return ir.Default(
                var=self.handle_var(node.var, node),
                value=node.value
            )

        def handle_var(self, node: ir.Var, parent):
            if node in self.to_keep:
                return node

            if node in self.renamed_vars:
                return self.renamed_vars[node]

            # Rename var
            result = ir.Var(
                type=node.type,
                name=f"{node.name}_{self.suffix}"
            )
            self.renamed_vars[node] = result
            return result

    def __init__(self):
        super().__init__()
        self.renamed_vars: dict[ir.Var, ir.Var] = {}

    def handle_logical(self, node: ir.Logical, parent: ir.Node):
        groups = group_tasks(node.body, {
            "aggregates_and_ranks": (ir.Aggregate, ir.Rank),
        })

        aggregates_and_ranks = groups["aggregates_and_ranks"]

        if not aggregates_and_ranks:
            return super().handle_logical(node, parent)

        # There should only be one, because the AggregationsRanksRewriter should have
        # separated them out.
        assert len(aggregates_and_ranks) == 1, "Multiple aggregate/ranks still found after rewriting dependencies"

        # Rename at this level, keeping some variables fixed, as they interact with the
        # outer scope.
        agg_rank = aggregates_and_ranks[0]
        if isinstance(agg_rank, ir.Aggregate):
            vars_to_keep = set(agg_rank.group)
            output_vars = helpers.get_agg_outputs(agg_rank)
            vars_to_keep.update(output_vars)

            result = self.RenameRewriter(vars_to_keep, 'agg').walk(node)
        else:
            assert isinstance(agg_rank, ir.Rank)
            vars_to_keep = set(agg_rank.group)
            output_vars = helpers.vars(agg_rank.args) + [agg_rank.result]
            vars_to_keep.update(output_vars)
            vars_to_keep.update(agg_rank.projection)

            result = self.RenameRewriter(vars_to_keep, 'rank').walk(node)

        # Process children
        result = super().handle_logical(result, parent)
        return result

class AggregationsRanksDeepCopy(Rewriter):
    class DeepCopyRewriter(Rewriter):
        def walk(self, node, parent=None):
            new_node = super().walk(node, parent)
            if isinstance(new_node, ir.Task):
                return new_node.clone()
            else:
                return new_node

    def __init__(self):
        super().__init__()

    def handle_logical(self, node: ir.Logical, parent: ir.Node):
        # First rewrite the children
        logical = super().handle_logical(node, parent)

        groups = group_tasks(logical.body, {
            "aggregates_and_ranks": (ir.Aggregate, ir.Rank),
        })

        if not groups["aggregates_and_ranks"]:
            # No aggregates or ranks, no need to deep copy
            return logical

        # Then make a deep copy of the logical to ensure all tasks have unique ids
        return self.DeepCopyRewriter().walk(logical)
