'''Quadrature nodes used in ``recursivenodes``.'''

import numpy as np
from scipy.special import gamma
from scipy.linalg import eigh_tridiagonal
from itertools import product
from recursivenodes.polynomials import jacobi, jacobider, jacobinorm2
from recursivenodes.utils import multiindex_equal, multiindex_up_to, npolys


def _jacobi_matrix(n, a=0., b=0.):
    ''' return as numpy arrays, the diagonals for the Jacobi recurrence relationship '''
    alpha = np.ndarray((n,))
    beta = np.ndarray((n,))
    for i in range(1, n+1):
        if i > 1:
            A = (2*i+a+b)*(2*i+a+b-1) / (2*i*(i+a+b))
            B = (a*a-b*b)*(2*i+a+b-1) / (2*i*(i+a+b)*(2*i+a+b-2))
        else:
            A = (a+b+2)/2
            B = (a-b)/2
        Ap = (2*(i+1)+a+b)*(2*(i+1)+a+b-1) / (2*(i+1)*(i+1+a+b))
        C = 2*((i+1)+a-1)*((i+1)+b-1)*(2*(i+1)+a+b) / (2*(i+1)*((i+1)+a+b)*(2*(i+1)+a+b-2))
        alpha[i-1] = -B / A
        beta[i-1] = (C / (A * Ap))
    return (alpha, beta)


def _monicjacobi(n, x, alpha, beta):
    ''' evaluate the monic Jacobi polynomial:
        - for the interval [-1,1]
        - with degree n
        - with the diagonals generated by _jacobi_matrix(n,a,b) for the Jacobi weight (1+x)^a (1-x)^b
        - at the points x (a numpy array)

        returns a numpy array with the same shape as x
    '''
    pm2 = np.ones(np.shape(x))
    pm1 = (x - alpha[0]) * pm2
    for i in range(1, n):
        pm2 *= -beta[i-1]
        pm2 += (x - alpha[i]) * pm1
        p = pm2
        pm2 = pm1
        pm1 = p
    return pm1


def gaussjacobi(n, a=0., b=0.):
    '''The `n`-point Gauss-Jacobi quadrature points and weights for weight function `(1+x)^a (1-x)^b` for the interval [-1,1],
    which exactly integrates polynomials with degree up to `2n-1`.

    Args:
        n (int): The number of points.
        a (float, optional): the left exponent of the weight function.
        b (float, optional): the right exponent of the weight function.

    Returns:
        (ndarray, ndarray): Each of shape (`n`,), the points (in ascending order) and weights of the quadrature rule.

    Reference:
        :cite:`GoWe69`
    '''
    alpha, beta = _jacobi_matrix(n, a, b)
    x, V = eigh_tridiagonal(alpha, (beta[:-1])**0.5)
    if (a == b):
        x = (x - x[::-1]) / 2.
        if (n % 2):
            x[n // 2] = 0.
    mu0 = 2**(a+b+1)*gamma(a+1)*gamma(b+1) / gamma(a+b+2)
    w = V[0, :]**2. * mu0
    if (a == b):
        w = (w + w[::-1]) / 2.
    return (x, w)


def lobattogaussjacobi(n, a=0., b=0.):
    '''The `n`-point Lobatto-Gauss-Jacobi quadrature points and weights for weight function `(1+x)^a (1-x)^b` for the interval [-1,1],
    which includes the endpoints `-1` and `1` and exactly integrates polynomials with degree up to `2n-3`.

    Args:
        n (int): The number of points.
        a (float, optional): the left exponent of the weight function.
        b (float, optional): the right exponent of the weight function.

    Returns:
        (ndarray, ndarray): Each of shape (`n`,), the points (in ascending order) and weights of the quadrature rule.

    Reference:
        :cite:`Golu73`
    '''
    mu0 = 2**(a+b+1)*gamma(a+1)*gamma(b+1) / gamma(a+b+2)
    if (n == 2):
        return (np.array([-1., 1.]), np.array([mu0/2., mu0/2.]))
    alpha, beta = _jacobi_matrix(n, a, b)
    ends = np.array([-1., 1.])
    p_nm2 = jacobi(n-2, ends, a, b)
    p_nm1 = jacobi(n-1, ends, a, b)
    p_nm2 = _monicjacobi(n-2, ends, alpha, beta)
    p_nm1 = _monicjacobi(n-1, ends, alpha, beta)
    rhs = np.array([-p_nm1[0], p_nm1[1]])
    Ainv = np.array([[p_nm2[1], -p_nm2[0]], [-p_nm1[1], p_nm1[0]]])
    Ainv /= (p_nm1[0] * p_nm2[1] - p_nm2[0] * p_nm1[1])
    sol = Ainv.dot(rhs)
    alpha[-1] = sol[0]
    beta[-2] = sol[1]
    x, V = eigh_tridiagonal(alpha, (beta[:-1])**0.5)
    if (a == b):
        x = (x - x[::-1]) / 2.
        if (n % 2):
            x[n // 2] = 0.
    w = V[0, :]**2. * mu0
    if (a == b):
        w = (w + w[::-1]) / 2.
    x[0] = -1.
    x[-1] = 1.
    return (x, w)


def simplexgausslegendre(d, n):
    '''An `n^d`-point Gaussian quadrature rule for the biunit simplex that exactly integrates polynomials with degree up to `2n-1`.

    Args:
        d (int): The spatial dimension.
        n (int): The points per coordinate direction

    Returns:
        (ndarray, ndarray): The points (shape (`n, \\dots, n, d`)) and weights (shape (`n, \\dots, n,`)) of the quadrature rule.
    '''
    if (d == 1):
        x, w = gaussjacobi(n, 0., 0.)
        return x.reshape((n, 1)), w
    xt, wt = simplexgausslegendre(d-1, n)
    xz, wz = gaussjacobi(n, d-1, 0.)
    x = np.ndarray((n,)*d+(d,))
    w = np.ndarray((n,)*d)
    for i in product(range(n), repeat=d-1):
        for j in range(n):
            x[i+(j,)][0:d-1] = (xt[i][0:d-1] + 1.) * ((1. - xz[j])/2.) - 1.
            x[i+(j,)][d-1] = xz[j]
            w[i+(j,)] = wt[i] * wz[j] / 2**(d-1)
    return x, w


