"""
nitrogen.autodiff.forward
-------------------------

This module implements a simple forward accumulation
model for automatic differentiation. Its main object
is the ``adarray`` class.

Mathematical functions implemented include

==============================================   ====================================================
**Arithmetic and powers**                        **Description**
----------------------------------------------   ----------------------------------------------------
:func:`~nitrogen.autodiff.forward.add`           Addition, :math:`x + y`.
:func:`~nitrogen.autodiff.forward.subtract`      Subtraction, :math:`x - y`.
:func:`~nitrogen.autodiff.forward.mul`           Multiplication, :math:`x * y`.
:func:`~nitrogen.autodiff.forward.div`           Division, :math:`x/y`.
:func:`~nitrogen.autodiff.forward.powi`          Integer powers, :math:`x^i`.
:func:`~nitrogen.autodiff.forward.powf`          Real (or complex) powers, :math:`x^p`.
:func:`~nitrogen.autodiff.forward.sqrt`          Square root, :math:`\\sqrt{x}`.
==============================================   ====================================================

==============================================   ====================================================
**Trigometric and hyperbolic**                   **Description**
----------------------------------------------   ----------------------------------------------------
:func:`~nitrogen.autodiff.forward.sin`           Sine, :math:`\\sin(x)`.
:func:`~nitrogen.autodiff.forward.cos`           Subtraction, :math:`\\cos(x)`.
:func:`~nitrogen.autodiff.forward.asin`          Inverse sine, :math:`\\arcsin(x)`.
:func:`~nitrogen.autodiff.forward.acos`          Inverse cosine, :math:`\\arccos(x)`.
:func:`~nitrogen.autodiff.forward.sinh`          Hyperbolic sine, :math:`\\sinh(x)`.
:func:`~nitrogen.autodiff.forward.cosh`          Hyperbolic cosine, :math:`\\cosh(x)`.
:func:`~nitrogen.autodiff.forward.tanh`          Hyperbolic tangent, :math:`\\tanh(x)`.
==============================================   ====================================================

==============================================   ====================================================
**Exponents and logarithms**                     **Description**
----------------------------------------------   ----------------------------------------------------
:func:`~nitrogen.autodiff.forward.exp`           Exponential, :math:`\\exp(x)`.
:func:`~nitrogen.autodiff.forward.log`           Natural logarithm, :math:`\\log(x)`.
==============================================   ====================================================

==============================================   ====================================================
**Linear algebra**                               **Description**
----------------------------------------------   ----------------------------------------------------
:func:`~nitrogen.autodiff.forward.chol_sp`       Cholesky decomposition (symmetric, packed).
:func:`~nitrogen.autodiff.forward.int_tp`        Triangular matrix inverse (packed).
:func:`~nitrogen.autodiff.forward.llt_tp`        :math:`L L^T` for triangular matrix (packed).
:func:`~nitrogen.autodiff.forward.ltl_tp`        :math:`L^T L` for triangular matrix (packed).
==============================================   ====================================================

"""

import numpy as np
import warnings

class adarray:

    """
    Forward-type automatic differentiation object.

    Attributes
    ----------
    k : int
        The maximum derivative order.
    ni : int
        The number of independent variables.
    nck : ndarray
        A binomial coefficient table.
    idx : ndarray
        The multi-index table for derivatives of
        `ni` variables to order `k`.
    nd : int
        The number of unique derivatives.
    d : ndarray
        The derivative array, whose first index is the
        generalized derivative index. (See Notes.)
    zlevel : int
        The highest non-zero derivative order. If -1,
        then this adarray is identically zero.
    zlevels : ndarray
        The zero-levels of individual variables.

    Notes
    -----
    The derivative information is stored in the :attr:`d` attribute,
    an ndarray whose first index is a generalized derivative
    index. ``d[0]`` is the the value of the base array and ``d[i]``
    with ``i`` > 0 are the derivatives stored in *lexical* order.
    This ordering sorts derivatives first by their total degree: zeroth
    derivatives (the value), then first derivatives, then second
    derivatives, and so on. Within a group of derivatives of a given
    order, they are sorted by the derivative order with respect to the
    first independent variable, then by the order of the second, and
    so on. This ordering is the same as that of :attr:`idx`

    The values of higher-order derivatives are stored by convention
    with a factor equal to the inverse of the multi-index factorial, i.e.
    a derivative with multi-index ``[2, 0, 1, 3]`` would be stored as the
    corresponding derivative divided by :math:`2!\\times 0!\\times 1!\\times 3!`

    The zero-level indicators `zlevel` and `zlevels` are independent of
    each other.

    The accounting arrays, `nck`, `idx`, and `zlevels` should be considered
    immutable. They might be shared by multiple adarrays, unlike `d`. In general,
    do not modify adarray attributes directly.

    """

    def __init__(self,base_shape,k,ni, nck = None, idx = None, dtype = np.float64,
                 d = None, zlevel = None, zlevels = None):
        """
        Create a new adarray object.

        Parameters
        ----------
        base_shape : tuple of int
            Shape of base array. adarray.d will have shape
            ``(nd,) + base_shape``. Shape may be ().
        k : int
            Maximum derivative degree. `k` must be greater
            or equal to 0.
        ni : int
            Number of independent variables. `n` must be
            greater than or equal to 1.
        nck : ndarray, optional
            Return value of ncktab(nmax,kmax) with
            ``nmax >= k + ni`` and ``kmax >= min(k, ni)``.
            If None, this will be calculated.
        idx : ndarray, optional
            Return value of ``idxtab(k,ni)``.
            If None, this will be calculated.
        dtype : data-type, optional
            Data-type of initialized derivative array
        d : ndarray, optional
            A pre-allocated derivative array. If provided,
            `dtype` will be ignored and `d` must have
            a shape equal to ``(nd,) + base_shape``
        zlevel : int, optional
            The zero level indicator. If None, this
            will be set to `k`. The default is None.
        zlevels : array_like of int, optional
            The zero-levels of each variables. If None,
            each will be set to `k`. The default is None.

        See Also
        --------
        ncktab : Binomial coefficient table
        idxtab : Multi-index table

        """

        self.k = k
        self.ni = ni

        if nck is None:
            self.nck = ncktab(k + ni, min(k,ni))
        else:
            self.nck = nck # no copy

        if idx is None:
            self.idx = idxtab(k,ni)
        else:
            self.idx = idx # no copy

        self.nd = self.nck[k + ni, min(ni,k)]

        if d is None:
            self.d = np.empty((self.nd,) + base_shape, dtype = dtype)
        else:
            if d.shape != (self.nd,) + base_shape:
                raise ValueError("d does not have the correct shape")
            self.d = d

        if zlevel is None:
            self.zlevel = k
        else:
            self.zlevel = zlevel

        if zlevels is None:
            self.zlevels = np.array([k] * ni)
        else:
            self.zlevels = np.array(zlevels)

        # In general, zlevel should be
        # between (inclusive) max(zlevels) and sum(zlevels)
        # but that will not be explicitly kept track of.
        #
        # The most exhaustive z-leveling is a boolean
        # mask for each derivative, but I generally think
        # that the combination of the total zlevel
        # and the zlevels for each variables covers
        # most scenarios. I am trying to reduce overhead
        # for small grid-size evaluations.

    def copy(self, out = None):
        """
        Copy an adarray object.

        Parameters
        ----------
        out : adarray
            Output location. If None, this will be created
            if None. The default is None.

        Returns
        -------
        adarray
            An adarray object, with ``d`` attribute
            copied via ``d.copy()``.

        """
        if out is None:
            out = adarray(self.d.shape[1:],self.k,self.ni,
                    nck = self.nck,
                    idx = self.idx,
                    dtype = self.d.dtype,
                    d = self.d.copy(), # Copy the derivative ndarray itself
                    zlevel = self.zlevel,
                    zlevels = self.zlevels)
        else:
            # We assume out has the right shape, data-type, and that
            # nck and idx are already provided
            #
            np.copyto(out.d, self.d)
            out.zlevel = self.zlevel
            out.zlevels = self.zlevels

        return out

    # Define binary operators: +, -, ...
    # ADDITION
    def __add__(self,other):
        """ z = self + other """
        if type(other) == type(self):
            z = add(self, other) # adarray add
        else:
            z = self.copy()
            z.d[0] += other # Attempt ndarray iadd
            z.zlevel = max(z.zlevel, 0) # Assume other is non-zero
            z.zlevels = np.maximum(z.zlevels, 0)
        return z
    def __radd__(self,other):
        """ z = other + self """
        return self.__add__(other)

    # MULTIPLICATION
    def __mul__(self,other):
        """ z = self * other """
        if type(other) == type(self):
            z = mul(self, other)
        else:
            z = self.copy()
            z.d *= other # Attempt ndarray imul
            # assuming other to be constant.
            # NumPy broadcasting is important here.
            # If other is a scalar, then everything is multiplied
            # by it. If other is an ndarray of the same base_shape
            # as z, then it will be broadcast over each
            # derivative, as desired.
            #
            # zlevel is unchanged by constant multiplication
            # zlevels is  "  "

        return z
    def __rmul__(self,other):
        """ z = other * self """
        return self.__mul__(other)

    # DIVISION
    def __truediv__(self, other):
        """ z = self / other """
        if type(other) == type(self):
            z = div(self,other)
        else:
            z = self.copy()
            z.d /= other # ndarray idiv, using broadcasting
            # zlevel is unchanged by constant division
            # zlevels " "  "
        return z
    def __rtruediv__(self, other):
        """ z = other / self """
        # z = (self**-1.0) * other
        return (powf(self,-1.0)).__mul__(other)

    # SUBTRACTION
    def __sub__(self,other):
        """ z = self - other """
        return self.__add__(-other)
    def __rsub__(self,other):
        """ z = other - self """
        return (-self).__add__(other)

    # UNARY OPERATIONS
    def __neg__(self):
        """ z = -self """
        z = self.copy()
        np.negative(z.d, out = z.d) #z.d = -z.d
        # zlevel is unchanged by negation
        # zlevels " " "
        return z
    def __pos__(self):
        """ z = +self"""
        return self.copy()

def array(d,k,ni,copyd = False,zlevel = None, zlevels = None):
    """
    Create an adarray object from a raw derivative array.

    Parameters
    ----------
    d : ndarray
        The derivative array with shape (nd,...)
    k : int
        The maximum derivative order.
    ni : int
        The number of independent variables.
    copyd : boolean, optional
        If True, a copy of `d` will be made for returned
        adarray. If False, the adarray will use
        the same reference. The default is False.
    zlevel : int, optional
        The zero-level of `d`. If None, this
        will be set safely to `k`. The default is None.
    zlevels : array_like of int, optional
        The zero-level of each variable. If None,
        this will be set safely to `k` for each. The
        default is None.

    Returns
    -------
    adarray

    """

    if type(d) != np.ndarray:
        raise TypeError("d must be an ndarray")

    base_shape = d.shape[1:]

    if copyd:
        dinit = d.copy()
    else:
        dinit = d # use reference

    # Check the length of the first dimension
    if np.ndim(d) < 1:
        return ValueError("d must have at least 1 dimension")

    nck = ncktab(k + ni, min(k,ni))
    if d.shape[0] != nck[k+ni,min(k,ni)]:
        return ValueError("The first dimension of d has an incorrect length")

    if zlevel is None:
        zlevel = k

    if zlevels is None:
        zlevels = np.array([k] * ni)
    else:
        zlevels = np.array(zlevels)

    return adarray(base_shape, k, ni, nck, None, d.dtype, dinit, zlevel, zlevels)

def ndize1(x):
    """
    Create a 1D ndarray of adarray objects from a
    single adarray by promoting the first index
    of the base shape.

    Parameters
    ----------
    x : adarray
        The original adarray.

    Returns
    -------
    ndarray
        The 1D array of adarrays.

    Notes
    -----
    The derivative arrays of the new adarray objects
    are views of `x.d`.

    """

    if np.ndim(x.d) < 2:
        raise ValueError("The base shape must have at least 1 dimension")

    n = x.d.shape[1]    # the size of the new ndarray

    X = np.ndarray((n,), dtype = adarray)
    for i in range(n):
        X[i] = array(x.d[:,i], x.k, x.ni,
                     copyd = False,     # keep references
                     zlevel = x.zlevel,
                     zlevels = x.zlevels)

    return X

def nck(n,k):
    """
    Calculate the binomical coefficient (`n` choose `k`).

    This function uses a simple recursive algorithm. Use of
    :func:`ncktab` may be significantly faster (i.e.
    ``ncktab(n,k)[n,k]``.)

    Parameters
    ----------
    n, k : int
        Arguments of the binomial coefficient.

    Returns
    -------
    np.uint64
        The binomial coefficient (`n` choose `k`).

    Examples
    --------
    >>> adf.nck(4,2)
    6

    >>> adf.nck(4,0)
    1

    See Also
    --------
    ncktab : Binomial coefficient table

    """

    if k > n or k < 0 or n < 0:
        return np.uint64(0)
    elif n == 0 :
        return np.uint64(1)
    elif k == 0 or k == n:
        return np.uint64(1)
    else:
        return nck(n-1,k-1) + nck(n-1,k)

def ncktab(nmax,kmax=None):
    """
    Binomial coefficient table

    Parameters
    ----------
    nmax : int
        Maximum value of n (first argument of coeff).
    kmax : int, optional
        Maximum value of k. If None, its value is
        set to nmax. The default is None.

    Returns
    -------
    tab : ndarray
        2D array of shape (`nmax` + 1, `kmax` + 1) containing
        the binomial coefficients, `tab[n,k]` =
        (`n` choose `k`). The values of invalid elements
        (i.e. `k` > `n`) are undefined. The data-type is
        ``np.uint64``.

    Examples
    --------
    >>> adf.ncktab(3)
    array([[1, 0, 0, 0],
           [1, 1, 0, 0],
           [1, 2, 1, 0],
           [1, 3, 3, 1]], dtype=uint64)

    """

    if kmax is None:
        kmax = nmax

    if nmax < 0:
        raise ValueError('nmax must be >= 0')
    if kmax < 0:
        raise ValueError('kmax must be >= 0')

    tab = np.zeros((nmax+1,kmax+1), dtype = np.uint64)

    for n in range(nmax+1):
        for k in range(kmax+1):

            if k > n:
                tab[n,k] = 0
            elif n < 1:
                tab[n,k] = 1
            elif k < 1:
                tab[n,k] = 1
            elif k == n:
                tab[n,k] = 1
            else:
                tab[n,k] = tab[n-1,k-1] + tab[n-1,k]

    return tab

def nderiv(deriv, nvar):
    """
    The number of derivatives up to order `deriv`,
    inclusively, in `nvar` variables. This equals
    the binomial coefficient (`deriv` + `nvar`, `nvar`).

    Parameters
    ----------
    deriv : int
        The maximum derivative order.
    nvar : int
        The number of independent variables.

    Returns
    -------
    np.uint64
        The number of derivatives.

    Examples
    --------
    >>> adf.nderiv(3,2)
    10

    >>> adf.nderiv(2,6)
    28

    """
    n = deriv + nvar
    k = min(deriv, nvar)
    return ncktab(n,k)[n,k]

def nckmulti(a,b,nck):
    """
    Calculate multi-index binomial coefficient (`a`, `b`)

    Parameters
    ----------
    a,b : ndarray
        1D multi-indices
    nck : ndarray
        Binomial coefficient table as returned by
        ``ncktab(nmax,kmax)``, with nmax >= max(`a`)
        and kmax >= max(`b`)

    Returns
    -------
    np.float64
        Multi-index binomial coefficient.

    See Also
    --------
    ncktab : Binomial coefficient table

    Examples
    --------
    >>> adf.nckmulti(np.array([4,2,3]), np.array([2,1,2]), adf.ncktab(4))
    36.0

    """

    ni = a.size # lenth of multi-index

    value = np.float64(1.0) # double float

    for i in range(ni):
        value *= nck[a[i],b[i]]

    return value

def idxtab(k,ni):

    """
    Multi-index table for derivatives up to order
    `k` of `ni` variables.

    The number of multi-indices (i.e. the number of
    unique derivatives) is `nd` = nck(`k` + `ni`, `ni`)

    Parameters
    ----------
    k : int
        Maximum derivative order.
    ni : int
        Number of variables.

    Returns
    -------
    ndarray
        Multi-index table with shape (`nd`, `ni`) and
        data-type ``np.uint32``.

    Examples
    --------
    >>> adf.idxtab(2,3)
    array([[0, 0, 0],
           [1, 0, 0],
           [0, 1, 0],
           [0, 0, 1],
           [2, 0, 0],
           [1, 1, 0],
           [1, 0, 1],
           [0, 2, 0],
           [0, 1, 1],
           [0, 0, 2]], dtype=uint32)

    """

    #nd = nck(k + ni, ni) # The number of derivatives
    nd = ncktab(k+ni,min(k,ni))[k+ni, min(k,ni)] # --this is faster than nck
    idx = np.zeros((nd,ni),dtype=np.uint32)

    # The first row is already all zeros -- no derivative
    # Compute next row by incrementing.

    for i in range(1,nd):

        # idxtab[i,:] equals the increment of
        # idxtab[i-1,:]

        # We will treat the second row in the table
        # as a special case. It is just [1, 0, 0, ...]
        if i == 1:
            idx[i,0] = 1
            continue

        # The current degree is >= 1

        # Look for the right-most non-zero degree
        for j in range(ni-1,-1,-1):

            if idx[i-1,j] > 0:
                # If it is not at the last position, shift
                # one degree right one position.
                if j < (ni-1):
                    idx[i,:] = idx[i-1,:]
                    idx[i,j] -= 1
                    idx[i,j+1] += 1
                    break
                else: # it is at the last position, j = nd-1
                    # Look for the *next* non-zero position
                    for p in range(ni-2,-1,-1):
                        if idx[i-1,p] > 0:
                            # We found another non-zero entry.
                            # Shift one from this entry to
                            # the right and add everything
                            # from the last position
                            idx[i,:] = idx[i-1,:]
                            idx[i,ni-1] = 0
                            idx[i,p] -= 1
                            idx[i,p+1] += 1 + idx[i-1,ni-1]
                            break
                    else: # We did not find another non-zero
                        # this means we have an index like
                        # 0,0,0,...,b
                        # so the next index is
                        # b+1,0,0,...,0
                            idx[i,:] = 0
                            idx[i,0] = 1 + idx[i-1,ni-1]
                    break

    return idx

def idxposk(a,nck):
    """
    Calculate the relative lexical position of multi-index `a`
    within the block of its degree ``k`` = sum(`a`).

    Parameters
    ----------
    a : ndarray
        1D multi-index
    nck : ndarray
        A binomial coefficent table as returned by
        ``ncktab(nmax,kmax)`` with ``nmax >= ni - 3 + k - a[0]``
        and ``kmax >= min(ni - 2, k - a[0] - 1)``. Simpler
        requirements that satisfy these are ``nmax >= ni + k - 1`` and
        ``kmax >= min(ni, k - 1)``.

    Returns
    -------
    posk : np.uint64
        Relative position of multi-index in its
        block of degree ``k`` = sum(`a`).

    Examples
    --------
    >>> adf.idxposk(np.array([0,0,0]),adf.ncktab(3,3))
    0

    >>> adf.idxposk(np.array([2,0,1,3]),adf.ncktab(5,2))
    33

    See Also
    --------
    ncktab : Binomial coefficient table

    """

    ni = a.size

    if ni == 0 or ni == 1:
        posk = np.uint64(0)
    else:
        k = np.sum(a)

        if k == 0:
            posk = np.uint64(0)
        else:
            a0 = a[0]

            # Because we assume that we will only index `nck` as
            # nck[n,min(k,n-k)], we cannot use the following simple sum
            #
            #  posk = np.sum( nck[(ni-2):(ni-2+(k-a0)), ni-2] ) + idxposk(a[1:], nck)
            #
            # Instead, we can explicitly loop over the sum
            #
            posk = np.uint64(0)
            for p in range(k-a0):
                posk += nck[p + ni - 2, min(p, ni-2)]
            posk += idxposk(a[1:], nck)

    return posk

def idxpos(a,nck):
    """
    Calculate the absolute lexical position of multi-index `a`.

    Parameters
    ----------
    a : ndarray
        1D multi-index
    nck : ndarray
        A binomial coefficient table as returned by
        ``ncktab(nmax,kmax)``, with ``nmax >= ni + k - 1``
        and ``kmax >= min(ni, k - 1)``, where ``ni = a.size``
        and ``k = sum(a)``.

    Returns
    -------
    pos : np.uint64
        The absolute lexical position of multi-index `a`

    Examples
    --------
    >>> adf.idxpos(np.array([0,0,0]),adf.ncktab(3,3))
    0

    >>> adf.idxpos(np.array([2,0,1,3]),adf.ncktab(9,4))
    159

    See Also
    --------
    ncktab: Binomial coefficient table

    """

    k = np.sum(a) # Degree of multi-index a
    ni = a.size   # Number of variables

    if k == 0:
        return np.uint64(0)

    else:

        offset = nck[ni + k - 1, min(ni,k-1)] # The number of multi-indices with degree
                                              # less than k
        posk = idxposk(a,nck)    # The position of this multi-index within
                                 # the block of multi-indices of the same degree k
        return offset + posk

def mvleibniz(X, Y, k, ni, nck, idx, out=None, Xzlevel = None, Yzlevel = None,
              Xzlevels = None, Yzlevels = None):
    """
    Multivariate Leibniz formula for derivative arrays.

    Parameters
    ----------
    X,Y : ndarray
        Derivative array factors (e.g. :attr:`adarray.d`). These
        must have matching shapes.
    k : int
        Maximum derivative order of `X` and `Y`
    ni : int
        Number of independent variables for `X` and `Y`
    nck : ndarray
        Binomial coefficient table, as returned
        by ``ncktab(nmax,kmax)`` with ``nmax`` >= `k` + `ni` - 1
        and ``kmax`` >= min(`ni`, `k` - 1)
    idx : ndarray
        Multi-index table, as returned by ``idxtab(k, ni)``
    out : ndarray, optional
        Output location. If None, a new ndarray is created with the
        result type of X and Y's data-types
    Xzlevel, Yzlevel : int, optional
        Zero-level for input. If None, this is assumed to be `k`.
        The default is None.
    Xzlevels, Yzlevels : ndarray, optional
        The zero-levels for each variable. If None, this is assumed
        to be `k` for all. The default is None.

    Returns
    -------
    out : ndarray
        The derivative array for `X` * `Y`

    See Also
    --------
    adarray
    ncktab : Binomial coefficient table
    idxtab : Multi-index table

    Notes
    -----
    :func:`mvleibniz` is low-level function applied directly to
    derivative arrays (usually :attr:`adarray.d`).
    Typically, multiplication should be applied at high-level
    with the * operator directly with :class:`adarray` objects.

    Xzlevel and Xzlevels are used independently to determine
    whether certain derivatives are to be skipped.

    """

    # Initialize result Z to zero
    res_type = np.result_type(X.dtype, Y.dtype)
    if out is None:
        out = np.zeros(X.shape, dtype = res_type)
    else:
        if out.dtype != res_type:
            raise TypeError("out data-type is incompatible with X * Y")
        out.fill(0)

    if Xzlevel is None:
        Xzlevel = k
    if Yzlevel is None:
        Yzlevel = k

    if Xzlevels is None:
        Xzlevels = np.array([k] * ni)
    if Yzlevels is None:
        Yzlevels = np.array([k] * ni)

    Z = out # Reference only

    nd,_ = idx.shape

    #
    # Sum through all derivatives of X and Y
    # For each pair, accumulate the result
    # into the appropriate derivative of Z
    #
    # Derivatives are ordered with lower degrees
    # first.

    for iX in range(nd):
        idxX = idx[iX,:]   # The derivative index of X
        kX = np.sum(idxX)  # The derivative degree
        if kX > Xzlevel:
            break # Skip remaining X derivatives. They are zero.
        if (idxX > Xzlevels).any():
            continue

        for iY in range(nd):
            idxY = idx[iY,:]  # The derivative index of Y
            kY = np.sum(idxY) # The derivative degree
            if kY > Yzlevel:
                break # Skip remaining Y derivatives. They are zero.
            if (idxY > Yzlevels).any():
                continue

            kZ = kX + kY   # The product's degree
            if kZ > k: # Total degree of this term is too high
                break # Skip remaining Y derivatives

            # the product of the X derivative and Y derivatives
            # contributes to the Z derivative with index:
            idxZ = idxX + idxY
            # Convert this multi-index to a 1D index
            iZ = idxpos(idxZ, nck)

            # adarrays are normalized to include (inverse) factorial
            # factors for each derivative. This convention removes
            # the binomial coefficient in the generalized Leibniz
            # formula. I.e. we do not need to calculate
            #    nckmulti(idxZ, idxX, nck)
            #

            # Add this term to Z
            # Use of '+=' with ndarrays will not create a new copy
            Z[iZ] += X[iX] * Y[iY]

    return Z

def mvchain(df,X,k,ni,nck,idx, out = None, Xzlevel = None, Xzlevels = None):
    """
    Multivariate chain rule Z = f(X) via Taylor series.

    Parameters
    ----------
    df : ndarray
        An array containing the derivatives of
        single-argument function f through order `k`.
        The shape of `df` is ``(k+1,) + X.shape[1:]``
    X : ndarray
        Derivative array in ``adarray.d`` format
    k : int
        Maximum derivative order
    ni : int
        Number of independent variables
    nck : ndarray
        Binomial coefficient table for `X` satisfying requirements
        for :func:`mvleibniz`.
    idx : ndarray
        Multi-index table for `X` satisfying requirements for
        :func:`mvleibniz`.
    out : ndarray, optional
        Output location. If None, a new ndarray is created with the
        same data-type as X.
    Xzlevel : int, optional
        The zlevel of the X derivative array. If None, this is
        assumed to be `k`.
    Xzlevels : ndarray, optional
        The zlevels of each variable. If None, this is
        assumed to be `k` for each.

    Returns
    -------
    out : ndarray
        The derivative array for f(`X`)

    See Also
    --------
    ncktab : Binomial coefficient table
    idxtab : Multi-index table
    mvleibniz : Generalized Leibniz product rule
    adchain : Chain rule for adarray objects

    Notes
    -----
    :func:`mvchain` is a low-level function that acts directly
    on derivative arrays (usually :attr:`adarray.d`). In most cases, the
    high-level function :func:`adchain` should be used directly with
    :class:`adarray` objects.
    """

    X0 = X[:1].copy() # Value of X
    X[:1].fill(0) # X now equals "X-X0"
    if Xzlevel is None:
        Xzlevel = k
    if Xzlevels is None:
        Xzlevels = np.array([k] * ni)

    # Initialize result to zero (and create if necessary)
    res_type = np.result_type(df, X)
    if out is None:
        out = np.ndarray(X.shape, dtype = res_type)

    if out.dtype != res_type:
        raise TypeError("out data-type is incompatible with f(X)")

    # Initialize result to zero
    out.fill(0)

    Z = out # Reference only

    for i in range(k+1):

        if i == 0:
            # Initialize Xi = X**i = X**0 to 1 (constant)
            Xi = np.zeros(X.shape,dtype = X.dtype)
            Xi[0] = 1.0
            Xi[:1].fill(1.0)
            Xizlevel = 0 # the zero-level is 0 (constant)
            Xizlevels = np.array([0] * ni)
            fact = 1.0 # Initialize factorial to 1
        else:
            Xi = mvleibniz(Xi,X,k,ni,nck,idx,
                           Xzlevel = Xizlevel, Yzlevel = Xzlevel,
                           Xzlevels = Xizlevels, Yzlevels = Xzlevels)
            # Xi <-- Xi * (X-X0)
            # The call to mvleibniz needs to create a new output buffer
            # because we cannot overwrite Xi while mvleibniz is operating
            # (this probably could be done more efficiently -- revisit)

            # Determine the zlevel of Xi
            if Xizlevel < 0 or Xzlevel < 0: # if either X**i or X is iden. 0
                Xizlevel = -1 # so is the next power
            else:
                Xizlevel = min(Xizlevel + Xzlevel, k)

            if (Xizlevels < 0).all() or (Xzlevels < 0).all():
                Xizlevels = np.array([-1] * ni)
            else:
                Xizlevels = np.minimum(Xizlevels + Xzlevels, k)

            fact = fact * i

        if (df[i] != 0 ).any() and Xizlevel >= 0 and (Xizlevels >= 0).all():
            # If both df is non-zero and
            # Xi is non-zero
            Z += (df[i] * (1.0/fact)) * Xi # note broadcast of df[i] over Xi

    # Restore the value of X
    np.copyto(X[:1], X0)

    return Z

def const(value,k,ni,nck = None, idx = None):
    """
    Create an :class:`adarray` object for a constant scalar or array.

    Parameters
    ----------
    value : array_like
        The constant value. If this is not an ndarray already, it
        will try to be converted.
    k : int
        The maximum derivative degree.
    ni : int
        The number of independent variables
    nck, idx : ndarray, optional
        See :class:`adarray` constructor.

    Returns
    -------
    adarray
        A constant adarray object.

    Examples
    --------
    >>> adf.const(1., 2, 2).d
    array([1., 0., 0., 0., 0., 0.])

    >>> adf.const(42j, 1, 3).d
    array([0.+42.j, 0. +0.j, 0. +0.j, 0. +0.j])

    >>> adf.const([2.0, 3.0], 2, 2).d
    array([[2., 3.],
           [0., 0.],
           [0., 0.],
           [0., 0.],
           [0., 0.],
           [0., 0.]])
    """

    if type(value) != np.ndarray:
        value = np.array(value)

    base_shape = value.shape
    nck = ncktab(k + ni, min(k,ni))
    nd = nck[k + ni, min(k,ni)]

    d = np.zeros((nd,)+base_shape, dtype = value.dtype)

    # Copy value to d[:1]
    np.copyto(d[:1], value)

    # Determine the zlevel
    if (value == 0).all():
        zlevel = -1 # Identically zero
        zlevels = np.array([-1] * ni)
    else:
        zlevel = 0 # A non-zero constant
        zlevels = np.array([0] * ni)

    c = array(d, k, ni, copyd = False, zlevel = zlevel, zlevels = zlevels)

    return c

def sym(value,i,k,ni,nck = None, idx = None):
    """
    Create an :class:`adarray` for a symbol (i.e. one
    the independent variables with respect to which derivatives
    are being taken).

    Parameters
    ----------
    value : array_like
        The value of the variable.
    i : int
        The variable index, `i` = 0, ... , `ni` - 1
    k, ni: int
        See :class:`adarray` constructor.
    nck,idx : ndarray, optional
        See :class:`adarray` constructor.

    Returns
    -------
    adarray
        An :class:`adarray` object for variable `i` and value `value`.

    Examples
    --------
    >>> adf.sym(2.0, 0, 2, 2).d
    array([2., 1., 0., 0., 0., 0.])

    >>> adf.sym(3.0, 1, 2, 2).d
    array([3., 0., 1., 0., 0., 0.])

    >>> adf.sym([3.0, 4.0j], 1, 2, 2).d
    array([[3.+0.j, 0.+4.j],
           [0.+0.j, 0.+0.j],
           [1.+0.j, 1.+0.j],
           [0.+0.j, 0.+0.j],
           [0.+0.j, 0.+0.j],
           [0.+0.j, 0.+0.j]])

    """

    if k < 0:
        raise ValueError("k be >= 0")
    if ni < 1:
        raise ValueError("ni must be >= 1 to construct symbol")

    if i > ni+1 or i < 0 :
        raise ValueError('Symbol index i must be 0, ..., ni-1')

    x = const(value, k, ni, nck, idx)
    # const returns x with zlevel = -1 or 0 depending on value

    if k > 0 :
        x.d[(i+1):(i+2)].fill(1)

        x.zlevel = 1

        x.zlevels = np.array([0] * ni)
        x.zlevels[i] = 1

    # If k == 0, then the zlevel/zlevels from const can be kept

    return x

def empty_like(x, dtype = None):
    """
    Create an *uninitialized* adarray with the same properties as `x`,
    including base array data-type. The zlevel will be maximum.

    Parameters
    ----------
    x : adarray
        Prototype object

    dtype : dtype, optional
        Base data-type. If None, then `x.d.dtype` is used.

    Returns
    -------
    adarray
        A new adarray with the same properties as `x`

    >>> empty_like(const([3.3, 2.1], 2, 2)).d.shape
    (6, 2)

    """
    if dtype is None:
        dtype = x.d.dtype

    return adarray(x.d.shape[1:], x.k, x.ni, x.nck, x.idx, dtype)

def const_like(value, x, dtype = None):
    """
    Create a constant adarray initialized to
    `value` with similar properties to `x`.

    Parameters
    ----------
    value : scalar or array_like
        Value.
    x : adarray
        Prototype.
    dtype : dtype, optional
        Data-type. If None, then `x.d.dtype` is used.

    Returns
    -------
    adarray

    """

    if dtype is None:
        dtype = x.d.dtype

    base_shape = x.d.shape[1:]
    if np.isscalar(value):
        value = np.full(base_shape, value, dtype=dtype)

    return const(value, x.k, x.ni, x.nck, x.idx)

def add(x, y, out = None):
    """
    Add x + y

    Parameters
    ----------
    x,y : adarray
        Input argument
    out : adarray, optional
        Output location. If None, this
        will be created. The default is None.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.0, 0, 2, 2)
    >>> y = adf.sym(3.0, 1, 2, 2)
    >>> adf.add(x,y).d
    array([4., 1., 1., 0., 0., 0.])

    """

    res_type = np.result_type(x.d, y.d)
    if out is None:
        out = empty_like(x, dtype = res_type)

    if res_type != out.d.dtype:
        raise TypeError("output data-type incompatible with x + y")

    np.add(x.d, y.d, out = out.d)
    out.zlevel = max(x.zlevel, y.zlevel)
    out.zlevels = np.maximum(x.zlevels, y.zlevels)

    return out

def subtract(x, y, out = None):
    """
    Subtract x - y

    Parameters
    ----------
    x,y : adarray
        Input argument
    out : adarray, optional
        Output location. If None, this
        will be created. The default is None.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.0, 0, 2, 2)
    >>> y = adf.sym(3.0, 1, 2, 2)
    >>> adf.subtract(x,y).d
    array([-2.,  1., -1.,  0.,  0.,  0.])

    """

    res_type = np.result_type(x.d, y.d)
    if out is None:
        out = empty_like(x, dtype = res_type)

    if res_type != out.d.dtype:
        raise TypeError("output data-type incompatible with x - y")

    np.subtract(x.d, y.d, out = out.d)
    out.zlevel = max(x.zlevel, y.zlevel)
    out.zlevels = np.maximum(x.zlevels, y.zlevels)

    return out

def mul(x, y, out = None):
    """
    Multiply x * y

    Parameters
    ----------
    x,y : adarray
        Input argument
    out : adarray, optional
        Output buffer. If None, this will be created.
        The default is None.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.0, 0, 2, 2)
    >>> y = adf.sym(3.0, 1, 2, 2)
    >>> adf.mul(x,y).d
    array([3., 3., 1., 0., 1., 0.])

    """

    res_type = np.result_type(x.d, y.d)
    if out is None:
        out = empty_like(x, dtype = res_type)

    if res_type != out.d.dtype:
        raise TypeError("output data-type incompatible with x * y")

    mvleibniz(x.d,y.d,x.k,x.ni,x.nck,x.idx, out = out.d,
              Xzlevel = x.zlevel, Yzlevel = y.zlevel,
              Xzlevels = x.zlevels, Yzlevels = y.zlevels)

    if x.zlevel < 0 or y.zlevel < 0 :
        # Either of the factors is identically zero
        # So is the result.
        out.zlevel = -1
    else:
        # Both zlevels are >= 0
        out.zlevel = min(x.zlevel + y.zlevel, x.k)

    if (x.zlevels < 0).any() or (y.zlevels < 0).any():
        out.zlevels = np.array([-1] * x.ni)
    else:
        out.zlevels = np.minimum(x.zlevels + y.zlevels, x.k)

    return out

def div(x, y, out = None):
    """
    Divide x / y

    Parameters
    ----------
    x,y : adarray
        Input argument
    out : adarray, optional
        Output buffer. If None, this will be created.
        The default is None.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.0, 0, 2, 2)
    >>> y = adf.sym(3.0, 1, 2, 2)
    >>> adf.div(x,y).d
    array([ 0.33333333,  0.33333333, -0.11111111,  0.        , -0.11111111,
            0.03703704])

    """

    res_type = np.result_type(x.d, y.d)
    if out is None:
        out = empty_like(x, dtype = res_type)

    if res_type != out.d.dtype:
        raise TypeError("output data-type incompatible with x / y")

    # Calculate 1 / y
    iy = powf(y, -1.0)
    # Multiply x * (1/y)
    return mul(x,iy, out = out)

def adchain(df,x, out=None):

    """
    Calculate f(x) via chain rule.

    Parameters
    ----------
    df : ndarray
        Single-variable derivatives of f(x) w.r.t to its argument
        up to order :attr:`x.k`

    x : adarray
        Argument of f(`x`)

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        The result f(`x`).

    """

    if out is None:
        out = empty_like(x)

    mvchain(df, x.d, x.k, x.ni, x.nck, x.idx, out = out.d,
            Xzlevel = x.zlevel, Xzlevels = x.zlevels)

    # Determine the zlevel of the result
    # We assume, safely, that df has a maximum zlevel
    if x.zlevel <= 0:
        # x is identically zero or constant
        # The result is also constant
        out.zlevel = 0
    else:
        # x has a zlevel >= 1
        out.zlevel = x.k

    # For zlevels, first assume full dependence
    out.zlevels = np.array([x.k]*x.ni)
    #
    # then for any variable i that x does not
    # depend on (x.zlevels[i] <= 0), neither
    # will the output depend on it
    #
    out.zlevels[x.zlevels <= 0] = 0

    return out

def sin(x, out=None):
    """
    Sine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input in radians.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.0, 0, 3, 1)
    >>> adf.sin(x).d
    array([ 0.84147098,  0.54030231, -0.42073549, -0.09005038])
    """

    xval = x.d[0] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    for i in range(k+1):

        if i == 0:
            df[i] = np.sin(xval)
        elif i == 1:
            df[i] = np.cos(xval)
        else:
            df[i] = -df[i-2]

    return adchain(df, x, out = out)

def cos(x, out=None):
    """
    Cosine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input in radians.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(2.0, 0, 3, 1)
    >>> adf.cos(x).d
    array([-0.41614684, -0.90929743,  0.20807342,  0.15154957])

    """

    xval = x.d[0] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    for i in range(k+1):

        if i == 0:
            df[i] = np.cos(xval)
        elif i == 1:
            df[i] = -np.sin(xval)
        else:
            df[i] = -df[i-2]

    return adchain(df, x, out = out)

def asin(x, out = None):
    """
    Arcsine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input, :math:`x \\in [-1,1]`
    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result. The real part lies in :math:`[-\\pi/2, \\pi/2]`.

    Examples
    --------
    >>> x = adf.sym(0.35, 0, 3, 1)
    >>> adf.asin(x).d
    array([0.3575711 , 1.06752103, 0.21289593, 0.28767379])

    """

    xval = x.d[:1] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    df[0] = np.arcsin(xval)

    #
    # The first derivative of arcsin(x) is 1/sqrt(1-x^2)
    #
    # We'll construct this now
    #
    y = sym(xval, 0, k, 1)
    z = powf(1.0 - y*y, -0.5)
    # z.d contains the **scaled** derivatives of 1/sqrt(1-x^2) w.r.t. x
    #
    scale = 1.0
    for i in range(1, k+1):
        df[i] = z.d[i-1] * scale
        scale *= i

    return adchain(df, x, out = out)

def acos(x, out = None):
    """
    Arccosine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input, :math:`x \\in [-1,1]`
    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result. The real part lies in :math:`[0,\\pi]`.

    Examples
    --------
    >>> x = adf.sym(0.35, 0, 3, 1)
    >>> adf.acos(x).d
    array([ 1.21322522, -1.06752103, -0.21289593, -0.28767379])

    """

    #
    # arccos(x) = pi/2 - arcsin(x)
    #

    out = asin(x, out = out)  # Put arcsin(x) into place
    # We are assuming that asin returns with
    # zlevel/zlevels all >= 0
    out.d *= -1               # Multiply by -1
    out.d[0] += np.pi/2       # Add pi/2 to value

    return out

def exp(x, out = None):
    """
    Exponential for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input argument.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.exp(x).d
    array([4.48168907, 4.48168907, 2.24084454, 0.74694818])

    """

    xval = x.d[:1] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    np.exp(xval, out = df[:1])
    for i in range(1,k+1):
        np.copyto(df[i:(i+1)], df[:1])

    return adchain(df, x, out = out)

def log(x, out = None):
    """
    Natural logarithm for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input argument.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(3.0, 0, 3, 1)
    >>> adf.log(x).d
    array([ 1.09861229,  0.33333333, -0.05555556,  0.01234568])

    """

    xval = x.d[:1] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    # Value of f: ln(x)
    np.log(xval, out = df[:1])

    # First derivative of f:
    # 1/x
    if k > 0:
        np.copyto(df[1:2], 1.0/xval)

    for i in range(2,k+1):
        # i^th derivative of f:
        # (1/x)**i = 1/x * (1/x)**(i-1)
        #
        np.multiply((-i + 1) * df[1:2], df[(i-1):i], out = df[i:(i+1)])

    return adchain(df, x, out = out)

def sinh(x, out = None):
    """
    Hyperbolic sine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input argument.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.sinh(x).d
    array([2.12927946, 2.35240962, 1.06463973, 0.39206827])

    """

    xval = x.d[:1] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    np.sinh(xval, out = df[:1])

    # The derivative of sinh is cosh
    # and the derivative of cosh is sinh
    if k >= 1:
        np.cosh(xval, out = df[1:2])

    for i in range(2,k+1):
        np.copyto(df[i:(i+1)], df[(i-2):(i-1)])

    return adchain(df, x, out = out)

def cosh(x, out = None):
    """
    Hyperbolic cosine for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input argument.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.cosh(x).d
    array([2.35240962, 2.12927946, 1.17620481, 0.35487991])

    """

    xval = x.d[:1] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,)+x.d.shape[1:], dtype = xval.dtype)

    np.cosh(xval, out = df[:1])

    # The derivative of sinh is cosh
    # and the derivative of cosh is sinh
    if k >= 1:
        np.sinh(xval, out = df[1:2])

    for i in range(2,k+1):
        np.copyto(df[i:(i+1)], df[(i-2):(i-1)])

    return adchain(df, x, out = out)

def tanh(x, out = None):
    """
    Hyperbolic tangent for :class:`adarray` objects.

    Parameters
    ----------
    x : adarray
        Input argument.

    out : adarray, optional
        Output location of result. `out` must have the same
        properties as x. If None, a new adarray is allocated and
        returned.

    Returns
    -------
    adarray
        Result.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.tanh(x).d
    array([ 0.90514825,  0.18070664, -0.1635663 ,  0.0878162 ])

    """

    sh = sinh(x)
    ch = cosh(x)

    # tanh = sinh / cosh

    return div(sh, ch, out = out)

def powi(x, i, out = None):
    """
    x**i for integer i

    Parameters
    ----------
    x : adarray
        Base argument.
    i : integer
        Integer exponent.
    out : adarray, optional
        Output location of result.

    Returns
    -------
    adarray
        Result.

    Notes
    -----
    If `i` == 0, then powi returns a constant 1 for
    any value of `x`. For negative `i`, `x` is inverted
    and then the positive power is applied to 1/`x`.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.powf(x, 3).d
    array([3.375, 6.75 , 4.5  , 1.   ])

    """

    if out is None:
        out = empty_like(x)

    # If exponent is negative,
    # then invert argument and
    # and sign of exponent
    if i < 0:
        x = powf(x, -1.0) # 1 / x
        i = -i

    # Check for identity, x**1 = x
    #
    if i == 1: # Just a copy
        np.copyto(out.d, x.d)
        out.zlevel = x.zlevel
        out.zlevels = x.zlevels
        return out

    # Otherwise,
    # initialize out to constant 1.
    res = const_like(1.0, x)

    # Do right-to-left binary
    # exponentiation
    a = x
    while True:
        if i % 2 == 1:
            res = res * a

        i = i // 2  # floor division
        if i == 0:
            break
        a = a * a

    np.copyto(out.d, res.d)
    out.zlevel = res.zlevel
    out.zlevels = res.zlevels

    return out

def powf(x, p, out = None):
    """
    x**p for general p

    Parameters
    ----------
    x : adarray
        Base argument.
    p : float or complex
        Exponent.
    out : adarray, optional
        Output location of result.

    Returns
    -------
    adarray
        Result.

    Notes
    -----
    powf uses the NumPy
    `float_power() <https://numpy.org/doc/stable/reference/generated/numpy.float_power.html#numpy.float_power>`_
    function to
    compute the value array. It inherits the branch-cut convention
    of this function.

    Examples
    --------
    >>> x = adf.sym(1.5, 0, 3, 1)
    >>> adf.powf(x, -2.5).d
    array([ 0.36288737, -0.60481228,  0.70561433, -0.70561433])

    """

    if np.result_type(x.d.dtype, p) != np.result_type(x.d.dtype):
        raise TypeError("Invalid type combination")

    xval = x.d[0] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,) + x.d.shape[1:], dtype = xval.dtype)

    df[0] = np.float_power(xval, p)
    for i in range(1, k+1):
        df[i] = (p-i+1) * (df[i-1] / xval)

    return adchain(df, x, out = out)

def sqrt(x, out = None):
    """
    sqrt(x)

    Parameters
    ----------
    x : adarray
        Input argument
    out : adarray, optional
        Output buffer. If None, this will be created.
        The default is None.

    Returns
    -------
    adarray
        Result.

    Notes
    -----
    The adarray sqrt function uses the NumPy
    `sqrt() <https://numpy.org/doc/stable/reference/generated/numpy.sqrt.html#numpy.sqrt>`_
    function as its underlying routine. Its branch-cut convention is inherited.

    Examples
    --------
    >>> x = adf.sym(2.5, 0, 3, 1)
    >>> adf.sqrt(x).d
    array([ 1.58113883,  0.31622777, -0.03162278,  0.00632456])
    """

    xval = x.d[0] # Value array of x
    k = x.k

    df = np.ndarray( (k+1,) + x.d.shape[1:], dtype = xval.dtype)

    df[0] = np.sqrt(xval) # Uses numpy branch cut

    if k >= 1:
        ixval = 1.0 / xval # The inverse value
    for i in range(1, k+1):
        df[i] = (1.5 - i) * (df[i-1] * ixval)

    return adchain(df, x, out = out)

def reduceOrder(F, i, k, ni, idx, out = None):
    """
    Reduce the derivative array for F with respect to
    variable i. The returned derivative array is
    that for the function :math:`\\partial_i F`.

    Parameters
    ----------
    F : ndarray
        The derivative array up to degree `k`
        in `ni` variables.
    i : int
        The variable index (0, ..., `ni` - 1) to
        reduce.
    k : int
        The initial derivative order of F.
    ni : int
        The number of independent variables.
    idx : ndarray
        The return value of idxtab with suitable parameters.
    out : ndarray, optional
        Output buffer. If None, this will be
        created. The default is None.

    Returns
    -------
    out : ndarray
        The derivative array with shape (`nd_reduced`, ...)
        corresponding to the new function :math:`\\partial_i F`

    """

    nd,_ = idx.shape # The number of derivatives

    if k <= 0:
        raise ValueError("Cannot reduce a derivative array with k = 0")

    if i < 0 or i >= ni :
        raise ValueError(f"Cannot reduce w.r.t variable index i = {i:d} with only ni = {ni:d} variables")

    # k and ni are now both > 0
    #
    # Calculate the number of derivatives for
    # order k - 1 in ni variables. We already know
    # that nd is the number of deriv for order k
    # in ni variables, so we can use the simple result:
    nd_reduced = (nd * k) // (k + ni)  # This should be an integer result always!

    if out is None:
        out = np.ndarray( (nd_reduced,) + F.shape[1:], dtype = F.dtype)
    G = out # reference only

    # The index table `idx` for F is already provided,
    # so we will loop through this instead of building
    # a new index table for the reduced derivative array
    # The derivatives we are interested in appear
    # in the same order in the F (unreduced) and G (reduced)
    # derivative arrays.
    # This is to say that if one has two multi-indices of F
    # a = [a0 a1 a2 ... ai ...] and b = [b0 b1 b2 ... bi ...]
    # a appears before b
    # then the corresponding reduced indices
    # a' = [a0 a1 a2 ... ai-1 ... ] and b' = [b0 b1 b2 ... bi-1 ...]
    # appear in the same order: a' before b'
    #
    #
    iG = 0 # running position in the reduced derivative array
    for iF in range(nd):

        idxF = idx[iF,:]  # The multi-index of the original function's derivatives

        if idxF[i] < 1:
            # This derivative is not part of the reduced array
            continue
        # else, idxF[i] >= 1. It contributes to the
        # derivative element of G corresponding to idxF[i]--
        #
        # Because we store the derivatives with a factor
        # equal to the inverse of the multi-index factorial,
        # we need to correct this in the reduced array
        #
        # The correct conversion factor is
        #  idxF! / idxG! = idxF[i]! / idxG[i]!
        #                = idxF[i]! / (idxF[i] - 1)!
        #                = idxF[i]
        #
        # Now we do this:
        # G[iG] = F[iF] * idxF[i]
        #
        # Depending on the shape of G, we will split this up
        if len(G.shape) == 1:
            G[iG] = F[iF] * idxF[i] # Scalar multiplication
        else:
            np.multiply(F[iF], idxF[i], out = G[iG]) # use np.multiply to save on temp memory

        iG += 1

    # Check that nothing went wrong in the book keeping
    assert iG == nd_reduced, "Derivative array reduction was not successful!"

    return G

def n2N(n):
    """
    Calculate the square matrix rank N
    for a packed storage size n

    Parameters
    ----------
    n : int
        The packed length.

    Returns
    -------
    N : np.uint64
        The matrix rank.

    Examples
    --------
    >>> adf.n2N(21)
    6

    """

    N = np.uint64((np.sqrt(8*n+1)-1)/2)

    return N

def chol_sp(H, out = None):
    """
    Cholesky decomposition of a symmetric matrix in packed format.
    If real symmetric, then H should be positive definite.
    If complex symmetric (*not* Hermitian), then H should have
    non-zero pivots.

    Parameters
    ----------
    H : ndarray of adarray
        H is stored in 1D packed format (see :mod:`nitrogen.linalg.packed`)
    out : ndarray of adarray
        Output buffer. If None, this will be created.
        If out = H, then in-place decomposition is performed

    Returns
    -------
    out : ndarray of adarray
        The lower triangle Cholesky decomposition L in packed storage.
        H = L @ L.T

    """

    if H.ndim != 1:
        raise ValueError('H must be 1-dimensional')

    if out is None:
        out = np.ndarray(H.size, dtype = adarray)
        for i in range(H.size):
            out[i] = empty_like(H[i])

    # Copy H to out for in-place routine
    if out is not H:
        for i in range(H.size):
            H[i].copy(out = out[i])

    _chol_sp_unblocked(out)

    return out

def _chol_sp_unblocked(H):
    """
    An unblocked, in-place implementation of Cholesky
    decomposition for symmetric matrices H.

    Parameters
    ----------
    H : ndarray of adarray
        H is a symmetric matrix in 1D packed storage.
        (Lower triangle row-packed)

    Returns
    -------
    ndarray of adarray
        The in-place result.

    Notes
    -----
    This routine uses a standard Cholesky algorithm
    for *real* symmetric matrices.

    """

    # H is a 1d ndarray of adarray objects
    # in packed format
    n = H.size
    N = n2N(n)

    L = np.ndarray((N,N), dtype = adarray)
    # Copy references to adarrays to the lower
    # triangle of a full "reference" matrix
    # References above the diagonal are undefined.
    k = 0
    for i in range(N):
        for j in range(i+1):
            L[i,j] = H[k]
            k += 1

    # Compute the Cholesky decomposition
    # with a simple unblocked algorithm
    tol = 1e-10
    pivmax = np.abs( np.sqrt(L[0,0].d[:1]) ) # for pivot threshold checking

    for j in range(N):
        r = L[j,:j]         # The j^th row, left of the diagonal

        # L[j,j] <-- sqrt(d - r @ r.T)
        sqrt( (L[j,j] - r @ r.T).copy(), out = L[j,j])

        # Check that the pivot is sufficiently large
        if (np.abs(L[j,j].d[:1]) / pivmax < tol).any() :
            warnings.warn("Small diagonal (less than rel. tol. = {:.4E} encountered in Cholesky decomposition".format(tol))

        # Store the new maximum pivot
        pivmax = np.maximum(pivmax, np.abs(L[j,j].d[:1]))

        # Calculate the column below the diagonal element j
        #B = L[j+1:,:j]      # The block between r and c
        #c = L[j+1:,j]       # The j^th column, below the diagonal
        for i in range(j+1,N):
            Bi = L[i,:j]    # An ndarray
            ci = L[i, j]    # An adarray
            #L[j+1:,j] = (c - B @ r.T) / L[j,j]
            div( (ci - Bi @ r.T).copy(), L[j,j], out = L[i,j] )

    return H

def inv_tp(L, out = None):
    """
    Invert a triangular matrix in lower row-packed (or upper column-packed)
    storage.

    Parameters
    ----------
    L : ndarray of adarray
        L is stored in 1D packed format (see :mod:`nitrogen.linalg.packed`)
    out : ndarray of adarray
        Output buffer. If None, this will be created.
        If out = L, then in-place inversion is performed

    Returns
    -------
    out : ndarray of adarray
        The inverse of the triangular matrix in packed storage.

    """

    if L.ndim != 1:
        raise ValueError('L must be 1-dimensional')

    if out is None:
        out = np.ndarray(L.size, dtype = adarray)
        for i in range(L.size):
            out[i] = empty_like(L[i])

    # Copy L to out for in-place routine
    if out is not L:
        for i in range(L.size):
            L[i].copy(out = out[i])

    # Now perform in-place on `out`
    _inv_tp_unblocked(out)

    return out

def _inv_tp_unblocked(L):
    """
    Invert a lower triangular matrix in row-packed storage.

    Parameters
    ----------
    L : ndarray of adarray
        L is a triangular matrix in 1D packed storage.
        (Row-packed for lower, column-packed for upper)

    Returns
    -------
    ndarray of adarray
        The in-place result.

    """

    n = L.size
    N = n2N(n)
    one = np.uint64(1)

    X = np.ndarray((N,N), dtype = adarray)
    # Copy references to adarrays to the lower
    # triangle of a full "reference" matrix
    # Elements above the diagonal are not defined!
    k = 0
    for i in range(N):
        for j in range(i+1):
            X[i,j] = L[k]
            k += 1

    # Compute the triangular inverse
    # with a simple in-place element by element algorithm
    abstol = 1e-15
    # In-place lower triangle inversion
    for j in range(N - one, -1,-1):

        # Compute j^th diagonal element
        if (np.abs(X[j,j].d[:1]) < abstol).any():
            warnings.warn(f"Small diagonal (less than abs. tol. = {abstol:.4E})" \
                          "encounted in triangle inversion")

        #X[j,j] = 1.0/X[j,j]
        powf(X[j,j].copy(), -1.0, out = X[j,j])

        # Compute j^th column, below diagonal
        for i in range(N - one, j, -1):
            mul( -X[j,j], X[i, j+1:i+1] @ X[j+1:i+1, j], out = X[i,j])

    return L

def llt_tp(L, out = None):
    """
    L @ L.T of a lower triangular matrix.

    Parameters
    ----------
    L : ndarray of adarray
        Lower triangular matrix in packed storage.
    out : ndarray of adarray
        Output buffer. If None, this will be created.
        If out = L, then in-place multiplication is performed

    Returns
    -------
    out : ndarray of adarray
        The result in packed storage.

    """

    if L.ndim != 1:
        raise ValueError('L must be 1-dimensional')

    if out is None:
        out = np.ndarray(L.size, dtype = adarray)
        for i in range(L.size):
            out[i] = empty_like(L[i])

    # Copy L to out for in-place routine
    if out is not L:
        for i in range(L.size):
            L[i].copy(out = out[i])

    # Perform in-place L @ L.T
    _llt_tp_unblocked(out)

    return out

def _llt_tp_unblocked(L):
    """
    An unblocked, in-place routine for multiplying
    L @ L.T where L is a lower triangular matrix
    in packed row-order storage.

    This is equivalent to U.T @ U where U is in
    upper triangular packed column-order storage.

    The resulting symmetric matrix is returned in
    packed storage.

    Parameters
    ----------
    L : ndarray of adarray
        A lower triangular matrix in 1D packed
        row-order storage.

    Returns
    -------
    ndarray of adarray
        The in-place result.

    """

    # Calculate matrix dimensions
    n = L.size
    N = n2N(n)
    one = np.uint64(1)

    A = np.ndarray((N,N), dtype = adarray)
    # Copy references to adarrays to the lower
    # triangle of a full "reference" matrix
    # References above the diagonal are undefined.
    k = 0
    for i in range(N):
        for j in range(i+1):
            A[i,j] = L[k]
            k += 1

    # This is similar to a "reverse Cholesky decomposition"
    # so we will work in the opposite direction as that

    for j in range(N-one, -1, -1):

        r = A[j,:j]         # The j^th row, left of the diagonal

        for i in range(N-one, j, -1):
            Bi = A[i,:j]    # An ndarray
            ci = A[i, j]    # An adarray

            # ci <-- Ljj * ci + Bi @ r.T
            ( A[j,j] * ci + Bi @ r.T ).copy(out = A[i,j])

        (A[j,j]*A[j,j] + r @ r.T).copy(out = A[j,j])

    return L

def ltl_tp(L, out = None):
    """
    L.T @ L with a lower triangular matrix.

    Parameters
    ----------
    L : ndarray of adarray
        Lower triangular matrix in packed storage.
    out : ndarray of adarray
        Output buffer. If None, this will be created.
        If out = L, then in-place multiplication is performed

    Returns
    -------
    out : ndarray of adarray
        The symmetric result in packed storage.

    """

    if L.ndim != 1:
        raise ValueError('L must be 1-dimensional')

    if out is None:
        out = np.ndarray(L.size, dtype = adarray)
        for i in range(L.size):
            out[i] = empty_like(L[i])

    # Copy L to out for in-place routine
    if out is not L:
        for i in range(L.size):
            L[i].copy(out = out[i])

    # Perform in-place L @ L.T
    _ltl_tp_unblocked(out)

    return out

def _ltl_tp_unblocked(L):
    """
    An unblocked, in-place routine for multiplying
    L.T @ L where L is a lower triangular matrix
    in packed row-order storage.

    This is equivalent to U @ U.T where U is in
    upper triangular packed column-order storage.

    The resulting symmetric matrix is returned in
    packed storage.

    Parameters
    ----------
    L : ndarray of adarray
        A lower triangular matrix in 1D packed
        row-order storage.

    Returns
    -------
    ndarray of adarray
        The in-place result.

    """

    # Calculate matrix dimensions
    n = L.size
    N = n2N(n)

    A = np.ndarray((N,N), dtype = adarray)
    # Copy references to adarrays to the lower
    # triangle of a full "reference" matrix
    # References above the diagonal are undefined.
    k = 0
    for i in range(N):
        for j in range(i+1):
            A[i,j] = L[k]
            k += 1

    # This is the "converse" of the llt routine
    # for L @ L.T
    for i in range(N):
        for j in range(i+1):
            # Compute A[i,j]
            # This is the dot product of the
            # i^th row of L.T and the
            # j^th column of L
            #
            # The i^th row of L.T is zero until
            # its i^th element
            #
            # The j^th column of L is zero until
            # its j^th element
            #
            # So the dot product need only begin
            # at the max(i,j)^th element
            #
            # By the loop ranges, j is always <= i
            # so max(i,j) = i, and we can begin
            # the dot product with the i^th element

            # The first factor is the
            # i^th row of L.T beginning at its i^th element
            # This is the transpose of the i^th column of
            # L beginning at its i^th element, which is in
            # the lower triangle, so A's reference is OK
            F1 = (A[i:,i]).T
            # The second factor is the j^th column
            # of L beginning at its i^th element, which is
            # also in the lower triangle, so OK
            F2 = A[i:,j]

            (F1 @ F2).copy(out = A[i,j])

    return L

def cost(k,ni, quiet = False):
    """
    Estimate the cost of adarray operations.

    Parameters
    ----------
    k : int
        The total degree, :math:`k \\geq 0`
    ni : int
        The number of variables, :math:`n_i \\geq 1`
    quiet : bool, optional
        Suppress printed output. The default is False.

    Returns
    -------
    None

    """

    nck = ncktab(k + 2*ni)
    k = np.uint32(k)
    two = np.uint32(2)
    one = np.uint32(1)

    nd = nck[k+ni,k] # The number of derivatives

    # First, calculate the memory cost. This equals
    # the number of derivatives stored
    mem_cost = nd

    # Estimate the number of floating point operations
    # required for some simple arithmetic. (Assuming real numbers.)
    #
    # Addition requires the addition of each derivative separately
    add_cost = nd
    #
    # Multiplication requires the generalized Leibniz formula,
    # "mvleibniz". This has
    # sum_(kp=0)^k  (kp + n - 1 choose kp) (k-kp+ni choose k-kp)
    #      =  (k + 2n choose k)
    # terms. Each term requires one multiplication and one addition
    # (to a given derivative of the result) for a total of 2 FLOPs
    # per term
    mul_cost = two * nck[k + 2*ni, k]

    #
    # Now calculate the cost of a call to mvchain
    # This assumes df is already calculated.
    #
    # The total includes
    #  (k+1) * mul_cost (for calculating powers) and
    #  (k+1) * (2*nd + 1) for FLOPs scaling by df and summing result.
    mvc_cost = (k + one) * (mul_cost + two*nd + one)
    #
    # Any actual usage of mvchain requires the calculation
    # of df, which will be variable depending on the
    # function itself. Often these go as ~ k primitive multiplies,
    # which is typically small compared to the remaining cost.
    #
    fun_cost = mvc_cost + k + one
    if not quiet:
        print(f"Memory scaling         = {mem_cost:d}")
        print(f"Addition FLOPs         = {add_cost:d}")
        print(f"Multiplication FLOPs   = {mul_cost:d}")
        print(f"Chain rule FLOPs       = {mvc_cost:d}")
        print(f"Typical function eval. = {fun_cost:d}")

    return


# coding: utf-8
import itertools
import random
import warnings
from operator import itemgetter

import numpy as np
import scipy
from scipy.optimize import fmin_l_bfgs_b
from scipy.sparse import coo_matrix
from scipy.spatial.distance import pdist, squareform
from sksparse.cholmod import cholesky

from gphypo.base_bo import BaseBO

def adj_metric(u, v):
    '''
    give this function to scipy.spatial.distance.dist
    :param u:
    :param v:
    :return: 1 if (u, v) is adj else 0
    '''
    if np.abs(u - v).sum() == 1:
        return 1
    else:
        return 0

def mat_flatten(x):
    return np.array(x).flatten()

def create_normalized_X_grid(meshgrid):
    def create_normalized_meshgrid(meshgrid):
        mesh_shape = meshgrid[0].shape
        fixed_shape = mesh_shape[:2][::-1] + mesh_shape[2:]
        new_mesh = np.meshgrid(*[np.arange(shape) for shape in fixed_shape])
        return new_mesh

    normalized_meshgrid = create_normalized_meshgrid(list(meshgrid))
    normalized_meshgrid = np.array(normalized_meshgrid)
    normalized_X_grid = normalized_meshgrid.reshape(normalized_meshgrid.shape[0], -1).T
    return normalized_X_grid

def create_adjacency_mat_using_pdist(dim_list):
    meshgrid = np.array(np.meshgrid(*[np.arange(ndim) for ndim in dim_list]))
    X_grid = meshgrid.reshape(meshgrid.shape[0], -1).T
    dist = pdist(X_grid, adj_metric)
    return squareform(dist)

def create_adjacency_mat(dim_list, calc_sparse=True):
    K_func = lambda x: scipy.sparse.csc_matrix(x) if calc_sparse else x

    xp = scipy.sparse if calc_sparse else np
    adj1d_list = [create_adjacency_mat_using_pdist([ndim]) for ndim in dim_list]
    if len(dim_list) == 1:
        return K_func(adj1d_list[0])

    K = xp.kron(adj1d_list[1], xp.eye(dim_list[0])) + xp.kron(xp.eye(dim_list[1]), adj1d_list[0])

    prod = dim_list[0] * dim_list[1]

    for i in range(2, len(dim_list)):
        K = xp.kron(K, xp.eye(dim_list[i])) + xp.kron(xp.eye(prod), adj1d_list[i])
        prod *= dim_list[i]

    return K_func(K)

class GMRF_BO(BaseBO):
    def __init__(self, bo_param_list, environment,
                 burnin=False, is_edge_normalized=False, n_early_stopping=None,
                 gt_available=False, normalize_output="zero_mean_unit_var", optimizer="fmin_l_bfgs_b",
                 update_hyperparam_func='pairwise_sampling',
                 acquisition_func='ucb', acquisition_param_dic={'beta': 5, 'pi': 0.01}, n_stop_pairwise_sampling=np.inf,
                 GAMMA=0.01, GAMMA0=0.01, GAMMA_Y=0.01, ALPHA=0.1, initial_k=1, initial_theta=1, n_ctr=None):

        '''
        meshgrid: Output from np.methgrid.
        e.g. np.meshgrid(np.arange(-1, 1, 0.1), np.arange(-1, 1, 0.1)) for 2D space
        with |x_i| < 1 constraint. s
        environment: Environment class which is equipped with sample() method to
        return observed value.
        beta (optional): Hyper-parameter to tune the exploration-exploitation
        balance. If beta is large, it emphasizes the variance of the unexplored
        solution solution (i.e. larger curiosity)
        '''

        super(GMRF_BO, self).__init__(bo_param_list, environment,
                                      n_early_stopping=n_early_stopping,
                                      gt_available=gt_available, normalize_output=normalize_output,
                                      acquisition_func=acquisition_func, acquisition_param_dic=acquisition_param_dic,
                                      optimizer=optimizer, n_ctr=n_ctr)

        self.n_stop_pairwise_sampling = n_stop_pairwise_sampling
        self.GAMMA = GAMMA
        self.GAMMA0 = GAMMA0
        self.GAMMA_Y = GAMMA_Y
        self.ALPHA = ALPHA

        self.does_pairwise_sampling = False
        assert update_hyperparam_func in [None, "pairwise_sampling", "simple_loglikelihood", "loglikelihood"] #TODO "simple_loglikelihood", "loglikelihood" does not work (2017/9/7)
        if update_hyperparam_func == 'pairwise_sampling':
            self.update_hyperparam = self.update_hyper_params_by_adj_idxes
            self.does_pairwise_sampling = True
        elif update_hyperparam_func == 'simple_loglikelihood':
            self.update_hyperparam = self.update_hyperparams_by_simple_loglikelihood
            self.normalized_X_grid = create_normalized_X_grid(self.meshgrid)
        elif update_hyperparam_func == 'loglikelihood':
            self.update_hyperparam = self.update_hyperparams_by_loglikelihood
        else:
            self.update_hyperparam = None

        tau0 = - create_adjacency_mat(self.ndim_list)
        row_sum_list = - mat_flatten(tau0.sum(axis=0))
        self.diff_list = row_sum_list.max() - row_sum_list
        print("Edge Num: %s" % np.count_nonzero(self.diff_list))
        print(tau0.toarray())
        self.is_edge_normalized = is_edge_normalized

        # # This changes the edge weight. This seems to be not so effective(Result does not change)
        # if is_edge_normalized:
        #     weight_arr = np.matrix(self.ndim * 2 / row_sum_list)
        #     weight_mat = weight_arr.T.dot(weight_arr)
        #     # weight_mat = np.sqrt(weight_arr.T.dot(weight_arr))
        #
        #     if type(tau0) == scipy.sparse.csc.csc_matrix:
        #         tau0 = tau0.multiply(weight_mat)
        #     else:
        #         tau0 *= weight_mat  # O.K. for dense mat but N.G. for sparse mat
        #
        #     tau0 = scipy.sparse.csc_matrix(tau0)
        #     print(tau0.toarray())

        self.baseTau0 = tau0
        row_idxes, col_idxes, _ = scipy.sparse.find(self.baseTau0 != 0)
        self.adj_pair_list = [(r, c) for r, c in zip(row_idxes, col_idxes)]

        # Grid-based BurnIn
        if burnin and not environment.reload:
            if self.n_ctr:
                self.do_all_point_burnin()
            else:
                self.do_grid_based_burnin()

        if environment.reload or burnin:
            if self.update_hyperparam is not None:
                self.update_hyperparam()

        self.update()
        self.total_clicked_ratio_list = []

    def do_all_point_burnin(self):
        n_sample_per_point = int(self.n_ctr / self.n_points)
        for i in range(self.n_points):
            self.sample(self.X_grid[i], n_sample_per_point)

    def do_grid_based_burnin(self):
        burnin = 2 ** self.ndim
        quarter_point_list = [[bo_param[len(bo_param) // 4], bo_param[len(bo_param) // 4 * 3]] for bo_param in
                              self.bo_param_list]

        for coodinate in itertools.product(*quarter_point_list):
            coodinate = np.array(coodinate)
            self.sample(coodinate, self.n_ctr)
            if self.does_pairwise_sampling:
                adj_idx = self.get_pairwise_idx(self.X_grid2idx_dic[tuple(coodinate)])
                self.sample(self.X_grid[adj_idx], self.n_ctr)

        X_seq, T_seq = self.point_info_manager.X_seq, self.point_info_manager.T_seq
        self.bestT = np.max(T_seq)
        self.bestX = X_seq[np.argmax(T_seq)]

        self.does_pairwise_sampling = False
        print('%d burins has finised!' % burnin)

    def calc_tau(self):
        gamma = self.GAMMA
        gamma_y = self.GAMMA_Y
        gamma0 = self.GAMMA0

        n_grid = self.point_info_manager.get_n_grid()
        sum_grid = self.point_info_manager.get_sum_grid()

        if self.is_edge_normalized:
            # corner will be treated as center nodes. This is explained as A -> A' in PPTX file.
            gamma_tilda = n_grid * gamma + gamma0 + gamma_y * self.diff_list.flatten()
        else:
            gamma_tilda = n_grid * gamma + gamma0  # Normal

        # print ('gamma_tilda: {}'.format(gamma_tilda) )
        mu_tilda = np.array([s * gamma + self.ALPHA * gamma0 for s in sum_grid]) / gamma_tilda

        tau1 = -scipy.sparse.diags(gamma_tilda)
        tau0 = self.baseTau0 * self.GAMMA_Y

        tau0 += scipy.sparse.diags(gamma_tilda - mat_flatten(tau0.sum(axis=0)))

        return tau0, -tau1, mu_tilda, gamma_tilda

    def get_adj_idxes(self, idx):
        return scipy.sparse.find(self.baseTau0[idx] != 0)[1]

    def get_pairwise_idx(self, idx):
        adj_idxes = scipy.sparse.find(self.baseTau0[idx] != 0)[1]
        return random.choice(adj_idxes)

    def learn(self):
        T = self.point_info_manager.get_T(excludes_none=True)

        grid_idx = np.argmax(self.acquisition_func.compute(self.mu, self.sigma, T))

        continue_flg = self.sample(self.X_grid[grid_idx])

        if not continue_flg:
            return False

        if self.does_pairwise_sampling:
            adj_idx = self.get_pairwise_idx(grid_idx)
            continue_flg = self.sample(self.X_grid[adj_idx])

            if not continue_flg:
                return False

        if self.point_info_manager.update_cnt > self.n_stop_pairwise_sampling:
            self.does_pairwise_sampling = False

        self.update()
        return True

    def learn_from_clicks(self, mu2ratio_dir='./mu2ratio', mu_sigma_csv_fn='./mu2ratio/mu_sigma.csv',
                          ratio_csv_fn='./mu2ratio/ratios.csv'):
        self.call_mu2ratio(mu2ratio_dir=mu2ratio_dir, mu_sigma_csv_fn=mu_sigma_csv_fn, ratio_csv_fn=ratio_csv_fn)
        self.sample_using_ratio_csv(ratio_csv_fn=ratio_csv_fn)
        self.update()

        self.save_mu_sigma_csv(mu_sigma_csv_fn)

        n_total_clicked = self.environment.result_df.output.values[-self.n_points:].astype(np.float64).sum()
        total_clicked_ratio = n_total_clicked / self.n_ctr
        self.total_clicked_ratio_list.append(total_clicked_ratio)
        print('ratio: %s' % total_clicked_ratio)

        return True

    def update(self, n_start_opt_hyper_param=0):
        A, B, mu_tilda, gamma_tilda = self.calc_tau()

        # start = time.time()
        # cov = scipy.sparse.linalg.inv(A) # This is slow

        factor = cholesky(A)
        cov = scipy.sparse.csc_matrix(factor(np.eye(A.shape[0])))
        # end = time.time() - start
        # print("one calc: %s sec" % end)

        self.mu = mat_flatten(cov.dot(B).dot(mu_tilda))
        self.sigma = mat_flatten(np.sqrt(cov[np.diag_indices_from(cov)]))

        # Update hyperparamters
        if self.update_hyperparam is not None and self.point_info_manager.update_cnt > n_start_opt_hyper_param:
            self.update_hyperparam()

    def update_hyper_params_by_adj_idxes(self):
        T = self.point_info_manager.get_T(excludes_none=False)
        var_list = []

        for i, j in self.adj_pair_list:
            t_i, t_j = T[i], T[j]
            if (t_i is not None) and (t_j is not None):
                var_list.append((t_i - t_j) ** 2)

        var_list = np.unique(var_list)
        var = var_list.mean()
        if self.n_ctr:
            self.GAMMA_Y = 1 / var
            self.GAMMA = self.GAMMA_Y * (2 * self.ndim)
            self.GAMMA0 = self.GAMMA * 0.01
        else:
            self.GAMMA_Y = 1 / var / (self.ndim ** 2)
            self.GAMMA = self.GAMMA_Y * (2 * self.ndim)
            self.GAMMA0 = self.GAMMA * 0.01
        self.ALPHA = np.mean(self.point_info_manager.get_T(excludes_none=True))

        print("New GammaY: %s" % self.GAMMA_Y)
        print("New Gamma: %s" % self.GAMMA)
        print("New ALPHA: %s" % self.ALPHA)

    # TODO this method does not work well and caliculation is heavy...
    def update_hyperparams_by_simple_loglikelihood(self):
        mean_grid = self.point_info_manager.get_T()

        unobserbed_idxes = np.isnan(mean_grid)
        obserbed_idxes = (np.isnan(mean_grid) == False)

        y_hat = mean_grid[obserbed_idxes]

        n_obserbed = len(y_hat)

        new_X_grid = np.concatenate([self.normalized_X_grid[obserbed_idxes], self.normalized_X_grid[unobserbed_idxes]],
                                    axis=0)
        dist = pdist(new_X_grid, metric=adj_metric)  # TODO heavy calculation

        A = np.zeros_like(dist)
        A[dist == 1 ** 2] = -1
        A = squareform(A)
        diag_arr = - A.sum(axis=0)
        A[np.diag_indices_from(A)] = diag_arr

        A00 = A[:n_obserbed, :n_obserbed]
        A01 = A[n_obserbed:, :n_obserbed]
        A11 = A[n_obserbed:, n_obserbed:]

        Lambda = A00 - (A01.T).dot(np.linalg.inv(A11).dot(A01))

        print(Lambda.shape, A.shape)

        t = y_hat.dot(Lambda.dot(y_hat[:, np.newaxis]))[0] / 2

        self.k += 1 / 2
        self.theta = self.theta / (t * self.theta + 1)
        self.GAMMA_Y = (self.k - 1) * self.theta
        self.GAMMA = self.GAMMA_Y * (2 * self.ndim)
        self.GAMMA0 = self.GAMMA * 0.01

        if self.normalize_output is None:
            self.ALPHA = np.mean(self.point_info_manager.get_T(excludes_none=True))
            print("New ALPHA: %s" % self.ALPHA)

        print("New GammaY: %s" % self.GAMMA_Y)
        print("New Gamma: %s" % self.GAMMA)
        print("t: %s" % t)
        print('k: %s, theta: %s' % (self.k, self.theta))

        log_likelihood = (y_hat.dot(Lambda.dot(y_hat[:, np.newaxis])) - np.log(np.linalg.det(Lambda) + 0.00001))[0]
        print('log_likelihood: %s' % log_likelihood)

    # TODO this method does not work well...
    def update_hyperparams_by_loglikelihood(self):
        # Update hyper-paramter of GMRF below
        def obj_func(eval_gradient=True):
            if eval_gradient:
                lml, grad = self.log_marginal_likelihood(eval_gradient=True)
                return -lml, -grad
            else:
                return -self.log_marginal_likelihood()

        # optima = [(self._constrained_optimization(obj_func, theta))] # Use when gradient can be calculated
        optima = [(self._constrained_optimization(lambda x: obj_func(x, eval_gradient=False),
                                                  gradExist=False))]  # Use when gradient cannot be calculated

        # Select result from run with minimal (negative) log-marginal likelihood
        lml_values = list(map(itemgetter(1), optima))

        self.GAMMA, self.GAMMA_Y = optima[np.argmin(lml_values)][0]
        self.GAMMA0 = 0.01 * self.GAMMA

        print("New GammaY: %s" % self.GAMMA_Y)
        print("New Gamma: %s" % self.GAMMA)

        self.ALPHA = self.point_info_manager.get_real_T(excludes_none=True).mean()
        print("New ALPHA: %s" % self.ALPHA)

        self.log_marginal_likelihood_value_ = -np.min(lml_values)
        print("log_marginal_likelihood: %s" % self.log_marginal_likelihood_value_)

    # TODO this method does not work well...
    def log_marginal_likelihood(self, eval_gradient=False):
        """Returns log-marginal likelihood of theta for training data.(theta is hyper-paramters)

        Parameters
        ----------
        eval_gradient : bool, default: False
            If True, the gradient of the log-marginal likelihood with respect
            to the kernel hyperparameters at position theta is returned
            additionally. If True, theta must not be None.

        Returns
        -------
        log_likelihood : float
            Log-marginal likelihood of theta for training data.

        log_likelihood_gradient : array, shape = (n_kernel_params,), optional
            Gradient of the log-marginal likelihood with respect to the kernel
            hyperparameters at position theta.
            Only returned when eval_gradient is True.
        """
        A, B, mu_tilda, gamma_tilda, r_grid, n_grid = self.calc_tau(return_all=True)

        gamma = self.GAMMA
        gamma_y = self.GAMMA_Y

        A1 = np.copy(self.baseTau0)
        A1[np.diag_indices_from(A1)] = 2 * self.ndim

        # A1 = (A - B) / gamma_y
        A2 = np.diag(n_grid)

        B_inv = np.linalg.inv(B)
        A_inv = np.linalg.inv(A)

        M_lambda = B - B.dot(np.linalg.inv(A).dot(B))

        # log_likelihood = mu_tilda.dot(M_lambda).dot(
        #     mu_tilda[:, np.newaxis]) - np.nan_to_num(np.log(np.linalg.det(M_lambda)))
        log_likelihood = mu_tilda.dot(M_lambda).dot(mu_tilda[:, np.newaxis])

        # log_likelihood = -log_likelihood
        print('log_likelihood: %f' % log_likelihood)

        Lambda = B - B.dot(A_inv.dot(B))
        Lambda_inv = np.linalg.inv(Lambda)

        M_gamma = A2 - A2.dot(2 * gamma * A_inv - (gamma ** 2) * A_inv.dot(A2).dot(A_inv)).dot(A2)

        mean_r_grid = np.nan_to_num(np.array([r.sum() for r in r_grid]) / n_grid)

        log_likelihood_grad_gamma = mean_r_grid.dot(M_gamma.dot(mean_r_grid[:, np.newaxis]))[0]

        tmpMat = A2.dot(A_inv).dot(A1).dot(A_inv).dot(A2)

        log_likelihood_grad_gamma_y = (gamma ** 2) * (
            np.trace(Lambda_inv.dot(tmpMat) + mean_r_grid.dot(tmpMat.dot(mean_r_grid[:, np.newaxis]))[0]))

        log_likelihood_gradient = np.array([log_likelihood_grad_gamma, log_likelihood_grad_gamma_y])
        print('log_likelihood_grad: %s' % log_likelihood_gradient)

        if eval_gradient:
            return log_likelihood, log_likelihood_gradient
        else:
            return log_likelihood

    def _constrained_optimization(self, obj_func, bounds=[(1e-5, 1e5), (1e-5, 1e5)], gradExist=True):
        initial_theta = [self.GAMMA, self.GAMMA_Y]
        if gradExist:
            if self.optimizer == "fmin_l_bfgs_b":
                theta_opt, func_min, convergence_dict = \
                    fmin_l_bfgs_b(obj_func, initial_theta, bounds=bounds)

                if convergence_dict["warnflag"] != 0:
                    warnings.warn("fmin_l_bfgs_b terminated abnormally with the "
                                  " state: %s" % convergence_dict)

            elif callable(self.optimizer):
                theta_opt, func_min = \
                    self.optimizer(obj_func, initial_theta, bounds=bounds)
            else:
                raise ValueError("Unknown optimizer %s." % self.optimizer)
        else:
            theta_opt, func_min, convergence_dict = \
                fmin_l_bfgs_b(obj_func, initial_theta, approx_grad=True, bounds=bounds)

        return theta_opt, func_min


# operator overloading

class time:
    def __init__(self, h=0, m=0):
        self.hours = h
        self.minutes = m

    def display(self):
        print(self.hours, "hour and", self.minutes, "minutes")

    def __add__(self, other):
        m1 = self.hours*60 + self.minutes
        m2 = other.hours*60 + other.minutes
        mtotal = m1+m2

        temp = time()
        temp.hours = mtotal//60
        temp.minutes = mtotal%60
        return temp

first = time(1, 35)
second = time(0, 50)


n = int(input('Введите количество фраз '))
phrases = []
for i in range(n):
    phrases.append(input())

search_phrase = input('Введите фразу для поиска ')

for phrase in phrases:
    if(search_phrase.lower() in phrase.lower()):
        print(phrase)


from BilibiliParser import BilibiliParser
from bs4 import BeautifulSoup

# Create variable lists
v_id = []
v_hide = []
v_title = []
v_link = []
v_desc = []
v_watchnum = []
v_textnum = []
v_uptime = []
v_upername = []
v_uperspace = []

# Create data tag name
titleinfo = ['avid', 'hide']
tagsinfo = ['.watch-num', '.hide', '.time', '']

# definition Beautifulsoup select method
'''
这里直接简单粗暴地使用了Beautifulsoup对象的select方法，
在 Chrome 中对想要定位到想要抓取的位置-右键-copy-copy selector
后面获取数据就会使用到这一方法，所以为了方便调用提前设定
'''

def Selector(selector_path):
    selector = soup.select(selector_path)

    return selector

# define url
'''
可以在B站的搜索页面上查看网页参数的变化，
然后修改BilibiliParser对象的Params属性
'''
url = 'https://search.bilibili.com/all?keyword=python&order=totalrank&duration=4&tids_1=36'

# use for-loop to get data
for page in range(1, 11):  # transform the number of pages you want get
    # parse html
    page_params = {'page': page}  # give params
    parser = BilibiliParser(url=url, page_params=page_params)
    html = parser.getUrl(parser.url, params=parser.page_params)
    soup = parser.htmlParser(html.text)

    # get info-data
    ## id & hide
    for type in titleinfo:
        selector_path = '#server-search-app > div.contain > div.body-contain > div > div.result-wrap.clearfix > ul > li > div > div.headline.clearfix > span.type.{}'
        if type == 'avid':
            v_selectorinfo = Selector(selector_path.format(type))
            for each in v_selectorinfo:
                v_id.append(each.get_text())
        else:
            v_SelectorInfo = Selector(selector_path.format(type))
            for each in v_SelectorInfo:
                v_hide.append(each.get_text())

    ## link & title
    for tag in soup.find_all(name='li', class_='video matrix'):
        v_link.append(tag.a.attrs['href'])
        v_title.append(tag.a.attrs['title'])

    ## desc
    desc_selector_path = '#server-search-app > div.contain > div.body-contain > div > div.result-wrap.clearfix > ul > li > div > div.des.hide'
    for desc in Selector(desc_selector_path):
        v_desc.append(desc.get_text())

    ## watchnum & textnum & uptime & uper-info
    for item in tagsinfo:
        selector_path = '#server-search-app > div.contain > div.body-contain > div > div.result-wrap.clearfix > ul > li > div > div.tags > span.so-icon{}'
        if item == '.watch-num':
            tagSelector = Selector(selector_path.format(item))
            for each in tagSelector:
                v_watchnum.append(each.get_text())
        elif item == '.hide':
            tagSelector = Selector(selector_path.format(item))
            for each in tagSelector:
                v_textnum.append(each.get_text())
        elif item == '.time':
            tagSelector = Selector(selector_path.format(item))
            for each in tagSelector:
                v_uptime.append(each.get_text())
        else:
            selector_path = '#server-search-app > div.contain > div.body-contain > div > div.result-wrap.clearfix > ul > li > div > div.tags > span > a'
            tagSelector = Selector(selector_path)
            for each in tagSelector:
                v_upername.append(each.get_text())
                v_uperspace.append(each['href'])

# Export Data
'''
将所有数据打包成一个字典后再用转成 DataFrame 格式
然后导出数据
'''
## use dict to store data
dict_data = {}
total_data = [
    v_id, v_hide, v_title, v_link, v_desc, v_watchnum, v_textnum, v_uptime,
    v_upername, v_uperspace
]
col_name = [
    'v_id', 'v_hide', 'v_title', 'v_link', 'v_desc', 'v_watchnum', 'v_textnum',
    'v_uptime', 'v_upername', 'v_uperspace'
]
for key, value in zip(col_name, total_data):
    dict_data[key] = value

## tidy & export data
data = parser.dataTidy(dict_data)
data.to_excel(r'testoutput.xlsx', index=False, encoding='utf-8')


import numpy as np
import matplotlib.pyplot as plt
import grid_plot as gp

def riemann():

    # grid info
    xmin = 0.0
    xmax = 1.0

    nzones = 2
    ng = 0

    gr = gp.FVGrid(nzones, xmin=xmin, xmax=xmax)

    #------------------------------------------------------------------------
    # plot a domain without ghostcells
    gr.draw_grid(emphasize_end=True)

    # label a few
    gr.label_center(0, r"$i$", fontsize="medium")
    gr.label_center(1, r"$i+1$", fontsize="medium")
    gr.label_edge(1, r"$i+\myhalf$", fontsize="medium")

    gr.mark_cell_left_state(1, r"$U_{i+\myhalf,R}^{n+\myhalf}$", fontsize="large",
                            color="b")
    gr.mark_cell_right_state(0, r"$U_{i+\myhalf,L}^{n+\myhalf}$", fontsize="large",
                             color="b")

    gr.label_cell_center(0, r"$U_i$")
    gr.label_cell_center(1, r"$U_{i+1}$")

    # flux
    plt.arrow(gr.xl[ng+nzones//2]-0.25*gr.dx, 1.05, 0.5*gr.dx, 0,
                shape='full', head_width=0.075, head_length=0.05,
                lw=1, width=0.03,
                edgecolor="none", facecolor="red",
                length_includes_head=True, zorder=100)

    plt.text(gr.xl[ng+nzones//2], 1.15, r"$F(U_{i+\myhalf}^{n+\myhalf})$", color="red",
               horizontalalignment="center")

    gr.clean_axes(padding=False)
    plt.ylim(-0.25, 1.25)

    plt.subplots_adjust(left=0.05,right=0.95,bottom=0.05,top=0.95)

    f = plt.gcf()
    f.set_size_inches(6.0,2.25)

    plt.tight_layout()

    plt.savefig("riemann_comp.pdf")

if __name__== "__main__":
    riemann()


#!/usr/bin/env python
# -*- coding:utf-8 -*-

import os
import time

import utils.globalvar as g
from utils.common import programEroor
from utils.common import makeDirectory
from utils.common import copyFile

try:
    from watchdog.observers import Observer
    from watchdog.events import FileSystemEventHandler
except ImportError:
    programEroor("[!] Please pip install watchdog")

class ArcHandler(FileSystemEventHandler):
    def on_created(self, event):
        backup_path = g.getBackupPath()
        day_time = time.strftime("%Y%m%d")
        src_path = event.src_path
        dst_path = os.path.join(backup_path, day_time ,src_path[1:])

        try:
            if event.is_directory:
                makeDirectory(dst_path)
            else:
                copyFile(src_path, dst_path)
        except IOError:
            pass

class ArcWarden(object):
    def __init__(self, watch_path):
        self.watch_path = watch_path
        self.event_handler = ArcHandler()
        self.observer = Observer()

    def start(self):
        self.observer.schedule(self.event_handler, path=self.watch_path, recursive=True)
        self.observer.start()

    def stop(self):
        self.observer.stop()

    def join(self):
        self.observer.join()


import urllib.request
import string

pagename = 'ocr'

#http://www.pythonchallenge.com/pc/def/ocr.html
print('{0}'.format("http://www.pythonchallenge.com/pc/def/" + pagename + ".html"))

with urllib.request.urlopen("http://www.pythonchallenge.com/pc/def/" + pagename + ".html") as URL:
	filebody = URL.read()

# After reading the html contents we see the text:
# "find rare characters in the mess below"

# \n\n<!-- mess -->\n\n
# parse the mess out of the filebody
sequenceBegin = '-->\n\n<!--'
sequenceEnd = '-->\n\n'
messBegin = filebody.index(sequenceBegin.encode())+len(sequenceBegin)
messEnd = len(filebody) - len(sequenceEnd)
mess = filebody[messBegin:messEnd]

# parse the mess character by character into a dictionary to record each character and the hit rate
dict = {}
for i in range(len(mess)):
	c = mess[i:i+1]
	if c in dict:
		dict[c] += 1
	else:
		dict[c] = 1

# replace each high occuring character with a '' to reduce the mess to a URL
for key, value in dict.items():
	if value > 100:
		mess = mess.replace(key,''.encode())

#http://www.pythonchallenge.com/pc/def/equality.html
print('{0}'.format("http://www.pythonchallenge.com/pc/def/" + mess.decode() + ".html"))


from django.conf.urls import url, include
from django.contrib import admin
from django.conf.urls.static import static
from django.conf import settings
from django.contrib.auth.views import logout

from nabes_app.views import (IndexView, NewsletterListView, CreateAccountView,
ProfileView, ProfileUpdateView, LeadershipListView, DirectoryListView, StreetListView,
directory_printed)

urlpatterns = [
    url(r'^admin/', admin.site.urls),
    url(r'^', include('django.contrib.auth.urls')),
    url(r'^logout/$', logout, name='logout'),
    url(r'^create_account/$', CreateAccountView.as_view(), name='create_account_view'),
    url(r'accounts/profile/$', ProfileView.as_view(), name='profile_view'),
    url(r'accounts/profile/update/$', ProfileUpdateView.as_view(), name='profile_update_view'),
    url(r'^$', IndexView.as_view(), name='index_view'),
    url(r'^newsletter/$', NewsletterListView.as_view(), name='newsletter_listview'),
    url(r'^leadership/$', LeadershipListView.as_view(), name='leadership_listview'),
    url(r'^directory/$', DirectoryListView.as_view(), name='directory_listview'),
    url(r'^streets/$', StreetListView.as_view(), name='street_listview'),
    url(r'^pdf/$', directory_printed, name='pdf_view'),
] + static(settings.MEDIA_URL, document_root=settings.MEDIA_ROOT)