# inquadpack.py ---

# Copyright (C) 2017 Omid Khanmohamadi <OmidLink@Gmail.com>

# Author: Omid Khanmohamadi <OmidLink@Gmail.com>

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.


import numpy as np
import autograd as ag
import autograd.numpy as anp
import scipy.integrate as spi
from scipy.optimize import fsolve
from functools import partial

norm2 = partial(np.linalg.norm, ord=2)

max_nsubdivs = 16
nsubdivs = 0
hgtfcn_steepness_const = np.sqrt(20.)

def chebpts(n, interval=np.array([-1., 1.]), kind=2):
    m = n-1

    theta = np.pi*np.arange(-m, m+1, 2)
    if kind == 1:
        # roots of T_n (no bdry pts) aka 1st-kind Chebyshev pts or
        # Chebyshev--Gauss pts
        x = np.sin(theta/(2*m+2))
    else:
        # extrema of T_{n-1} (include bdry pts) aka 2nd-kind Chebyshev pts
        # or Chebyshev--Gauss--Lobatto pts
        x = np.sin(theta/(2*m))

    return scalenodes(x, interval)

def scalenodes(x, interval):
    if np.allclose(interval, np.array([-1., 1.])):
        return x
    return interval[1]*(x + 1)/2 + interval[0]*(1 - x)/2


def sgn_conds_satisfied(lvlsets, sgns, xc):
    # zip returns an iterator (in Python 3; returns a list in
    # Python 2, which is less efficient. Use itertools.izip in 2.)
    # of tuples, where the i-th tuple contains the i-th element
    # from each of the argument sequences or iterables.
    return np.all(np.array([s*psi(xc) >= 0. for psi, s in zip(lvlsets, sgns)]))


def volume(hyprect):
    return np.prod(np.diff(hyprect, axis=0))


def sgnUL(m, s, is_surf, sigma):
    return sigma*m if (m == sigma*s or is_surf) else 0


sgnL = partial(sgnUL, sigma=-1)
sgnU = partial(sgnUL, sigma=1)


# def psiUL(x, psi, idx, xUL_idx):
#     return psi(np.insert(x, idx, xUL_idx))


def sides_except(hyprect, idx):
    """Sides of hyprect except the side corresponding to index idx."""
    return np.c_[hyprect[:, :idx], hyprect[:, idx+1:]]


def halves(hyprect):
    """Halve hyprect along its largest extent.
    """
    # In case of multiple occurrences of the maximum values, the
    # indices corresponding to the /first/ occurrence are returned by
    # np.argmax. This works for us, as we only need to halve the
    # hyprect, so it suffices to pick only one of the identical
    # maximum extents (including the first).
    idx = np.argmax(halfwidths(hyprect))

    # items before column idx.
    hyprect_before_idx = hyprect[:, :idx]

    # items after column idx.
    hyprect_after_idx = hyprect[:, idx+1:]

    # midpoint of the edge in column idx.
    hyprect_idx_mid = np.mean(hyprect[:, idx])

    # hyprect1 and hyprect2 differ only in column idx.
    hyprect1_at_idx = np.array([hyprect[0, idx],
                                hyprect_idx_mid])

    hyprect2_at_idx = np.array([hyprect_idx_mid,
                                hyprect[1, idx]])

    hyprect1 = np.c_[hyprect_before_idx, hyprect1_at_idx, hyprect_after_idx]

    hyprect2 = np.c_[hyprect_before_idx, hyprect2_at_idx, hyprect_after_idx]

    return hyprect1, hyprect2


def halfwidths(hyprect):
    """Halfwidth of hyprectangle hyprect.

    hyprect is a 2xd np.array whose 1st row contains the lower bounds of the
    hyprectangle (xL_i) and whose 2nd row contains the upper bounds
    (xU_i), i.e.,
    [xL;
     xU].
    """
    return np.diff(hyprect, axis=0)/2


def centerpt(hyprect):
    # axis=0: find the mean of each column. That is do
    #
    # (sum_{i0} (hyprect)_{i0,i1})/(size(hyprect, axis=0))
    return np.mean(hyprect, axis=0)


def bound_fx_minus_fxc(f, centerpt_hyprect, l2_halfwidths_hyprect):
    """Linear approximation to sup |f(x)-f(xc)|, for all x in hyprect.
    """
    xc = centerpt_hyprect
    Df = ag.jacobian(f)
    Df_xc = Df(xc)

    # If the derivative Df_xc is a vector (ndim == 1, aka gradient^T),
    # take its 2-norm (which results in a scalar) times the 2-norm of
    # the halfwidths_hyprect vector, which is also a scalar. The
    # returned result is a scalar.
    if np.ndim(Df_xc) == 1:
        return norm2(Df(xc)) * l2_halfwidths_hyprect

    # If the derivative Df_xc is a matrix (ndim == 2, aka Jacobian),
    # return the product of the vector containing the 2-norm of each
    # row (axis=1) of this matrix and the 2-norm of the
    # halfwidths_hyprect vector. This vector-scalar product is done by
    # numpy broadcasting the scalar to a vector behind the scenes and
    # then doing an elementwise product of two vectors. The returned
    # result is a vector.
    else:
        return norm2(Df(xc), axis=1) * l2_halfwidths_hyprect


def roots_chebpoly(x, idx, lvlsets, interval):
    nchebpts = 5
    chebpdeg = nchebpts-1
    xcheb = chebpts(nchebpts, interval)
    # TODO: Preallocate np.array or change to list (for possible
    # performance gain).
    R = np.array([])
    for psi in lvlsets:
        # np.insert(x, k, y) does Saye's "x+ye_k"; it makes a copy of
        # the array x and inserts y before the given index k of x. The
        # returned array has y at its index k. In our special use
        # case, np.insert(x, k, y) is equivalent to
        # np_r[x[:k], y, x[k:]].
        psi_xcheb = np.array([psi(anp.insert(x, idx, xch)) for xch in xcheb])
        chebp_coeffs = np.polynomial.chebyshev.chebfit(
            xcheb, psi_xcheb, chebpdeg)
        chebp = np.polynomial.chebyshev.Chebyshev(chebp_coeffs)
        chebp_rts = chebp.roots()
        chebp_rts_real_mask = np.isreal(np.real_if_close(chebp_rts))
        chebp_rts_real = np.unique(np.real(chebp_rts[chebp_rts_real_mask]))
        interval_mask = np.logical_and(chebp_rts_real > interval[0],
                                       chebp_rts_real < interval[1])
        chebp_rts_real_interval = chebp_rts_real[interval_mask]
        R = np.r_[R, chebp_rts_real_interval]
    return R


def roots_fsolve(x, idx, lvlsets, interval):
    interval_midpt = np.mean(interval)

    # TODO: Preallocate np.array or change to list (for possible
    # performance gain).
    R = np.array([])
    for psi in lvlsets:
        # np.insert(x, k, y) does Saye's "x+ye_k"; it makes a copy of
        # the array x and inserts y before the given index k of x,
        # returning [ x[:k], y, x[k:] ].
        r_full = fsolve(lambda y: psi(anp.insert(x, idx, y)),
                        x0=interval_midpt, full_output=True)
        # r_full[2] is an integer flag, set to 1 if a solution was
        # found, otherwise refer to `mesg` for more information.
        iter_converged = (r_full[2] == 1)
        if iter_converged:
            R = np.r_[R, r_full[0]]
    return R


# Algorithm 1
def eval_integrand(x, idx, integrand, lvlsets, sgns, interval):
    """
    x is in R^{d-1}
    idx is k
    integrand is f: R^{d} -> R
    lvlsets is a list with psi_i: R^{d} -> R
    sgns is a list with s_i: +1, -1, or 0
    interval is the pair (x_1, x_2)
    """
    # Convert interval to a contiguous /flattened/ (i.e., 1D) array.
    # This results in a view, unless a copy is needed.
    interval = np.ravel(interval)
    rts = roots_chebpoly(x, idx, lvlsets, interval)

    # np.unique returns the SORTED, unique elements of its input array.
    R = np.unique(np.r_[interval, rts])

    I = 0.
    for j in np.arange(len(R)-1):
        xc = anp.insert(x, idx, (R[j+1]+R[j])/2)

        if sgn_conds_satisfied(lvlsets, sgns, xc):
            # quad (fixed_quad) returns in index 1 an estimate of the
            # absolute error in the result (the statically returned
            # value of None). The trailing [0] is to catch only the
            # index 0 of the result, which is the value of the
            # integral.
            I += spi.quad(lambda y: integrand(anp.insert(x, idx, y)),
                          R[j], R[j+1])[0]
    return I


# Algorithm 2
def eval_surf_integrand(x, idx, integrand, lvlset, interval):
    interval = np.ravel(interval)
    r = roots_chebpoly(x, idx, lvlset, interval)

    if len(r) == 1:
        evalpt = anp.insert(x, idx, r)
        phi = lvlset[0]
        Dphi_evalpt = ag.jacobian(phi)(evalpt)
        return integrand(evalpt)*norm2(Dphi_evalpt)/np.abs(Dphi_evalpt[idx])
    else:
        return 0.


# Algorithm 3
def inquad(integrand, lvlsets, sgns, hyprect, is_surf):
    """
    hyprect is U
    is_surf is S
    """

    global nsubdivs

    # Base case of recursion.
    if np.size(hyprect, axis=1) == 1:
        return eval_integrand(np.array([]), 0,
                              integrand, lvlsets, sgns, hyprect)

    xc = centerpt(hyprect)
    l2_halfwidths_hyprect = norm2(halfwidths(hyprect))

    # Pruning.
    #
    # Iterating over a list and removing elements skips every other
    # element if done in forward direction. For example,
    #
    # abc = ['a', 'b', 'c']
    # for l in abc:
    #     abc.remove(l)
    # print abc
    #
    # gives ['b']. When you remove 'a' from abc, the remaining
    # elements slide down the list. The list is now ['b', 'c'], 'b' is
    # at index 0 and 'c' is at index 1. Since the next iteration is
    # going to look at index 1 (which is the 'c' element), 'b' gets
    # skipped entirely. One solution, which we use below, is to
    # iterate over the list in reverse. Alternatively, if we were
    # dealing with a single list, we could consider using in-place
    # (slice) list comprehension [https://stackoverflow.com/a/1208792].
    for ilvlset in reversed(range(len(lvlsets))):
        psi = lvlsets[ilvlset]
        s = sgns[ilvlset]

        bnd_psi_x_minus_psi_xc = bound_fx_minus_fxc(psi,
                                                    xc,
                                                    l2_halfwidths_hyprect)

        psi_xc = psi(xc)
        if np.abs(psi_xc) >= bnd_psi_x_minus_psi_xc:
            if s*psi_xc >= 0.:
                del lvlsets[ilvlset]
                del sgns[ilvlset]
            else:
                # The domain of integration is empty.
                return 0.

    # Tensor-product Gaussian quadrature.
    #
    # Use when the previous pruning steps leave an empty list of
    # level-set functions, which occurs when the bound-evaluation step
    # in pruning determines that the domain of integration is the
    # entire hyprectangle.
    if len(lvlsets) == 0:
        # nquad expects each element of ranges (e.g., ranges[0]) to
        # be/return a sequence of two numbers. ranges has in element r
        # the bounds of integration w.r.t to x_r, whereas hyprect has
        # that bound in its column c, hence the need for transpose,
        # ranges=hyprect.T. For example, to get
        #
        # ranges = [[1, 2], [3, 4], [5, 6]]
        #
        # we must transpose
        #
        # hyprect = np.array([[1, 3, 5],
        #                     [2, 4, 6]])
        return spi.nquad(integrand, ranges=hyprect.T)[0]

    psi0 = lvlsets[0]
    Dpsi0 = ag.jacobian(psi0)
    Dpsi0_xc = Dpsi0(xc)

    # Propose a candidate height fcn direction (its suitability is
    # checked below).
    idx = np.argmax(np.abs(Dpsi0_xc))

    lvlsets_new = []
    sgns_new = []

    for psi, s in zip(lvlsets, sgns):
        Dpsi = ag.jacobian(psi)
        Dpsi_xc = Dpsi(xc)
        bnd_Dpsi_x_minus_Dpsi_xc = bound_fx_minus_fxc(
            Dpsi, xc, l2_halfwidths_hyprect)
        Dpsi_idx_xc = Dpsi_xc[idx]
        bnd_Dpsi_idx_x_minus_Dpsi_idx_xc = bnd_Dpsi_x_minus_Dpsi_xc[idx]

        # Cond. (a) (no sign change in the derivative, with a safety
        # margin of bnd_Dpsi_x_minus_Dpsi_xc to be strictly away from
        # zero) guarantees the existence of the height function as
        # well as the monotonicity of the level-set function psi in
        # the height function direction idx. Cond (b) (small bound on
        # derivative Dpsi of psi relative to its partial derivative
        # Dpsi_idx in the hyprect) concerns the curved-surface
        # "Jacobian factor" that arises explicitly in the surface
        # integral and implicitly in the volume integrals (as the
        # derivative of the integrands). It is related to how "steep"
        # the height function is and it needs to be small (<
        # hgtfcn_steepness_const, which is an empirical const) for the
        # approximation of these integrals using Gaussian quadrature
        # to be reasonably accurate.
        if (np.abs(Dpsi_idx_xc) > bnd_Dpsi_idx_x_minus_Dpsi_idx_xc  # (a)
            and
            norm2((Dpsi_xc     + bnd_Dpsi_x_minus_Dpsi_xc))  # (b)
            /     (Dpsi_idx_xc - bnd_Dpsi_idx_x_minus_Dpsi_idx_xc)
            < hgtfcn_steepness_const):

            xL_idx = hyprect[0, idx]
            xU_idx = hyprect[1, idx]

            # When Python doesn't find a name in the local namespace,
            # it looks it up in the enclosing namespace (and global
            # and built-in namespaces afterwards), and it uses late
            # binding, which means it binds the name to the object at
            # the time the function is /called/, rather than at the
            # time the function is /defined/. One approach to forcing
            # early binding semantics is the "default argument hack",
            # which binds the variable as a default argument at
            # function definition time. (N.B: this changes the
            # function signature.) Other approaches are using a
            # "function maker" (closure) or using partial, from
            # functools. See, e.g.,
            # https://stackoverflow.com/q/5190146 and
            # https://stackoverflow.com/q/23400785, for more info.
            #
            # psiL = partial(psiUL, psi=psi, idx=idx, xUL_idx=xL_idx)
            # psiU = partial(psiUL, psi=psi, idx=idx, xUL_idx=xU_idx)
            def psiL(x, idx=idx, xL_idx=xL_idx):
                return psi(np.insert(x, idx, xL_idx))

            def psiU(x, idx=idx, xU_idx=xU_idx):
                return psi(np.insert(x, idx, xU_idx))

            sgn_Dpsi_idx_xc = np.sign(Dpsi_idx_xc)
            sL = sgnL(sgn_Dpsi_idx_xc, s, is_surf)
            sU = sgnU(sgn_Dpsi_idx_xc, s, is_surf)

            lvlsets_new += [psiL, psiU]
            sgns_new += [sL, sU]

        # Height function direction is not suitable for this psi.
        else:
            # max_nsubdivs is only referenced (but not assigned to
            # (bound to a new object) and, since it is not found in
            # the local (i.e., this function's) namespace, it is
            # looked up in the global (module) namespace. max_nsubdivs
            # is in /scope/ here, meaning that it can be seen here
            # (think telescope), although it is not in the local
            # namespace. (N.B.: A scope refers to a region of a
            # program from where a namespace can be accessed without a
            # prefix.) nsubdivs, on the other hand, is assigned to
            # (bound to nsubdivs + 1). This makes it a local variable.
            # But it is also referenced (in nsubdivs < max_nsubdivs),
            # and the referencing happens before the assignment. This
            # causes a "UnboundLocalError: local variable 'nsubdivs'
            # referenced before assignment", unless we use 'global
            # nsubdivs' before either referencing or assignment
            # occurs, as we have at the beginning of this function.
            if nsubdivs < max_nsubdivs:
                # Split hyprect into two halves along its largest extent
                hyprect1, hyprect2 = halves(hyprect)
                nsubdivs += 1

                # Upon recursion to compute inquad_hyprect1, the
                # parameter name lvlsets is created in the local
                # namespace of the inner inquad() and is bound to the
                # SAME OBJECT the argument name lvlsets is bound to in
                # the local namespace of the outer inquad(). Since the
                # object lvlsets is bound to is mutable (it is a
                # list), the pruning step
                #
                # del lvlsets[ilvlset]
                # del sgns[ilvlset]
                #
                # (if run) will in fact modify that object, and NOT a
                # copy of it. As a consequence, the recursion to
                # compute inquad_hyprect2 will use the MODIFIED
                # object, and NOT the original, leading to a nasty
                # bug! lvlsets.copy() (shallow copy) creates a new
                # container in which it copies the references to the
                # objects in lvlsets. This solves the modification
                # problem above, as del will delete the references in
                # the new list (container) but the references in the
                # original list are untouched. Then, the same happens
                # for computing inquad_hyprect2. Similarly for sgns.
                # http://python.net/~goodger/projects/pycon/2007/idiomatic/handout.html#other-languages-have-variables
                # https://stackoverflow.com/questions/2612802/how-to-clone-or-copy-a-list-in-python
                inquad_hyprect1 = inquad(integrand,
                                         lvlsets.copy(), sgns.copy(),
                                         hyprect1, is_surf)

                inquad_hyprect2 = inquad(integrand,
                                         lvlsets.copy(), sgns.copy(),
                                         hyprect2, is_surf)

                return inquad_hyprect1 + inquad_hyprect2
            else:
                # We have exceeded the maximum number of subdivisions
                # allowed. Switch to a low order method.
                if not is_surf and sgn_conds_satisfied(lvlsets, sgns, xc):
                    return volume(hyprect)*integrand(xc)
                else:
                    return 0

    # def integrand_new(x):
    #
    # for lvlsets embedded in R^{>3} (but not R^2) leads to
    #
    # return _quadpack._qagse(func,a,b,args,full_output,epsabs,epsrel,limit)
    # TypeError: integrand_new() takes 1 positional argument but 2 were given
    #
    # since nquad (_qagse) expects func's signature to be ``func(x0,
    # x1, ..., xn, t0, t1, ..., tm)``. Integration over ``x0`` is the
    # innermost integral, and ``xn`` is the outermost.
    #
    # def integrand_new(x, idx=idx, integrand=integrand,
    # lvlsets=lvlsets, sgns=sgns):
    #
    # leads to
    #
    # File "inquadpack.py", line 534, in ellipsoid
    #   return x[0]**2 + 4.*x[1]**2 + 9.*x[2]**2 - 1.
    # IndexError: index 2 is out of bounds for axis 0 with size 2
    #
    # which is much harder to catch. Since nquad expects func to have
    # two positional arguments (one for each integration dimension;
    # recall that we have recursed once---resulting in one dimension
    # reduction---when we first run integrand_new) it is binding the
    # first argument of func to the value of x0 and its second
    # argument (which is our idx) to that of x1. A disaster! The index
    # idx, which is supposed to be an integer, is bound to a float, and
    # the only clue is the warning
    #
    # inquadpack.py:465: VisibleDeprecationWarning: using a
    # non-integer number instead of an integer will result in an error
    # in the future
    # interval = hyprect[:, idx]
    #
    # Using `early-binding hack' (via default argument assignment)
    # considered harmful! It changes the signature.
    #
    # One can call * a gather parameter (when used in function
    # definition) or a scatter operator (when used at function
    # invocation).
    #
    # https://stackoverflow.com/questions/5239856/foggy-on-asterisk-in-python
    # https://stackoverflow.com/questions/2322355/proper-name-for-python-operator
    def integrand_new(*x):
        interval = hyprect[:, idx]
        if is_surf:
            return eval_surf_integrand(x, idx, integrand,
                                       lvlsets, interval)
        else:
            return eval_integrand(x, idx, integrand,
                                  lvlsets, sgns, interval)

    hyprect_new = sides_except(hyprect, idx)
    return inquad(integrand_new, lvlsets_new.copy(), sgns_new.copy(),
                  hyprect_new, is_surf=False)


if __name__ == '__main__':
    from tests.scratch_work import demo_calls

    demo_calls()
