from __future__ import annotations

from v0.relationalai.semantics.metamodel import ir, helpers
from v0.relationalai.semantics.metamodel.visitor import Rewriter
from v0.relationalai.semantics.metamodel.compiler import Pass, group_tasks
from v0.relationalai.semantics.metamodel.util import OrderedSet, ordered_set
from v0.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 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(model.root)

        r = AggregationsRanksRewriter(dep_info)
        result = r.walk(model)

        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 = dep_info
        self.rewritten: dict[int, ir.Node] = {}

    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 = []
        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
                agg_body.add(dep)

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

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

            # Construct the final Logical holding the aggregate/rank contents.
            # Output variables need to be hoisted
            if isinstance(agg_rank, ir.Aggregate):
                output_vars = [v for v in helpers.vars(agg_rank.args) if not helpers.is_aggregate_input(v, 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 = ir.Logical(
                engine=node.engine,
                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)

        # Make a deep copy so that each task has a unique id. This is important for later
        # rewrite passes (namely QuantifyVars) which identify tasks by id.
        result = DeepCopyRewriter().walk(result)
        return result

# 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] = {}

        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
        agg_rank = aggregates_and_ranks[0]
        if isinstance(agg_rank, ir.Aggregate):
            vars_to_keep = set(agg_rank.group)
            output_vars = [v for v in helpers.vars(agg_rank.args) if not helpers.is_aggregate_input(v, 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 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
