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) * 2 * np.pi / (order + 1)
        else:
            return (np.array(range(0, order)) + 1) * np.pi / (2 * (order + 1))