# 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

from typing import Any, Callable

import numpy as np
import sympy

from gt4py.eve import codegen
from gt4py.eve.codegen import FormatTemplate as as_fmt
from gt4py.next.iterator import builtins, ir as gtir
from gt4py.next.iterator.ir_utils import common_pattern_matcher as cpm


MATH_BUILTINS_MAPPING = {
    "abs": "abs({})",
    "neg": "(- {})",
    "sin": "math.sin({})",
    "cos": "math.cos({})",
    "tan": "math.tan({})",
    "arcsin": "asin({})",
    "arccos": "acos({})",
    "arctan": "atan({})",
    "sinh": "math.sinh({})",
    "cosh": "math.cosh({})",
    "tanh": "math.tanh({})",
    "arcsinh": "asinh({})",
    "arccosh": "acosh({})",
    "arctanh": "atanh({})",
    "sqrt": "math.sqrt({})",
    "exp": "math.exp({})",
    "log": "math.log({})",
    "gamma": "tgamma({})",
    "cbrt": "cbrt({})",
    "isfinite": "isfinite({})",
    "isinf": "isinf({})",
    "isnan": "isnan({})",
    "floor": "math.ifloor({})",
    "ceil": "ceil({})",
    "trunc": "trunc({})",
    "minimum": "min({}, {})",
    "maximum": "max({}, {})",
    "fmod": "fmod({}, {})",
    "power": "math.pow({}, {})",
    "float": "dace.float64({})",
    "float32": "dace.float32({})",
    "float64": "dace.float64({})",
    "int": "dace.int32({})" if np.dtype(int).itemsize == 4 else "dace.int64({})",
    "int32": "dace.int32({})",
    "int64": "dace.int64({})",
    "bool": "dace.bool_({})",
    "plus": "({} + {})",
    "minus": "({} - {})",
    "multiplies": "({} * {})",
    "divides": "({} / {})",
    "floordiv": "({} // {})",
    "eq": "({} == {})",
    "not_eq": "({} != {})",
    "less": "({} < {})",
    "less_equal": "({} <= {})",
    "greater": "({} > {})",
    "greater_equal": "({} >= {})",
    "and_": "({} and {})",
    "or_": "({} or {})",
    "xor_": "({} != {})",
    "mod": "({} % {})",
    "not_": "(not {})",
}


def _builtin_cast(val: str, target_type: str) -> str:
    assert target_type in builtins.TYPE_BUILTINS
    return MATH_BUILTINS_MAPPING[target_type].format(val)


def _builtin_get_domain_range(field: str, axis: str) -> str:
    # The builtin function returns a tuple of two values, the range start and stop.
    # Here we return part of the symbol name: the full name also contains an
    # additional suffix '_0' for range start or '_1' for stop, which is obtained
    # from to the `tuple_get` index.
    return f"__{field}_{axis}_range"


def _builtin_if(cond: str, true_val: str, false_val: str) -> str:
    return f"{true_val} if {cond} else {false_val}"


def _builtin_tuple_get(index: str, tuple_name: str) -> str:
    return f"{tuple_name}_{index}"


def _builtin_make_const_list(arg: str) -> str:
    """
    Takes a single scalar argument and broadcasts this value on the local dimension
    of map expression. In a dataflow, we represent it as a tasklet that writes
    a value to a scalar node.
    """
    return arg


GENERAL_BUILTIN_MAPPING: dict[str, Callable[..., str]] = {
    "cast_": _builtin_cast,
    "get_domain_range": _builtin_get_domain_range,
    "if_": _builtin_if,
    "make_const_list": _builtin_make_const_list,
    "tuple_get": _builtin_tuple_get,
}


def format_builtin(builtin: str, *args: Any) -> str:
    if builtin in MATH_BUILTINS_MAPPING:
        fmt = MATH_BUILTINS_MAPPING[builtin]
        return fmt.format(*args)
    elif builtin in GENERAL_BUILTIN_MAPPING:
        expr_func = GENERAL_BUILTIN_MAPPING[builtin]
        return expr_func(*args)
    else:
        raise NotImplementedError(f"'{builtin}' not implemented.")


class PythonCodegen(codegen.TemplatedGenerator):
    """Helper class to visit a symbolic expression and translate it to Python code.

    The generated Python code can be use either as the body of a tasklet node or,
    as in the case of field domain definitions, for sybolic array shape and map range.
    """

    Literal = as_fmt("{value}")

    def visit_AxisLiteral(self, node: gtir.AxisLiteral, **kwargs: Any) -> str:
        return node.value

    def visit_FunCall(self, node: gtir.FunCall, args_map: dict[str, gtir.Node]) -> str:
        if isinstance(node.fun, gtir.Lambda):
            # update the mapping from lambda parameters to corresponding argument expressions
            lambda_args_map = args_map | {
                p.id: arg for p, arg in zip(node.fun.params, node.args, strict=True)
            }
            return self.visit(node.fun.expr, args_map=lambda_args_map)
        elif cpm.is_call_to(node, "deref"):
            assert len(node.args) == 1
            if not isinstance(node.args[0], gtir.SymRef):
                # shift expressions are not expected in this visitor context
                raise NotImplementedError(f"Unexpected deref with arg type '{type(node.args[0])}'.")
            return self.visit(node.args[0], args_map=args_map)
        elif isinstance(node.fun, gtir.SymRef):
            args = self.visit(node.args, args_map=args_map)
            builtin_name = str(node.fun.id)
            return format_builtin(builtin_name, *args)
        raise NotImplementedError(f"Unexpected 'FunCall' node ({node}).")

    def visit_InfinityLiteral(self, node: gtir.InfinityLiteral, **kwargs: Any) -> str:
        return str(sympy.oo) if node == gtir.InfinityLiteral.POSITIVE else str(-sympy.oo)

    def visit_SymRef(self, node: gtir.SymRef, args_map: dict[str, gtir.Node]) -> str:
        symbol = str(node.id)
        if symbol in args_map:
            return self.visit(args_map[symbol], args_map=args_map)
        return symbol


def get_source(node: gtir.Node) -> str:
    """
    Specialized visit method for symbolic expressions.

    The visitor uses `args_map` to map lambda parameters to the corresponding argument expressions.

    Returns:
        A string containing the Python code corresponding to a symbolic expression
    """
    return PythonCodegen.apply(node, args_map={})
