from __future__ import division

import numpy
from dolo.numeric.misc import cartesian

# Credits : both routines below are ported from the Compecon Toolbox
# by Paul L Fackler and Mario J. Miranda.
# It is downloadable at http://www4.ncsu.edu/~pfackler/compecon/toolbox.html
def hermgauss(n):

    from numpy import pi, fix, zeros, sqrt

    maxit = 100
    pim4 = 1 / pi ** 0.25
    m = int(fix((n + 1) / 2))

    x = zeros(n)
    w = zeros(n)
    # reasonable starting values
    for i in range(m):
        if i == 0:
            z = sqrt(2 * n + 1) - 1.85575 * ((2 * n + 1) ** (-1 / 6))
        elif i == 1:
            z = z - 1.14 * (n ** 0.426) / z
        elif i == 2:
            z = 1.86 * z + 0.86 * x[0]
        elif i == 3:
            z = 1.91 * z + 0.91 * x[1]
        else:
            z = 2 * z + x[i - 2]
        # root finding iterations
        its = 0
        while its < maxit:
            its += 1
            p1 = pim4
            p2 = 0
            for j in range(n):
                p3 = p2
                p2 = p1
                p1 = z * sqrt(2 / (j + 1)) * p2 - sqrt(j / (j + 1)) * p3
            pp = sqrt(2 * n) * p2
            z1 = z
            z = z1 - p1 / pp
            if abs(z - z1) < 1e-14:
                break
        if its >= maxit:
            raise Exception("Failure to converge")
        x[n - i - 1] = z
        x[i] = -z
        w[i] = 2 / pp ** 2
        w[n - i - 1] = w[i]

    return [x, w]


def gauss_hermite_nodes(orders, sigma, mu=None):
    """
    Computes the weights and nodes for Gauss Hermite quadrature.

    Parameters
    ----------
    orders : int, list, array
        The order of integration used in the quadrature routine
    sigma : array-like
        If one dimensional, the variance of the normal distribution being
        approximated. If multidimensional, the variance-covariance matrix of
        the multivariate normal process being approximated.

    Returns
    -------
    x : array
        Quadrature nodes
    w : array
        Quadrature weights
    """
    if isinstance(orders, int):
        orders = [orders]

    import numpy

    if mu is None:
        mu = numpy.array([0] * sigma.shape[0])

    herms = [hermgauss(i) for i in orders]

    points = [h[0] * numpy.sqrt(2) for h in herms]
    weights = [h[1] / numpy.sqrt(numpy.pi) for h in herms]

    if len(orders) == 1:
        # Note: if sigma is 2D, x will always be 2D, even if sigma is only 1x1.
        # print(points.shape)
        x = numpy.array(points[0]) * numpy.sqrt(float(sigma))
        if sigma.ndim == 2:
            x = x[:, None]
        w = weights[0]
        return [x, w]

    else:
        x = cartesian(points).T

        from functools import reduce

        w = reduce(numpy.kron, weights)

        zero_columns = numpy.where(sigma.sum(axis=0) == 0)[0]
        for i in zero_columns:
            sigma[i, i] = 1.0

        C = numpy.linalg.cholesky(sigma)

        x = numpy.dot(C, x) + mu[:, numpy.newaxis]

        x = numpy.ascontiguousarray(x.T)

        for i in zero_columns:
            x[:, i] = 0

        return [x, w]


# from numpy.polynomial.hermite import hermgauss


if __name__ == "__main__":

    orders = [8, 8]
    mu = numpy.array([0.05, 0.01])
    sigma = numpy.array(
        [
            [0.1, 0.015],
            [0.015, 0.1],
        ]
    )

    from numpy.polynomial.hermite import hermgauss as hermgauss_numpy

    [xg, wg] = hermgauss_numpy(10)
    [x, w] = hermgauss(10)
    print(w - wg)
    print(x - xg)
