"""
Rust code printer

The `RustCodePrinter` converts SymPy expressions into Rust expressions.

A complete code generator, which uses `rust_code` extensively, can be found
in `sympy.utilities.codegen`. The `codegen` module can be used to generate
complete source code files.

"""

# Possible Improvement
#
# * make sure we follow Rust Style Guidelines_
# * make use of pattern matching
# * better support for reference
# * generate generic code and use trait to make sure they have specific methods
# * use crates_ to get more math support
#     - num_
#         + BigInt_, BigUint_
#         + Complex_
#         + Rational64_, Rational32_, BigRational_
#
# .. _crates: https://crates.io/
# .. _Guidelines: https://github.com/rust-lang/rust/tree/master/src/doc/style
# .. _num: http://rust-num.github.io/num/num/
# .. _BigInt: http://rust-num.github.io/num/num/bigint/struct.BigInt.html
# .. _BigUint: http://rust-num.github.io/num/num/bigint/struct.BigUint.html
# .. _Complex: http://rust-num.github.io/num/num/complex/struct.Complex.html
# .. _Rational32: http://rust-num.github.io/num/num/rational/type.Rational32.html
# .. _Rational64: http://rust-num.github.io/num/num/rational/type.Rational64.html
# .. _BigRational: http://rust-num.github.io/num/num/rational/type.BigRational.html

from __future__ import annotations
from functools import reduce
import operator
from typing import Any

from sympy.codegen.ast import (
    float32, float64, int32,
    real, integer,  bool_
)
from sympy.core import S, Rational, Float, Lambda
from sympy.core.expr import Expr
from sympy.core.numbers import equal_valued
from sympy.functions.elementary.integers import ceiling, floor
from sympy.printing.codeprinter import CodePrinter
from sympy.printing.precedence import PRECEDENCE

# Rust's methods for integer and float can be found at here :
#
# * `Rust - Primitive Type f64 <https://doc.rust-lang.org/std/primitive.f64.html>`_
# * `Rust - Primitive Type i64 <https://doc.rust-lang.org/std/primitive.i64.html>`_
#
# Function Style :
#
# 1. args[0].func(args[1:]), method with arguments
# 2. args[0].func(), method without arguments
# 3. args[1].func(), method without arguments (e.g. (e, x) => x.exp())
# 4. func(args), function with arguments

# dictionary mapping SymPy function to (argument_conditions, Rust_function).
# Used in RustCodePrinter._print_Function(self)

class float_floor(floor):
    """
    Same as `sympy.floor`, but mimics the Rust behavior of returning a float rather than an integer
    """
    def _eval_is_integer(self):
        return False

class float_ceiling(ceiling):
    """
    Same as `sympy.ceiling`, but mimics the Rust behavior of returning a float rather than an integer
    """
    def _eval_is_integer(self):
        return False


function_overrides = {
    "floor": (floor, float_floor),
    "ceiling": (ceiling, float_ceiling),
}

# f64 method in Rust
known_functions = {
    # "": "is_nan",
    # "": "is_infinite",
    # "": "is_finite",
    # "": "is_normal",
    # "": "classify",
    "float_floor": "floor",
    "float_ceiling": "ceil",
    # "": "round",
    # "": "trunc",
    # "": "fract",
    "Abs": "abs",
    # "": "signum",
    # "": "is_sign_positive",
    # "": "is_sign_negative",
    # "": "mul_add",
    "Pow": [(lambda base, exp: equal_valued(exp, -1), "recip", 2),   # 1.0/x
            (lambda base, exp: equal_valued(exp, 0.5), "sqrt", 2),   # x ** 0.5
            (lambda base, exp: equal_valued(exp, -0.5), "sqrt().recip", 2),   # 1/(x ** 0.5)
            (lambda base, exp: exp == Rational(1, 3), "cbrt", 2),    # x ** (1/3)
            (lambda base, exp: equal_valued(base, 2), "exp2", 3),    # 2 ** x
            (lambda base, exp: exp.is_integer, "powi", 1),           # x ** y, for i32
            (lambda base, exp: not exp.is_integer, "powf", 1)],      # x ** y, for f64
    "exp": [(lambda exp: True, "exp", 2)],   # e ** x
    "log": "ln",
    # "": "log",          # number.log(base)
    # "": "log2",
    # "": "log10",
    # "": "to_degrees",
    # "": "to_radians",
    "Max": "max",
    "Min": "min",
    # "": "hypot",        # (x**2 + y**2) ** 0.5
    "sin": "sin",
    "cos": "cos",
    "tan": "tan",
    "asin": "asin",
    "acos": "acos",
    "atan": "atan",
    "atan2": "atan2",
    # "": "sin_cos",
    # "": "exp_m1",       # e ** x - 1
    # "": "ln_1p",        # ln(1 + x)
    "sinh": "sinh",
    "cosh": "cosh",
    "tanh": "tanh",
    "asinh": "asinh",
    "acosh": "acosh",
    "atanh": "atanh",
    "sqrt": "sqrt",  # To enable automatic rewrites
}

# i64 method in Rust
# known_functions_i64 = {
#     "": "min_value",
#     "": "max_value",
#     "": "from_str_radix",
#     "": "count_ones",
#     "": "count_zeros",
#     "": "leading_zeros",
#     "": "trainling_zeros",
#     "": "rotate_left",
#     "": "rotate_right",
#     "": "swap_bytes",
#     "": "from_be",
#     "": "from_le",
#     "": "to_be",    # to big endian
#     "": "to_le",    # to little endian
#     "": "checked_add",
#     "": "checked_sub",
#     "": "checked_mul",
#     "": "checked_div",
#     "": "checked_rem",
#     "": "checked_neg",
#     "": "checked_shl",
#     "": "checked_shr",
#     "": "checked_abs",
#     "": "saturating_add",
#     "": "saturating_sub",
#     "": "saturating_mul",
#     "": "wrapping_add",
#     "": "wrapping_sub",
#     "": "wrapping_mul",
#     "": "wrapping_div",
#     "": "wrapping_rem",
#     "": "wrapping_neg",
#     "": "wrapping_shl",
#     "": "wrapping_shr",
#     "": "wrapping_abs",
#     "": "overflowing_add",
#     "": "overflowing_sub",
#     "": "overflowing_mul",
#     "": "overflowing_div",
#     "": "overflowing_rem",
#     "": "overflowing_neg",
#     "": "overflowing_shl",
#     "": "overflowing_shr",
#     "": "overflowing_abs",
#     "Pow": "pow",
#     "Abs": "abs",
#     "sign": "signum",
#     "": "is_positive",
#     "": "is_negnative",
# }

# These are the core reserved words in the Rust language. Taken from:
# https://doc.rust-lang.org/reference/keywords.html

reserved_words = ['abstract',
                  'as',
                  'async',
                  'await',
                  'become',
                  'box',
                  'break',
                  'const',
                  'continue',
                  'crate',
                  'do',
                  'dyn',
                  'else',
                  'enum',
                  'extern',
                  'false',
                  'final',
                  'fn',
                  'for',
                  'gen',
                  'if',
                  'impl',
                  'in',
                  'let',
                  'loop',
                  'macro',
                  'match',
                  'mod',
                  'move',
                  'mut',
                  'override',
                  'priv',
                  'pub',
                  'ref',
                  'return',
                  'Self',
                  'self',
                  'static',
                  'struct',
                  'super',
                  'trait',
                  'true',
                  'try',
                  'type',
                  'typeof',
                  'unsafe',
                  'unsized',
                  'use',
                  'virtual',
                  'where',
                  'while',
                  'yield']


class TypeCast(Expr):
    """
    The type casting operator of the Rust language.
    """

    def __init__(self, expr, type_) -> None:
        super().__init__()
        self.explicit = expr.is_integer and type_ is not integer
        self._assumptions = expr._assumptions
        if self.explicit:
            setattr(self, 'precedence', PRECEDENCE["Func"] + 10)

    @property
    def expr(self):
        return self.args[0]

    @property
    def type_(self):
        return self.args[1]

    def sort_key(self, order=None):
        return self.args[0].sort_key(order=order)


class RustCodePrinter(CodePrinter):
    """A printer to convert SymPy expressions to strings of Rust code"""
    printmethod = "_rust_code"
    language = "Rust"

    type_aliases = {
        integer: int32,
        real: float64,
    }

    type_mappings = {
        int32: 'i32',
        float32: 'f32',
        float64: 'f64',
        bool_: 'bool'
    }

    _default_settings: dict[str, Any] = dict(CodePrinter._default_settings, **{
        'precision': 17,
        'user_functions': {},
        'contract': True,
        'dereference': set(),
    })

    def __init__(self, settings={}):
        CodePrinter.__init__(self, settings)
        self.known_functions = dict(known_functions)
        userfuncs = settings.get('user_functions', {})
        self.known_functions.update(userfuncs)
        self._dereference = set(settings.get('dereference', []))
        self.reserved_words = set(reserved_words)
        self.function_overrides = function_overrides

    def _rate_index_position(self, p):
        return p*5

    def _get_statement(self, codestring):
        return "%s;" % codestring

    def _get_comment(self, text):
        return "// %s" % text

    def _declare_number_const(self, name, value):
        type_ = self.type_mappings[self.type_aliases[real]]
        return "const %s: %s = %s;" % (name, type_, value)

    def _format_code(self, lines):
        return self.indent_code(lines)

    def _traverse_matrix_indices(self, mat):
        rows, cols = mat.shape
        return ((i, j) for i in range(rows) for j in range(cols))

    def _get_loop_opening_ending(self, indices):
        open_lines = []
        close_lines = []
        loopstart = "for %(var)s in %(start)s..%(end)s {"
        for i in indices:
            # Rust arrays start at 0 and end at dimension-1
            open_lines.append(loopstart % {
                'var': self._print(i),
                'start': self._print(i.lower),
                'end': self._print(i.upper + 1)})
            close_lines.append("}")
        return open_lines, close_lines

    def _print_caller_var(self, expr):
        if len(expr.args) > 1:
            # for something like `sin(x + y + z)`,
            # make sure we can get '(x + y + z).sin()'
            # instead of 'x + y + z.sin()'
            return '(' + self._print(expr) + ')'
        elif expr.is_number:
            return self._print(expr, _type=True)
        else:
            return self._print(expr)

    def _print_Function(self, expr):
        """
        basic function for printing `Function`

        Function Style :

        1. args[0].func(args[1:]), method with arguments
        2. args[0].func(), method without arguments
        3. args[1].func(), method without arguments (e.g. (e, x) => x.exp())
        4. func(args), function with arguments
        """

        if expr.func.__name__ in self.known_functions:
            cond_func = self.known_functions[expr.func.__name__]
            func = None
            style = 1
            if isinstance(cond_func, str):
                func = cond_func
            else:
                for cond, func, style in cond_func:
                    if cond(*expr.args):
                        break
            if func is not None:
                if style == 1:
                    ret = "%(var)s.%(method)s(%(args)s)" % {
                        'var': self._print_caller_var(expr.args[0]),
                        'method': func,
                        'args': self.stringify(expr.args[1:], ", ") if len(expr.args) > 1 else ''
                    }
                elif style == 2:
                    ret = "%(var)s.%(method)s()" % {
                        'var': self._print_caller_var(expr.args[0]),
                        'method': func,
                    }
                elif style == 3:
                    ret = "%(var)s.%(method)s()" % {
                        'var': self._print_caller_var(expr.args[1]),
                        'method': func,
                    }
                else:
                    ret = "%(func)s(%(args)s)" % {
                        'func': func,
                        'args': self.stringify(expr.args, ", "),
                    }
                return ret
        elif hasattr(expr, '_imp_') and isinstance(expr._imp_, Lambda):
            # inlined function
            return self._print(expr._imp_(*expr.args))
        else:
            return self._print_not_supported(expr)

    def _print_Mul(self, expr):
        contains_floats = any(arg.is_real and not arg.is_integer for arg in expr.args)
        if contains_floats:
            expr = reduce(operator.mul,(self._cast_to_float(arg) if arg != -1 else arg for arg in expr.args))

        return super()._print_Mul(expr)

    def _print_Add(self, expr, order=None):
        contains_floats = any(arg.is_real and not arg.is_integer for arg in expr.args)
        if contains_floats:
            expr = reduce(operator.add, (self._cast_to_float(arg) for arg in expr.args))

        return super()._print_Add(expr, order)

    def _print_Pow(self, expr):
        if expr.base.is_integer and not expr.exp.is_integer:
            expr = type(expr)(Float(expr.base), expr.exp)
            return self._print(expr)
        return self._print_Function(expr)

    def _print_TypeCast(self, expr):
        if not expr.explicit:
            return self._print(expr.expr)
        else:
            return self._print(expr.expr) + ' as %s' % self.type_mappings[self.type_aliases[expr.type_]]

    def _print_Float(self, expr, _type=False):
        ret = super()._print_Float(expr)
        if _type:
            return ret + '_%s' % self.type_mappings[self.type_aliases[real]]
        else:
            return ret

    def _print_Integer(self, expr, _type=False):
        ret = super()._print_Integer(expr)
        if _type:
            return ret + '_%s' % self.type_mappings[self.type_aliases[integer]]
        else:
            return ret

    def _print_Rational(self, expr):
        p, q = int(expr.p), int(expr.q)
        float_suffix = self.type_mappings[self.type_aliases[real]]
        return '%d_%s/%d.0' % (p, float_suffix, q)

    def _print_Relational(self, expr):
        if (expr.lhs.is_integer and not expr.rhs.is_integer) or (expr.rhs.is_integer and not expr.lhs.is_integer):
            lhs = self._cast_to_float(expr.lhs)
            rhs = self._cast_to_float(expr.rhs)
        else:
            lhs = expr.lhs
            rhs = expr.rhs
        lhs_code = self._print(lhs)
        rhs_code = self._print(rhs)
        op = expr.rel_op
        return "{} {} {}".format(lhs_code, op, rhs_code)

    def _print_Indexed(self, expr):
        # calculate index for 1d array
        dims = expr.shape
        elem = S.Zero
        offset = S.One
        for i in reversed(range(expr.rank)):
            elem += expr.indices[i]*offset
            offset *= dims[i]
        return "%s[%s]" % (self._print(expr.base.label), self._print(elem))

    def _print_Idx(self, expr):
        return expr.label.name

    def _print_Dummy(self, expr):
        return expr.name

    def _print_Exp1(self, expr, _type=False):
        return "E"

    def _print_Pi(self, expr, _type=False):
        return 'PI'

    def _print_Infinity(self, expr, _type=False):
        return 'INFINITY'

    def _print_NegativeInfinity(self, expr, _type=False):
        return 'NEG_INFINITY'

    def _print_BooleanTrue(self, expr, _type=False):
        return "true"

    def _print_BooleanFalse(self, expr, _type=False):
        return "false"

    def _print_bool(self, expr, _type=False):
        return str(expr).lower()

    def _print_NaN(self, expr, _type=False):
        return "NAN"

    def _print_Piecewise(self, expr):
        if expr.args[-1].cond != True:
            # We need the last conditional to be a True, otherwise the resulting
            # function may not return a result.
            raise ValueError("All Piecewise expressions must contain an "
                             "(expr, True) statement to be used as a default "
                             "condition. Without one, the generated "
                             "expression may not evaluate to anything under "
                             "some condition.")
        lines = []

        for i, (e, c) in enumerate(expr.args):
            if i == 0:
                lines.append("if (%s) {" % self._print(c))
            elif i == len(expr.args) - 1 and c == True:
                lines[-1] += " else {"
            else:
                lines[-1] += " else if (%s) {" % self._print(c)
            code0 = self._print(e)
            lines.append(code0)
            lines.append("}")

        if self._settings['inline']:
            return " ".join(lines)
        else:
            return "\n".join(lines)

    def _print_ITE(self, expr):
        from sympy.functions import Piecewise
        return self._print(expr.rewrite(Piecewise, deep=False))

    def _print_MatrixBase(self, A):
        if A.cols == 1:
            return "[%s]" % ", ".join(self._print(a) for a in A)
        else:
            raise ValueError("Full Matrix Support in Rust need Crates (https://crates.io/keywords/matrix).")

    def _print_SparseRepMatrix(self, mat):
        # do not allow sparse matrices to be made dense
        return self._print_not_supported(mat)

    def _print_MatrixElement(self, expr):
        return "%s[%s]" % (expr.parent,
                           expr.j + expr.i*expr.parent.shape[1])

    def _print_Symbol(self, expr):

        name = super()._print_Symbol(expr)

        if expr in self._dereference:
            return '(*%s)' % name
        else:
            return name

    def _print_Assignment(self, expr):
        from sympy.tensor.indexed import IndexedBase
        lhs = expr.lhs
        rhs = expr.rhs
        if self._settings["contract"] and (lhs.has(IndexedBase) or
                rhs.has(IndexedBase)):
            # Here we check if there is looping to be done, and if so
            # print the required loops.
            return self._doprint_loops(rhs, lhs)
        else:
            lhs_code = self._print(lhs)
            rhs_code = self._print(rhs)
            return self._get_statement("%s = %s" % (lhs_code, rhs_code))

    def _print_sign(self, expr):
        arg = self._print(expr.args[0])
        return "(if (%s == 0.0) { 0.0 } else { (%s).signum() })" % (arg, arg)

    def _cast_to_float(self, expr):
        if not expr.is_number:
            return TypeCast(expr, real)
        elif expr.is_integer:
            return Float(expr)
        return expr

    def _can_print(self, name):
        """ Check if function ``name`` is either a known function or has its own
            printing method. Used to check if rewriting is possible."""

        # since the whole point of function_overrides is to enable proper printing,
        # we presume they all are printable

        return name in self.known_functions or name in function_overrides or getattr(self, '_print_{}'.format(name), False)

    def _collect_functions(self, expr):
        functions = set()
        if isinstance(expr, Expr):
            if expr.is_Function:
                functions.add(expr.func)
            for arg in expr.args:
                functions = functions.union(self._collect_functions(arg))
        return functions

    def _rewrite_known_functions(self, expr):
        if not isinstance(expr, Expr):
            return expr

        expression_functions = self._collect_functions(expr)
        rewriteable_functions = {
            name: (target_f, required_fs)
            for name, (target_f, required_fs) in self._rewriteable_functions.items()
            if self._can_print(target_f)
            and all(self._can_print(f) for f in required_fs)
        }
        for func in expression_functions:
            target_f, _ = rewriteable_functions.get(func.__name__, (None, None))
            if target_f:
                expr = expr.rewrite(target_f)
        return expr

    def indent_code(self, code):
        """Accepts a string of code or a list of code lines"""

        if isinstance(code, str):
            code_lines = self.indent_code(code.splitlines(True))
            return ''.join(code_lines)

        tab = "    "
        inc_token = ('{', '(', '{\n', '(\n')
        dec_token = ('}', ')')

        code = [ line.lstrip(' \t') for line in code ]

        increase = [ int(any(map(line.endswith, inc_token))) for line in code ]
        decrease = [ int(any(map(line.startswith, dec_token)))
                     for line in code ]

        pretty = []
        level = 0
        for n, line in enumerate(code):
            if line in ('', '\n'):
                pretty.append(line)
                continue
            level -= decrease[n]
            pretty.append("%s%s" % (tab*level, line))
            level += increase[n]
        return pretty