import abc
import math
import numpy as np

class ExpansionBasis1D(object):
    """
    A family of one-dimensional functions that make up a basis.
    """

    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def eval(self, idx, coord, diff_order = 0):
        """
        Evaluate the diff_order-th derivative of the idx-th basis function at
        the specified coordinates (where zeroth derivative means evaluating the
        function itself).
        """
        pass

    @abc.abstractmethod
    def colloc_grid(self, order):
        """
        Get the coordinates of the optimal collocation grid of the specified
        order (i.e. exactly order coordinates will be returned).
        """
        pass

class CosBasis(ExpansionBasis1D):
    PARITY_NONE = 0
    PARITY_EVEN = 1
    PARITY_ODD  = 2

    _diff_fact = [(1, np.cos), (-1, np.sin), (-1, np.cos), (1, np.sin)]
    _parity = None

    def __init__(self, parity = PARITY_NONE):
        self._parity = parity

    def eval(self, idx, coord, diff_order = 0):
        fact, f = self._diff_fact[diff_order % 4]

        if self._parity == self.PARITY_EVEN:
            idx *= 2
        elif self._parity == self.PARITY_ODD:
            idx = 2 * idx + 1

        fact *= idx ** diff_order
        return fact * f(idx * coord)

    def colloc_grid(self, order):
        if self._parity == self.PARITY_NONE:
            return (np.array(range(0, order)) + 1) * np.pi / (order + 1)
            return (np.array(range(0, order))) * np.pi / (order - 1)
        else:
            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))

class ChebBasis(ExpansionBasis1D):

    _scale = None

    def __init__(self, scale = 1.0):
        self._scale = scale

    def eval(self, idx, coord, diff_order = 0):
        x = 2.0 * coord / self._scale - 1.0
        t = np.arccos(x)
        st = np.sin(t)

        #idx += 1

        if diff_order == 0:
            return np.cos(idx * t);
        elif diff_order == 1:
            return (2.0 / self._scale) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx, idx * np.sin(idx * t) / st)
        elif diff_order == 2:
            return ((2.0 / self._scale) ** 2) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx * (idx * idx - 1.0) / 3., -(idx ** 2) * np.cos(idx * t) / (st * st) + idx * np.cos(t) * np.sin(idx * t) / (st * st * st))
        else:
            raise NotImplementedError

    def colloc_grid(self, order):
        if self._parity == self.PARITY_NONE:
            return (np.array(range(0, order)) + 1) * 2 * np.pi / (order + 1)
        else:
            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))

class ChebEvenBasis(ExpansionBasis1D):

    _scale = None

    def __init__(self, scale = 1.0):
        self._scale = scale

    def eval(self, idx, coord, diff_order = 0):
        x = coord / self._scale
        t = np.arccos(x)
        st = np.sin(t)

        idx *= 2

        if diff_order == 0:
            return np.cos(idx * t);
        elif diff_order == 1:
            return (1.0 / self._scale) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx, idx * np.sin(idx * t) / st)
        elif diff_order == 2:
            return ((1.0 / self._scale) ** 2) * np.where(np.fabs(np.fabs(x) - 1.0) < 1e-10, idx * idx * (idx * idx - 1.0) / 3., -(idx ** 2) * np.cos(idx * t) / (st * st) + idx * np.cos(t) * np.sin(idx * t) / (st * st * st))
        else:
            raise NotImplementedError

    def colloc_grid(self, order):
        if self._parity == self.PARITY_NONE:
            return (np.array(range(0, order)) + 1) * 2 * np.pi / (order + 1)
        else:
            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))


class ChebTLBasis(ExpansionBasis1D):

    _scale = None

    def __init__(self, scale = 1.0):
        self._scale = scale

    def eval(self, idx, coord, diff_order = 0):
        y   = np.abs(coord)
        y12 = np.sqrt(y)
        t = 2 * np.arctan2(np.sqrt(self._scale), y12)
        #idx += 1

        if diff_order == 0:
            return np.cos(idx * t)
        elif diff_order == 1:
            return np.where(np.abs(y) < 1e-10,
                    -2.0 * idx * idx * ((-1.0)**idx) / self._scale,
                    idx * np.sqrt(self._scale) * np.sin(idx * t) / (y12 * (y12 * y12 + self._scale)))
        elif diff_order == 2:
            return np.where(np.abs(y) < 1e-10,
                    4.0 * idx * idx * (idx * idx + 2) * ((-1.0) ** idx) / (3 * (self._scale**2)),
                    -((self._scale * idx * idx) * np.cos(idx * t) / (y * ((y + self._scale)**2)) + np.sqrt(self._scale) * idx * np.sin(idx * t) / (y12 * ((y + self._scale)**2)) + np.sqrt(self._scale) * idx * np.sin(idx * t) / (2 * y * y12 * (y +
                        self._scale))))
        else:
            raise NotImplementedError

    def colloc_grid(self, order):
        if self._parity == self.PARITY_NONE:
            return (np.array(range(0, order)) + 1) * 2 * np.pi / (order + 1)
        else:
            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))

class ChebSBBasis(ExpansionBasis1D):

    _scale = None

    def __init__(self, scale = 1.0):
        self._scale = scale

    def eval(self, idx, coord, diff_order = 0):
        idx *= 2
        val = np.arctan2(self._scale, np.abs(coord))

        if diff_order == 0:
            return np.sin((idx + 1) * val)
        elif diff_order == 1:
            return -self._scale * (idx + 1) * np.sign(coord) * np.cos((idx + 1) * val) / ((self._scale ** 2) + (coord ** 2))
        elif diff_order == 2:
            return self._scale * (idx + 1) * (2 * np.abs(coord) * np.cos((1 + idx) * val) - self._scale * (1 + idx) * np.sin((1 + idx) * val)) / (((self._scale ** 2) + (coord ** 2)) ** 2)
        else:
            raise NotImplementedError

    def colloc_grid(self, order):
        idx = np.array(range(0, order))
        t = (idx + 2) * np.pi / (2 * order + 2)
        return self._scale / np.tan(t)

class ChebTBBasis(ExpansionBasis1D):

    _scale = None

    def __init__(self, scale = 1.0):
        self._scale = scale

    def eval(self, idx, coord, diff_order = 0):
        idx *= 2
        val = np.arctan2(self._scale, np.abs(coord))

        if diff_order == 0:
            return np.cos(idx * val)
        elif diff_order == 1:
            return -self._scale * (idx + 1) * np.sign(coord) * np.cos((idx + 1) * val) / ((self._scale ** 2) + (coord ** 2))
        elif diff_order == 2:
            return self._scale * (idx + 1) * (2 * np.abs(coord) * np.cos((1 + idx) * val) - self._scale * (1 + idx) * np.sin((1 + idx) * val)) / (((self._scale ** 2) + (coord ** 2)) ** 2)
        else:
            raise NotImplementedError

    def colloc_grid(self, order):
        idx = np.array(range(0, order))
        t = (idx + 2) * np.pi / (2 * order + 4)
        return self._scale / np.tan(t)