#  Copyright (c) 2020, Apple Inc. All rights reserved.
#
#  Use of this source code is governed by a BSD-3-clause license that can be
#  found in the LICENSE.txt file or at https://opensource.org/licenses/BSD-3-Clause
import math
import numpy as np

from coremltools.converters.mil.mil import types
from coremltools.converters.mil.mil.operation import Operation, precondition, SYMBOL, VALUE
from coremltools.converters.mil.mil.types.symbolic import is_symbolic
from coremltools.converters.mil.mil.ops.defs._op_reqs import register_op
from coremltools.converters.mil.mil.input_type import (
    DefaultInputs,
    FloatInputType,
    InputSpec,
    ScalarOrTensorInputType,
    StringInputType
)

class elementwise_unary(Operation):
    """
    Elementwise Unary Op Superclass
    """
    input_spec = InputSpec(x=ScalarOrTensorInputType(),)

    def __init__(self, **kwargs):
        super(elementwise_unary, self).__init__(**kwargs)

    def type_inference(self):
        return self.x.sym_type

"""
Elementwise unary op implementation(s)
"""

@register_op(doc_str="")
class abs(elementwise_unary):
    """
    Return the absolute values of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(abs, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.abs(self.x.val)


@register_op(doc_str="")
class acos(elementwise_unary):
    """
    Return the inverse cosine values of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(acos, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.arccos(self.x.val)


@register_op(doc_str="")
class asin(elementwise_unary):
    """
    Return the inverse sine of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(asin, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.arcsin(self.x.val)


@register_op(doc_str="")
class atan(elementwise_unary):
    """
    Return the inverse tangent of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(atan, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.arctan(self.x.val)


@register_op(doc_str="")
class atanh(elementwise_unary):
    """
    Return the inverse hyperbolic tangent values of the input
    ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(atanh, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.arctanh(self.x.val)


@register_op(doc_str="")
class ceil(elementwise_unary):
    """
    Return the ceil values of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(ceil, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.ceil(self.x.val)


@register_op(doc_str="")
class clip(Operation):
    """
    Clip the values in the input ``x`` to ``[alpha, beta]``, element-wise.
    Any values less than ``alpha`` are set to ``alpha``, and any values greater
    than ``beta`` are set to ``beta``.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)
    alpha: const f32 (Required)
    beta: const f32 (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    input_spec = InputSpec(
        x=ScalarOrTensorInputType(),
        alpha=FloatInputType(const=True),
        beta=FloatInputType(const=True),
    )

    def __init__(self, **kwargs):
        super(clip, self).__init__(**kwargs)

    def type_inference(self):
        return self.x.sym_type

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.minimum(np.maximum(self.x.val, self.alpha.val), self.beta.val)


@register_op(doc_str="")
class cos(elementwise_unary):
    """
    Return cosine of ``x`` element-wise. Input domain is ``(-inf, inf)`` and
    output range is ``[-1,1]``.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], T>

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(cos, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.cos(self.x.val)


@register_op(doc_str="")
class cosh(elementwise_unary):
    """
    Return hyperbolic cosine of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], T>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """

    def __init__(self, **kwargs):
        super(cosh, self).__init__(**kwargs)

    @precondition(allow=VALUE)
    def value_inference(self):
        return np.cosh(self.x.val)


@register_op(doc_str="")
class erf(elementwise_unary):
    """
    Return the gauss error function of the input ``x``, element-wise.

    Parameters
    ----------
    x: tensor<[\*d], T> (Required)

    Returns
    -------
    tensor<[\*d], f32>
        * A tensor of the same shape as ``x``.

    Attributes
    ----------
    T: fp32
    """