import numpy as np

from mygrad.operation_base import BroadcastableOp, Operation

__all__ = [
    "Sin",
    "Sinc",
    "Cos",
    "Tan",
    "Csc",
    "Sec",
    "Cot",
    "Arcsin",
    "Arccos",
    "Arctan",
    "Arccsc",
    "Arcsec",
    "Arccot",
    "Arctan2",
]


class Sin(Operation):
    """ f(a) -> sin(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.sin(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad * np.cos(a.data)


def _dsinc(x):
    x = x * np.pi
    return (x * np.cos(x) - np.sin(x)) / x ** 2


class Sinc(Operation):
    """ f(a) -> sin(pi*a)/(pi*a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.sinc(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        x = a.data
        return np.pi * grad * np.piecewise(x, [x == 0, x != 0], [np.zeros_like, _dsinc])


class Cos(Operation):
    """ f(a) -> cos(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.cos(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad * -np.sin(a.data)


class Tan(Operation):
    """ f(a) -> tan(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.tan(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad / np.cos(a.data) ** 2


class Csc(Operation):
    """ f(a) -> csc(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return 1 / np.sin(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad * -np.cos(a.data) / np.sin(a.data) ** 2


class Sec(Operation):
    """ f(a) -> sec(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return 1 / np.cos(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad * np.sin(a.data) / np.cos(a.data) ** 2


class Cot(Operation):
    """ f(a) -> cot(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return 1 / np.tan(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return -grad / np.sin(a.data) ** 2


class Arcsin(Operation):
    """ f(a) -> arcsin(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.arcsin(a.data)

    def backward_var(self, grad, index, **kwargs):
        # d arcsin / dx at x = -1, 1 returns 0, not NaN
        a = self.variables[index]
        return np.select([np.abs(a.data) != 1], [grad / np.sqrt(1 - a.data ** 2)])


class Arccos(Operation):
    """ f(a) -> arccos(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.arccos(a.data)

    def backward_var(self, grad, index, **kwargs):
        # d arccos / dx at x = -1, 1 returns 0, not NaN
        a = self.variables[index]
        return np.select([np.abs(a.data) != 1], [-grad / np.sqrt(1 - a.data ** 2)])


class Arctan(Operation):
    """ f(a) -> arctan(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.arctan(a.data)

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return grad / (1 + a.data ** 2)


class Arccsc(Operation):
    """ f(a) -> arccsc(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.arcsin(1 / a.data)

    def backward_var(self, grad, index, **kwargs):
        # d arccsc / dx at x = -1, 1 returns 0, not NaN
        a = self.variables[index]
        return np.select(
            [np.abs(a.data) != 1], [-grad / (np.abs(a.data) * np.sqrt(a.data ** 2 - 1))]
        )


class Arcsec(Operation):
    """ f(a) -> arcsec(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.arccos(1 / a.data)

    def backward_var(self, grad, index, **kwargs):
        # d arcsec / dx at x = -1, 1 returns 0, not NaN
        a = self.variables[index]
        return np.select(
            [np.abs(a.data) != 1], [grad / (np.abs(a.data) * np.sqrt(a.data ** 2 - 1))]
        )


class Arccot(Operation):
    """ f(a) -> arccot(a)"""

    def __call__(self, a):
        self.variables = (a,)
        return np.piecewise(
            a.data, [a.data == 0, a.data != 0], [np.pi / 2, lambda x: np.arctan(1 / x)]
        )

    def backward_var(self, grad, index, **kwargs):
        a = self.variables[index]
        return -grad / (1 + a.data ** 2)


class Arctan2(BroadcastableOp):
    """ f(a, b) -> arctan(a/b)"""

    def __call__(self, a, b):
        self.variables = (a, b)
        return np.arctan2(a.data, b.data)

    def backward_var(self, grad, index, **kwargs):
        a, b = self.variables
        if index == 0:
            return grad * b.data / (a.data ** 2 + b.data ** 2)
        else:
            return -1.0 * grad * a.data / (a.data ** 2 + b.data ** 2)
