# Source code for ufl.mathfunctions

    """This module provides basic mathematical functions."""
    # Copyright (C) 2008-2016 Martin Sandve Alnæs
    #
    # This file is part of UFL (https://www.fenicsproject.org)
    #
    # SPDX-License-Identifier:    LGPL-3.0-or-later
    #
    # Modified by Anders Logg, 2008
    # Modified by Kristian B. Oelgaard, 2011

    import cmath
    import math
    import numbers
    import warnings

    from ufl.constantvalue import (
        ComplexValue,
        ConstantValue,
        FloatValue,
        IntValue,
        RealValue,
        Zero,
        as_ufl,
        is_true_ufl_scalar,
    )
    from ufl.core.operator import Operator
    from ufl.core.ufl_type import ufl_type

    """
    TODO: Include additional functions available in <cmath> (need derivatives as well):

    Exponential and logarithmic functions:
    log10    Compute common logarithm (function)

    TODO: Any other useful special functions?

    About bessel functions:
    http://en.wikipedia.org/wiki/Bessel_function

    Portable implementations of bessel functions:
    http://www.boost.org/doc/libs/1_47_0/libs/math/doc/sf_and_dist/html/math_toolkit/main_overview/tr1.html

    Implementation in C++ std::tr1:: or boost::math::tr1::
    - BesselK: cyl_bessel_k(nu, x)
    - BesselI: cyl_bessel_i(nu, x)
    - BesselJ: cyl_bessel_j(nu, x)
    - BesselY: cyl_neumann(nu, x)
    """

    # --- Function representations ---

    @ufl_type(is_abstract=True, is_scalar=True, num_ops=1)
    class MathFunction(Operator):
        """Base class for all unary scalar math functions."""

        # Freeze member variables for objects in this class
        __slots__ = ("_name",)

        def __init__(self, name, argument):
            """Initialise."""
            Operator.__init__(self, (argument,))
            if not is_true_ufl_scalar(argument):
                raise ValueError("Expecting scalar argument.")
            self._name = name

        def evaluate(self, x, mapping, component, index_values):
            """Evaluate."""
            a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
            try:
                if isinstance(a, numbers.Real):
                    res = getattr(math, self._name)(a)
                else:
                    res = getattr(cmath, self._name)(a)
            except ValueError:
                warnings.warn(f"Value error in evaluation of function {self._name} with argument {a}.")
                raise
            return res

        def __str__(self):
            """Format as a string."""
            return f"{self._name}({self.ufl_operands[0]})"

    @ufl_type()
    class Sqrt(MathFunction):
        """Square root."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Sqrt."""
            if isinstance(argument, RealValue | Zero | numbers.Real):
                if float(argument) < 0:
                    return ComplexValue(cmath.sqrt(complex(argument)))
                else:
                    return FloatValue(math.sqrt(float(argument)))
            if isinstance(argument, ComplexValue | complex):
                return ComplexValue(cmath.sqrt(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "sqrt", argument)

    @ufl_type()
    class Exp(MathFunction):
        """Exponentiation.."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Exp."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.exp(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.exp(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "exp", argument)

    @ufl_type()
    class Ln(MathFunction):
        """Natural logarithm."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Ln."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.log(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.log(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "ln", argument)

        def evaluate(self, x, mapping, component, index_values):
            """Evaluate."""
            a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
            try:
                return math.log(a)
            except TypeError:
                return cmath.log(a)

    @ufl_type()
    class Cos(MathFunction):
        """Cosine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Cos."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.cos(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.cos(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "cos", argument)

    @ufl_type()
    class Sin(MathFunction):
        """Sine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Sin."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.sin(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.sin(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "sin", argument)

    @ufl_type()
    class Tan(MathFunction):
        """Tangent."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Tan."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.tan(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.tan(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "tan", argument)

    @ufl_type()
    class Cosh(MathFunction):
        """Hyperbolic cosine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Cosh."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.cosh(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.cosh(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "cosh", argument)

    @ufl_type()
    class Sinh(MathFunction):
        """Hyperbolic sine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Sinh."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.sinh(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.sinh(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "sinh", argument)

    @ufl_type()
    class Tanh(MathFunction):
        """Hyperbolic tangent."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Tanh."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.tanh(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.tanh(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "tanh", argument)

    @ufl_type()
    class Acos(MathFunction):
        """Inverse cosine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Acos."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.acos(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.acos(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "acos", argument)

    @ufl_type()
    class Asin(MathFunction):
        """Inverse sine."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Asin."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.asin(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.asin(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "asin", argument)

    @ufl_type()
    class Atan(MathFunction):
        """Inverse tangent."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Atan."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.atan(float(argument)))
            if isinstance(argument, (ComplexValue)):
                return ComplexValue(cmath.atan(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "atan", argument)

    @ufl_type(is_scalar=True, num_ops=2)
    class Atan2(Operator):
        """Inverse tangent with two inputs."""

        __slots__ = ()

        def __new__(cls, arg1, arg2):
            """Create a new Atan2."""
            if isinstance(arg1, RealValue | Zero) and isinstance(arg2, RealValue | Zero):
                return FloatValue(math.atan2(float(arg1), float(arg2)))
            if isinstance(arg1, (ComplexValue)) or isinstance(arg2, (ComplexValue)):
                raise TypeError("Atan2 does not support complex numbers.")
            return Operator.__new__(cls)

        def __init__(self, arg1, arg2):
            """Initialise."""
            Operator.__init__(self, (arg1, arg2))
            if isinstance(arg1, ComplexValue | complex) or isinstance(arg2, ComplexValue | complex):
                raise TypeError("Atan2 does not support complex numbers.")
            if not is_true_ufl_scalar(arg1):
                raise ValueError("Expecting scalar argument 1.")
            if not is_true_ufl_scalar(arg2):
                raise ValueError("Expecting scalar argument 2.")

        def evaluate(self, x, mapping, component, index_values):
            """Evaluate."""
            a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
            b = self.ufl_operands[1].evaluate(x, mapping, component, index_values)
            try:
                res = math.atan2(a, b)
            except TypeError:
                raise ValueError("Atan2 does not support complex numbers.")
            except ValueError:
                warnings.warn(f"Value error in evaluation of function atan2 with arguments {a}, {b}.")
                raise
            return res

        def __str__(self):
            """Format as a string."""
            return f"atan2({self.ufl_operands[0]},{self.ufl_operands[1]})"

    @ufl_type()
    class Erf(MathFunction):
        """Erf function."""

        __slots__ = ()

        def __new__(cls, argument):
            """Create a new Erf."""
            if isinstance(argument, RealValue | Zero):
                return FloatValue(math.erf(float(argument)))
            if isinstance(argument, (ConstantValue)):
                return ComplexValue(math.erf(complex(argument)))
            return MathFunction.__new__(cls)

        def __init__(self, argument):
            """Initialise."""
            MathFunction.__init__(self, "erf", argument)

        def evaluate(self, x, mapping, component, index_values):
            """Evaluate."""
            a = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
            return math.erf(a)

    @ufl_type(is_abstract=True, is_scalar=True, num_ops=2)
    class BesselFunction(Operator):
        """Base class for all bessel functions."""

        __slots__ = "_name"

        def __init__(self, name, nu, argument):
            """Initialise."""
            if not is_true_ufl_scalar(nu):
                raise ValueError("Expecting scalar nu.")
            if not is_true_ufl_scalar(argument):
                raise ValueError("Expecting scalar argument.")

            # Use integer representation if suitable
            fnu = float(nu)
            inu = int(nu)
            if fnu == inu:
                nu = as_ufl(inu)
            else:
                nu = as_ufl(fnu)

            Operator.__init__(self, (nu, argument))

            self._name = name

        def evaluate(self, x, mapping, component, index_values):
            """Evaluate."""
            a = self.ufl_operands[1].evaluate(x, mapping, component, index_values)
            try:
                import scipy.special  # type: ignore
            except ImportError:
                raise ValueError(
                    "You must have scipy installed to evaluate bessel functions in python."
                )
            name = self._name[-1]
            if isinstance(self.ufl_operands[0], IntValue):
                nu = int(self.ufl_operands[0])
                functype = "n" if name != "i" else "v"
            else:
                nu = self.ufl_operands[0].evaluate(x, mapping, component, index_values)
                functype = "v"
            func = getattr(scipy.special, name + functype)
            return func(nu, a)

        def __str__(self):
            """Format as a string."""
            return f"{self._name}({self.ufl_operands[0]}, {self.ufl_operands[1]})"

    @ufl_type()
    class BesselJ(BesselFunction):
        """Bessel J function."""

        __slots__ = ()

        def __init__(self, nu, argument):
            """Initialise."""
            BesselFunction.__init__(self, "cyl_bessel_j", nu, argument)

    @ufl_type()
    class BesselY(BesselFunction):
        """Bessel Y function."""

        __slots__ = ()

        def __init__(self, nu, argument):
            """Initialise."""
            BesselFunction.__init__(self, "cyl_bessel_y", nu, argument)

    @ufl_type()
    class BesselI(BesselFunction):
        """Bessel I function."""

        __slots__ = ()

        def __init__(self, nu, argument):
            """Initialise."""
            BesselFunction.__init__(self, "cyl_bessel_i", nu, argument)

    @ufl_type()
    class BesselK(BesselFunction):
        """Bessel K function."""

        __slots__ = ()

        def __init__(self, nu, argument):
            """Initialise."""
            BesselFunction.__init__(self, "cyl_bessel_k", nu, argument)