#include <cassert>
#include <cmath>
#include <string>
#include <sstream>
#include <iostream>
#include <fstream>
#include <algorithm>
#include <random>
#include <ctime>              // required for "struct timezone" and
                              // "gettimeofday()" used to set the randomseed
using namespace std;

#ifndef DISABLE_OPENMP
#include <omp.h>       // (OpenMP-specific)
#endif

//#include <boost/math/special_functions/gamma.hpp>

#include <mrc_simple.hpp>
#include <visfd.hpp>
using namespace visfd;

#include "err.hpp"
#include "settings.hpp"

int main(int argc, char **argv) {
  try {

    Settings settings; // parse the command-line argument list from the shell
    settings.ParseArgs(argc, argv);

    #ifndef DISABLE_OPENMP
    #pragma omp parallel
    {
      int rank, nthr;
      rank = omp_get_thread_num();
      //cerr << "rank=" << rank << endl;
      if (rank == 0) {
        nthr = omp_get_num_threads();
        cerr << "   (Using " << nthr << " threads (cpu cores).  You can change this using the \"-np n\"\n"
             << "    argument, or by setting the OMP_NUM_THREADS environment variable.)" << endl;
      }
    }
    #else
    cerr << " (Serial version)" << endl;
    #endif //#ifndef DISABLE_OPENMP

    // Read the input tomogram
    MrcSimple tomo_in;

    if (settings.in_file_name != "") {
      cerr << "Reading tomogram \""<<settings.in_file_name<<"\"" << endl;
      tomo_in.Read(settings.in_file_name, false);
      // (Note: You can also use "tomo_in.Read(cin);" or "cin >> tomo;")
      tomo_in.PrintStats(cerr);      //Optional (display the tomogram size & format)
      WarnMRCSignedBytes(tomo_in, settings.in_file_name, cerr);
    }
    else {
      assert((settings.image_size[0] > 0) &&
             (settings.image_size[1] > 0) &&
             (settings.image_size[2] > 0));
      tomo_in.Resize(settings.image_size);
    }

    // ---- mask ----

    // Optional: if there is a "mask", read that too
    MrcSimple mask;
    if (settings.mask_file_name != "") {
      cerr << "Reading mask \""<<settings.mask_file_name<<"\"" << endl;
      mask.Read(settings.mask_file_name, false);
      if ((mask.header.nvoxels[0] != tomo_in.header.nvoxels[0]) ||
          (mask.header.nvoxels[1] != tomo_in.header.nvoxels[1]) ||
          (mask.header.nvoxels[2] != tomo_in.header.nvoxels[2]))
        throw InputErr("Error: The size of the mask image does not match the size of the input image.\n");
    }

    // ---- Voxel size? ----

    float voxel_width[3] = {1.0, 1.0, 1.0};
    if (settings.voxel_width > 0.0) {
      // Did the user manually specify the width of each voxel?
      voxel_width[0] = settings.voxel_width;
      voxel_width[1] = settings.voxel_width;
      voxel_width[2] = settings.voxel_width;
    }
    else {
      // Otherwise, infer it from the header of the MRC file
      voxel_width[0] = tomo_in.header.cellA[0]/tomo_in.header.nvoxels[0];
      voxel_width[1] = tomo_in.header.cellA[1]/tomo_in.header.nvoxels[1];
      voxel_width[2] = tomo_in.header.cellA[2]/tomo_in.header.nvoxels[2];
      if (settings.voxel_width_divide_by_10) {
        voxel_width[0] *= 0.1;
        voxel_width[1] *= 0.1;
        voxel_width[2] *= 0.1;
      }
      cerr << "voxel width in physical units = ("
           << voxel_width[0] << ", "
           << voxel_width[1] << ", "
           << voxel_width[2] << ")\n";
    }

    if ((voxel_width[0] <= 0.0) ||
        (voxel_width[1] <= 0.0) ||
        (voxel_width[2] <= 0.0))
      throw InputErr("Error in tomogram header: Invalid voxel width(s).\n"
                     "Use the -w argument to specify the voxel width.");

    if ((voxel_width[0] != voxel_width[1]) ||
        (voxel_width[1] != voxel_width[2]))
      throw InputErr("Error in tomogram header: Unequal voxel widths in the x, y and directions.\n"
                     "Use the -w argument to specify the voxel width.");

    if (settings.in_coords_file_name != "") {
      for (int iz = 0; iz < tomo_in.header.nvoxels[2]; iz++)
        for (int iy = 0; iy < tomo_in.header.nvoxels[1]; iy++)
          for (int ix = 0; ix < tomo_in.header.nvoxels[0]; ix++)
            tomo_in.aaafI[iz][iy][ix] = 0.0;
      fstream coords_file;
      coords_file.open(settings.in_coords_file_name, ios::in);
      if (! coords_file)
        throw InputErr("Error: unable to open \""+
                       settings.in_coords_file_name +"\" for reading.\n");
      while (coords_file) {
        float x, y, z;
        coords_file >> x;
        coords_file >> y;
        coords_file >> z;
        int ix, iy, iz;
        ix = static_cast<int>(x / voxel_width[0]);
        iy = static_cast<int>(y / voxel_width[1]);
        iz = static_cast<int>(z / voxel_width[2]);
        if (((0 <= ix) && (ix < tomo_in.header.nvoxels[0])) &&
            ((0 <= iy) && (iy < tomo_in.header.nvoxels[1])) &&
            ((0 <= iz) && (iz < tomo_in.header.nvoxels[2])))
          tomo_in.aaafI[iz][iy][ix] = 1.0;
      }
    } //if (settings.in_coords_file_name != "") {

    double vol_total = settings.compartment_volume;
    if (settings.compartment_volume < 0) {
      if (mask.aaafI) {
        for (int iz=0; iz<tomo_in.header.nvoxels[2]; iz++)
          for (int iy=0; iy<tomo_in.header.nvoxels[1]; iy++)
            for (int ix=0; ix<tomo_in.header.nvoxels[0]; ix++)
              vol_total += mask.aaafI[iz][iy][ix];
      }
      else {
        vol_total = (tomo_in.header.nvoxels[0] *
                     tomo_in.header.nvoxels[1] *
                     tomo_in.header.nvoxels[2]);
      }
      vol_total *= (voxel_width[0]*voxel_width[1]*voxel_width[2]);
    }

    if (settings.num_particles < 0) {
      // If the user did not explicitly specify the number of particles
      // (ie "ribosomes") present in the original image, we can attempt
      // to infer it from the image itself.  To do this, I assume that
      // the image is filled with voxels containing 0s everywhere except
      // at the center of where each particle is located, where the voxel is 1.
      // Consequently, the number of non-zero voxels (not excluded by the mask)
      // should equal the sum of all of the voxel entries.
      // (This should also equal the number of non-zero voxels.  However
      //  sometimes these images are themselves slightly blurred to make it
      //  easier to see each particle.  This blurring should not effect the sum)
      double sum = 0.0;
      for (int iz=0; iz<tomo_in.header.nvoxels[2]; iz++) {
        for (int iy=0; iy<tomo_in.header.nvoxels[1]; iy++) {
          for (int ix=0; ix<tomo_in.header.nvoxels[0]; ix++) {
            if (mask.aaafI)
              sum += tomo_in.aaafI[iz][iy][ix] * mask.aaafI[iz][iy][ix];
            else
              sum += tomo_in.aaafI[iz][iy][ix];
          }
        }
      }
      settings.num_particles = sum;
    }

    // DEBUG: The next if-then statement is for debugging only.
    if (settings.randomize_input_image) {
      size_t nparticles = floor(settings.num_particles);
      size_t nvoxels;
      if (mask.aaafI) {
        for (int iz=0; iz<tomo_in.header.nvoxels[2]; iz++)
          for (int iy=0; iy<tomo_in.header.nvoxels[1]; iy++)
            for (int ix=0; ix<tomo_in.header.nvoxels[0]; ix++)
              if (mask.aaafI[iz][iy][ix] != 0.0)
                nvoxels += 1;
      }
      else {
        nvoxels = (tomo_in.header.nvoxels[0] *
                   tomo_in.header.nvoxels[1] *
                   tomo_in.header.nvoxels[2]);
      }
      vector<bool> random_bit_list(nvoxels, false);
      for (size_t i=0; i < nparticles; i++)
        random_bit_list[i] = true;

      long random_seed = settings.random_seed;
      if (random_seed <= 0) {
        random_seed = time(nullptr);
        cerr << "(random_seed = " << random_seed << ")" << endl;
      }
      shuffle(random_bit_list.begin(), random_bit_list.end(),
              default_random_engine(random_seed));
      size_t i=0;
      for (int iz = 0; iz < tomo_in.header.nvoxels[2]; iz++) {
        for (int iy = 0; iy < tomo_in.header.nvoxels[1]; iy++) {
          for (int ix = 0; ix < tomo_in.header.nvoxels[0]; ix++) {
            tomo_in.aaafI[iz][iy][ix] = 0.0;
            if ((! mask.aaafI) || mask.aaafI[iz][iy][ix] != 0.0) {
              tomo_in.aaafI[iz][iy][ix] = random_bit_list[i];
              i++;
            }
          }
        }
      }
    } //if (settings.randomize_input_image)

    // First we must allocate a new array to store the filtered image.
    MrcSimple tomo_out = tomo_in;

    for (int i_sig = 0; i_sig < settings.vfSigma.size(); ++i_sig) {

      // ---- main calculation: -----

      //A Gaussian filter is a (weighted) average over a volume of nearby voxels
      //Below we estimate the effective volume of those nearby voxels, by taking
      //advantage of the fact that Gaussian is normalized (it's integral = 1).
      //Normalization effectively reduces the height of the filter by the number
      //of voxels which effectively belong to the filter.
      // Example: a cube filter
      // If we were averaging the brightness of the 27 voxels
      // belonging to a 3*3*3 cube surrounding the voxel of interest
      // then the effective volume of the filter is is 3*3*3=27, and
      // weight of each voxel's contribution to the average is 1/27.  Similarly,
      // for a more general filter, we can interpret the peak height of the
      // filter as 1/volume. Below, we calculate the height of the 3D Gaussian's
      // peak after normalization.  Normally the 3D Gaussian is:
      //   (2*pi*sigma^2)^(-3/2) * exp(-0.5*(x^2+y^2+z^2)/sigma^2)
      // and the peak height is (2*pi*sigma^2)^(-3/2),
      // However this formula does not work for small sigma << 1.
      // For small sigma, you have to normalize by computing a discrete sum
      // (over a finite range).  We do this below (in a non-intuitive way).

      float sigma = settings.vfSigma[i_sig];

      sigma /= voxel_width[0]; //scale by voxel_width (if specified)

      // How wide is the filter window (over what size window do we integrate)?
      if (settings.filter_truncate_ratio <= 0) {
        assert(settings.filter_truncate_threshold > 0.0);
        settings.filter_truncate_ratio = sqrt(-2*log(settings.filter_truncate_threshold));
      }
      int filter_truncate_halfwidth = floor(sigma *
                                            settings.filter_truncate_ratio);

      float gauss_peak_height_3D;

      {
        // GenFilterGauss1D generates a 1D discrete normalized Gaussian (located
        // in the afH[]) member. afH[0] is the height at the central peak.
        // However, for now I just use them to estimate the height of the
        // Gaussian peak.  (See below)
        Filter1D<float, int> filter1D =
          GenFilterGauss1D(sigma, filter_truncate_halfwidth);

        // A 3D Gaussian blur can be performed by expoiting the fact that
        // multidimensional Gaussians are "seperable" filters:  You can blur
        // in the X direction, then in the Y direcion, then in the Z direction.
        // The result is the same that you would get from a 3D filter Gaussian
        // which is the product of Gaussians in the X,Y,Z directions.
        // So the peak height is the product of the 3 1D Gaussian peak heights.
        gauss_peak_height_3D =
          (filter1D.afH[0] *
           filter1D.afH[0] *
           filter1D.afH[0]);
      }

      // The Gaussian peak's height is 1/volume.

      float volume_gaussian_bin = 1.0 / gauss_peak_height_3D;
      volume_gaussian_bin *= (voxel_width[0]*voxel_width[1]*voxel_width[2]);

      float num_bins = vol_total / volume_gaussian_bin;

      // Now blur the original image (if the user did not do it already)
      if ((! settings.precomputed_gaussian_blur) ||
          (settings.vfSigma.size() > 1)) {

        ApplyGauss(tomo_in.header.nvoxels,
                   tomo_in.aaafI,
                   tomo_out.aaafI, //<-store resulting image here
                   mask.aaafI,
                   sigma,
                   filter_truncate_halfwidth,
                   true,
                   &cerr);

        // Densities everywhere are assumed to be in physical units
        // (ie. of 1/Angstroms^3),  NOT   1/voxels^3
        // The Gaussian blur computes densities in 1/voxels^3
        // To compensate for this, divide the densities by voxel_width^3
        MultiplyScalarArr(static_cast<float>(
                          1.0/(voxel_width[0]*voxel_width[1]*voxel_width[2])),
                          tomo_out.header.nvoxels,
                          tomo_out.aaafI);
      }

      // After bluring the image, find the lowest density:

      float extreme_density;
      int afXextreme[3] = {-1, -1, -1};
      float global_minima;
      float global_maxima;

      // Find the voxels with the minima and maxima intensities.
      // Discard global minima or maxima which lie on the boundary
      // The safe, careful way to do this is to only consider
      // voxels which are surrounded by other voxels and are
      // local minima or maxima.
      for (int iz=0; iz < tomo_out.header.nvoxels[2]; iz++) {
        for (int iy=0; iy < tomo_out.header.nvoxels[1]; iy++) {
          for (int ix=0; ix < tomo_out.header.nvoxels[0]; ix++) {

            bool is_local_minima = true;
            bool is_local_maxima = true;

            if (! settings.extrema_on_boundary) {
              if ((ix == 0) || (ix ==  tomo_out.header.nvoxels[0]-1) ||
                  (ix == 0) || (ix ==  tomo_out.header.nvoxels[0]-1) ||
                  (ix == 0) || (ix ==  tomo_out.header.nvoxels[0]-1)) {
                is_local_minima = false;
                is_local_maxima = false;
              }
            }

            float center_val = tomo_out.aaafI[iz][iy][ix];
            for (int jz=-1; jz<=1; jz++) {
              for (int jy=-1; jy<=1; jy++) {
                for (int jx=-1; jx<=1; jx++) {
                  if (! settings.extrema_on_boundary) {
                    if (mask.aaafI && (mask.aaafI[iz+jz][iy+jy][ix+jx] == 0.0))
                    {
                      is_local_minima = false;
                      is_local_maxima = false;
                    }
                  }
                  if (tomo_out.aaafI[iz+jz][iy+jy][ix+jx] <= center_val)
                    is_local_minima = false;
                  if (tomo_out.aaafI[iz+jz][iy+jy][ix+jx] >= center_val)
                    is_local_maxima = false;
                }
              }
            }
            if ((center_val < global_minima) || (afXextreme[0] == -1)) {
              global_minima = center_val;
              if (settings.use_min_density) {
                extreme_density = global_minima;
                afXextreme[0] = ix;
                afXextreme[1] = iy;
                afXextreme[2] = iz;
              }
            }
            if ((center_val > global_maxima) || (afXextreme[0] == -1)) {
              global_maxima = center_val;
              if (! settings.use_min_density) {
                extreme_density = global_maxima;
                afXextreme[0] = ix;
                afXextreme[1] = iy;
                afXextreme[2] = iz;
              }
            }
          } //for (int ix=1; ix < tomo_out.header.nvoxels[0]-1; ix++)
        } //for (int iy=1; iy < tomo_out.header.nvoxels[1]-1; iy++)
      } //for (int iz=1; iz < tomo_out.header.nvoxels[2]-1; iz++)

      // Did we fail to find any local minima or maxima densities?
      if (afXextreme[0] == -1) {
        string extrema_type = "minina";
        if (! settings.use_min_density)
          extrema_type = "maxima";
        stringstream msg_ss;
        msg_ss << "Error: There are no local density " << extrema_type << "\n"
               << "       in the image (at scale sigma = " << sigma << ")\n"
               << "       Aborting...\n";
        throw InputErr(msg_ss.str());
      }

      float ave_density = settings.num_particles / vol_total;

      // Suppose you divide the volume of the region (ie. cell, compartment...)
      // into equal size "bins" of equal size ("volume_gaussian_bin").
      // A fixed number of particles lie within this volume.
      // ...What is the probability that the number of particles in a bin
      //    does not exceed "extreme_density"?
      //  This is given by the (cumulative) Poisson distribution:
      //     prob = lambda^k * exp(-lambda) / k!
      //          = lambda^k * exp(-lambda) / Gamma(k+1)  <-continuous version
      // "k" is the number of particles in this Gaussian-shaped "bin"
      long double k = extreme_density * volume_gaussian_bin;
      // "lambda" is the expected number of particles in this "bin"
      long double lambda = ave_density * volume_gaussian_bin;

      long double prob_cdf_obvious_way = 0.0;
      if (settings.use_min_density) {
        // Calculate the culmulative probability of seeing this many
        // particles in the bin or less.  First try it the obvious way:

        for (long i=0; i <= floor(k); i++)
          prob_cdf_obvious_way += pow(lambda, i) * exp(-lambda) / tgamma(i+1.0);
      }
      else {
      // Calculate the culmulative probability of seeing this many
      // particles in the bin or more.  First try it the obvious way:
        for (long i=0; i < floor(k); i++)
          prob_cdf_obvious_way += pow(lambda, i) * exp(-lambda) / tgamma(i+1.0);
        prob_cdf_obvious_way = 1.0 - prob_cdf_obvious_way;
      }

      cerr << "####################################" << endl;
      cerr << "## DEBUG MESSAGES (PLEASE IGNORE) ##" << endl;
      cerr << "## location (in voxels) = ("
           << afXextreme[0] << ","
           << afXextreme[1] << ","
           << afXextreme[2] << ")" << endl;
      cerr << "## num_in_bin = " << k << endl;
      cerr << "## num_expected_in_bin = " << lambda << endl;
      cerr << "## prob_cdf_obvious_way = " << prob_cdf_obvious_way << endl;
      cerr << "####################################" << endl;

      long double prob_cdf;

      // However, in our case the number of particles in this "bin" is not (k)
      // necessarily an integer.  In that case use the continuous version of
      // (the cumulative distribution) of the Poisson distributio.
      // COMMENTING OUT: This requires the BOOST libraries
      // COMMENTING OUT: if (use_min_density)
      // COMMENTING OUT:   prob_cdf = gamma_q(k+1,lambda);
      // COMMENTING OUT: else
      // COMMENTING OUT:   prob_cdf = 1.0 - gamma_q(k,lambda);
      //    (gamma_q() is defined in boost)

      // The cumulative distribution of the Poisson distribution
      // is also given by the upper incomplete Gamma function
      // COMMENTING OUT: USE THE BOOST LIBRARY
      //double prob_cdf = gamma_q(floor(k)+1, lambda);
      // Unfortunately the  BOOST library can not be easily distributed with
      // this code.  (Even if I throw away most ofthe BOOST code using the
      // "bcp" utility, the remaining BOOST code exceeds the size of this
      // the code for this project by a factor of 50.)
      // Hence we will use the integer approximation ("obvious_way"):

      prob_cdf = prob_cdf_obvious_way;

      // ...Now what is the probability that the number of particles
      //    in ANY OF THE BINS does not exceed "extreme_density"?

      long double prob_total = 1.0 - pow((1.0 - prob_cdf), num_bins);

      double effective_bin_size = (pow(volume_gaussian_bin, 1.0/3) *
                                   voxel_width[0]);
      cout << prob_total
           << " " << extreme_density
           << " " << afXextreme[0]
           << " " << afXextreme[1]
           << " " << afXextreme[2]
           << " " << effective_bin_size
           << endl;

      // Discussion:
      //      What do we expect this number to be?
      // What is the expected value of "prob_total" for a randomly distributed
      // particles in the same number of bins.
      //
      // Answer: 0.5
      //
      // The proof is not specific to the kind of distribution we are using.
      // (a Poisson distribution)
      //
      // Proof:
      //
      // Let "c" be the cumulative probability distrubution for the entire
      // system (in this "prob_total", but it could also be "prob_cdf")
      //
      // Let p(x) be the probability density of measuring x.
      //
      // Let C(x) be the cumulative distribution of measuring something <= x.
      // (in this case the probability that none of the bins have
      //  a density exceeding x)  This means that:
      //
      // C(X) = \int_{-\infty}^X  dx  p(x)
      //
      // Let C^{-1}(c)  denote the inverse of C(x)
      // ==> C^{-1}(c) = x   if   C(x) = c
      //
      // What is the average value of c?
      //
      // <c> = \int_{-infty}^{\infty} * c(x) * p(x) * dx
      //
      //     = \int_0^1         c *  p(x(c)) * dx(c)/dc * dc
      //
      //     = \int_0^1 c * p( C^{-1}(c) ) * (d/dc) C^{-1}(c)  * dc
      //
      //     = \int_0^1 c * p( C^{-1}(c) ) * ( 1 / (d/dx) C( C^{-1}(c) ) ) * dc
      //
      //     = \int_0^1 c * p( C^{-1}(c) ) * ( 1 / p( C^{-1}(c) ) )  * dc
      //
      //     = \int_0^1 c * dc
      //
      //     = 0.5
      //
      // However this number will be lower if you repeat this multiple times
      // (using different bin-sizes ("sigma" values), for example)
      // because by repeating, you give it more opportunities to find
      // a more extreme value.  (Here we scan over a range of "sigma" values.)

    } // for (int i_sig = 0; i_sig < settings.vfSigma.size(); i_sig++) {...

    if ((settings.out_file_name != "") && (settings.vfSigma.size() == 1))
    {
      // Write the image file containing the filtered version for that sigma
      cerr << "writing a tomogram containing the most recently calculated density cloud"
           << endl;

      tomo_out.Write(settings.out_file_name);
      // (You can also use "file_stream << tomo_out;")
    }

  } //try {
  catch (const std::exception& e) {
    cerr << "\n" << e.what() << endl;
    exit(1);
  }

} // main()


// Copyright Nick Thompson 2017.
// Use, modification and distribution are subject to the
// Boost Software License, Version 1.0.
// (See accompanying file LICENSE_1_0.txt
// or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_MATH_SPECIAL_LEGENDRE_STIELTJES_HPP
#define BOOST_MATH_SPECIAL_LEGENDRE_STIELTJES_HPP

/*
 * Constructs the Legendre-Stieltjes polynomial of degree m.
 * The Legendre-Stieltjes polynomials are used to create extensions for Gaussian quadratures,
 * commonly called "Gauss-Konrod" quadratures.
 *
 * References:
 * Patterson, TNL. "The optimum addition of points to quadrature formulae." Mathematics of Computation 22.104 (1968): 847-856.
 */

#include <iostream>
#include <vector>
#include <boost/math/tools/roots.hpp>
#include <boost/math/special_functions/legendre.hpp>

namespace boost{
namespace math{

template<class Real>
class legendre_stieltjes
{
public:
    legendre_stieltjes(size_t m)
    {
        if (m == 0)
        {
           throw std::domain_error("The Legendre-Stieltjes polynomial is defined for order m > 0.\n");
        }
        m_m = static_cast<int>(m);
        std::ptrdiff_t n = m - 1;
        std::ptrdiff_t q;
        std::ptrdiff_t r;
        bool odd = n & 1;
        if (odd)
        {
           q = 1;
           r = (n-1)/2 + 2;
        }
        else
        {
           q = 0;
           r = n/2 + 1;
        }
        m_a.resize(r + 1);
        // We'll keep the ones-based indexing at the cost of storing a superfluous element
        // so that we can follow Patterson's notation exactly.
        m_a[r] = static_cast<Real>(1);
        // Make sure using the zero index is a bug:
        m_a[0] = std::numeric_limits<Real>::quiet_NaN();

        for (std::ptrdiff_t k = 1; k < r; ++k)
        {
            Real ratio = 1;
            m_a[r - k] = 0;
            for (std::ptrdiff_t i = r + 1 - k; i <= r; ++i)
            {
                // See Patterson, equation 12
                std::ptrdiff_t num = (n - q + 2*(i + k - 1))*(n + q + 2*(k - i + 1))*(n-1-q+2*(i-k))*(2*(k+i-1) -1 -q -n);
                std::ptrdiff_t den = (n - q + 2*(i - k))*(2*(k + i - 1) - q - n)*(n + 1 + q + 2*(k - i))*(n - 1 - q + 2*(i + k));
                ratio *= static_cast<Real>(num)/static_cast<Real>(den);
                m_a[r - k] -= ratio*m_a[i];
            }
        }
    }

    Real norm_sq() const
    {
        Real t = 0;
        bool odd = m_m & 1;
        for (size_t i = 1; i < m_a.size(); ++i)
        {
            if(odd)
            {
                t += 2*m_a[i]*m_a[i]/static_cast<Real>(4*i-1);
            }
            else
            {
                t += 2*m_a[i]*m_a[i]/static_cast<Real>(4*i-3);
            }
        }
        return t;
    }

    Real operator()(Real x) const
    {
        // Trivial implementation:
        // Em += m_a[i]*legendre_p(2*i - 1, x);  m odd
        // Em += m_a[i]*legendre_p(2*i - 2, x);  m even
        size_t r = m_a.size() - 1;
        Real p0 = 1;
        Real p1 = x;

        Real Em;
        bool odd = m_m & 1;
        if (odd)
        {
            Em = m_a[1]*p1;
        }
        else
        {
            Em = m_a[1]*p0;
        }

        unsigned n = 1;
        for (size_t i = 2; i <= r; ++i)
        {
            std::swap(p0, p1);
            p1 = boost::math::legendre_next(n, x, p0, p1);
            ++n;
            if (!odd)
            {
               Em += m_a[i]*p1;
            }
            std::swap(p0, p1);
            p1 = boost::math::legendre_next(n, x, p0, p1);
            ++n;
            if(odd)
            {
                Em += m_a[i]*p1;
            }
        }
        return Em;
    }

    Real prime(Real x) const
    {
        Real Em_prime = 0;

        for (size_t i = 1; i < m_a.size(); ++i)
        {
            if(m_m & 1)
            {
                Em_prime += m_a[i]*detail::legendre_p_prime_imp(static_cast<unsigned>(2*i - 1), x, policies::policy<>());
            }
            else
            {
                Em_prime += m_a[i]*detail::legendre_p_prime_imp(static_cast<unsigned>(2*i - 2), x, policies::policy<>());
            }
        }
        return Em_prime;
    }

    std::vector<Real> zeros() const
    {
        using boost::math::constants::half;

        std::vector<Real> stieltjes_zeros;
        std::vector<Real> legendre_zeros = legendre_p_zeros<Real>(m_m - 1);
        int k;
        if (m_m & 1)
        {
            stieltjes_zeros.resize(legendre_zeros.size() + 1, std::numeric_limits<Real>::quiet_NaN());
            stieltjes_zeros[0] = 0;
            k = 1;
        }
        else
        {
            stieltjes_zeros.resize(legendre_zeros.size(), std::numeric_limits<Real>::quiet_NaN());
            k = 0;
        }

        while (k < (int)stieltjes_zeros.size())
        {
            Real lower_bound;
            Real upper_bound;
            if (m_m & 1)
            {
                lower_bound = legendre_zeros[k - 1];
                if (k == (int)legendre_zeros.size())
                {
                    upper_bound = 1;
                }
                else
                {
                    upper_bound = legendre_zeros[k];
                }
            }
            else
            {
                lower_bound = legendre_zeros[k];
                if (k == (int)legendre_zeros.size() - 1)
                {
                    upper_bound = 1;
                }
                else
                {
                    upper_bound = legendre_zeros[k+1];
                }
            }

            // The root bracketing is not very tight; to keep weird stuff from happening
            // in the Newton's method, let's tighten up the tolerance using a few bisections.
            boost::math::tools::eps_tolerance<Real> tol(6);
            auto g = [&](Real t) { return this->operator()(t); };
            auto p = boost::math::tools::bisect(g, lower_bound, upper_bound, tol);

            Real x_nk_guess = p.first + (p.second - p.first)*half<Real>();
            boost::uintmax_t number_of_iterations = 500;

            auto f = [&] (Real x) { Real Pn = this->operator()(x);
                                    Real Pn_prime = this->prime(x);
                                    return std::pair<Real, Real>(Pn, Pn_prime); };

            const Real x_nk = boost::math::tools::newton_raphson_iterate(f, x_nk_guess,
                                                  p.first, p.second,
                                                  2*std::numeric_limits<Real>::digits10,
                                                  number_of_iterations);

            BOOST_ASSERT(p.first < x_nk);
            BOOST_ASSERT(x_nk < p.second);
            stieltjes_zeros[k] = x_nk;
            ++k;
        }
        return stieltjes_zeros;
    }

private:
    // Coefficients of Legendre expansion
    std::vector<Real> m_a;
    int m_m;
};

}}
#endif


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2013-2015 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef ERROREST_HH
#define ERROREST_HH

#include <vector>

#include <boost/fusion/include/for_each.hpp>
#include <boost/utility.hpp>
#include <boost/multi_array.hpp>

namespace Kaskade
{
  /**
   * \ingroup refcrit
   * \brief Base class for refinement criteria.
   */
  class RefinementCriterion
  {
  public:
    /**
     * \brief Computes the thresholds for refinement.
     * \param normalizedErrors a twodimensional array containing the normalized error contribution for each cell and
     *                         each variable, i.e. normalizedErrors[i][j] contains the error contribution of cell i
     *                         to the error in variable j, divided by the tolerance, such that the sum over i
     *                         should not exceed one (for acceptance)
     * \returns an array containing a threshold value such that any cell i for which normalizedErrors[i][j] > return[j]
     *          for any j is to be refined, or an empty array in case all variables are accurate enough
     */
    std::vector<double> threshold(boost::multi_array<double,2> const& normalizedErrors) const;

  private:
    virtual double computeThreshold(std::vector<double>& normalizedErrors, double totalError, int j) const = 0;
  };

  /**
   * \ingroup refcrit
   * \brief Refines the grid according to the normalized errors and the given thresholds.
   */
  template <class GridManager, class GridView>
  void markAndRefine(GridManager& gridManager, GridView const& gridView, boost::multi_array<double,2> const& normalizedErrors,
                     std::vector<double> const& threshold)
  {
    auto end = gridView.template end<0>();
    for (auto ci=gridView.template begin<0>(); ci!=end; ++ci)
    {
      auto idx = gridView.indexSet().index(*ci);
      for (int j=0; j<threshold.size(); ++j)
        if (normalizedErrors[idx][j] >= threshold[j])
        {
          gridManager.mark(1,*ci);
          break;
        }
    }
    gridManager.adaptAtOnce();
  }

  /**
   * \ingroup refcrit
   * \brief Fixed fraction refinement criterion.
   * Determines the refinement thresholds such that at least the specified fraction of the cells are
   * marked for refinement. This ensures that not too few cells are marked, which could lead to very many
   * refinement iterations, but may lead to erratic refinement if there are indeed only very local refinements
   * necessary.
   *
   * In very rare circumstances (namely when there are many cells with exactly the same local error estimate),
   * much more cells than the required fraction can be marked for refinement.
   */
  class FixedFractionCriterion: public RefinementCriterion
  {
  public:
    FixedFractionCriterion(double fraction = 0.2);

  private:
    double fraction;
    virtual double computeThreshold(std::vector<double>& normalizedErrors, double totalError, int j) const;
  };

  /**
   * \ingroup refcrit
   * \brief Bulk refinement criterion.
   * Determines the refinement thresholds such that approximately the specified fraction of the total error
   * is removed by the refinement (under the unrealistically optimistic assumption that refinement eliminates
   * the error completely in that cell - in reality, it's only reduced by a certain factor, so the total
   * error is reduced somewhat less).
   */
  class BulkCriterion: public RefinementCriterion
  {
  public:
    BulkCriterion(double fraction = 0.2);

  private:
    double fraction;
    virtual double computeThreshold(std::vector<double>& normalizedErrors, double totalError, int j) const;
  };

  /**
   * \ingroup refcrit
   * \brief Max value refinement criterion.
   * Determines the refinement thresholds such that all cells with an error contribution exceeding a certain
   * fraction of the maximum error contribution are refined.
   */
  class MaxValueCriterion: public RefinementCriterion
  {
  public:
    MaxValueCriterion(double fraction = 0.2);

  private:
    double fraction;
    virtual double computeThreshold(std::vector<double>& normalizedErrors, double totalError, int j) const;
  };

  /**
   * \ingroup refcrit
   * \brief Babuska-Rheinboldt refinement criterion.
   * Determines the refinement thresholds such that all cells with an error contribution exceeding the
   * expected error contribution of the worst cell \em after refinement will be marked. This tends to
   * equilibrate the error contributions quite fast.
   *
   * Note that the local convergence order should be estimated (important in the vicinity of singularities)
   * but is here assumed to be the fixed specified order.
   */
  class BabuskaRheinboldtCriterion: public MaxValueCriterion
  {
  public:
    BabuskaRheinboldtCriterion(std::vector<int> const& order);

  private:
    std::vector<int> order;
    virtual double computeThreshold(std::vector<double>& normalizedErrors, double totalError, int j) const;
  };

  //---------------------------------------------------------------------------------------------

  namespace ErrorestDetail
  {
    template <class GroupByCell>
    struct GroupedSummationCollector
    {
      GroupedSummationCollector(GroupByCell const& group_):
        group(group_)
      {}

      template <class Cell>
      int integrationOrder(Cell const& , int shapeFunctionOrder) const
      {
        return shapeFunctionOrder;
      }

      // need to define this because boost::multi_array does not perform assignment if the arrays have
      // different shape.
      GroupedSummationCollector<GroupByCell>& operator=(GroupedSummationCollector<GroupByCell> const& c)
      {
        auto shape = c.sums.shape();

        sums.resize(boost::extents[shape[0]][shape[1]]);
        sums = c.sums;
        group = c.group;

        return *this;
      }

      template <class CellPointer, class Index, class Sequence>
      void operator()(CellPointer const& ci, Index idx, double weight, Sequence const& x)
      {
        if (sums.num_elements()==0)
          sums.resize(boost::extents[group.nGroups][boost::fusion::size(x)]);
        int i = 0; // for_each needs constant functor...
        boost::fusion::for_each(x,Add(sums,group[idx],i,weight));
      }

      void join(GroupedSummationCollector<GroupByCell> const& c)
      {
        auto shape = c.sums.shape();
        auto myshape = sums.shape();

        if(c.sums.num_elements()==0) return;
        if (sums.num_elements()==0 || myshape[0]!=shape[0] || myshape[1]!=shape[1]) // not yet initialized
        {
          sums.resize(boost::extents[shape[0]][shape[1]]);
          sums = c.sums;                                  // -> do a simple copy
        }
        else
          for (size_t i=0; i<shape[0]; ++i)               // otherwise add up
            for (size_t j=0; j<shape[1]; ++j)
              sums[i][j] += c.sums[i][j];
      }

      boost::multi_array<double,2> sums;

    private:
      GroupByCell group;

      struct Add
      {
        Add(boost::multi_array<double,2>& sums_, int idx_, int& i_, double w_):
          sums(sums_), idx(idx_), i(i_) , w(w_)
        {}

        template <class T>
        void operator()(T const& t) const { sums[idx][i++] += w*t; }

      private:
        boost::multi_array<double,2>& sums;
        int idx;
        int& i;
        double w;
      };
    };

  }
}

#endif


#include "distances.h"
#include <EMD_wrapper.h>

#define EIGEN_NO_DEBUG
#define EIGEN_DONT_PARALLELIZE
#include <Eigen/Dense>

#include <cmath> // for std::isnormal
#include <limits>
#include <vector>

#if defined(WITH_ARRAYFIRE)
#include <cuda_runtime.h>
#if WITH_GPU_PROFILING
#include <nvtx3/nvToolsExt.h>
#endif
#define AF_DEFINE_CUDA_TYPES
#include <af/cuda.h>
#include <arrayfire.h>

#include "lp_distance.cuh"

using af::array;
#endif

using MatrixXd = Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>;

std::vector<Metric> expand_metrics(const ManifoldGenerator& generator, int E, Distance distance,
                                   const std::vector<Metric>& metrics)
{
  // Expand the 'metrics' vector now that we know the value of E.
  std::vector<Metric> expandedMetrics;

  // For the Wasserstein distance, it's more convenient to have one 'metric' for each variable (before taking lags).
  // However, for the L^1 / L^2 distances, it's more convenient to have one 'metric' for each individual
  // point of each observations, so metrics.size() == M.E_actual().
  if (distance == Distance::Wasserstein) {
    // Add a metric for the main variable and for the dt variable.
    // These are always treated as a continuous values (though perhaps in the future this will change).
    expandedMetrics.push_back(Metric::Diff);
    if (generator.E_dt(E) > 0) {
      expandedMetrics.push_back(Metric::Diff);
    }

    // Add in the metrics for the 'extra' variables as they were supplied to us.
    for (int k = 0; k < generator.numExtras(); k++) {
      expandedMetrics.push_back(metrics[k]);
    }
  } else {
    // Add metrics for the main variable and the dt variable and their lags.
    // These are always treated as a continuous values (though perhaps in the future this will change).
    for (int lagNum = 0; lagNum < E + generator.E_dt(E); lagNum++) {
      expandedMetrics.push_back(Metric::Diff);
    }

    // The user specified how to treat the extra variables.
    for (int k = 0; k < generator.numExtras(); k++) {
      int numLags = (k < generator.numExtrasLagged()) ? E : 1;
      for (int lagNum = 0; lagNum < numLags; lagNum++) {
        expandedMetrics.push_back(metrics[k]);
      }
    }
  }

  return expandedMetrics;
}

DistanceIndexPairs lazy_lp_distances(int Mp_i, const Options& opts, const Manifold& M, const Manifold& Mp)
{
  std::vector<int> inds;
  std::vector<double> dists;

  inds.reserve(M.numPoints());
  dists.reserve(M.numPoints());

  // We'll store the points we are comparing in the following two arrays.
  auto x = std::unique_ptr<double[]>(new double[M.E_actual()], std::default_delete<double[]>());
  auto y = std::unique_ptr<double[]>(new double[M.E_actual()], std::default_delete<double[]>());

  Mp.lazy_fill_in_point(Mp_i, y.get());

  const bool skipOtherPanels = opts.panelMode && (opts.idw < 0);

  // Compare every observation in the M manifold to the
  // Mp_i'th observation in the Mp manifold.
  for (int i = 0; i < M.numPoints(); i++) {

    if (skipOtherPanels && (M.panel(i) != Mp.panel(Mp_i))) {
      continue;
    }

    // Calculate the distance between M[i] and Mp[Mp_i]
    double dist_i = 0.0;

    M.lazy_fill_in_point(i, x.get());

    // If we have panel data and the M[i] / Mp[Mp_j] observations come from different panels
    // then add the user-supplied penalty/distance for the mismatch.
    if (opts.panelMode && opts.idw > 0) {
      dist_i += opts.idw * (M.panel(i) != Mp.panel(Mp_i));
    }

    for (int j = 0; j < M.E_actual(); j++) {
      // Get the sub-distance between M[i,j] and Mp[Mp_i, j]
      double dist_ij;

      // If either of these values is missing, the distance from
      // M[i,j] to Mp[Mp_i, j] is opts.missingdistance.
      // However, if the user doesn't specify this, then the entire
      // M[i] to Mp[Mp_i] distance is set as missing.
      if ((x[j] == MISSING_D) || (y[j] == MISSING_D)) {
        if (opts.missingdistance == 0) {
          dist_i = MISSING_D;
          break;
        } else {
          dist_ij = opts.missingdistance;
        }
      } else { // Neither M[i,j] nor Mp[Mp_i, j] is missing.
        // How do we compare them? Do we treat them like continuous values and subtract them,
        // or treat them like unordered categorical variables and just check if they're the same?
        if (opts.metrics[j] == Metric::Diff) {
          dist_ij = x[j] - y[j];
        } else { // Metric::CheckSame
          dist_ij = (x[j] != y[j]);
        }
      }

      if (opts.distance == Distance::MeanAbsoluteError) {
        dist_i += abs(dist_ij) / M.E_actual();
      } else { // Distance::Euclidean
        dist_i += dist_ij * dist_ij;
      }
    }

    if (dist_i != 0 && dist_i != MISSING_D) {
      if (opts.distance == Distance::MeanAbsoluteError) {
        dists.push_back(dist_i);
      } else { // Distance::Euclidean
        dists.push_back(sqrt(dist_i));
      }
      inds.push_back(i);
    }
  }

  return { inds, dists };
}

DistanceIndexPairs eager_lp_distances(int Mp_i, const Options& opts, const Manifold& M, const Manifold& Mp)
{
  std::vector<int> inds;
  std::vector<double> dists;

  inds.reserve(M.numPoints());
  dists.reserve(M.numPoints());

  const bool skipOtherPanels = opts.panelMode && (opts.idw < 0);

  // Compare every observation in the M manifold to the
  // Mp_i'th observation in the Mp manifold.
  for (int i = 0; i < M.numPoints(); i++) {

    if (skipOtherPanels && (M.panel(i) != Mp.panel(Mp_i))) {
      continue;
    }

    // Calculate the distance between M[i] and Mp[Mp_i]
    double dist_i = 0.0;

    // If we have panel data and the M[i] / Mp[Mp_j] observations come from different panels
    // then add the user-supplied penalty/distance for the mismatch.
    if (opts.panelMode && opts.idw > 0) {
      dist_i += opts.idw * (M.panel(i) != Mp.panel(Mp_i));
    }

    for (int j = 0; j < M.E_actual(); j++) {
      // Get the sub-distance between M[i,j] and Mp[Mp_i, j]
      double dist_ij;

      // If either of these values is missing, the distance from
      // M[i,j] to Mp[Mp_i, j] is opts.missingdistance.
      // However, if the user doesn't specify this, then the entire
      // M[i] to Mp[Mp_i] distance is set as missing.
      if ((M(i, j) == MISSING_D) || (Mp(Mp_i, j) == MISSING_D)) {
        if (opts.missingdistance == 0) {
          dist_i = MISSING_D;
          break;
        } else {
          dist_ij = opts.missingdistance;
        }
      } else { // Neither M[i,j] nor Mp[Mp_i, j] is missing.
        // How do we compare them? Do we treat them like continuous values and subtract them,
        // or treat them like unordered categorical variables and just check if they're the same?
        if (opts.metrics[j] == Metric::Diff) {
          dist_ij = M(i, j) - Mp(Mp_i, j);
        } else { // Metric::CheckSame
          dist_ij = (M(i, j) != Mp(Mp_i, j));
        }
      }

      if (opts.distance == Distance::MeanAbsoluteError) {
        dist_i += abs(dist_ij) / M.E_actual();
      } else { // Distance::Euclidean
        dist_i += dist_ij * dist_ij;
      }
    }

    if (dist_i != 0 && dist_i != MISSING_D) {
      if (opts.distance == Distance::MeanAbsoluteError) {
        dists.push_back(dist_i);
      } else { // Distance::Euclidean
        dists.push_back(sqrt(dist_i));
      }
      inds.push_back(i);
    }
  }

  return { inds, dists };
}

// This function compares the M(i,.) multivariate time series to the Mp(j,.) multivariate time series.
// The M(i,.) observation has data for E consecutive time points (e.g. time(i), time(i+1), ..., time(i+E-1)) and
// the Mp(j,.) observation corresponds to E consecutive time points (e.g. time(j), time(j+1), ..., time(j+E-1)).
// At each time instant we observe n >= 1 pieces of data.
// These may either be continuous data or unordered categorical data.
//
// The Wasserstein distance (using the 'curve-matching' strategy) is equivalent to the (minimum) cost of turning
// the first time series into the second time series. In a simple example, say E = 2 and n = 1, and
//         M(i,.) = [ 1, 2 ] and Mp(j,.) = [ 2, 2 ].
// To turn M(i,.) into Mp(j,.) the first element needs to be increased by 1, so the overall cost is
//         Wasserstein( M(i,.), Mp(j,.) ) = 1.
// The distance can also reorder the points, so for example say
//         M(i,.) = [ 1, 100 ] and Mp(j,.) = [ 100, 1 ].
// If we just change the 1 to 100 and the 100 to 1 then the cost of each is 99 + 99 = 198.
// However, Wasserstein can instead reorder these points at a cost of
//         Wasserstein( M(i,.), Mp(j,.) ) = 2 * gamma * (time(1)-time(2))
// so if the observations occur on a regular grid so time(i) = i then the distance will just be 2 * gamma.
//
// The return value of this function is a matrix which shows the pairwise costs associated to each
// potential Wasserstein solution. E.g. the (n,m) element of the returned matrix shows the cost
// of turning the individual point M(i, n) into Mp(j, m).
//
// When there are missing values in one or other observation, we can either ignore this time period
// and compute the Wasserstein for the mismatched regime where M(i,.) is of size len_i and Mp(j,.) is
// of size len_j, where len_i != len_j is possible. Alternatively, we can fill in the affected elements
// of the cost matrix with some user-supplied 'missingDistance' value and then len_i == len_j is upheld.
std::unique_ptr<double[]> wasserstein_cost_matrix(const Manifold& M, const Manifold& Mp, int i, int j,
                                                  const Options& opts, int& len_i, int& len_j)
{
  // The M(i,.) observation will be stored as one flat vector of length M.E_actual():
  // - the first M.E() observations will the lagged version of the main time series
  // - the next M.E() observations will be the lagged 'dt' time series (if it is included, i.e., if M.E_dt() > 0)
  // - the next n * M.E() observations will be the n lagged extra variables,
  //   so in total that is M.E_lagged_extras() = n * M.E() observations
  // - the remaining M.E_actual() - M.E() - M.E_dt() - M.E_lagged_extras() are the unlagged extras and the distance
  //   between those two vectors forms a kind of minimum distance which is added to the time-series curve matching
  //   Wasserstein distance.

  bool skipMissing = (opts.missingdistance == 0);

  // We'll store the points we are comparing in the following two arrays.
  auto x = std::unique_ptr<double[]>(new double[M.E_actual()], std::default_delete<double[]>());
  auto y = std::unique_ptr<double[]>(new double[M.E_actual()], std::default_delete<double[]>());

  if (opts.lowMemoryMode) {
    M.lazy_fill_in_point(i, x.get());
    Mp.lazy_fill_in_point(j, y.get());
  } else {
    M.eager_fill_in_point(i, x.get());
    Mp.eager_fill_in_point(j, y.get());
  }

  int numLaggedExtras = M.E_lagged_extras() / M.E();

  auto M_i = Eigen::Map<MatrixXd>(x.get(), 1 + (M.E_dt() > 0) + numLaggedExtras, M.E());
  auto Mp_j = Eigen::Map<MatrixXd>(y.get(), 1 + (M.E_dt() > 0) + numLaggedExtras, M.E());

  auto M_i_missing = (M_i.array() == M.missing()).colwise().any();
  auto Mp_j_missing = (Mp_j.array() == Mp.missing()).colwise().any();

  if (skipMissing) {
    // N.B. Can't .sum() a vector of bools to count them in Eigen.
    len_i = M.E() - M_i_missing.count();
    len_j = Mp.E() - Mp_j_missing.count();
  } else {
    len_i = M.E();
    len_j = Mp.E();
  }

  double gamma = 1.0;
  if (M.E_dt() > 0) {
    // Imagine the M_i time series as a plot, and calculate the
    // aspect ratio of this plot, so we can rescale the time variable
    // to get the user-supplied aspect ratio.
    double minData = std::numeric_limits<double>::max();
    double maxData = std::numeric_limits<double>::min();
    double maxTime = 0.0;
    for (int t = 0; t < M_i.cols(); t++) {
      if (M_i(0, t) != MISSING_D) {
        if (M_i(0, t) < minData) {
          minData = M_i(0, t);
        }
        if (M_i(0, t) > maxData) {
          maxData = M_i(0, t);
        }
      }
      if (M_i(1, t) != MISSING_D && M_i(1, t) > maxTime) {
        maxTime = M_i(1, t);
      }
    }

    double epsilon = 1e-6; // Some small number in case the following ratio gets wildly large/small
    gamma = opts.aspectRatio * (maxData - minData + epsilon) / (maxTime + epsilon);
  }

  int timeSeriesDim = M_i.rows();

  double unlaggedDist = 0.0;
  int numUnlaggedExtras = M.E_extras() - M.E_lagged_extras();
  for (int e = 0; e < numUnlaggedExtras; e++) {
    double x_e = x[M_i.size() + e];
    double y_e = y[Mp_j.size() + e];

    bool eitherMissing = (x_e == M.missing()) || (y_e == M.missing());

    if (eitherMissing) {
      unlaggedDist += opts.missingdistance;
    } else {
      if (opts.metrics[timeSeriesDim + e] == Metric::Diff) {
        unlaggedDist += abs(x_e - y_e);
      } else {
        unlaggedDist += (x_e != y_e);
      }
    }
  }

  // If we have panel data and the M[i] / Mp[j] observations come from different panels
  // then add the user-supplied penalty/distance for the mismatch.
  if (opts.panelMode && opts.idw > 0) {
    unlaggedDist += opts.idw * (M.panel(i) != Mp.panel(j));
  }

  auto flatCostMatrix = std::make_unique<double[]>(len_i * len_j);
  std::fill_n(flatCostMatrix.get(), len_i * len_j, unlaggedDist);
  Eigen::Map<MatrixXd> costMatrix(flatCostMatrix.get(), len_i, len_j);

  for (int k = 0; k < timeSeriesDim; k++) {
    int n = 0;
    for (int nn = 0; nn < M_i.cols(); nn++) {
      if (skipMissing && M_i_missing[nn]) {
        continue;
      }

      int m = 0;

      for (int mm = 0; mm < Mp_j.cols(); mm++) {
        if (skipMissing && Mp_j_missing[mm]) {
          continue;
        }
        double dist;
        bool eitherMissing = M_i_missing[nn] || Mp_j_missing[mm];

        if (eitherMissing) {
          dist = opts.missingdistance;
        } else {
          if (opts.metrics[k] == Metric::Diff) {
            dist = abs(M_i(k, nn) - Mp_j(k, mm));
          } else {
            dist = M_i(k, nn) != Mp_j(k, mm);
          }
        }

        // For the time data, we add in the 'gamma' scaling factor calculated earlier
        if ((M.E_dt() > 0) && (k == 1)) {
          dist *= gamma;
        }

        costMatrix(n, m) += dist;

        m += 1;
      }

      n += 1;
    }
  }

  return flatCostMatrix;
}

// TODO: Subtract the D(x,x) and D(y,y) parts from this.
double approx_wasserstein(double* C, int len_i, int len_j, double eps, double stopErr)
{
  Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> costMatrix(C, len_i, len_j);

  double r = 1.0 / len_i;
  double c = 1.0 / len_j;

  Eigen::MatrixXd K = Eigen::exp(-costMatrix.array() / eps);
  Eigen::MatrixXd Kp = len_i * K.array();

  Eigen::VectorXd u = Eigen::VectorXd::Ones(len_i) / len_i;
  Eigen::VectorXd v = Eigen::VectorXd::Ones(len_j) / len_j;

  int maxIter = 10000;
  for (int iter = 0; iter < maxIter; iter++) {

    v = c / (K.transpose() * u).array();
    u = 1.0 / (Kp * v).array();

    if (iter % 10 == 0) {
      // Compute right marginal (diag(u) K diag(v))^T1
      Eigen::VectorXd tempColSums = (u.asDiagonal() * K * v.asDiagonal()).colwise().sum();
      double LInfErr = (tempColSums.array() - c).abs().maxCoeff();
      if (LInfErr < stopErr) {
        break;
      }
    }
  }

  Eigen::MatrixXd transportPlan = u.asDiagonal() * K * v.asDiagonal();
  double dist = (transportPlan.array() * costMatrix.array()).sum();
  return dist;
}

double wasserstein(double* C, int len_i, int len_j)
{
  // Create vectors which are just 1/len_i and 1/len_j of length len_i and len_j.
  auto w_1 = std::make_unique<double[]>(len_i);
  std::fill_n(w_1.get(), len_i, 1.0 / len_i);
  auto w_2 = std::make_unique<double[]>(len_j);
  std::fill_n(w_2.get(), len_j, 1.0 / len_j);

  int maxIter = 10000;
  double cost;
  EMD_wrap(len_i, len_j, w_1.get(), w_2.get(), C, &cost, maxIter);
  return cost;
}

DistanceIndexPairs wasserstein_distances(int Mp_i, const Options& opts, const Manifold& M, const Manifold& Mp)
{
  std::vector<int> inds;
  std::vector<double> dists;

  inds.reserve(M.numPoints());
  dists.reserve(M.numPoints());

  const bool skipOtherPanels = opts.panelMode && (opts.idw < 0);

  // Compare every observation in the M manifold to the
  // Mp_i'th observation in the Mp manifold.
  for (int i = 0; i < M.numPoints(); i++) {

    if (skipOtherPanels && (M.panel(i) != Mp.panel(Mp_i))) {
      continue;
    }

    int len_i, len_j;
    auto C = wasserstein_cost_matrix(M, Mp, i, Mp_i, opts, len_i, len_j);

    if (len_i > 0 && len_j > 0) {
      double dist_i = wasserstein(C.get(), len_i, len_j);

      // Alternatively, the approximate version based on Sinkhorn's algorithm can be called with something like:
      // double dist_i = approx_wasserstein(C.get(), len_i, len_j, 0.1, 0.1)
      // In that case, the "std::isnormal" is really needed on the next line, as some
      // instability gives us some 'nan' distances using that method.

      if (dist_i != 0 && std::isnormal(dist_i)) {
        dists.push_back(dist_i);
        inds.push_back(i);
      }
    }
  }

  return { inds, dists };
}

/////////////////////////////////////////////////////////////// ArrayFire PORTED versions BEGIN HERE

#if defined(WITH_ARRAYFIRE)
DistanceIndexPairsOnGPU afLPDistances(const int numPredictions, const Options& opts, const ManifoldOnGPU& M,
                                      const ManifoldOnGPU& Mp, const af::array& metricOpts)
{
  constexpr bool useCustomKernel = true;

#if WITH_GPU_PROFILING
  auto range = nvtxRangeStartA(__FUNCTION__);
#endif

  const af_dtype cType = M.mdata.type();

  if (useCustomKernel) {
    af::array valids(M.numPoints, numPredictions, b8);
    af::array dists(M.numPoints, numPredictions, cType);
    if (cType == f64) {
      cuLPDistances(valids.device<char>(), dists.device<double>(), numPredictions,
                    opts.distance == Distance::MeanAbsoluteError, opts.panelMode, opts.idw, opts.missingdistance,
                    M.E_actual, M.numPoints, M.mdata.device<double>(), M.panel.device<int>(), Mp.mdata.device<double>(),
                    Mp.panel.device<int>(), metricOpts.device<char>(), afcu::getStream(0));
    } else if (cType == f32) {
      cuLPDistances(valids.device<char>(), dists.device<float>(), numPredictions,
                    opts.distance == Distance::MeanAbsoluteError, opts.panelMode, opts.idw, opts.missingdistance,
                    M.E_actual, M.numPoints, M.mdata.device<float>(), M.panel.device<int>(), Mp.mdata.device<float>(),
                    Mp.panel.device<int>(), metricOpts.device<char>(), afcu::getStream(0));
    }
    valids.unlock();
    dists.unlock();
    M.mdata.unlock();
    M.panel.unlock();
    Mp.mdata.unlock();
    Mp.panel.unlock();
    metricOpts.unlock();

#if WITH_GPU_PROFILING
    nvtxRangeEnd(range);
#endif
    return { valids, dists };
  } else {
    using af::array;
    using af::moddims;
    using af::select;
    using af::seq;
    using af::span;
    using af::sum;
    using af::tile;

    // Mp_i goes from 0 to numPredictions - 1
    // All coloumns of Manifold M are considered valid for batch operation

    const bool imdoZero = opts.missingdistance == 0;
    const int numLibraryPoints = M.numPoints;
    const int eacts = M.E_actual;

    array anyMCols, distsMat;
    {
      array predsM = tile(M.mdata, 1, 1, numPredictions);
      array predsMp = tile(moddims(Mp.mdata(span, seq(numPredictions)), eacts, 1, numPredictions), 1, numLibraryPoints);
      array diffMMp = predsM - predsMp;
      array compMMp = (predsM != predsMp).as(cType);
      array distMMp = select(tile(metricOpts, 1, numLibraryPoints, numPredictions), diffMMp, compMMp);
      array missing = predsM == MISSING_D || predsMp == MISSING_D;

      distsMat = (imdoZero ? distMMp : select(missing, opts.missingdistance, distMMp));
      anyMCols = anyTrue(missing, 0);

      if (opts.distance == Distance::MeanAbsoluteError) {
        distsMat = af::abs(distsMat) / double(eacts);
      } else {
        distsMat = distsMat * distsMat;
      }
    }
    if (opts.panelMode && opts.idw > 0) {
      array npPanelMp = tile(Mp.panel(seq(numPredictions)).T(), numLibraryPoints);
      array npPanelM = tile(M.panel, 1, numPredictions);
      array penalty = (opts.idw * (npPanelM != npPanelMp));
      array penalties = tile(moddims(penalty, 1, numLibraryPoints, numPredictions), eacts);

      distsMat += penalties;
    }
    array accDists = sum(distsMat, 0);
    array distances = select(anyMCols * imdoZero, double(MISSING_D), accDists);
    array valids = (distances != 0.0 && distances != double(MISSING_D));
    array dists = (opts.distance == Distance::MeanAbsoluteError ? distances : af::sqrt(distances));

    valids = moddims(valids, numLibraryPoints, numPredictions);
    dists = moddims(dists, numLibraryPoints, numPredictions);

#if WITH_GPU_PROFILING
    nvtxRangeEnd(range);
#endif
    return { valids, dists };
  }
}

array afWassersteinCostMatrix(const bool& skipMissing, const Options& opts, const array& metricOpts,
                              const ManifoldOnGPU& M, const array& M_i, const array& M_i_missing, const array& x,
                              const int& len_i, const array& Mp_j, const array& Mp_j_missing, const array& y,
                              const int& len_j, const bool arePanelIdsSame)
{
  using af::array;
  using af::constant;
  using af::dim4;
  using af::moddims;
  using af::seq;
  using af::span;
  using af::sum;
  using af::tile;
  using af::where;

  const af_dtype cType = M.mdata.type();

  double gamma = 1.0;
  if (M.E_dt > 0) {
    array firstColumn = M_i(span, 0);
    array nonMissings = firstColumn != MISSING_D;
    array validIndexs = where(nonMissings);
    array validValues = firstColumn(validIndexs);
    double minData = af::min<double>(validValues);
    double maxData = af::max<double>(validValues);
    double maxTime = af::max<double>(M_i(span, 1));

    // Some small number in case the following ratio gets wildly large/small
    constexpr double epsilon = 1e-6;

    gamma = opts.aspectRatio * (maxData - minData + epsilon) / (maxTime + epsilon);
  }

  const int timeSeriesDim = M_i.dims(1);
  double unlaggedDist = 0.0;
  {
    const int numUnlaggedExtras = M.E_extras - M.E_lagged_extras;

    array eitherMissing = (x == M.missing || y == M.missing);

    array cond = metricOpts(seq(timeSeriesDim, numUnlaggedExtras + timeSeriesDim - 1));
    array ulDists = (cond * af::abs(x - y) + (1 - cond) * (x != y).as(cType));
    array dists = (eitherMissing * opts.missingdistance + (1 - eitherMissing) * ulDists);
    unlaggedDist = sum<double>(dists);
  }

  if (opts.panelMode && opts.idw > 0) {
    unlaggedDist += opts.idw * arePanelIdsSame;
  }

  const seq timeSeries(timeSeriesDim);
  array costMatrix = af::constant(unlaggedDist, len_j, len_i, cType);

  array cmIsMetricDiff = metricOpts(timeSeries);
  cmIsMetricDiff = tile(moddims(cmIsMetricDiff, 1, 1, timeSeriesDim), len_j, len_i);

  if (skipMissing) {
    // In this case: Unless both entries are available, no need to process anything else
    array idxM_i_missing = where(!M_i_missing);
    array idxMp_j_missing = where(!Mp_j_missing);

    array Mp_j_k = moddims(Mp_j(idxMp_j_missing, timeSeries), len_j, 1, timeSeriesDim);
    array M_i_k = moddims(M_i(idxM_i_missing, timeSeries), 1, len_i, timeSeriesDim);
    array cmMp_j = tile(Mp_j_k, 1, len_i);
    array cmM_i = tile(M_i_k, len_j);
    array cmDiff = select(cmIsMetricDiff, af::abs(cmMp_j - cmM_i), (cmM_i != cmMp_j).as(cType));

    if (M.E_dt > 0) {
      // For time series k = 1, scale by gamma
      cmDiff(span, span, 1) *= gamma;
    }
    costMatrix += sum(cmDiff, 2); // Add results to unlaggedDist
  } else {
    array Mp_j_missingT = tile(Mp_j_missing, 1, M_i_missing.dims(0)); // cost matrix shape
    array M_i_missingT = tile(M_i_missing.T(), Mp_j_missing.dims(0)); // cost matrix shape
    array eitherMissing = (M_i_missingT || Mp_j_missingT);            // one of the entries missing
    array eitherMissingT = tile(eitherMissing, 1, 1, timeSeriesDim);  // [len_j len_i timeSeriesDim 1]

    array Mp_j_k = moddims(Mp_j(span, timeSeries), len_j, 1, timeSeriesDim);
    array M_i_k = moddims(M_i(span, timeSeries), 1, len_i, timeSeriesDim);
    array cmMp_j = tile(Mp_j_k, 1, len_i);
    array cmM_i = tile(M_i_k, len_j);
    array cmDiff = select(cmIsMetricDiff, af::abs(cmMp_j - cmM_i), (cmM_i != cmMp_j).as(cType));
    array cmDist = select(eitherMissingT, opts.missingdistance, cmDiff);

    if (M.E_dt > 0) {
      // For time series k = 1, scale by gamma
      cmDist(span, span, 1) *= gamma;
    }
    costMatrix += sum(cmDist, 2); // Add results to unlaggedDist
  }

  return costMatrix;
}

DistanceIndexPairs afWassersteinDistances(int Mp_i, const Options& opts, const Manifold& hostM, const Manifold& hostMp,
                                          const ManifoldOnGPU& M, const ManifoldOnGPU& Mp,
                                          const std::vector<int>& inpInds, const af::array& metricOpts)
{
  using af::anyTrue;
  using af::seq;

  const bool skipMissing = (opts.missingdistance == 0);

  // Precompute values that are iteration invariant
  const seq mpjRange0(Mp.E_x);
  const seq mpjRange1(Mp_i, Mp_i + 1 + (Mp.E_dt > 0) + Mp.E_lagged_extras / Mp.E_x);

  const array Mp_j = Mp.mdata(mpjRange0, mpjRange1);

  const array Mp_j_missing = anyTrue(Mp_j == Mp.missing, 1);

  const int numUnlaggedExtrasEnd = M.E_extras - M.E_lagged_extras;

  const seq xseq(M.E_x + M.E_dt + M.E_lagged_extras, M.E_x + M.E_dt + M.E_lagged_extras + numUnlaggedExtrasEnd - 1);
  const seq yseq(Mp.E_x + Mp.E_dt + Mp.E_lagged_extras,
                 Mp.E_x + Mp.E_dt + Mp.E_lagged_extras + numUnlaggedExtrasEnd - 1);

  const array y = Mp.mdata(yseq, Mp_i);
  const int len_j = (skipMissing ? Mp.E_x - af::sum<int>(Mp_j_missing) : Mp.E_x);

  // Since both len_i and len_j should be greater than zero
  // to collect valid indices and respective distances, just return empty handed
  if (len_j <= 0) {
    return {};
  }

  // Return Items
  std::vector<int> inds;
  std::vector<double> dists;

  // Compare every observation in the M manifold to the Mp_i'th observation in the Mp manifold.
  for (int i : inpInds) {
    const seq miRange0(M.E_x);
    const seq miRange1(i, i + 1 + (M.E_dt > 0) + M.E_lagged_extras / M.E_x);

    array M_i = M.mdata(miRange0, miRange1);

    array M_i_missing = anyTrue(M_i == M.missing, 1);

    array x = M.mdata(xseq, i);

    const int len_i = (skipMissing ? M.E_x - af::sum<int>(M_i_missing) : M.E_x);

    if (len_i > 0) { // Short-check for len_j already passed
      // TODO I think afWassersteinCostMatrix can be further vectorized to run for all i's
      array cm = afWassersteinCostMatrix(skipMissing, opts, metricOpts, M, M_i, M_i_missing, x, len_i, Mp_j,
                                         Mp_j_missing, y, len_j, (hostM.panel(i) != hostMp.panel(Mp_i)));
      std::vector<double> C(cm.elements());
      cm.host(C.data());
      double dist_i = wasserstein(C.data(), len_i, len_j);

      // Alternative: approximate version based on Sinkhorn's algorithm
      // double dist_i = approx_wasserstein(C.get(), len_i, len_j, 0.1, 0.1)
      // In that case, the "std::isnormal" is really needed on the next line, as some
      // instability gives us some 'nan' distances using that method.
      if (dist_i != 0 && std::isnormal(dist_i)) {
        dists.push_back(dist_i);
        inds.push_back(i);
      }
    }
  }

  return { inds, dists };
}
#endif


/**
 * ****************************************************************************
 * Copyright (c) 2015, Robert Lukierski.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 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.
 *
 * 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.
 *
 * ****************************************************************************
 * Simple Kalman Filter.
 * ****************************************************************************
 */

#ifndef VISIONCORE_MATH_KALMAN_HPP
#define VISIONCORE_MATH_KALMAN_HPP

#include <VisionCore/Platform.hpp>

#include <Eigen/LU>

namespace vc
{

namespace math
{

#if 0 // TODO FIXME
template<int dimension, class StateType>
class ConstantProcess {
public:
    typedef Eigen::Matrix<double, StateType::DIM, 1> VecState;
    typedef Eigen::Matrix<double, StateType::DIM, StateType::DIM> MatStateState;

    /// Noise per second, per component
    VecState sigma;
    MatStateState noise;
    MatStateState jacobian;

    inline ConstantProcess() :
    sigma(VecState::Zero()),
    noise(MatStateState::Zero()),
    jacobian(MatStateState::Identity())
    {}

    MatStateState const& getJacobian(StateType const& /*state*/, const double /*dt*/) {
        /// No change over time due to process model
        return jacobian;
    }

    void updateState(StateType const& /*state*/, const double /*dt*/) {
        /// no-op - no change due to process model
    }

    MatStateState const& getNoiseCovariance(const double dt) {
        noise = (dt * sigma).asDiagonal();
        return noise;
    }

    void updateFromMeasurement(StateType & state, VecState const & innovation) {
        state.x += innovation;
    }

};

template<class StateType>
class AbsoluteMeasurement {
public:
    static const int DIM = StateType::DIM;
    typedef Eigen::Matrix<double, DIM, 1> VecMeas;
    typedef Eigen::Matrix<double, DIM, 1> VecState;
    typedef Eigen::Matrix<double, DIM, DIM> MatMeasMeas;
    typedef Eigen::Matrix<double, DIM, DIM> MatMeasState;

    VecMeas measurement;

    MatMeasState jacobian;
    MatMeasMeas covariance;

    AbsoluteMeasurement() :
    measurement(VecMeas::Zero()),
    jacobian(MatMeasState::Identity()),
    covariance(MatMeasMeas::Identity()) {}

    MatMeasState const& getJacobian(StateType const& /*state*/) {
        return jacobian;
    }

    /// Measurement noise covariance, aka uncertainty
    /// in measurement
    MatMeasMeas const& getCovariance(StateType const& /*state*/) {
        return covariance;
    }

    VecState const getInnovation(StateType const& state) {
        return measurement - state.x;
    }

};

template<typename T, typename TimeT = T>
class KalmanFilter
{
public:
    typedef T Scalar;
    typedef TimeT TimeType;

    template<unsigned int dim>
    class StateT
    {
    public:
        static const int Dimension = dim;
        typedef Eigen::Matrix<Scalar, Dimension, 1> StateVectorT;
        typedef Eigen::Matrix<Scalar, Dimension, Dimension> CovarianceMatrixT;

        StateVectorT StateVector;
        CovarianceMatrixT Covariance;

        StateT() : StateVector(StateVectorT::Zero()), Covariance(CovarianceMatrixT::Identity())
        {

        }
    };

    template<unsigned int sdim, unsigned int mdim>
    class MeasurementT
    {
    public:
        typedef StateT<sdim> StateType;
        static const int Dimension = mdim;
        typedef Eigen::Matrix<Scalar, Dimension, 1> MeasurementVectorT;
        typedef Eigen::Matrix<Scalar, StateType::Dimension, 1> StateVectorT;
        typedef Eigen::Matrix<Scalar, Dimension, Dimension> JacobianMatrixT;
        typedef Eigen::Matrix<Scalar, Dimension, StateType::Dimension> HMatrixT;

        MeasurementT()
        {

        }
    };

    template<unsigned int dim>
    void predict(StateT<dim>& state, TimeType dt)
    {
        const StateT<dim>::CovarianceMatrixT A(processModel.getJacobian(state, dt));
        state.Covariance = A * state.Covariance * A.transpose() + processModel.getNoiseCovariance(dt);
        /// @todo symmetrize?
        processModel.updateState(state, dt);
    }

    template<class MeasurementType>
    void correct(MeasurementType & m)
    {
        typedef Eigen::Matrix<Scalar, MeasurementType::DIM, StateType::DIM> MatMeasState;
        typedef Eigen::Matrix<Scalar, MeasurementType::DIM, MeasurementType::DIM> MatMeasMeas;
        typedef Eigen::Matrix<Scalar, MeasurementType::DIM, 1> VecMeas;
        /// @todo implement
        const MatMeasState & H = m.getJacobian(state);
        const MatMeasMeas & R = m.getCovariance(state);
        const VecState innovation = m.getInnovation(state);
        const MatMeasMeas S = H * state.covariance * H.transpose() + R;

        Eigen::eikfLUType<MatMeasMeas> luOfS = S.eikfLUFunc();
        MatStateState K = state.covariance * H.transpose() * luOfS.inverse();
        processModel.updateFromMeasurement(state, K * innovation);
        state.covariance = (MatStateState::Identity() - K * H) * state.covariance;

    }

};
#endif

}

}

#endif // VISIONCORE_MATH_KALMAN_HPP


//=======================================================================
// Copyright (c)
//
// Distributed under the Boost Software License, Version 1.0. (See
// accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)
//=======================================================================
/**
 * @file scheduling_jobs.hpp
 * @brief
 * @author Robert Rosolek
 * @version 1.0
 * @date 2013-11-19
 */
#ifndef PAAL_SCHEDULING_JOBS_HPP
#define PAAL_SCHEDULING_JOBS_HPP

#define BOOST_RESULT_OF_USE_DECLTYPE

#include "paal/data_structures/fraction.hpp"
#include "paal/utils/functors.hpp"
#include "paal/utils/type_functions.hpp"
#include "paal/utils/irange.hpp"
#include "paal/utils/assign_updates.hpp"

#include <boost/range/algorithm/copy.hpp>
#include <boost/range/algorithm/sort.hpp>
#include <boost/range/counting_range.hpp>
#include <boost/range/iterator_range.hpp>
#include <boost/range/numeric.hpp>

#include <algorithm>
#include <cassert>
#include <iterator>
#include <numeric>
#include <random>
#include <utility>
#include <vector>

namespace paal {
namespace greedy {

namespace detail {

template <class MachineIterator, class JobIterator, class GetSpeed,
          class GetLoad>
struct sched_traits {
    typedef typename std::iterator_traits<MachineIterator>::reference
        machine_reference;
    typedef typename std::iterator_traits<JobIterator>::reference job_reference;
    typedef pure_result_of_t<GetSpeed(machine_reference)> speed_t;
    typedef pure_result_of_t<GetLoad(job_reference)> load_t;
    typedef data_structures::fraction<load_t, speed_t> frac_t;
};

template <class MachineIterator, class JobIterator, class GetSpeed,
          class GetLoad, class Traits = sched_traits<
                             MachineIterator, JobIterator, GetSpeed, GetLoad>>
typename Traits::frac_t calculate_bound(const MachineIterator mfirst,
                                        const MachineIterator mlast,
                                        const JobIterator jfirst,
                                        const JobIterator jlast,
                                        GetSpeed get_speed, GetLoad get_load) {
    typedef typename Traits::speed_t Speed;
    typedef typename Traits::load_t Load;
    typedef typename Traits::frac_t Frac;

    auto jobs_num = jlast - jfirst;
    auto machines_num = mlast - mfirst;

    std::vector<Speed> speed_sum(machines_num);
    std::transform(mfirst, mlast, speed_sum.begin(), get_speed);
    boost::partial_sum(speed_sum, speed_sum.begin());

    std::vector<Load> load_sum(jobs_num);
    std::transform(jfirst, jlast, load_sum.begin(), get_load);
    boost::partial_sum(load_sum, load_sum.begin());

    typedef decltype(machines_num) MachinesNumType;
    assert(jobs_num > 0 && machines_num > 0);
    Frac result(get_load(*jfirst), get_speed(*mfirst));
    for (auto jobID : irange(jobs_num)) {
        Load load = get_load(jfirst[jobID]);
        auto get_single = [ = ](MachinesNumType i) {
            return Frac(load, get_speed(mfirst[i]));
        };
        auto get_summed = [&](MachinesNumType i) {
            return Frac(load_sum[jobID], speed_sum[i]);
        };
        auto condition = [ = ](MachinesNumType i) {
            return get_summed(i) >= get_single(i);
        };
        auto machines_ids = boost::counting_range(
            static_cast<MachinesNumType>(0), machines_num);
        // current range based version in boost is broken
        // should be replaced when released
        // https://github.com/boostorg/algorithm/pull/4
        auto it = std::partition_point(machines_ids.begin(), machines_ids.end(),
                                       condition);
        MachinesNumType machineID =
            (it != machines_ids.end()) ? *it : machines_num - 1;
        auto getMax = [ = ](MachinesNumType i) {
            return std::max(get_single(i), get_summed(i));
        };

        Frac candidate = getMax(machineID);
        if (machineID != 0) {
            assign_min(candidate, getMax(machineID - 1));
        }
        assign_max(result, candidate);
    }
    return result;
}

template <class MachineIterator, class JobIterator, class OutputIterator,
          class GetSpeed, class GetLoad, class RoundFun>
void schedule(MachineIterator mfirst, MachineIterator mlast, JobIterator jfirst,
              JobIterator jlast, OutputIterator result, GetSpeed get_speed,
              GetLoad get_load, RoundFun round) {
    typedef sched_traits<MachineIterator, JobIterator, GetSpeed, GetLoad>
        Traits;
    typedef typename Traits::speed_t Speed;
    typedef typename Traits::load_t Load;

    if (mfirst == mlast || jfirst == jlast) {
        return;
    }

    std::vector<MachineIterator> machines;
    boost::copy(boost::counting_range(mfirst, mlast),
                std::back_inserter(machines));
    auto get_speed_from_iterator = utils::make_lift_iterator_functor(get_speed);
    boost::sort(machines, utils::make_functor_to_comparator(
                              get_speed_from_iterator, utils::greater{}));

    std::vector<JobIterator> jobs;
    boost::copy(boost::counting_range(jfirst, jlast), std::back_inserter(jobs));
    auto get_load_from_iterator = utils::make_lift_iterator_functor(get_load);
    boost::sort(jobs, utils::make_functor_to_comparator(get_load_from_iterator,
                                                        utils::greater{}));

    auto bound = detail::calculate_bound(
        machines.begin(), machines.end(), jobs.begin(), jobs.end(),
        get_speed_from_iterator, get_load_from_iterator);
    Load bound_load = bound.num;
    Speed bound_speed = bound.den;
    Load current_load{};
    auto emit = [&result](MachineIterator miter, JobIterator jiter) {
        *result = std::make_pair(miter, jiter);
        ++result;
    };
    auto job_iter = jobs.begin();
    for (auto machine_iter = machines.begin(); machine_iter != machines.end();
         ++machine_iter) {
        auto &&machine = *(*machine_iter);
        Speed speed = get_speed(machine);
        while (job_iter != jobs.end()) {
            auto &&job = *(*job_iter);
            Load job_load = get_load(job) * bound_speed,
                 new_load = current_load + job_load;
            assert(new_load <= bound_load * (2 * speed));
            if (bound_load * speed < new_load) {
                Load frac_load = bound_load * speed - current_load;
                if (round(frac_load, job_load)) {
                    emit(*machine_iter, *job_iter);
                } else {
                    auto next_machine_iter = std::next(machine_iter);
                    assert(next_machine_iter != machines.end());
                    emit(*next_machine_iter, *job_iter);
                }
                ++job_iter;
                current_load = job_load - frac_load;
                break;
            }
            emit(*machine_iter, *job_iter);
            ++job_iter;
            current_load = new_load;
        }
    }
    assert(job_iter == jobs.end());
}
} //!detail

/*
 * @brief This is deterministic solve scheduling jobs on machines with different
 * speeds problem and return schedule
 *
 * Example:
 *  \snippet scheduling_jobs_example.cpp Scheduling Jobs Example
 *
 * example file is scheduling_jobs_example.cpp
 *
 * @param mfirst
 * @param mlast
 * @param jfirst
 * @param jlast
 * @param result
 * @param get_speed
 * @param get_load
 * @tparam MachineIterator
 * @tparam JobIterator
 * @tparam OutputIterator
 * @tparam GetSpeed
 * @tparam GetLoad
 */
template <class MachineIterator, class JobIterator, class OutputIterator,
          class GetSpeed, class GetLoad>
void schedule_deterministic(const MachineIterator mfirst,
                            const MachineIterator mlast,
                            const JobIterator jfirst, const JobIterator jlast,
                            OutputIterator result, GetSpeed get_speed,
                            GetLoad get_load) {
    detail::schedule(mfirst, mlast, jfirst, jlast, result, get_speed, get_load,
                     utils::always_true{});
}

/*
 * @brief This is randomized solve scheduling jobs on machines with different
 * speeds problem and return schedule.
 *
 * Example:
 *  \snippet scheduling_jobs_example.cpp Scheduling Jobs Example
 *
 * example file is scheduling_jobs_example.cpp
 *
 * @param mfirst
 * @param mlast
 * @param jfirst
 * @param jlast
 * @param result
 * @param get_speed
 * @param get_load
 * @param gen
 * @tparam MachineIterator
 * @tparam JobIterator
 * @tparam OutputIterator
 * @tparam GetSpeed
 * @tparam GetLoad
 * @tparam RandomNumberGenerator
 */
template <class MachineIterator, class JobIterator, class OutputIterator,
          class GetSpeed, class GetLoad,
          class RandomNumberGenerator = std::default_random_engine>
void schedule_randomized(const MachineIterator mfirst,
                         const MachineIterator mlast, const JobIterator jfirst,
                         const JobIterator jlast, OutputIterator result,
                         GetSpeed get_speed, GetLoad get_load,
                         RandomNumberGenerator &&gen =
                             std::default_random_engine(97345631u)) {
    typedef typename detail::sched_traits<MachineIterator, JobIterator,
                                          GetSpeed, GetLoad> Traits;
    double alpha = std::uniform_real_distribution<double>()(gen);
    auto round = [alpha](typename Traits::load_t fractional_load,
                         typename Traits::load_t total_load) {
        return total_load * alpha < fractional_load;
    };
    detail::schedule(mfirst, mlast, jfirst, jlast, result, get_speed, get_load,
                     round);
}

} //!greedy
} //!paal

#endif // PAAL_SCHEDULING_JOBS_HPP