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

"""
Python wrapper around the C extension for the theoretical 3-D
real-space correlation function, :math:`\\xi(r)`. Corresponding
C routines are in ``theory/xi/``, python interface is
:py:mod:`Corrfunc.theory.xi`.
"""

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)

__author__ = ('Manodeep Sinha')
__all__ = ('xi',)

def xi(boxsize, nthreads, binfile, X, Y, Z,
       weights=None, weight_type=None, verbose=False, output_ravg=False,
       xbin_refine_factor=2, ybin_refine_factor=2,
       zbin_refine_factor=1, max_cells_per_dim=100,
       copy_particles=True, enable_min_sep_opt=True,
       c_api_timer=False, isa='fastest', bin_type='custom', attrs_pair_weights=None):
    """
    Function to compute the projected correlation function in a
    periodic cosmological box. Pairs which are separated by less
    than the ``r`` bins (specified in ``binfile``) in 3-D real space.

    If ``weights`` are provided, the resulting correlation function
    is weighted.  The weighting scheme depends on ``weight_type``.

    .. note:: Pairs are double-counted. And if ``rmin`` is set to
        0.0, then all the self-pairs (i'th particle with itself) are
        added to the first bin => minimum number of pairs in the first bin
        is the total number of particles.

    Parameters
    ----------
    boxsize : double
        A double-precision value for the boxsize of the simulation
        in same units as the particle positions and the ``r`` bins.

    nthreads : integer
        Number of threads to use.

    binfile : string or an list/array of floats
        For string input: filename specifying the ``r`` bins for
        ``xi``. The file should contain white-space separated values
        of (rmin, rmax)  for each ``r`` wanted. The bins need to be
        contiguous and sorted in increasing order (smallest bins come first).

        For array-like input: A sequence of ``r`` values that provides the
        bin-edges. For example,
        ``np.logspace(np.log10(0.1), np.log10(10.0), 15)`` is a valid
        input specifying **14** (logarithmic) bins between 0.1 and 10.0. This
        array does not need to be sorted.

    X/Y/Z : arraytype, real (float/double)
        Particle positions in the 3 axes. Must be within [0, boxsize]
        and specified in the same units as ``rp_bins`` and boxsize. All
        3 arrays must be of the same floating-point type.

        Calculations will be done in the same precision as these arrays,
        i.e., calculations will be in floating point if XYZ are single
        precision arrays (C float type); or in double-precision if XYZ
        are double precision arrays (C double type).

    weights : array_like, real (float/double), optional
        A scalar, or an array of weights of shape (n_weights, n_positions) or
        (n_positions,). ``weight_type`` specifies how these weights are used;
        results are returned in the ``weightavg`` field.

    verbose : boolean (default false)
        Boolean flag to control output of informational messages

    output_ravg : boolean (default false)
        Boolean flag to output the average ``r`` for each bin. Code will
        run slower if you set this flag.

        Note: If you are calculating in single-precision, ``rpavg`` will
        suffer from numerical loss of precision and can not be trusted. If
        you need accurate ``rpavg`` values, then pass in double precision
        arrays for the particle positions.

    (xyz)bin_refine_factor : integer, default is (2,2,1); typically within [1-3]
        Controls the refinement on the cell sizes. Can have up to a 20% impact
        on runtime.

    max_cells_per_dim : integer, default is 100, typical values in [50-300]
        Controls the maximum number of cells per dimension. Total number of
        cells can be up to (max_cells_per_dim)^3. Only increase if ``rmax`` is
        too small relative to the boxsize (and increasing helps the runtime).

    copy_particles : boolean (default True)
        Boolean flag to make a copy of the particle positions
        If set to False, the particles will be re-ordered in-place

        .. versionadded:: 2.3.0

    enable_min_sep_opt : boolean (default true)
        Boolean flag to allow optimizations based on min. separation between
        pairs of cells. Here to allow for comparison studies.

        .. versionadded:: 2.3.0

    c_api_timer : boolean (default false)
        Boolean flag to measure actual time spent in the C libraries. Here
        to allow for benchmarking and scaling studies.

    isa : string (default ``fastest``)
        Controls the runtime dispatch for the instruction set to use. Options
        are: [``fastest``, ``avx512f``, ``avx``, ``sse42``, ``fallback``]

        Setting isa to ``fastest`` will pick the fastest available instruction
        set on the current computer. However, if you set ``isa`` to, say,
        ``avx`` and ``avx`` is not available on the computer, then the code
        will revert to using ``fallback`` (even though ``sse42`` might be
        available).  Unless you are benchmarking the different instruction
        sets, you should always leave ``isa`` to the default value. And if
        you *are* benchmarking, then the string supplied here gets translated
        into an ``enum`` for the instruction set defined in ``utils/defs.h``.

    weight_type : string, optional, Default: None.
        The type of weighting to apply. One of ["pair_product", None].

    bin_type : string, case-insensitive (default ``custom``)
        Set to ``lin`` for speed-up in case of linearly-spaced bins.
        In this case, the bin number for a pair separated by ``r`` is given by
        ``(r - binfile[0])/(binfile[-1] - binfile[0])*(len(binfile) - 1)``,
        i.e. only the first and last bins of input ``binfile`` are considered.
        Then setting ``output_ravg`` is virtually costless.
        For non-linear binning, set to ``custom``.
        In the vast majority of cases, bin_type='linear' will yield identical
        results to custom linear binning but with higher performance.
        In a few rare cases where a pair falls on a bin boundary,
        'linear' and custom linear may disagree on which bin the pair falls into
        due to finite floating point precision.
        ``auto`` will choose linear binning if input ``binfile`` is within
        ``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
        of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.

    pair_weights : array-like, optional. Default: None.
        Array of pair weights.

    sep_pair_weights : array-like, optional. Default: None.
        Array of separations corresponding to ``pair_weights``.

    attrs_pair_weights : tuple. Default: None.
        Attributes for pair weights; in case ``weight_type`` is "inverse_bitwise",
        the tuple of (offset to be added to the bitwise counts,
        default weight value if denominator is zero).

    Returns
    -------
    results : Numpy structured array
        A numpy structured array containing [rmin, rmax, ravg, xi, npairs,
        weightavg] for each radial specified in the ``binfile``. If
        ``output_ravg`` is not set then ``ravg`` will be set to 0.0 for all
        bins; similarly for ``weightavg``. ``xi`` contains the correlation
        function while ``npairs`` contains the number of pairs in that bin.
        If using weights, ``xi`` will be weighted while ``npairs`` will not be.

    api_time : float, optional
        Only returned if ``c_api_timer`` is set.  ``api_time`` measures only
        the time spent within the C library and ignores all python overhead.

    Example
    -------
    >>> from __future__ import print_function
    >>> import numpy as np
    >>> from os.path import dirname, abspath, join as pjoin
    >>> import Corrfunc
    >>> from Corrfunc.theory.xi import xi
    >>> binfile = pjoin(dirname(abspath(Corrfunc.__file__)),
    ...                 "../theory/tests/", "bins")
    >>> N = 100000
    >>> boxsize = 420.0
    >>> nthreads = 4
    >>> seed = 42
    >>> np.random.seed(seed)
    >>> X = np.random.uniform(0, boxsize, N)
    >>> Y = np.random.uniform(0, boxsize, N)
    >>> Z = np.random.uniform(0, boxsize, N)
    >>> weights = np.ones_like(X)
    >>> results = xi(boxsize, nthreads, binfile, X, Y, Z, weights=weights, weight_type='pair_product', output_ravg=True)
    >>> for r in results: print("{0:10.6f} {1:10.6f} {2:10.6f} {3:10.6f} {4:10d} {5:10.6f}"
    ...                         .format(r['rmin'], r['rmax'],
    ...                         r['ravg'], r['xi'], r['npairs'], r['weightavg']))
    ...                   # doctest: +NORMALIZE_WHITESPACE
      0.167536   0.238755   0.226592  -0.205733          4   1.000000
      0.238755   0.340251   0.289277  -0.176729         12   1.000000
      0.340251   0.484892   0.426819  -0.051829         40   1.000000
      0.484892   0.691021   0.596187  -0.131853        106   1.000000
      0.691021   0.984777   0.850100  -0.049207        336   1.000000
      0.984777   1.403410   1.225112   0.028543       1052   1.000000
      1.403410   2.000000   1.737153   0.011403       2994   1.000000
      2.000000   2.850200   2.474588   0.005405       8614   1.000000
      2.850200   4.061840   3.532018  -0.014098      24448   1.000000
      4.061840   5.788530   5.022241  -0.010784      70996   1.000000
      5.788530   8.249250   7.160648  -0.001588     207392   1.000000
      8.249250  11.756000  10.207213  -0.000323     601002   1.000000
     11.756000  16.753600  14.541171   0.000007    1740084   1.000000
     16.753600  23.875500  20.728773  -0.001595    5028058   1.000000

    """

    try:
        from Corrfunc._countpairs import countpairs_xi as xi_extn
    except ImportError:
        msg = "Could not import the C extension for the projected "\
              "correlation function."
        raise ImportError(msg)

    import numpy as np
    from future.utils import bytes_to_native_str
    from Corrfunc.utils import translate_isa_string_to_enum, translate_bin_type_string_to_enum,\
        get_edges, convert_to_native_endian,\
        sys_pipes, process_weights

    weights, _ = process_weights(weights, None, X, None, weight_type, autocorr=True)

    # Ensure all input arrays are native endian
    X, Y, Z = [convert_to_native_endian(arr, warn=False)
                        for arr in [X, Y, Z]]

    if weights is not None:
        weights = [convert_to_native_endian(arr, warn=False) for arr in weights]

    # Passing None parameters breaks the parsing code, so avoid this
    kwargs = {}
    for k in ['weights', 'weight_type', 'attrs_pair_weights']:
        v = locals()[k]
        if v is not None:
            kwargs[k] = v

    integer_isa = translate_isa_string_to_enum(isa)
    integer_bin_type = translate_bin_type_string_to_enum(bin_type)
    rbinfile = get_edges(binfile)
    with sys_pipes():
      extn_results = xi_extn(boxsize, nthreads, rbinfile,
                             X, Y, Z,
                             verbose=verbose,
                             output_ravg=output_ravg,
                             xbin_refine_factor=xbin_refine_factor,
                             ybin_refine_factor=ybin_refine_factor,
                             zbin_refine_factor=zbin_refine_factor,
                             max_cells_per_dim=max_cells_per_dim,
                             copy_particles=copy_particles,
                             enable_min_sep_opt=enable_min_sep_opt,
                             c_api_timer=c_api_timer,
                             isa=integer_isa,
                             bin_type=integer_bin_type, **kwargs)
    if extn_results is None:
        msg = "RuntimeError occurred"
        raise RuntimeError(msg)
    else:
        extn_results, api_time = extn_results

    results_dtype = np.dtype([(bytes_to_native_str(b'rmin'), np.float64),
                              (bytes_to_native_str(b'rmax'), np.float64),
                              (bytes_to_native_str(b'ravg'), np.float64),
                              (bytes_to_native_str(b'xi'), np.float64),
                              (bytes_to_native_str(b'npairs'), np.uint64),
                              (bytes_to_native_str(b'weightavg'), np.float64)])
    results = np.array(extn_results, dtype=results_dtype)

    if not c_api_timer:
        return results
    else:
        return results, api_time

if __name__ == '__main__':
    import doctest
    doctest.testmod()


# Copyright (c) 2015, Andrew Delong and Babak Alipanahi All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation and/or
# other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors
# may be used to endorse or promote products derived from this software without
# specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Author's note:
#     This file was distributed as part of the Nature Biotechnology
#     supplementary software release for DeepBind. Users of DeepBind
#     are encouraged to instead use the latest source code and binaries
#     for scoring sequences at
#        http://tools.genes.toronto.edu/deepbind/
#
import sys, os
import numpy as np
import argparse
from smat import *

#os.environ["PYTHONUNBUFFERED"] = "1"  # Disable output buffering
sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)

parser = argparse.ArgumentParser(description="Train a convolutional neural net on MNIST and print out the error rates.")
parser.add_argument("--device",type=int,default=None,help="Device # to use (which GPU). Default is 0.")
parser.add_argument("--show_filters",action="store_true",default=False,help="Plot the filters after each update, showing in a popup window.")
parser.add_argument("--f64",action="store_true",default=False,help="Use float64, if supported by the GPU. Default is float32.")
args = parser.parse_args()

if args.device is not None:
    set_backend_options(device=args.device)

dt = float64 if args.f64 else float32
set_default_dtype(dt)

print "Using device", get_backend_info().device
print "Checking cuDNN ...",
try:    cudnn_dll()
except: quit("Failed to load cuDNN shared library. Quitting.");
print "OK."

##############################################################################
#  Functions for loading DATA
##############################################################################

# Load a data as of pairs (X,Y) where:
#    X is a (batchsize x inputsize) matrix of inputs
#    Y is a (batchsize x outputsize) matrix of corresponding targets
#  MNIST has 60000 training examples total.
with np.load("data/mnist/mnist_train.npz") as mnist_file:
    inputs_train  = asarray(mnist_file['X'], dtype=dt)/255   # Load 60000 x 784 matrix of training inputs, scaled to range [0,1]
    targets_train = asarray(mnist_file['Y'], dtype=dt)       # Load 60000 x 10 matrix of training targets

with np.load("data/mnist/mnist_test.npz") as mnist_file:
    inputs_test  = asarray(mnist_file['X'], dtype=dt)/255    # Load 10000 x 784 matrix of testing inputs, scaled to range [0,1]
    targets_test = asarray(mnist_file['Y'], dtype=dt)        # Load 10000 x 10 matrix of testing targets

# Generate minibatches out of the full dataset.
trainsize = len(inputs_train)
testsize  = len(inputs_test)
batchsize = 100            ; assert trainsize % batchsize == 0 # Make sure we use all training examples
batches_train = [(inputs_train[i:i+batchsize], targets_train[i:i+batchsize]) for i in range(0, trainsize, batchsize)]
batches_test  = [(inputs_test [i:i+batchsize], targets_test [i:i+batchsize]) for i in range(0, testsize,  batchsize)]

##############################################################################
#  CONVNET CONFIGURATION
##############################################################################

# Configuration of neural net layers (number of filters, hidden units, etc)
input_w = 28          # Width  of MNIST image
input_h = 28          # Height of MNIST image
input_c = 1           # Number of input channels for MNIST

# Layer 1 is convolution, so call it C1
C1_filter_c = 32  # Number of filters in C1
C1_filter_w = 5   # Width  of filters in C1
C1_filter_h = 5   # Height of filters in C1
C1_stride   = 1   # Stride of C1
C1_w = (input_w-C1_filter_w)//C1_stride + 1  # Width  of C1 output featuremap
C1_h = (input_h-C1_filter_h)//C1_stride + 1  # Height of C1 output featuremap

# Layer 2 is pooling, so call it P1
P1_mode = "max"   # Pooling type ("max" or "avg")
P1_window_w = 3   # Width  of pooling windows in P1
P1_window_h = 3   # Height of pooling windows in P1
P1_stride   = 2   # Stride of P1
P1_w = (C1_w-P1_window_w)//P1_stride + 1  # Width  of P1 output featuremap
P1_h = (C1_h-P1_window_h)//P1_stride + 1  # Height of P1 output featuremap

# Layer 3 is fully connected, so call it F1
F1_size = 1000         # Number of neurons in F1
F1_dropout_rate = 0.5  # Dropout rate for F1

# Layer 4 is fully connected softmax, so call it F2
F2_size = 10      # 10 classes representing digits 0..9

##############################################################################
#  CONVNET PARAMETER ALLOCATION and INITIALIZATION
##############################################################################

# Count how many parameters there are in total for each later,
# so that we can allocate one big vector P for all parameters
num_params = { 'C1_weights' : C1_filter_c*C1_filter_w*C1_filter_h * (input_c),
               'C1_bias'    : C1_filter_c,
               'F1_weights' : F1_size * (C1_filter_c*P1_w*P1_h),
               'F1_bias'    : F1_size,
               'F2_weights' : F2_size * (F1_size),
               'F2_bias'    : F2_size }
num_params_total = sum(num_params.values())
P = zeros(num_params_total, dt)
P_grad = zeros_like(P)

# Slice the param vector P into parameters for each layer.
_ = 0  # temp counter
C1_weights = P[_:_+num_params['C1_weights']].reshape((C1_filter_c,  input_c*C1_filter_h*C1_filter_w)); _ += C1_weights.size;
C1_bias    = P[_:_+num_params['C1_bias'   ]];                                                          _ += C1_bias.size;
F1_weights = P[_:_+num_params['F1_weights']].reshape((C1_filter_c*P1_w*P1_h, F1_size));                _ += F1_weights.size;
F1_bias    = P[_:_+num_params['F1_bias'   ]].reshape((1,         F1_size));                            _ += F1_bias.size;
F2_weights = P[_:_+num_params['F2_weights']].reshape((F1_size,   F2_size));                            _ += F2_weights.size;
F2_bias    = P[_:_+num_params['F2_bias'   ]].reshape((1,         F2_size));                            _ += F2_bias.size;
assert _ == num_params_total

# Slice the gradient vector P_grad into parameters for each layer.
_ = 0
C1_weights_grad = P_grad[_:_+num_params['C1_weights']].reshape((C1_filter_c,  input_c*C1_filter_h*C1_filter_w)); _ += C1_weights_grad.size;
C1_bias_grad    = P_grad[_:_+num_params['C1_bias'   ]];                                                          _ += C1_bias_grad.size;
F1_weights_grad = P_grad[_:_+num_params['F1_weights']].reshape((C1_filter_c*P1_w*P1_h, F1_size));                _ += F1_weights_grad.size;
F1_bias_grad    = P_grad[_:_+num_params['F1_bias'   ]].reshape((1,         F1_size));                            _ += F1_bias_grad.size;
F2_weights_grad = P_grad[_:_+num_params['F2_weights']].reshape((F1_size,   F2_size));                            _ += F2_weights_grad.size;
F2_bias_grad    = P_grad[_:_+num_params['F2_bias'   ]].reshape((1,         F2_size));                            _ += F2_bias_grad.size;
assert _ == num_params_total

# Initialize parameters in P for each layer with a different random scale.
def set_rand(target, scale):
    target.ravel()[:] = randn(target.size) * scale
set_rand(C1_weights, 0.01);    C1_bias += 0.0001;    # Initialize biases as small positive values
set_rand(F1_weights, 0.01);    F1_bias += 0.0001;
set_rand(F2_weights, 0.01);

##############################################################################
#  CONVNET FORWARDPROP / BACKPROP
##############################################################################

# Send input mini-batch through our convnet.
# Returns final outputs and, if targets are given, returns gradients as well.
def eval_convnet(inputs, targets=None):

    # Network parameters stored in global variables, for simplicity of this demo
    global C1_weights, C1_bias
    global F1_weights, F1_bias
    global F2_weights, F2_bias

    # Gradients of parameters also stored in global variables, for simplicity
    global C1_weights_grad, C1_bias_grad
    global F1_weights_grad, F1_bias_grad
    global F2_weights_grad, F2_bias_grad

    # Forward propagate C1
    C1_hidden = relu(conv2(inputs, input_w, input_h, C1_weights, C1_filter_w, C1_filter_h, bias=C1_bias, stride=C1_stride))

    # Forward propagate P1
    P1_hidden = pool2(P1_mode, C1_hidden, C1_w, C1_h, P1_window_w, P1_window_h, stride=P1_stride)

    # Forward propagate F1
    F1_hidden = relu(dot(P1_hidden, F1_weights) + F1_bias)
    F1_hidden, F1_mask = dropout(F1_hidden, F1_dropout_rate, test_mode=(targets is None))

    # Forward propagate F2
    F2_hidden = softmax(dot(F1_hidden, F2_weights) + F2_bias)

    # If no targets provided (no gradient requested), just return the predictions of final layer
    if targets is None:
        return F2_hidden

    # Compute residuals
    F2_delta = F2_hidden - targets

    # Backward propagate F2_delta
    F2_bias_grad[:] = sum(F2_delta, axis=0)
    F2_weights_grad[:] = dot_tn(F1_hidden, F2_delta)
    F1_delta = dot_nt(F2_delta, F2_weights)

    # Backward propagate F1_delta
    F1_delta = dropout_grad(F1_delta, F1_mask)  # Backprop through dropout after F1 layer
    F1_delta *= relu_grad(F1_hidden)            # Backprop through relu after F1 layer
    F1_bias_grad[:] = sum(F1_delta, axis=0)
    F1_weights_grad[:] = dot_tn(P1_hidden, F1_delta)
    P1_delta = dot_nt(F1_delta, F1_weights)

    # Backward propagate P1_delta
    C1_delta = pool2_grad(P1_mode, C1_hidden, C1_w, C1_h, P1_window_w, P1_window_h, P1_hidden, P1_delta, stride=P1_stride)

    # Backward propagate C1_delta
    C1_delta *= relu_grad(C1_hidden)  # Backprop through relu after C1 layer
    conv2_biasgrad(C1_bias, C1_delta, C1_bias_grad)
    conv2_filtersgrad(inputs, input_w, input_h, C1_filter_w, C1_filter_h, C1_delta, C1_weights_grad, stride=C1_stride)

##############################################################################
#  Functions for PRINTING
##############################################################################

def make_filter_grid(filter_weights, filter_w, filter_h, max_cols=8):

    # Determine the range [-vmin, vmax] to map to [0, 255]
    vmin = float(filter_weights.min())
    vmax = float(filter_weights.max())
    vmax, vmin = max(vmax, -vmin), min(-vmax, vmin)

    # Scale all filters to range [0, 255] and reshape to a single (n, width, height) array
    images = ((filter_weights.asnumpy() - vmin) / (vmax - vmin) * 255).astype(np.uint8).reshape((-1, filter_h, filter_w))
    n = len(images)

    # Create a big image to store the filters, then copy each filter into its slot in the grid
    num_cols = min(n, max_cols)
    num_rows = (n + num_cols - 1) // num_cols
    grid = np.zeros((num_rows*(filter_h+1)+1, num_cols*(filter_w+1)+1), np.uint8)
    for col in range(num_cols):
        for row in range(num_rows):
            i = row*num_cols + col
            if i < len(images):
                grid[1+row*(filter_h+1):(row+1)*(filter_h+1),
                     1+col*(filter_w+1):(col+1)*(filter_w+1)] = images[i]

    return grid

def error_rate(batches):
    predictions = np.vstack([eval_convnet(inputs).asnumpy() for inputs, targets in batches])
    targets     = np.vstack([targets.asnumpy()              for inputs, targets in batches])
    num_errors  = np.sum( predictions[np.where(targets==1)] != predictions.max(axis=1) )
    return 100.*num_errors/len(predictions)

filter_plot_img = None   # Global variable to hold reference to the filter image currently being plotted in a pyplot window
def plot_filter_grid():
    global filter_plot_img
    import matplotlib.pyplot as plt
    import matplotlib.cm as cm

    filter_grid = make_filter_grid(C1_weights, C1_filter_w, C1_filter_h)
    if filter_plot_img is None:
        # The first time we plot the filters, use imshow and pop up a new window
        filter_plot_img = plt.imshow(filter_grid, cmap=cm.Greys_r, interpolation="NEAREST")
        plt.show(block=False)
    else:
        # When we want to just update the filters, replace the data and give pyplot event loop a chance to draw
        filter_plot_img.set_data(filter_grid)
        plt.pause(0.001)

def print_status(epoch=None):
    update_interval = 5
    if epoch is not None:
        print ".",
    if epoch is None or (epoch+1) % update_interval == 0:  # Only print status every 5 epochs.
        time_per_epoch = toc() / update_interval
        train_error = error_rate(batches_train)
        test_error  = error_rate(batches_test)
        status_msg = "start" if epoch is None else ("epoch[%d]"% (epoch+1))
        time_msg   = "(%.2fs/epoch)" % time_per_epoch if epoch is not None else ""
        print "\n%s: %.2f%% train err, %.2f%% test err %s " % (status_msg, train_error, test_error, time_msg),
        if args.show_filters:
            plot_filter_grid()
        tic()

tic()
print_status()

##############################################################################
#  SGD TRAINING LOOP
##############################################################################

# Parameters of SGD training
num_epoch  = 50
learn_rate = 0.02
momentum   = 0.90

# Allocate array to store momentum of every parameter; updated during training
P_momentum = zeros_like(P)

tic("training time")

# Start training!
for epoch in range(num_epoch):

    for inputs, targets in batches_train:

        # Generate compute per-layer gradient based on targets
        eval_convnet(inputs, targets)

        # Gradient step with very basic momentum
        P_grad *= -learn_rate/batchsize
        P_momentum *= momentum
        P_momentum += P_grad
        P += P_momentum

    # Print current classification error on training data
    print_status(epoch)

print "\nTotal training time = %.1fs" % toc("training time")


# Copyright 2022 The TEMPO Collaboration
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Module for the process tensor time evolving matrix product operator algorithm
(PT-TEMPO). This module is based on [Strathearn2018], [Pollock2018],
[Jorgensen2019], and [Fux2021].

**[Strathearn2018]**
A. Strathearn, P. Kirton, D. Kilda, J. Keeling and
B. W. Lovett,  *Efficient non-Markovian quantum dynamics using
time-evolving matrix product operators*, Nat. Commun. 9, 3322 (2018).

**[Pollock2018]**
F.  A.  Pollock,  C.  Rodriguez-Rosario,  T.  Frauenheim,
M. Paternostro, and K. Modi, *Non-Markovian quantumprocesses: Complete
framework and efficient characterization*, Phys. Rev. A 97, 012127 (2018).

**[Jorgensen2019]**
M. R. Jørgensen and F. A. Pollock, *Exploiting the causal tensor network
structure of quantum processes to efficiently simulate non-markovian path
integrals*, Phys. Rev. Lett. 123, 240602 (2019)

**[Fux2021]**
G. E. Fux, E. Butler, P. R. Eastham, B. W. Lovett, and
J. Keeling, *Efficient exploration of Hamiltonian parameter space for
optimal control of non-Markovian open quantum systems*, Phys. Rev. Lett. 126,
200401 (2021).
"""

from typing import Dict, Optional, Text, Union
from copy import copy

import numpy as np

from oqupy.base_api import BaseAPIClass
from oqupy.bath import Bath
from oqupy.config import PT_DEFAULT_TOLERANCE
from oqupy.config import PT_TEMPO_BACKEND_CONFIG
from oqupy.process_tensor import BaseProcessTensor
from oqupy.process_tensor import SimpleProcessTensor
from oqupy.process_tensor import FileProcessTensor
from oqupy.tempo.backends.pt_tempo_backend import PtTempoBackend
from oqupy.tempo.tempo import TempoParameters
from oqupy.tempo.tempo import guess_tempo_parameters
from oqupy.tempo.tempo import influence_matrix
from oqupy.operators import commutator, acommutator
from oqupy.operators import left_right_super
from oqupy.util import get_progress

PT_CLASS = {"simple": SimpleProcessTensor}

class PtTempo(BaseAPIClass):
    """
    Class to facilitate a PT-TEMPO computation.

    Parameters
    ----------
    bath: Bath
        The Bath (includes the coupling operator to the system).
    parameters: TempoParameters
        The parameters for the PT-TEMPO computation.
    start_time: float
        The start time.
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    def __init__(
            self,
            bath: Bath,
            start_time: float,
            end_time: float,
            parameters: TempoParameters,
            process_tensor_file: Optional[Union[Text, bool]] = None,
            overwrite: Optional[bool] = False,
            backend_config: Optional[Dict] = None,
            name: Optional[Text] = None,
            description: Optional[Text] = None) -> None:
        """Create a PtTempo object. """
        assert isinstance(bath, Bath), \
            "Argument 'bath' must be an instance of Bath."
        self._bath = bath
        self._dimension = self._bath.dimension
        self._correlations = self._bath.correlations

        try:
            tmp_start_time = float(start_time)
        except Exception as e:
            raise AssertionError("Start time must be a float.") from e
        self._start_time = tmp_start_time

        try:
            tmp_end_time = float(end_time)
        except Exception as e:
            raise AssertionError("End time must be a float.") from e
        self._end_time = tmp_end_time

        assert isinstance(parameters, TempoParameters), \
            "Argument 'parameters' must be an instance of TempoParameters."
        self._parameters = parameters

        self._process_tensor = None
        if process_tensor_file or isinstance(process_tensor_file, Text):
            if isinstance(process_tensor_file, Text):
                filename = process_tensor_file
            else:
                filename = None
            self._init_file_process_tensor(filename, overwrite)
        else:
            self._init_simple_process_tensor()

        if backend_config is None:
            self._backend_config = PT_TEMPO_BACKEND_CONFIG
        else:
            self._backend_config = backend_config

        super().__init__(name, description)

        tmp_coupling_comm = commutator(self._bath._coupling_operator)
        tmp_coupling_acomm = acommutator(self._bath._coupling_operator)
        self._coupling_comm = tmp_coupling_comm.diagonal()
        self._coupling_acomm = tmp_coupling_acomm.diagonal()

        tmp_num_steps = int((end_time - self._start_time)/self._parameters.dt)
        assert tmp_num_steps >= 2, \
            "Parameter `end_time` must be more than two times steps " \
            + "larger than the parameter `start_time`!"
        self._num_steps = tmp_num_steps

        self._backend_instance = None
        self._init_pt_tempo_backend()

    def _init_simple_process_tensor(self):
        """ToDo. """
        unitary = self._bath.unitary_transform
        if not np.allclose(unitary, np.identity(self._dimension)):
            transform_in = left_right_super(unitary.conjugate().T,
                                            unitary).T
            transform_out = left_right_super(unitary,
                                             unitary.conjugate().T).T
        else:
            transform_in = None
            transform_out = None
        self._process_tensor = SimpleProcessTensor(
            hilbert_space_dimension=self._dimension,
            dt=self._parameters.dt,
            transform_in=transform_in,
            transform_out=transform_out)

    def _init_file_process_tensor(self, filename, overwrite):
        """ToDo. """
        unitary = self._bath.unitary_transform
        if not np.allclose(unitary, np.identity(self._dimension)):
            transform_in = left_right_super(unitary.conjugate().T,
                                            unitary).T
            transform_out = left_right_super(unitary,
                                             unitary.conjugate().T).T
        else:
            transform_in = None
            transform_out = None

        if overwrite:
            mode = "overwrite"
        else:
            mode = "write"
        self._process_tensor = FileProcessTensor(
            mode=mode,
            filename=filename,
            hilbert_space_dimension=self._dimension,
            dt=self._parameters.dt,
            transform_in=transform_in,
            transform_out=transform_out)

    def _init_pt_tempo_backend(self):
        """Create and initialize the pt-tempo backend. """
        sum_north = np.array([1.0]*(self._dimension**2))
        sum_west = np.array([1.0]*(self._dimension**2))
        dkmax = self._parameters.dkmax
        if dkmax is None:
            dkmax = self._num_steps
        self._backend_instance = PtTempoBackend(
                dimension=self._dimension,
                influence=self._influence,
                process_tensor=self._process_tensor,
                sum_north=sum_north,
                sum_west=sum_west,
                num_steps=self._num_steps,
                dkmax=dkmax,
                epsrel=self._parameters.epsrel,
                config=self._backend_config)

    def _influence(self, dk: int):
        """Create the influence functional matrix for a time step distance
        of dk. """
        return influence_matrix(
            dk,
            parameters=self._parameters,
            correlations=self._correlations,
            coupling_acomm=self._coupling_acomm,
            coupling_comm=self._coupling_comm)

    @property
    def dimension(self) -> np.ndarray:
        """Hilbert space dimension. """
        return copy(self._dimension)

    def compute(self, progress_type: Optional[Text] = None) -> None:
        """
        Propagate (or continue to propagate) the TEMPO tensor network to
        time `end_time`.

        Parameters
        ----------
        progress_type: str (default = None)
            The progress report type during the computation. Types are:
            {``silent``, ``simple``, ``bar``}. If `None` then
            the default progress type is used.
        """
        if self._backend_instance.step is None:
            self._backend_instance.initialize()

        progress = get_progress(progress_type)
        with progress(self._backend_instance.num_steps) as prog_bar:
            while self._backend_instance.compute_step():
                prog_bar.update(self._backend_instance.step)
            prog_bar.update(self._backend_instance.step)

    def get_process_tensor(
            self,
            progress_type: Optional[Text] = None) -> BaseProcessTensor:
        """
        Returns a the computed process tensor. It performs the computation if
        it hasn't been already done.

        Parameters
        ----------
        progress_type: str (default = None)
            The progress report type during the computation. Types are:
            {``silent``, ``simple``, ``bar``}. If `None` then
            the default progress type is used.

        Returns
        -------
        process_tensor: SimpleProcessTensor
            The computed process tensor.
        """
        if self._backend_instance.step is None or \
            self._backend_instance.step < self._backend_instance.num_steps:
            self.compute(progress_type=progress_type)

        if len(self._process_tensor) < self._backend_instance.num_steps:
            self._backend_instance.update_process_tensor()

        return self._process_tensor

def pt_tempo_compute(
        bath: Bath,
        start_time: float,
        end_time: float,
        parameters: Optional[TempoParameters] = None,
        tolerance: Optional[float] = PT_DEFAULT_TOLERANCE,
        process_tensor_file: Optional[Union[Text, bool]] = None,
        overwrite: Optional[bool] = False,
        backend_config: Optional[Dict] = None,
        progress_type: Optional[Text] = None,
        name: Optional[Text] = None,
        description: Optional[Text] = None) -> BaseProcessTensor:
    """
    Shortcut for creating a process tensor by performing a PT-TEMPO
    computation.

    Parameters
    ----------
    bath: Bath
        The Bath (includes the coupling operator to the system).
    start_time: float
        The start time.
    end_time: float
        The time to which the PT-TEMPO should be computed.
    parameters: TempoParameters
        The parameters for the PT-TEMPO computation.
    tolerance: float
        Tolerance for the parameter estimation (only applicable if
        `parameters` is None).
    backend_config: dict (default = None)
        The configuration of the backend. If `backend_config` is
        ``None`` then the default backend configuration is used.
    progress_type: str (default = None)
        The progress report type during the computation. Types are:
        {``'silent'``, ``'simple'``, ``'bar'``}.  If `None` then
        the default progress type is used.
    name: str (default = None)
        An optional name for the tempo object.
    description: str (default = None)
        An optional description of the tempo object.
    """
    if parameters is None:
        assert tolerance is not None, \
            "If 'parameters' is 'None' then 'tolerance' must be " \
            + "a positive float."
        parameters = guess_tempo_parameters(bath=bath,
                                            start_time=start_time,
                                            end_time=end_time,
                                            tolerance=tolerance)
    ptt = PtTempo(bath,
                  start_time,
                  end_time,
                  parameters,
                  process_tensor_file,
                  overwrite,
                  backend_config,
                  name,
                  description)
    ptt.compute(progress_type=progress_type)
    return ptt.get_process_tensor()


"""PriNCe configuration module."""

import os.path as path
import platform
import sys
import warnings

import numpy as np

base_path = path.dirname(path.abspath(__file__))

#: Debug flag for verbose printing, 0 silences PriNCe entirely
debug_level = 1
#: Printout debug info only for functions in this list (just give the name,
#: "get_solution" for instance) Warning, this option slows down initialization
#: by a lot. Use only when needed.
override_debug_fcn = []
#: Override debug printout for debug levels < value for the functions above
override_max_level = 10
#: Print module name in debug output
print_module = False

# =================================================================
# Paths and library locations
# =================================================================

#: Directory where the data files for the calculation are stored
data_dir = path.join(base_path, 'data')

#: PrinceDB file name
db_fname = 'prince_db_05.h5'

#: Model file for redistribution functions (from SOPHIA or similar)
redist_fname = 'sophia_redistribution_logbins.npy'

#=========================================================================
# Physics configuration
#=========================================================================

#: Cosmological parameters

#: Hubble constant
H_0 = 70.5  #km s^-1 Mpc^-1
H_0s = 2.28475e-18  #s^-1

#: Omega_m
Omega_m = 0.27

#: Omega_Lambda
Omega_Lambda = 0.73

#: CMB energy kB*T0 [GeV]
E_CMB = 2.34823e-13

#===========================================================================
# Grids
#===========================================================================

#: Cosmic ray energy grid (defines system size for solver)
#: Number of bins in multiples of 4 recommended for maximal vectorization
#: efficiency for 256 bit AVX or similar
#: Format (log10(E_min), log10(E_max), nbins/decade of energy)
cosmic_ray_grid = (3, 14, 8)
#: Photon grid of target field, only for calculation of rates
photon_grid = (-15, -6, 8)

#: Scale of the energy grid
#:'E': logarithmic in energy E_i = E_min * (Delta)^i
#:'logE': linear grid in x = log_10(E): x_i = x_min + i * Delta
grid_scale ='E'

#: Order of semi-lagrangian for energy derivative
semi_lagr_method ='5th_order'

#===========================================================================
# Model options
#===========================================================================

#: Threshold lifetime value for explicit transport of particles of this type. It
#: means that if a particle is unstable with lifetime smaller than this threshold,
#: it will be decayed until all final state particles of this chain are stable.
#: In other words: short intermediate states will be integrated out
tau_dec_threshold = np.inf # All unstable particles decay
# tau_dec_threshold = 0.  # None unstable particles decay
# tau_dec_threshold = 850. # This value is for stable neutrons

#: Particle ID for which redistribution functions are needed to be taken into
#: account. The default value is 101 (proton). All particles with smaller
#: IDs, i.e. neutrinos, pions, muons etc., will have energy redistributions.
#: For larger IDs (nuclei) the boost conservation is employed.
redist_threshold_ID = 101

#: Cut on energy redistribution functions
#: Resitribution below this x value are set to 0.
#: "x_cut" : 0.,
#: "x_cut_proton" : 0.,
x_cut = 1e-4
x_cut_proton = 1e-1

#: cut on photon energy, cross section above y = E_cr e_ph / m_cr does not contribute
y_cut = np.inf

# Build equation system up to a maximal nuclear mass of
max_mass = np.inf

# Include secondaries like photons and neutrinos
secondaries = True
# List of specific particles to ignore
ignore_particles = [20,21] # (we ignore photons and electrons, as their physics is not fully implemented)

#===========================================================================
# Parameters of numerical integration
#===========================================================================

# Update rates at not more frequently than this value in z
update_rates_z_threshold = 0.01

# #Number of MKL threads (for sparse matrix multiplication the performance
# #advantage from using more than a few threads is limited by memory bandwidth)
MKL_threads = 4

# Sparse matrix-vector product from "CUPY"|"MKL"|"scipy"
linear_algebra_backend = "MKL"

# Check for CUPY library for GPU support
try:
    import cupy
    has_cupy = True
    mempool = cupy.get_default_memory_pool()
    mempool.free_all_blocks()
except ModuleNotFoundError:
    print('CUPY not found for GPU support. Degrading to MKL.')
    if linear_algebra_backend == 'cupy':
        linear_algebra_backend = 'MKL'
    has_cupy = False

#: determine shared library extension and MKL path
pf = platform.platform()

if 'Linux' in pf:
    mkl_path = path.join(sys.prefix, 'lib', 'libmkl_rt.so')
elif 'Darwin' in pf:
    mkl_path = path.join(sys.prefix, 'lib', 'libmkl_rt.dylib')
else:
    # Windows case
    mkl_path = path.join(sys.prefix, 'Library', 'bin', 'mkl_rt.dll')

# mkl library handler
mkl = None

# Check if MKL library found
if path.isfile(mkl_path):
    has_mkl = True
else:
    has_mkl = False

def set_mkl_threads(nthreads):
    global mkl, MKL_threads
    from ctypes import cdll, byref, c_int
    mkl = cdll.LoadLibrary(mkl_path)
    # Set number of threads
    MKL_threads = nthreads
    mkl.mkl_set_num_threads(byref(c_int(nthreads)))
    if debug_level >= 5:
        print('MKL threads limited to {0}'.format(nthreads))

if has_mkl:
    set_mkl_threads(MKL_threads)

if not has_mkl and linear_algebra_backend.lower() == 'mkl':
    print('MKL runtime not found. Degrading to scipy.')
    linear_algebra_backend = 'scipy'

def _download_file(url, outfile):
    """Downloads the PriNCe database from github release binaries."""

    from tqdm import tqdm
    import requests
    import math

    # Streaming, so we can iterate over the response.
    r = requests.get(url, stream=True)

    # Total size in bytes.
    total_size = int(r.headers.get('content-length', 0))
    block_size = 1024 * 1024
    wrote = 0
    with open(outfile, 'wb') as f:
        for data in tqdm(r.iter_content(block_size), total=math.ceil(total_size // block_size),
                         unit='MB', unit_scale=True):
            wrote = wrote + len(data)
            f.write(data)
    if total_size != 0 and wrote != total_size:
        raise Exception("ERROR, something went wrong")

# Download database file from github
base_url = 'https://github.com/joheinze/PriNCe/releases/download/'
release_tag = 'v0.5_alpha_release/'
url = base_url + release_tag + db_fname
if not path.isfile(path.join(data_dir, db_fname)):
    print('Downloading for PriNCe database file {0}.'.format(db_fname))
    if debug_level >= 2:
        print(url)
    _download_file(url, path.join(data_dir, db_fname))
else:
    import h5py
    try:
        with h5py.File(path.join(data_dir, db_fname), 'r') as prince_db:
            db_version = (prince_db.attrs['version'])
    except:
        print(f'Database file {db_fname} corrupted. Retrying download.')
        _download_file(url, path.join(data_dir, db_fname))
    finally:
        with h5py.File(path.join(data_dir, db_fname), 'r') as prince_db:
            db_version = (prince_db.attrs['version'])
        if debug_level >= 2:
            print(f'Using database file version {db_version}.')

# if path.isfile(path.join(data_dir, '...previous db name...')):
#     import os
#     print('Removing previous database {0}.'.format('...previous db name...'))
#     os.unlink(path.join(data_dir, '...previous db name...'))


# ----------------------------------------------------------------------------
# Copyright (c) 2016-2017, UniFrac development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file LICENSE, distributed with this software.
# ----------------------------------------------------------------------------
from warnings import warn
from functools import reduce
from operator import or_

import numpy as np
import skbio

import unifrac as qsu
from unifrac._meta import CONSOLIDATIONS

def is_biom_v210(f):
    import h5py
    if not h5py.is_hdf5(f):
        return False
    with h5py.File(f, 'r') as fp:
        if 'format-version' not in fp.attrs:
            return False

        version = fp.attrs.get('format-version', None)

        if version is None:
            return False

        if tuple(version) != (2, 1):
            return False

    return True

def is_newick(f):
    sniffer = skbio.io.format.newick.newick.sniffer_function
    return sniffer(f)[0]

def _validate(table, phylogeny):
    if not is_biom_v210(table):
        raise ValueError("Table does not appear to be a BIOM-Format v2.1")
    if not is_newick(phylogeny):
        raise ValueError("The phylogeny does not appear to be newick")

def _validate_meta(tables, phylogenies):
    for idx, (table, phylogeny) in enumerate(zip(tables, phylogenies)):
        if not is_biom_v210(table):
            raise ValueError(f"Table at position {idx} does not appear to be a"
                             " BIOM-Format v2.1")
        if not is_newick(phylogeny):
            raise ValueError(f"The phylogeny at position {idx} does not appear"
                             " to be newick")

def unweighted(table: str,
               phylogeny: str,
               threads: int = 1,
               variance_adjusted: bool = False,
               bypass_tips: bool = False) -> skbio.DistanceMatrix:
    """Compute Unweighted UniFrac

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default of 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Unweighted UniFrac was originally described in [1]_. Variance Adjusted
    UniFrac was originally described in [2]_, and while its application to
    Unweighted UniFrac was not described, factoring in the variance adjustment
    is still feasible and so it is exposed.

    References
    ----------
    .. [1] Lozupone, C. & Knight, R. UniFrac: a new phylogenetic method for
       comparing microbial communities. Appl. Environ. Microbiol. 71, 8228-8235
       (2005).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(table, phylogeny, 'unweighted',
                   variance_adjusted, 1.0, bypass_tips, threads)

def unweighted_fp32(table: str,
                    phylogeny: str,
                    threads: int = 1,
                    variance_adjusted: bool = False,
                    bypass_tips: bool = False) -> skbio.DistanceMatrix:
    """Compute Unweighted UniFrac using fp32 math

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default of 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Unweighted UniFrac was originally described in [1]_. Variance Adjusted
    UniFrac was originally described in [2]_, and while its application to
    Unweighted UniFrac was not described, factoring in the variance adjustment
    is still feasible and so it is exposed.

    References
    ----------
    .. [1] Lozupone, C. & Knight, R. UniFrac: a new phylogenetic method for
       comparing microbial communities. Appl. Environ. Microbiol. 71, 8228-8235
       (2005).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(table, phylogeny, 'unweighted_fp32',
                   variance_adjusted, 1.0, bypass_tips, threads)

def weighted_normalized(table: str,
                        phylogeny: str,
                        threads: int = 1,
                        variance_adjusted: bool = False,
                        bypass_tips: bool = False) -> skbio.DistanceMatrix:
    """Compute weighted normalized UniFrac

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default of 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Weighted UniFrac was originally described in [1]_. Variance Adjusted
    Weighted UniFrac was originally described in [2]_.

    References
    ----------
    .. [1] Lozupone, C. A., Hamady, M., Kelley, S. T. & Knight, R. Quantitative
       and qualitative beta diversity measures lead to different insights into
       factors that structure microbial communities. Appl. Environ. Microbiol.
       73, 1576-1585 (2007).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(str(table), str(phylogeny), 'weighted_normalized',
                   variance_adjusted, 1.0, bypass_tips, threads)

def weighted_normalized_fp32(table: str,
                             phylogeny: str,
                             threads: int = 1,
                             variance_adjusted: bool = False,
                             bypass_tips: bool = False
                             ) -> skbio.DistanceMatrix:
    """Compute weighted normalized UniFrac using fp32 math

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default of 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Weighted UniFrac was originally described in [1]_. Variance Adjusted
    Weighted UniFrac was originally described in [2]_.

    References
    ----------
    .. [1] Lozupone, C. A., Hamady, M., Kelley, S. T. & Knight, R. Quantitative
       and qualitative beta diversity measures lead to different insights into
       factors that structure microbial communities. Appl. Environ. Microbiol.
       73, 1576-1585 (2007).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(str(table), str(phylogeny), 'weighted_normalized_fp32',
                   variance_adjusted, 1.0, bypass_tips, threads)

def weighted_unnormalized(table: str,
                          phylogeny: str,
                          threads: int = 1,
                          variance_adjusted: bool = False,
                          bypass_tips: bool = False) -> skbio.DistanceMatrix:
    # noqa
    """Compute weighted unnormalized UniFrac

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default is 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Weighted UniFrac was originally described in [1]_. Variance Adjusted
    Weighted UniFrac was originally described in [2]_.

    References
    ----------
    .. [1] Lozupone, C. A., Hamady, M., Kelley, S. T. & Knight, R. Quantitative
       and qualitative beta diversity measures lead to different insights into
       factors that structure microbial communities. Appl. Environ. Microbiol.
       73, 1576-1585 (2007).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(str(table), str(phylogeny), 'weighted_unnormalized',
                   variance_adjusted, 1.0, bypass_tips, threads)

def weighted_unnormalized_fp32(table: str,
                               phylogeny: str,
                               threads: int = 1,
                               variance_adjusted: bool = False,
                               bypass_tips: bool = False
                               ) -> skbio.DistanceMatrix:
    # noqa
    """Compute weighted unnormalized UniFrac using fp32 math

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default is 1.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Weighted UniFrac was originally described in [1]_. Variance Adjusted
    Weighted UniFrac was originally described in [2]_.

    References
    ----------
    .. [1] Lozupone, C. A., Hamady, M., Kelley, S. T. & Knight, R. Quantitative
       and qualitative beta diversity measures lead to different insights into
       factors that structure microbial communities. Appl. Environ. Microbiol.
       73, 1576-1585 (2007).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    return qsu.ssu(str(table), str(phylogeny), 'weighted_unnormalized_fp32',
                   variance_adjusted, 1.0, bypass_tips, threads)

def generalized(table: str,
                phylogeny: str,
                threads: int = 1,
                alpha: float = 1.0,
                variance_adjusted: bool = False,
                bypass_tips: bool = False) -> skbio.DistanceMatrix:
    """Compute Generalized UniFrac

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default is 1
    alpha : float, optional
        The level of contribution of high abundance branches. Higher alpha
        increases the contribution of from high abundance branches while lower
        alpha reduces the contribution. Alpha was originally defined over the
        range [0, 1]. Default is 1.0.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Generalized UniFrac was originally described in [1]_. Variance Adjusted
    UniFrac was originally described in [2]_, but was not described in as
    applied to Generalized UniFrac. It is feasible to do, so it is exposed
    here.

    An alpha of 1.0 is Weighted normalized UniFrac. An alpha of 0.0 is
    approximately Unweighted UniFrac, and is if the proportions are
    dichotomized.

    References
    ----------
    .. [1] Chen, J., Bittinger, K., Charlson, E. S., Hoffmann C., Lewis, J.,
       Wu, G. D., Collman R. G., Bushman, F. D. & Hongzhe L. Associating
       microbiome composition with environmental covariates using generalized
       UniFrac distances. Bioinformatics 28(16), 2106–2113 (2012).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    if alpha == 1.0:
        warn("alpha of 1.0 is weighted-normalized UniFrac. "
             "Weighted-normalized is being used instead as it is more "
             "optimized.",
             Warning)
        return weighted_normalized(table, phylogeny, threads,
                                   variance_adjusted)
    else:
        return qsu.ssu(str(table), str(phylogeny), 'generalized',
                       variance_adjusted, alpha, bypass_tips, threads)

def generalized_fp32(table: str,
                     phylogeny: str,
                     threads: int = 1,
                     alpha: float = 1.0,
                     variance_adjusted: bool = False,
                     bypass_tips: bool = False) -> skbio.DistanceMatrix:
    """Compute Generalized UniFrac using fp32 math

    Parameters
    ----------
    table : str
        A filepath to a BIOM-Format 2.1 file.
    phylogeny : str
        A filepath to a Newick formatted tree.
    threads : int, optional
        The number of threads to use. Default is 1
    alpha : float, optional
        The level of contribution of high abundance branches. Higher alpha
        increases the contribution of from high abundance branches while lower
        alpha reduces the contribution. Alpha was originally defined over the
        range [0, 1]. Default is 1.0.
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    Generalized UniFrac was originally described in [1]_. Variance Adjusted
    UniFrac was originally described in [2]_, but was not described in as
    applied to Generalized UniFrac. It is feasible to do, so it is exposed
    here.

    An alpha of 1.0 is Weighted normalized UniFrac. An alpha of 0.0 is
    approximately Unweighted UniFrac, and is if the proportions are
    dichotomized.

    References
    ----------
    .. [1] Chen, J., Bittinger, K., Charlson, E. S., Hoffmann C., Lewis, J.,
       Wu, G. D., Collman R. G., Bushman, F. D. & Hongzhe L. Associating
       microbiome composition with environmental covariates using generalized
       UniFrac distances. Bioinformatics 28(16), 2106–2113 (2012).
    .. [2] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    _validate(table, phylogeny)
    if alpha == 1.0:
        warn("alpha of 1.0 is weighted-normalized UniFrac. "
             "Weighted-normalized is being used instead as it is more "
             "optimized.",
             Warning)
        return weighted_normalized_fp32(table, phylogeny, threads,
                                        variance_adjusted)
    else:
        return qsu.ssu(str(table), str(phylogeny), 'generalized_fp32',
                       variance_adjusted, alpha, bypass_tips, threads)

METHODS = {'unweighted': unweighted,
           'weighted_normalized': weighted_normalized,
           'weighted_unnormalized': weighted_unnormalized,
           'generalized': generalized,
           'unweighted_fp32': unweighted_fp32,
           'weighted_normalized_fp32': weighted_normalized_fp32,
           'weighted_unnormalized_fp32': weighted_unnormalized_fp32,
           'generalized_fp32': generalized_fp32}

def meta(tables: tuple, phylogenies: tuple, weights: tuple = None,
         consolidation: str = None, method: str = None,
         threads: int = 1, variance_adjusted: bool = False,
         alpha: float = None, bypass_tips: bool = False) -> \
         skbio.DistanceMatrix:
    """Compute meta UniFrac

    Parameters
    ----------
    tables : tuple of str
        Filepaths to BIOM-Format 2.1 files. This tuple is expected to be in
        index order with phylogenies.
    phylogenies : tuple of str
        Filepaths to Newick formatted trees. This tuple is expected to be in
        index order with tables.
    weights : tuple of float, optional
        The weight applied to each tree/table pair. This tuple is expected to
        be in index order with tables and phylogenies. Default is to weight
        each tree/table pair evenly.
    consolidation : str, optional
        The matrix consolidation method. The available choices are:
        'skipping_missing_matrices', 'missing_zero', 'missing_one',
        'skipping_missing_values'. The default is 'skipping_missing_values'.
    method : str
        The UniFrac method to use. The available choices are:
        'unweighted', 'weighted_unnormalized', 'weighted_normalized', and
        'generalized'.
    threads : int, optional
        The number of threads to use. Default is 1
    bypass_tips : bool
        Bypass the tips of the tree in the computation. This reduces compute
        by about 50%, but is an approximation.
    alpha : float, optional
        The level of contribution of high abundance branches. Higher alpha
        increases the contribution of from high abundance branches while lower
        alpha reduces the contribution. Alpha was originally defined over the
        range [0, 1]. Default is 1.0
    variance_adjusted : bool, optional
        Adjust for varianace or not. Default is False.

    Returns
    -------
    skbio.DistanceMatrix
        The resulting distance matrix.

    Raises
    ------
    IOError
        If the tree file is not found
        If the table is not found
    ValueError
        If the table does not appear to be BIOM-Format v2.1.
        If the phylogeny does not appear to be in Newick format.

    Notes
    -----
    UniFrac can be adapted to account for multiple genes, as originally
    done in [1]_.

    Generalized UniFrac was originally described in [2]_. Variance Adjusted
    UniFrac was originally described in [3]_, but was not described in as
    applied to Generalized UniFrac. It is feasible to do, so it is exposed
    here.

    References
    ----------
    .. [1] Lozupone C. A., Hamady M., Cantarel B. L., Coutinho P. M.,
       Henrissat B., Gordon J. I. & Knight R. The convergence of carbohydrate
       active gene repertoires in human gut microbes. PNAS 105(39):15076-81
       (2008).
    .. [2] Chen, J., Bittinger, K., Charlson, E. S., Hoffmann C., Lewis, J.,
       Wu, G. D., Collman R. G., Bushman, F. D. & Hongzhe L. Associating
       microbiome composition with environmental covariates using generalized
       UniFrac distances. Bioinformatics 28(16), 2106–2113 (2012).
    .. [3] Chang, Q., Luan, Y. & Sun, F. Variance adjusted weighted UniFrac: a
       powerful beta diversity measure for comparing communities based on
       phylogeny. BMC Bioinformatics 12:118 (2011).
    """
    if not len(tables):
        raise ValueError("No tables specified.")

    if not len(phylogenies):
        raise ValueError("No trees specified.")

    if len(tables) != len(phylogenies):
        raise ValueError("Number of trees and tables must be the same.")

    if weights is None:
        weights = tuple(1 for _ in phylogenies)
    else:
        if len(weights) != len(phylogenies):
            raise ValueError("Number of weights does not match number of "
                             "trees and tables.")

    if method is None:
        raise ValueError("No method specified.")
    method_ = METHODS.get(method.replace('-', '_'))
    if method_ is None:
        raise ValueError("Method (%s) unrecognized. Available methods are: %s"
                         % (method, ', '.join(METHODS.keys())))

    if consolidation is None:
        consolidation = 'skipping_missing_values'
    consolidation_ = CONSOLIDATIONS.get(consolidation.replace('-', '_'))
    if consolidation_ is None:
        raise ValueError("Consolidation (%s) unrecognized. Available "
                         "consolidations are: %s"
                         % (consolidation, ', '.join(CONSOLIDATIONS.keys())))

    if alpha is not None and method is not generalized:
        raise ValueError("The alpha parameter can only be set when the method "
                         "is set as 'generalized', the selected method is "
                         "'%s'." % method)

    _validate_meta(tables, phylogenies)

    kwargs = {'threads': threads,
              'bypass_tips': bypass_tips,
              'variance_adjusted': variance_adjusted}
    if alpha is not None:
        kwargs['alpha'] = alpha

    weights = np.array(weights, float)/sum(weights)
    dms = [method_(table, tree, **kwargs) for table, tree in zip(tables,
                                                                 phylogenies)]
    all_ids = sorted(reduce(or_, [set(dm.ids) for dm in dms]))
    dm = consolidation_(dms, [dm.ids for dm in dms], weights, all_ids)

    return skbio.DistanceMatrix(dm, ids=all_ids)