# -*- coding: utf-8 -*-
from abc import abstractmethod
from functools import reduce
from operator import mul

import numpy as nm

from sfepy.base.ioutils import InDir
from sfepy.base.base import Struct
from sfepy.discrete import PolySpace
from scipy.special import jacobi as scp_jacobi
from scipy.special import eval_jacobi as scp_eval_jacobi


def iter_by_order(order, dim, extended=False):
    """Iterates over all combinations of basis functions indexes
    needed to create multidimensional basis in a way that creates hierarchical
    basis

    Parameters
    ----------
    order : int
        desired order of multidimensional basis
    dim : int
        dimension of the basis
    extended : bool
        iterate over extended tensor product basis
        (Default value = False)

    Yields
    ------
    idx : tuple
        containing basis function indexes, used in
        ``_combine_polyvals`` and ``_combine_polyvals_der``
    """

    # nth(iter(map(lambda x: x + (order - reduce(add,x),)), range(order)), dim)
    # nth(dim, iterate(map(lambda x: x + (order - reduce(add,x),)),
    #                  map(tuple, range(order))))
    # nth(2, iterate(map(lambda x: x + (order - reduce(add,x),)),
    #                map(lambda x: (x,), range(order))))
    porder = order + 1
    if dim == 1:
        for i in range(porder):
            yield i,
        return
    elif dim == 2:
        for k in range(porder):
            for i in range(k + 1):
                yield i, k - i
        if not extended: return
        for k in range(1, porder):
            for i in range(1, porder):
                if i + k <= porder - 1:
                    continue
                yield i, k

    elif dim == 3:
        for k in range(porder):
            for j in range(k + 1):
                for i in range(j + 1):
                    yield i, j - i, k - j
        return


def get_n_el_nod(order, dim, extended=False):
    r"""Number of nodes per element for discontinuous legendre basis, i.e.
    number of iterations yielded by iter_by_order

    When extended is False

    .. math::
        N_p =  \frac{(n + 1) \cdot (n + 2) \cdot ... \cdot (n + d)}{d!}
    

    where `n` is the order and `d` the dimension.
    When extended is True

    .. math::
        N_p = (n + 1) ^ d

    where `n` is the order and `d` the dimension.

    Parameters
    ----------
    order : int
        desired order of multidimensional basis
    dim : int
        dimension of the basis
    extended : bool
        iterate over extended tensor product basis
        (Default value = False)


    Returns
    -------
    n_el_nod : int
        number of basis functions in basis
    """
    if extended:
        return (order + 1) ** dim
    return int(reduce(mul, map(lambda i: order + i + 1, range(dim))) /
               reduce(mul, range(1, dim + 1)))


class LegendrePolySpace(PolySpace):
    """Legendre hierarchical polynomials basis, over [0, 1] domain."""

    def __init__(self, name, geometry, order, extended):
        """

        Parameters
        ----------
        name
        geometry
        order : int
            approximation order, 0 for constant functions, 1 for linear etc.
        extended : bool
            for extended tensor product space
        """
        PolySpace.__init__(self, name, geometry, order)
        self.extended = extended  # only tensor product polyspace is extended
        self.n_v = geometry.n_vertex,
        self.dim = geometry.dim
        self.n_nod = get_n_el_nod(self.order, self.dim, self.extended)

        self.coefM = None
        self.expoM = None

    def _eval_base(self, coors, diff=0, ori=None,
                   suppress_errors=False, eps=1e-15):
        """Calls _combine_polyvals or _combine_polyvals_diff to build
        multidimensional basis implemented in subclasses to get basis for
        different geometries expects coors to be in shape
            (..., dim),
        Returns values of the basis in shape
            (coors.shape[:-1], 1, n_el_nod)
        or values of the gradient in shape
            (1, coors.shape[:-1], dim, n_el_nod)

        Parameters
        ----------
        coors : array_like
        ori : unused
            we do not need ori, because the basis is discontinous across
            elements this is treated in DGField(Default value = None)
        suppress_errors : has no effect
        diff : int or truthy
             (Default value = 0)
        eps : unused
             (Default value = 1e-15)

        Returns
        -------
        values : ndarray
        """
        coors = 2 * coors - 1  # transofrm from [0, 1] to [-1, 1]
        porder = self.order + 1
        if diff:
            diff = int(diff)
            values = nm.zeros((1,) + coors.shape[:-1] +
                              # (1,) for dummy axis used throughout sfepy
                              (self.dim,) * diff +
                              # (dim,) * diff order is shape of derivation tensor,
                              # we support only first derivative
                              (self.n_nod,))
            polyvals = nm.zeros(coors.shape + (porder,) + (diff + 1,))
            # diff + 1 is number of values of one dimensional base
            polyvals[..., 0] = self.legendreP(coors)
            polyvals[..., 1] = self.gradlegendreP(coors)

            for m, idx in enumerate(
                    iter_by_order(self.order, self.dim, self.extended)):
                for d in range(self.dim):
                    values[..., d, m] = 2 * self._combine_polyvals_diff(coors,
                                                                        polyvals,
                                                                        d, idx)
                    # 2 is due to transformation from [0,1] to [-1,1]
        else:
            values = nm.zeros(coors.shape[:-1] + (1, self.n_nod,))
            # 1, because no matter the dimension functions have only one value
            polyvals = self.legendreP(coors)
            for m, idx in enumerate(
                    iter_by_order(self.order, self.dim, self.extended)):
                values[..., 0, m] = self._combine_polyvals(coors, polyvals, idx)

        return values

    def get_interpol_scheme(self):
        r"""For dim > 1 returns F and P matrices according to gmsh basis
        specification [1]_:
        Let us assume that the approximation of the view's value over an element
        is written as a linear  combination of d basis functions
        :math:`f_i, i=0, ..., n-1` (the coefficients being stored
        in list-of-values).
        
        Defining

        .. math ::
            f_i = \sum\limits_{j=0}^{d-1} F_{ij}\cdot p_j,

        with

        :math:`p_j(u, v, w) = u^{P_j^{(0)}}\cdot v^{P_j^{(1)}}\cdot w^{P_j^{(2)}}`
        (`u`, `v` and `w` being the
        coordinates in the element's parameter space), then val-coef-matrix
        denotes the n x n matrix F and val-exp-matrix denotes the n x 3 matrix P
        where n is number of basis functions as calculated by ``get_n_el_nod``.
        
        Expects matrices to be saved in atributes coefM and expoM!

        .. [1] Remacle, J.-F., Chevaugeon, N., Marchandise, E., & Geuzaine, C.
           (2007). Efficient visualization of high-order finite elements.
           International Journal for Numerical Methods in Engineering, 69(4),
           750-771. https://doi.org/10.1002/nme.1787

        Returns
        -------
        interp_scheme_struct : Struct
            Struct with name of the scheme, geometry desc and P and F
        """
        interp_scheme_struct = None
        if self.dim > 1:
            interp_scheme_struct = Struct(name=self.name + "_interpol_scheme",
                                          F=self.coefM, P=self.expoM,
                                          desc=self.geometry.name)
        return interp_scheme_struct

    @abstractmethod
    def _combine_polyvals(self, coors, polyvals, idx):
        """Combines Legendre or Jacobi polynomials to get muldidimensional
        basis values according to element topolgy

        Parameters
        ----------
        coors : array_like
            coordinates of evaluation points
        polyvals : array_like
            values of legendre polynomials precomputed in _eval_base
        idx : tuple
            function index, for tensor-product correspond to orders of polynomials in variables

        Returns
        -------
        values : ndarray
        """

    @abstractmethod
    def _combine_polyvals_diff(self, coors, polyvals, der, idx):
        """Combines Legendre or Jacobi polynomials to get muldidimensional
        basis derivative values according to element topolgy.

        Parameters
        ----------
        coors : array_like
            coordinates of evaluation points
        polyvals : array_like
            values of legendre polynomials precomputed in _eval_base
        der : int
            derivative variable
        idx : tuple
            function index, for tensor-product correspond to orders of polynomials in variables

        Returns
        -------
        values : ndarray
        """

    # --------------------- #
    # 1D legendre polyspace #
    # --------------------- #
    def legendreP(self, coors):
        """
        Parameters
        ----------
        coors : array_like
            coordinates, preferably in interval [-1, 1] for which this basis is
            intented

        Returns
        -------
        values : ndarray
            values at coors of all the legendre polynomials up to self.order

        """
        return self.jacobiP(coors, alpha=0, beta=0)

    def gradlegendreP(self, coors, diff=1):
        """
        Parameters
        ----------
        diff : int
            default 1
        coors : array_like
            coordinates, preferably in interval [-1, 1] for which this basis is
            intented

        Returns
        -------
        values : ndarray
            values at coors of all the legendre polynomials up to self.order

        """
        return self.gradjacobiP(coors, 0, 0, diff=diff)

    """
    Explict legendre polynomials up to order 5
    """
    legendre_funs = [lambda x: 0 * x + 1,
                     # we need constant preserving shape and type of x
                     lambda x: 2 * x - 1,
                     lambda x: (6 * x ** 2 - 6 * x + 1),
                     lambda x: (20 * x ** 3 - 30 * x ** 2 + 12 * x - 1),
                     lambda x: (70 * x ** 4 - 140 * x ** 3
                                + 90 * x ** 2 - 20 * x + 1),
                     lambda x: (252 * x ** 5 - 630 * x ** 4 + 560 * x ** 3
                                - 210 * x ** 2 + 30 * x - 1)
                     ]

    def get_nth_fun(self, n):
        """Uses shifted Legendre polynomials formula on interval [0, 1].

        Convenience function for testing

        Parameters
        ----------
        n : int

        Returns
        -------
        fun : callable
            n-th function of the legendre basis
        """

        if n < 6:
            return self.legendre_funs[n]
        else:
            from scipy.special import comb as comb

            def fun(x):
                val = 0.0
                for k in range(n + 1):
                    val = val + comb(n, k) * comb(n + k, k) * (-x) ** k
                val *= -1 if n % 2 else 1
                return val

            return fun

    def get_nth_fun_der(self, n, diff=1):
        """Returns diff derivative of nth function. Uses shifted legendre
        polynomials formula on interval [0, 1].

        Useful for testing.

        Parameters
        ----------
        n : int
        diff : int
             (Default value = 1)

        Returns
        -------
        fun : callable
            derivative of n-th function of the 1D legendre basis

        """

        def dfun(x):
            """

            Parameters
            ----------
            x :
                

            Returns
            -------

            """
            from scipy.special import comb as comb, factorial
            val = x[:] * 0.0
            for k in range(diff, n + 1):
                val += comb(n, k) * comb(n + k, k) * factorial(k) / \
                       factorial(k - diff) * (-x) ** (k - diff)
            val *= -1 if (n - diff) % 2 else 1
            return val

        return dfun
    # --------------------- #

    # --------------------- #
    #  1D jacobi polyspace  #
    # --------------------- #
    def jacobiP(self, coors, alpha, beta):
        """Values of the jacobi polynomials on interval [-1, 1]
        up to self.order + 1 at coors

        Parameters
        ----------
        coors : array_like
        beta : float
        alpha : float
            

        Returns
        -------
        values : ndarray
            output shape is shape(coor) + (self.order + 1,)

        """
        if not isinstance(coors, nm.ndarray):
            sh = (1,)
        else:
            sh = nm.shape(coors)
        values = nm.ones(sh + (self.order + 1,))

        for i in range(self.order + 1):
            values[..., i] = scp_eval_jacobi(i, alpha, beta, coors)
            # for some reason eval_jacobi consumes last dimension if it is one,
            # when called with order array

        return values

    def gradjacobiP(self, coors, alpha, beta, diff=1):
        """diff derivative of the jacobi polynomials on interval [-1, 1]
        up to self.order + 1 at coors

        Parameters
        ----------
        coors :

        alpha : float

        beta : float

        diff : int
             (Default value = 1)

        Returns
        -------
        values : ndarray
            output shape is shape(coor) + (self.order + 1,)

        """
        if isinstance(coors, (int, float)):
            sh = (1,)
        else:
            sh = nm.shape(coors)

        values = nm.zeros(sh + (self.order + 1,))
        for i in range(self.order + 1):
            jacob_poly = scp_jacobi(i, alpha, beta)
            values[..., i] = jacob_poly.deriv(m=diff)(coors)
            # Warning
            # Computing values of high-order polynomials (around order > 20)
            # using polynomial coefficients is numerically unstable. To evaluate
            # polynomial values, the eval_* functions should be used instead.
            # On that note jacob_poly.deriv seems to use stable version.

        return values
    # --------------------- #


class LegendreTensorProductPolySpace(LegendrePolySpace):
    name = "legendre_tensor_product"


    def __init__(self, name, geometry, order, ):
        super().__init__(name, geometry, order, extended=True)
        self.n_nod = get_n_el_nod(self.order, self.dim, self.extended)
        if self.dim > 1:
            self.coefM, self.expoM = self.build_interpol_scheme()

    def _combine_polyvals(self, coors, polyvals, idx):
        return nm.prod(polyvals[..., range(len(idx)), idx], axis=-1)

    def _combine_polyvals_diff(self, coors, polyvals, dvar, idx):
        dimz = range(len(idx))
        derz = nm.zeros(len(idx), dtype=nm.int32)
        derz[dvar] = 1
        return nm.prod(polyvals[..., dimz, idx, derz], axis=-1)

    def build_interpol_scheme(self):
        """Builds F and P matrices returned by self.get_interpol_scheme.

        Note that this function returns coeficients according to gmsh
        parametrization of Quadrangle i.e. [-1, 1] x [-1, 1] and hence the form
        of basis function is not the same as exhibited by the
        LegendreTensorProductPolySpace object which acts on parametrization
        [0, 1] x [0, 1].


        Returns
        -------
        F : ndarray
            coefficient matrix
        P : ndarray
            exponent matrix
        """
        P = nm.zeros((self.n_nod, 3), dtype=nm.int32)
        for m, idx in enumerate(
                iter_by_order(self.order, self.dim, self.extended)):
            P[m, :self.dim] = idx

        F = nm.zeros((self.n_nod, self.n_nod))
        for m, idx in enumerate(
                iter_by_order(self.order, self.dim, self.extended)):
            xcoefs = list(scp_jacobi(idx[0], 0, 0).coef)[::-1]
            xcoefs = nm.array(xcoefs + [0] * (self.order + 1 - len(xcoefs)))
            ycoefs = list(scp_jacobi(idx[1], 0, 0).coef)[::-1]
            ycoefs = nm.array(ycoefs + [0] * (self.order + 1 - len(ycoefs)))
            coef_mat = nm.outer(xcoefs, ycoefs)
            F[m, :] = [coef_mat[idx] for idx in
                       iter_by_order(self.order, self.dim, self.extended)]
        return F, P


class LegendreSimplexPolySpace(LegendrePolySpace):
    name = "legendre_simplex"


    def __init__(self, name, geometry, order, extended=False):
        super().__init__(name, geometry, order, extended)
        if self.dim == 1:
            return
        if order <= 5:
            self.coefM = simplex_coefM5[:self.n_nod, :self.n_nod]
            self.expoM = simplex_expoM5[:self.n_nod, :]
            return

        indir = InDir(__file__)
        try:
            self.coefM = nm.loadtxt(
                    indir("legendre2D_simplex_coefs.txt")
            )[:self.n_nod, :self.n_nod]
            self.expoM = nm.loadtxt(
                    indir("legendre2D_simplex_expos.txt")
            )[:self.n_nod, :]
        except IOError as e:
            raise IOError(
                ("File {} not found, run gen_legendre_simplex_base.py"
                 + " to generate it. This is needed approx order > 5.")
                    .format(e.args[0]))

    def _combine_polyvals(self, coors, polyvals, idx):

        from scipy.special import eval_jacobi
        if len(idx) == 1:  # 1D
            return nm.prod(polyvals[..., range(len(idx)), idx], axis=-1)
        elif len(idx) == 2:  # 2D
            r = coors[..., 0]
            s = coors[..., 1]
            a = 2 * (1 + r) / (1 - s) - 1
            a[nm.isnan(a)] = -1.
            b = s
            return eval_jacobi(idx[0], 0, 0, a) \
                   * eval_jacobi(idx[1],2 * idx[0] + 1, 0,b) * (1 - b) ** idx[0]
        elif len(idx) == 3:  # 3D
            r = coors[..., 0]
            s = coors[..., 1]
            t = coors[..., 2]
            a = -2 * (1 + r) / (s + t) - 1
            b = 2 * (1 + s) / (1 - t) - t
            c = t
            a[nm.isnan(a)] = -1.
            b[nm.isnan(b)] = -1.
            return eval_jacobi(idx[0], 0, 0, a) * \
                   eval_jacobi(idx[1], 2 * idx[0] + 1, 0, 0, b) * \
                   eval_jacobi(idx[2], 2 * idx[0] + 2 * idx[1] + 2, 0, c) * \
                                                    (1 - c) ** (idx[0] + idx[1])

    def _combine_polyvals_diff(self, coors, polyvals, dvar, idx):

        if len(idx) == 1:  # 1D
            dimz = range(len(idx))
            derz = nm.zeros(len(idx), dtype=nm.int32)
            derz[dvar] = 1
            return nm.prod(polyvals[..., dimz, idx, derz], axis=-1)
        elif len(idx) == 2:  # 2D
            r = coors[..., 0]
            s = coors[..., 1]
            a = 2 * (1 + r) / (1 - s) - 1
            b = s
            a[nm.isnan(a)] = -1.
            di = idx[0]
            dj = idx[1]

            fa = self.jacobiP(a, 0, 0)[..., di]
            dfa = self.gradjacobiP(a, 0, 0)[..., di]
            gb = self.jacobiP(b, 2 * di + 1, 0)[..., dj]
            dgb = self.gradjacobiP(b, 2 * di + 1, 0)[..., dj]

            if dvar == 0:  # d/dr
                # r - derivative
                # d / dr = da / dr * d / da + db / dr * d / db
                #        = (2 / (1 - s)) d / da = (2 / (1 - b)) d / da
                dmodedr = dfa * gb
                dmodedr = dmodedr * (
                            (0.5 * (1 - b)) ** (di - 1)) if di > 0 else dmodedr
                return 2 ** di * dmodedr

            elif dvar == 1:  # d/ds
                # s - derivative
                # d / ds = ((1 + a) / 2) / ((1 - b) / 2) d / da + d / db
                dmodeds = dfa * (gb * (0.5 * (1 + a)))
                dmodeds = dmodeds * (
                            (0.5 * (1 - b)) ** (di - 1)) if di > 0 else dmodeds

                tmp = dgb * ((0.5 * (1 - b)) ** di)
                tmp = tmp - 0.5 * di * gb * (
                            (0.5 * (1 - b)) ** (di - 1)) if di > 0 else tmp
                dmodeds = dmodeds + fa * tmp
                return 2 ** di * dmodeds
        elif len(idx) == 3:  # 3D
            #  UNTESTED!
            r = coors[..., 0]
            s = coors[..., 1]
            t = coors[..., 2]
            a = -2 * (1 + r) / (s + t) - 1
            b = 2 * (1 + s) / (1 - t) - t
            c = t
            a[nm.isnan(a)] = -1.
            b[nm.isnan(b)] = -1.
            di = idx[0]
            dj = idx[1]
            dk = idx[2]
            fa = self.jacobiP(a, 0, 0)[..., di]
            dfa = self.gradjacobiP(a, 0, 0)[..., di]
            gb = self.jacobiP(b, 2 * di + 1, 0)[..., dj]
            dgb = self.gradjacobiP(b, 2 * di + 1, 0)[..., dj]
            hc = self.jacobiP(c, 2 * (di + dj) + 2, 0, dk)
            dhc = self.gradjacobiP(c, 2 * (di + dj) + 2, 0, dk)

            V3Dr = dfa * (gb * hc)
            V3Dr = V3Dr * ((0.5 * (1 - b)) ** (di - 1)) if di > 0 else V3Dr
            V3Dr = V3Dr * ((0.5 * (1 - c)) ** (
                        di + dj - 1)) if di + dj > 0 else V3Dr

            tmp = dgb * ((0.5 * (1 - b)) ** di)
            tmp = tmp + (-0.5 * di) * (
                        gb * (0.5 * (1 - b)) ** (di - 1)) if di > 0 else tmp
            tmp = tmp * ((0.5 * (1 - c)) ** (
                        di + dj - 1)) if di + dj > 0 else tmp
            tmp = fa * (tmp * hc)

            if dvar == 0:
                return V3Dr * (2 ** (2 * di + dj + 1))
            elif dvar == 1:
                V3Ds = 0.5 * (1 + a) * V3Dr

                V3Ds = V3Ds + tmp
                return V3Ds * (2 ** (2 * di + dj + 1))
            elif dvar == 2:
                V3Dt = 0.5 * (1 + a) * V3Dr + 0.5 * (1 + b) * tmp
                tmp = dhc * ((0.5 * (1 - c)) ** (di + dj))
                tmp = tmp - 0.5 * (di + dj) * (hc * ((0.5 * (1 - c)) ** (
                            di + dj - 1))) if di + dj > 0 else tmp
                tmp = fa * (gb * tmp)
                tmp = tmp * ((0.5 * (1 - b)) ** di)
                V3Dt = V3Dt + tmp
                return V3Dt * (2 ** (2 * di + dj + 1))


simplex_coefM5 = nm.array([
    1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    -1,3,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    -2,2,4,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    1,-8,0,10,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    2,-12,-4,10,20,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    4,-8,-24,4,24,24,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    -1,15,0,-45,0,0,35,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
    -2,26,4,-66,-48,0,42,84,0,0,0,0,0,0,0,0,0,0,0,0,0,
    -4,36,24,-60,-192,-24,28,168,168,0,0,0,0,0,0,0,0,0,0,0,0,
    -8,24,96,-24,-192,-240,8,96,240,160,0,0,0,0,0,0,0,0,0,0,0,
    1,-24,0,126,0,0,-224,0,0,0,126,0,0,0,0,0,0,0,0,0,0,
    2,-44,-4,210,84,0,-336,-336,0,0,168,336,0,0,0,0,0,0,0,0,0,
    4,-72,-24,276,408,24,-352,-1248,-384,0,144,864,864,0,0,0,0,0,0,0,0,
    8,-96,-96,240,1056,240,-224,-1824,-2400,-160,72,864,2160,1440,0,0,0,0,0,0,0,
    16,-64,-320,96,960,1440,-64,-960,-2880,-2240,16,320,1440,2240,1120,0,0,0,0,0,0,
    -1,35,0,-280,0,0,840,0,0,0,-1050,0,0,0,0,462,0,0,0,0,0,
    -2,66,4,-496,-128,0,1392,864,0,0,-1620,-1920,0,0,0,660,1320,0,0,0,0,
    -4,116,24,-760,-672,-24,1848,3888,648,0,-1860,-7200,-3240,0,0,660,3960,3960,0,0,0,
    -8,184,96,-944,-2112,-240,1808,9216,5040,160,-1480,-12480,-18000,-3200,0,440,5280,13200,8800,0,0,
    -16,240,320,-800,-4480,-1440,1120,11520,18720,2240,-720,-10880,-33120,-26880,-1120,176,3520,15840,24640,12320,0,
    -32,160,960,-320,-3840,-6720,320,5760,20160,17920,-160,-3840,-20160,-35840,-20160,32,960,6720,17920,20160,8064],
    dtype=nm.int64).reshape(21, 21)

simplex_expoM5 = nm.array([
    0,0,0,
    0,1,0,
    1,0,0,
    0,2,0,
    1,1,0,
    2,0,0,
    0,3,0,
    1,2,0,
    2,1,0,
    3,0,0,
    0,4,0,
    1,3,0,
    2,2,0,
    3,1,0,
    4,0,0,
    0,5,0,
    1,4,0,
    2,3,0,
    3,2,0,
    4,1,0,
    5,0,0],
    dtype=nm.int64).reshape(21, 3)
