# GT4Py - GridTools Framework
#
# Copyright (c) 2014-2024, ETH Zurich
# All rights reserved.
#
# Please, refer to the LICENSE file in the root directory.
# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations

import collections
import dataclasses
import enum
import functools
import operator
from typing import Optional

from gt4py import eve
from gt4py.next import common, utils
from gt4py.next.iterator import ir as itir
from gt4py.next.iterator.ir_utils import (
    common_pattern_matcher as cpm,
    domain_utils,
    ir_makers as im,
    misc as ir_misc,
)
from gt4py.next.iterator.transforms import (
    cse,
    fixed_point_transformation,
    inline_center_deref_lift_vars,
    inline_lambdas,
    inline_lifts,
    merge_let,
    trace_shifts,
)
from gt4py.next.iterator.type_system import inference as type_inference
from gt4py.next.type_system import type_info, type_specifications as ts


def _merge_arguments(
    args1: dict[str, itir.Expr], args2: dict[str, itir.Expr]
) -> dict[str, itir.Expr]:
    new_args = {**args1}
    for stencil_param, stencil_arg in args2.items():
        assert stencil_param not in new_args
        new_args[stencil_param] = stencil_arg
    return new_args


def _is_tuple_expr_of_literals(expr: itir.Expr):
    if cpm.is_call_to(expr, "make_tuple"):
        return all(_is_tuple_expr_of_literals(arg) for arg in expr.args)
    if cpm.is_call_to(expr, "tuple_get"):
        return _is_tuple_expr_of_literals(expr.args[1])
    return isinstance(expr, itir.Literal)


def _inline_as_fieldop_arg(
    arg: itir.Expr, *, uids: utils.IDGeneratorPool
) -> tuple[itir.Expr, dict[str, itir.Expr]]:
    assert cpm.is_applied_as_fieldop(arg)
    arg = ir_misc.canonicalize_as_fieldop(arg)

    stencil, *_ = arg.fun.args  # type: ignore[attr-defined]  # ensured by `is_applied_as_fieldop`
    inner_args: list[itir.Expr] = arg.args
    extracted_args: dict[str, itir.Expr] = {}  # mapping from outer-stencil param to arg

    stencil_params: list[itir.Sym] = []
    stencil_body: itir.Expr = stencil.expr

    for inner_param, inner_arg in zip(stencil.params, inner_args, strict=True):
        if isinstance(inner_arg, itir.Literal):
            # note: only literals, not all scalar expressions are required as it doesn't make sense
            # for them to be computed per grid point.
            stencil_body = im.let(inner_param, im.promote_to_const_iterator(inner_arg))(
                stencil_body
            )
        else:
            stencil_params.append(inner_param)
            new_outer_stencil_param = next(uids["__iasfop"])
            extracted_args[new_outer_stencil_param] = inner_arg

    return im.lift(im.lambda_(*stencil_params)(stencil_body))(
        *extracted_args.keys()
    ), extracted_args


def _deduplicate_as_fieldop_args(
    args: dict[str, itir.Expr], stencil_body: itir.Expr
) -> tuple[dict[str, itir.Expr], itir.Expr]:
    new_args_inverted: dict[itir.Expr, list[str]] = collections.defaultdict(list)
    for name, arg in args.items():
        new_args_inverted[arg].append(name)

    new_args: dict[str, itir.Expr] = {}
    new_stencil_body = stencil_body
    for arg, names in new_args_inverted.items():
        # put internal names at the end
        new_name, *aliases = sorted(names, key=lambda s: (s.startswith("_"), s))
        new_args[new_name] = arg
        if aliases:
            new_stencil_body = im.let(*((alias, new_name) for alias in aliases))(new_stencil_body)

    return new_args, new_stencil_body


def _prettify_as_fieldop_args(
    args: dict[str, itir.Expr], stencil_body: itir.Expr
) -> tuple[dict[str, itir.Expr], itir.Expr]:
    new_args: dict[str, itir.Expr] = {}
    remap_table: dict[str, str] = {}
    for name, arg in args.items():
        if isinstance(arg, itir.SymRef):
            assert arg.id not in new_args  # ensured by deduplication
            new_args[arg.id] = arg
            remap_table[name] = arg.id
        else:
            new_args[name] = arg

    return new_args, im.let(*remap_table.items())(stencil_body)


def fuse_as_fieldop(
    expr: itir.Expr,
    eligible_args: list[bool],
    *,
    offset_provider_type: common.OffsetProviderType,
    enable_cse: bool,
    uids: utils.IDGeneratorPool,
) -> itir.Expr:
    assert cpm.is_applied_as_fieldop(expr)

    assert isinstance(expr.fun.args[0], itir.Lambda) or cpm.is_call_to(expr.fun.args[0], "scan")
    stencil, restore_scan = ir_misc.unwrap_scan(expr.fun.args[0])

    domain = expr.fun.args[1] if len(expr.fun.args) > 1 else None

    args: list[itir.Expr] = expr.args

    new_args: dict[str, itir.Expr] = {}
    new_stencil_body: itir.Expr = stencil.expr

    for eligible, stencil_param, arg in zip(eligible_args, stencil.params, args, strict=True):
        if eligible:
            if cpm.is_applied_as_fieldop(arg):
                pass
            elif cpm.is_call_to(arg, "if_"):
                # transform scalar `if` into per-grid-point `if`
                # TODO(tehrengruber): revisit if we want to inline if_
                arg = im.op_as_fieldop("if_")(*arg.args)
            elif _is_tuple_expr_of_literals(arg):
                arg = im.op_as_fieldop(im.lambda_()(arg))()
            else:
                raise NotImplementedError()

            inline_expr, extracted_args = _inline_as_fieldop_arg(arg, uids=uids)

            new_stencil_body = im.let(stencil_param, inline_expr)(new_stencil_body)

            new_args = _merge_arguments(new_args, extracted_args)
        else:
            # just a safety check if typing information is available
            type_inference.reinfer(arg)
            if arg.type and not isinstance(arg.type, ts.DeferredType):
                assert isinstance(arg.type, ts.TypeSpec)

                assert not isinstance(
                    type_info.apply_to_primitive_constituents(type_info.extract_dtype, arg.type),
                    ts.ListType,
                )
            new_args = _merge_arguments(new_args, {stencil_param.id: arg})

    new_args, new_stencil_body = _deduplicate_as_fieldop_args(new_args, new_stencil_body)
    new_args, new_stencil_body = _prettify_as_fieldop_args(new_args, new_stencil_body)

    stencil = im.lambda_(*new_args.keys())(new_stencil_body)
    new_stencil = restore_scan(stencil)

    # simplify stencil directly to keep the tree small
    new_stencil = inline_lambdas.InlineLambdas.apply(
        new_stencil, opcount_preserving=True, force_inline_lift_args=False
    )
    new_stencil = inline_center_deref_lift_vars.InlineCenterDerefLiftVars.apply(
        new_stencil, is_stencil=True, uids=uids
    )  # to keep the tree small
    new_stencil = merge_let.MergeLet().visit(new_stencil)
    new_stencil = inline_lambdas.InlineLambdas.apply(
        new_stencil, opcount_preserving=True, force_inline_lift_args=True
    )
    new_stencil = inline_lifts.InlineLifts().visit(new_stencil)

    new_node = im.as_fieldop(new_stencil, domain)(*new_args.values())
    if enable_cse:
        # TODO(havogt): We should investigate how to keep the tree small without having to run CSE.
        new_node = cse.CommonSubexpressionElimination.apply(
            new_node, within_stencil=False, uids=uids, offset_provider_type=offset_provider_type
        )

    return new_node


def _arg_inline_predicate(node: itir.Expr, shifts: set[tuple[itir.OffsetLiteral, ...]]) -> bool:
    if _is_tuple_expr_of_literals(node):
        return True

    if (
        is_applied_fieldop := cpm.is_applied_as_fieldop(node)
        and not cpm.is_call_to(node.fun.args[0], "scan")
    ) or cpm.is_call_to(node, "if_"):
        # always inline arg if it is an applied fieldop with only a single arg
        if is_applied_fieldop and len(node.args) == 1:
            return True
        # argument is never used, will be removed when inlined
        if len(shifts) == 0:
            return True
        # applied fieldop with list return type must always be inlined as no backend supports this
        type_inference.reinfer(node)
        assert isinstance(node.type, ts.TypeSpec)
        dtype = type_info.apply_to_primitive_constituents(type_info.extract_dtype, node.type)
        if isinstance(dtype, ts.ListType):
            return True
        # only accessed at the center location
        if shifts in [set(), {()}]:
            return True
        # TODO(tehrengruber): Disabled as the InlineCenterDerefLiftVars does not support this yet
        #  and it would increase the size of the tree otherwise.
        # if len(shifts) == 1 and not any(
        #     trace_shifts.Sentinel.ALL_NEIGHBORS in access for access in shifts
        # ):
        #     return True  # noqa: ERA001 [commented-out-code]

    return False


def _make_tuple_element_inline_predicate(node: itir.Expr):
    if cpm.is_applied_as_fieldop(node):  # field, or tuple of fields
        return True
    if isinstance(node.type, ts.FieldType) and isinstance(node, itir.SymRef):
        return True
    return False


@dataclasses.dataclass(frozen=True, kw_only=True)
class FuseAsFieldOp(
    fixed_point_transformation.CombinedFixedPointTransform, eve.PreserveLocationVisitor
):
    """
    Merge multiple `as_fieldop` calls into one.

    >>> from gt4py import next as gtx
    >>> from gt4py.next import utils
    >>> from gt4py.next.iterator.ir_utils import ir_makers as im
    >>> IDim = gtx.Dimension("IDim")
    >>> field_type = ts.FieldType(dims=[IDim], dtype=ts.ScalarType(kind=ts.ScalarKind.INT32))
    >>> d = im.domain("cartesian_domain", {IDim: (0, 1)})
    >>> nested_as_fieldop = im.op_as_fieldop("plus", d)(
    ...     im.op_as_fieldop("multiplies", d)(
    ...         im.ref("inp1", field_type), im.ref("inp2", field_type)
    ...     ),
    ...     im.ref("inp3", field_type),
    ... )
    >>> print(nested_as_fieldop)
    as_fieldop(λ(__arg0, __arg1) → ·__arg0 + ·__arg1, c⟨ IDimₕ: [0, 1[ ⟩)(
      as_fieldop(λ(__arg0, __arg1) → ·__arg0 × ·__arg1, c⟨ IDimₕ: [0, 1[ ⟩)(inp1, inp2), inp3
    )
    >>> print(
    ...     FuseAsFieldOp.apply(
    ...         nested_as_fieldop,
    ...         offset_provider_type={},
    ...         allow_undeclared_symbols=True,
    ...         uids=utils.IDGeneratorPool(),
    ...     )
    ... )
    as_fieldop(λ(inp1, inp2, inp3) → ·inp1 × ·inp2 + ·inp3, c⟨ IDimₕ: [0, 1[ ⟩)(inp1, inp2, inp3)
    """  # noqa: RUF002  # ignore ambiguous multiplication character

    class Transformation(enum.Flag):
        #: Let `f_expr` be an expression with list dtype then
        #: `let(f, f_expr) -> as_fieldop(...)(f)` -> `as_fieldop(...)(f_expr)`
        FUSE_MAKE_TUPLE = enum.auto()
        #: `as_fieldop(...)(as_fieldop(...)(a, b), c)`
        #: -> as_fieldop(fused_stencil)(a, b, c)
        FUSE_AS_FIELDOP = enum.auto()
        INLINE_LET_VARS_OPCOUNT_PRESERVING = enum.auto()

        @classmethod
        def all(self) -> FuseAsFieldOp.Transformation:
            return functools.reduce(operator.or_, self.__members__.values())

    PRESERVED_ANNEX_ATTRS = ("domain",)
    REINFER_TYPES = True

    enabled_transformations = Transformation.all()

    uids: utils.IDGeneratorPool
    offset_provider_type: common.OffsetProviderType
    enable_cse: bool  # option to disable is mainly for testing purposes

    @classmethod
    def apply(
        cls,
        node: itir.Program,
        *,
        offset_provider_type: common.OffsetProviderType,
        uids: utils.IDGeneratorPool,
        allow_undeclared_symbols=False,
        within_set_at_expr: Optional[bool] = None,
        enabled_transformations: Optional[Transformation] = None,
        enable_cse: bool = True,
    ):
        enabled_transformations = enabled_transformations or cls.enabled_transformations

        node = type_inference.infer(
            node,
            offset_provider_type=offset_provider_type,
            allow_undeclared_symbols=allow_undeclared_symbols,
        )

        if within_set_at_expr is None:
            within_set_at_expr = not isinstance(node, itir.Program)

        new_node = cls(
            uids=uids,
            enabled_transformations=enabled_transformations,
            offset_provider_type=offset_provider_type,
            enable_cse=enable_cse,
        ).visit(node, within_set_at_expr=within_set_at_expr)
        # The `FuseAsFieldOp` pass does not fully preserve the type information yet. In particular
        # for the generated lifts this is tricky and error-prone. For simplicity, we just reinfer
        # everything here ensuring later passes can use the information.
        new_node = type_inference.infer(
            new_node,
            offset_provider_type=offset_provider_type,
            allow_undeclared_symbols=allow_undeclared_symbols,
        )
        return new_node

    def transform_fuse_make_tuple(self, node: itir.Node, **kwargs):
        if not cpm.is_call_to(node, "make_tuple"):
            return None

        for arg in node.args:
            type_inference.reinfer(arg)
            assert not isinstance(arg.type, ts.FieldType) or (
                hasattr(arg.annex, "domain")
                and isinstance(arg.annex.domain, domain_utils.SymbolicDomain)
            )

        eligible_els = [_make_tuple_element_inline_predicate(arg) for arg in node.args]
        field_args = [arg for i, arg in enumerate(node.args) if eligible_els[i]]
        distinct_domains = set(arg.annex.domain.as_expr() for arg in field_args)
        if len(distinct_domains) != len(field_args):
            new_els: list[itir.Expr | None] = [None for _ in node.args]
            field_args_by_domain: dict[itir.FunCall, list[tuple[int, itir.Expr]]] = {}
            for i, arg in enumerate(node.args):
                if eligible_els[i]:
                    assert isinstance(arg.annex.domain, domain_utils.SymbolicDomain)
                    domain = arg.annex.domain.as_expr()
                    field_args_by_domain.setdefault(domain, [])
                    field_args_by_domain[domain].append((i, arg))
                else:
                    new_els[i] = arg  # keep as is

            if len(field_args_by_domain) == 1 and all(eligible_els):
                # if we only have a single domain covering all args we don't need to create an
                # unnecessary let
                ((domain, inner_field_args),) = field_args_by_domain.items()
                new_node = im.op_as_fieldop(lambda *args: im.make_tuple(*args), domain)(
                    *(arg for _, arg in inner_field_args)
                )
                new_node = self.visit(new_node, **{**kwargs, "recurse": False})
            else:
                let_vars = {}
                for domain, inner_field_args in field_args_by_domain.items():
                    if len(inner_field_args) > 1:
                        var = next(self.uids["__fasfop"])
                        fused_args = im.op_as_fieldop(lambda *args: im.make_tuple(*args), domain)(
                            *(arg for _, arg in inner_field_args)
                        )
                        type_inference.reinfer(arg)
                        # don't recurse into nested args, but only consider newly created `as_fieldop`
                        # note: this will always inline (as we inline center accessed)
                        let_vars[var] = self.visit(fused_args, **{**kwargs, "recurse": False})
                        for outer_tuple_idx, (inner_tuple_idx, _) in enumerate(inner_field_args):
                            new_el = im.tuple_get(outer_tuple_idx, var)
                            new_el.annex.domain = domain_utils.SymbolicDomain.from_expr(domain)
                            new_els[inner_tuple_idx] = new_el
                    else:
                        i, arg = inner_field_args[0]
                        new_els[i] = arg
                assert not any(el is None for el in new_els)
                assert let_vars
                new_node = im.let(*let_vars.items())(im.make_tuple(*new_els))
                new_node = inline_lambdas.inline_lambda(new_node, opcount_preserving=True)
            return new_node
        return None

    def transform_fuse_as_fieldop(self, node: itir.Node, **kwargs):
        if cpm.is_applied_as_fieldop(node):
            node = ir_misc.canonicalize_as_fieldop(node)
            stencil = node.fun.args[0]  # type: ignore[attr-defined]  # ensure cpm.is_applied_as_fieldop
            assert isinstance(stencil, itir.Lambda) or cpm.is_call_to(stencil, "scan")
            args: list[itir.Expr] = node.args
            shifts = trace_shifts.trace_stencil(stencil, num_args=len(args))

            eligible_els = [
                _arg_inline_predicate(arg, arg_shifts)
                for arg, arg_shifts in zip(args, shifts, strict=True)
            ]
            if any(eligible_els):
                return self.visit(
                    fuse_as_fieldop(
                        node,
                        eligible_els,
                        uids=self.uids,
                        offset_provider_type=self.offset_provider_type,
                        enable_cse=self.enable_cse,
                    ),
                    **{**kwargs, "recurse": False},
                )
        return None

    def transform_inline_let_vars_opcount_preserving(self, node: itir.Node, **kwargs):
        # when multiple `as_fieldop` calls are fused that use the same argument, this argument
        # might become referenced once only. In order to be able to continue fusing such arguments
        # try inlining here.
        if cpm.is_let(node):
            new_node = inline_lambdas.inline_lambda(node, opcount_preserving=True)
            if new_node is not node:  # nothing has been inlined
                return self.visit(new_node, **kwargs)

        return None

    def generic_visit(self, node, **kwargs):
        if cpm.is_applied_as_fieldop(node):  # don't descend in stencil
            return im.as_fieldop(*node.fun.args)(*self.visit(node.args, **kwargs))

        # TODO(tehrengruber): This is a common pattern that should be absorbed in
        #  `FixedPointTransformation`.
        if kwargs.get("recurse", True):
            return super().generic_visit(node, **kwargs)
        else:
            return node

    def visit(self, node, **kwargs):
        if isinstance(node, itir.SetAt):
            return itir.SetAt(
                expr=self.visit(node.expr, **kwargs | {"within_set_at_expr": True}),
                # rest doesn't need to be visited
                domain=node.domain,
                target=node.target,
            )

        # don't execute transformations unless inside `SetAt` node
        if not kwargs.get("within_set_at_expr"):
            return self.generic_visit(node, **kwargs)

        # inline all fields with list dtype. This needs to happen before the children are visited
        # such that the `as_fieldop` can be fused.
        # TODO(tehrengruber): what should we do in case the field with list dtype is a let itself?
        #  This could duplicate other expressions which we did not intend to duplicate.
        # TODO(tehrengruber): This should be moved into a `transform_` method, but
        #  `FixedPointTransformation` does not support pre-order transformations yet.
        if cpm.is_let(node):
            for arg in node.args:
                type_inference.reinfer(arg)
            eligible_els = [
                isinstance(arg.type, ts.FieldType) and isinstance(arg.type.dtype, ts.ListType)
                for arg in node.args
            ]
            if any(eligible_els):
                node = inline_lambdas.inline_lambda(node, eligible_params=eligible_els)
                return self.visit(node, **kwargs)

        node = super().visit(node, **kwargs)

        return node
