// The contents of this file are in the public domain. See LICENSE_FOR_EXAMPLE_PROGRAMS.txt
/*
    This is an example illustrating the use of the tools in dlib for doing distribution
    estimation or detecting anomalies using one-class support vector machines.

    Unlike regular classifiers, these tools take unlabeled points and try to learn what
    parts of the feature space normally contain data samples and which do not.  Typically
    you use these tools when you are interested in finding outliers or otherwise
    identifying "unusual" data samples.

    In this example, we will sample points from the sinc() function to generate our set of
    "typical looking" points.  Then we will train some one-class classifiers and use them
    to predict if new points are unusual or not.  In this case, unusual means a point is
    not from the sinc() curve.
*/

#include <iostream>
#include <vector>
#include <dlib/svm.h>
#include <dlib/gui_widgets.h>
#include <dlib/array2d.h>
#include <dlib/image_transforms.h>

using namespace std;
using namespace dlib;

// Here is the sinc function we will be trying to learn with the one-class SVMs
double sinc(double x)
{
    if (x == 0)
        return 2;
    return 2*sin(x)/x;
}

int main()
{
    // We will use column vectors to store our points.  Here we make a convenient typedef
    // for the kind of vector we will use.
    typedef matrix<double,0,1> sample_type;

    // Then we select the kernel we want to use.  For our present problem the radial basis
    // kernel is quite effective.
    typedef radial_basis_kernel<sample_type> kernel_type;

    // Now make the object responsible for training one-class SVMs.
    svm_one_class_trainer<kernel_type> trainer;
    // Here we set the width of the radial basis kernel to 4.0.  Larger values make the
    // width smaller and give the radial basis kernel more resolution.  If you play with
    // the value and observe the program output you will get a more intuitive feel for what
    // that means.
    trainer.set_kernel(kernel_type(4.0));

    // Now sample some 2D points.  The points will be located on the curve defined by the
    // sinc() function.
    std::vector<sample_type> samples;
    sample_type m(2);
    for (double x = -15; x <= 8; x += 0.3)
    {
        m(0) = x;
        m(1) = sinc(x);
        samples.push_back(m);
    }

    // Now train a one-class SVM.  The result is a function, df(), that outputs large
    // values for points from the sinc() curve and smaller values for points that are
    // anomalous (i.e. not on the sinc() curve in our case).
    decision_function<kernel_type> df = trainer.train(samples);

    // So for example, let's look at the output from some points on the sinc() curve.
    cout << "Points that are on the sinc function:\n";
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -0;   m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -0.5; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -4.1; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;
    m(0) = -0.5; m(1) = sinc(m(0)); cout << "   " << df(m) << endl;

    cout << endl;
    // Now look at some outputs for points not on the sinc() curve.  You will see that
    // these values are all notably smaller.
    cout << "Points that are NOT on the sinc function:\n";
    m(0) = -1.5; m(1) = sinc(m(0))+4;   cout << "   " << df(m) << endl;
    m(0) = -1.5; m(1) = sinc(m(0))+3;   cout << "   " << df(m) << endl;
    m(0) = -0;   m(1) = -sinc(m(0));    cout << "   " << df(m) << endl;
    m(0) = -0.5; m(1) = -sinc(m(0));    cout << "   " << df(m) << endl;
    m(0) = -4.1; m(1) = sinc(m(0))+2;   cout << "   " << df(m) << endl;
    m(0) = -1.5; m(1) = sinc(m(0))+0.9; cout << "   " << df(m) << endl;
    m(0) = -0.5; m(1) = sinc(m(0))+1;   cout << "   " << df(m) << endl;

    // The output is as follows:
    /*
    Points that are on the sinc function:
        0.000389691
        0.000389691
        -0.000239037
        -0.000179978
        -0.000178491
        0.000389691
        -0.000179978

    Points that are NOT on the sinc function:
        -0.269389
        -0.269389
        -0.269389
        -0.269389
        -0.269389
        -0.239954
        -0.264318
    */

    // So we can see that in this example the one-class SVM correctly indicates that
    // the non-sinc points are definitely not points from the sinc() curve.

    // It should be noted that the svm_one_class_trainer becomes very slow when you have
    // more than 10 or 20 thousand training points.  However, dlib comes with very fast SVM
    // tools which you can use instead at the cost of a little more setup.  In particular,
    // it is possible to use one of dlib's very fast linear SVM solvers to train a one
    // class SVM.  This is what we do below.  We will train on 115,000 points and it only
    // takes a few seconds with this tool!
    //
    // The first step is constructing a feature space that is appropriate for use with a
    // linear SVM.  In general, this is quite problem dependent.  However, if you have
    // under about a hundred dimensions in your vectors then it can often be quite
    // effective to use the empirical_kernel_map as we do below (see the
    // empirical_kernel_map documentation and example program for an extended discussion of
    // what it does).
    //
    // But putting the empirical_kernel_map aside, the most important step in turning a
    // linear SVM into a one-class SVM is the following.  We append a -1 value onto the end
    // of each feature vector and then tell the trainer to force the weight for this
    // feature to 1.  This means that if the linear SVM assigned all other weights a value
    // of 0 then the output from a learned decision function would always be -1.  The
    // second step is that we ask the SVM to label each training sample with +1.  This
    // causes the SVM to set the other feature weights such that the training samples have
    // positive outputs from the learned decision function.  But the starting bias for all
    // the points in the whole feature space is -1.  The result is that points outside our
    // training set will not be affected, so their outputs from the decision function will
    // remain close to -1.

    empirical_kernel_map<kernel_type> ekm;
    ekm.load(trainer.get_kernel(),samples);

    samples.clear();
    std::vector<double> labels;
    // make a vector with just 1 element in it equal to -1.
    sample_type bias(1);
    bias = -1;
    sample_type augmented;
    // This time sample 115,000 points from the sinc() function.
    for (double x = -15; x <= 8; x += 0.0002)
    {
        m(0) = x;
        m(1) = sinc(x);
        // Apply the empirical_kernel_map transformation and then append the -1 value
        augmented = join_cols(ekm.project(m), bias);
        samples.push_back(augmented);
        labels.push_back(+1);
    }
    cout << "samples.size(): "<< samples.size() << endl;

    // The svm_c_linear_dcd_trainer is a very fast SVM solver which only works with the
    // linear_kernel.  It has the nice feature of supporting this "force_last_weight_to_1"
    // mode we discussed above.
    svm_c_linear_dcd_trainer<linear_kernel<sample_type> > linear_trainer;
    linear_trainer.force_last_weight_to_1(true);

    // Train the SVM
    decision_function<linear_kernel<sample_type> > df2 = linear_trainer.train(samples, labels);

    // Here we test it as before, again we note that points from the sinc() curve have
    // large outputs from the decision function.  Note also that we must remember to
    // transform the points in exactly the same manner used to construct the training set
    // before giving them to df2() or the code will not work.
    cout << "Points that are on the sinc function:\n";
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0;   m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0.5; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -4.1; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -1.5; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0.5; m(1) = sinc(m(0)); cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;

    cout << endl;
    // Again, we see here that points not on the sinc() function have small values.
    cout << "Points that are NOT on the sinc function:\n";
    m(0) = -1.5; m(1) = sinc(m(0))+4;   cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -1.5; m(1) = sinc(m(0))+3;   cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0;   m(1) = -sinc(m(0));    cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0.5; m(1) = -sinc(m(0));    cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -4.1; m(1) = sinc(m(0))+2;   cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -1.5; m(1) = sinc(m(0))+0.9; cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;
    m(0) = -0.5; m(1) = sinc(m(0))+1;   cout << "   " << df2(join_cols(ekm.project(m),bias)) << endl;

    // The output is as follows:
    /*
    Points that are on the sinc function:
        1.00454
        1.00454
        1.00022
        1.00007
        1.00371
        1.00454
        1.00007

    Points that are NOT on the sinc function:
        -1
        -1
        -1
        -1
        -0.999998
        -0.781231
        -0.96242
    */

    // Finally, to help you visualize what is happening here we are going to plot the
    // response of the one-class classifiers on the screen.  The code below creates two
    // heatmap images which show the response.  In these images you can clearly see where
    // the algorithms have identified the sinc() curve.  The hotter the pixel looks, the
    // larger the value coming out of the decision function and therefore the more "normal"
    // it is according to the classifier.
    const long size = 500;
    array2d<double> img1(size,size);
    array2d<double> img2(size,size);
    for (long r = 0; r < img1.nr(); ++r)
    {
        for (long c = 0; c < img1.nc(); ++c)
        {
            double x = 30.0*c/size - 19;
            double y = 8.0*r/size - 4;
            m(0) = x;
            m(1) = y;
            img1[r][c] = df(m);
            img2[r][c] = df2(join_cols(ekm.project(m),bias));
        }
    }
    image_window win1(heatmap(img1), "svm_one_class_trainer");
    image_window win2(heatmap(img2), "svm_c_linear_dcd_trainer");
    win1.wait_until_closed();
}


#include "anova.h"

// STL
#include <vector>
#include <algorithm>

// debug
#include <iostream>
#include <iomanip>

#include <boost/math/distributions/fisher_f.hpp>

template<typename T>
static T *transform(std::vector<T> &v)
{
    int size = v.size();
    T *result = new T[size];
    std::copy(v.begin(), v.end(), result);
    return result;
}

namespace stats
{

class MeanEnumer
{
public:
    MeanEnumer() = default;

    void operator()(const Sample &s)
    {
        _n += s.getCount();
        for (auto &el: s.getElements())
            _sum += el;
    }

    inline double getMean(){ return _sum / _n; }
    inline int getN(){ return _n; }
private:
    int _n = 0;
    double _sum = 0;
};

double Anova::FTest(const Samples &samples)
{
    assert(samples.size() > 1);
    if (samples.size()  <= 1)
        throw std::domain_error("count of samples must be more than 1");

    MeanEnumer enumer = std::for_each(samples.begin(), samples.end(), MeanEnumer());
    double generalMean = enumer.getMean();
    int N = enumer.getN();

    double Qa = 0;

    std::for_each(samples.begin(), samples.end(), [&Qa, &generalMean](const Sample &s){
        Qa += s.getCount() * pow(s.getMean() - generalMean, 2);
    });

    double Qoct = 0;
    std::for_each(samples.begin(), samples.end(), [&Qoct](const Sample &s){
        for (auto &el: s.getElements())
            Qoct += pow(el - s.getMean(), 2);
    });

    int sampleSize = samples.size();
    _v1 = sampleSize - 1;
    _v2 = N - sampleSize;
    _F = (Qa / _v1) / (Qoct / _v2);
    return _F;
}

double Anova::FCriticalTest(double alpha) const
{
    using namespace boost::math;
    fisher_f dist(_v1, _v2);
    return quantile(complement(dist, alpha));
}

//---------Two factor ANOVA---------------------------------------
TwoFactorAnova::Result TwoFactorAnova::FTest(const Samples &rowSamples)
{
    int m = rowSamples.size();
    assert(m != 0);
    if (m  == 0)
        throw std::domain_error("count of samples must be more than 1");

    int r = rowSamples[0].getCount();
    Samples columnSamples(r);
    Sample commonSample;
    for (int i = 0; i < m; i++)
    {
        const Sample &rowSample = rowSamples[i];
        int j = 0;
        for (auto &el: rowSample.getElements())
        {
            columnSamples[j].add(el);
            commonSample.add(el);
            j++;
        }
    }

    commonSample.refresh();
    for (auto &sample: columnSamples)
    {
        sample.refresh();
    }

    double Qa = 0.0, Qb = 0.0, Qoct = 0.0;

    for (const Sample &sample: columnSamples)
    {
        Qa += r*pow(sample.getMean() - commonSample.getMean(), 2);
    }

    double Da = Qa / (m - 1);

    for (const Sample &sample: rowSamples)
    {
        Qb += m*pow(sample.getMean() - commonSample.getMean(), 2);
    }

    double Db = Qb / (r - 1);

    for (int i = 0; i < m; i++)
    {
        for (int j = 0; j < r; j++)
        {
            Qoct += pow((rowSamples[i].getElements())[j]
                        + commonSample.getMean() - (rowSamples[i].getMean() + columnSamples[j].getMean()), 2);
        }
    }
    double Doct = Qoct / ((m - 1)*(r - 1));

    _Va = m - 1;
    _Vb = r - 1;
    _V2 = (m - 1) * (r - 1);

    Result res;
    res.FA = _F.FA =  Da / Doct;
    res.FB = _F.FB =  Db / Doct;
    res.FAB = _F.FAB = std::max(res.FA, res.FB) / std::min(res.FA, res.FB);

    return res;
}

double TwoFactorAnova::FcriticalTest(double alpha, int v1, int v2) const
{
    using namespace boost::math;
    fisher_f dist(v1, v2);
    return quantile(complement(dist, alpha));
}

double TwoFactorAnova::FCriticalTestA(double alpha) const
{
    return FcriticalTest(alpha, _Va, _V2);
}

double TwoFactorAnova::FCriticalTestB(double alpha) const
{
    return FcriticalTest(alpha, _Vb, _V2);
}

double TwoFactorAnova::FCriticalTestAB(double alpha) const
{
    return FcriticalTest(alpha, _V2, _V2);
}

}// stats


/* Copyright (c) 2020 C. Pattison
 * 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.
 *
 * 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 OWNER 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.
 */

#pragma once
#include "syk.hpp"
#include "syk_types.hpp"
#include "cuda_diagonalize.hpp"

#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <exception>
#include <algorithm>
#include <complex>

namespace syk {
using namespace std::complex_literals;

std::vector<double> cpu_hamiltonian_eigenvals(const MatrixType& hamiltonian) {
    auto eigensolver = Eigen::SelfAdjointEigenSolver<MatrixType>();
    eigensolver.compute(hamiltonian, false);

    if(eigensolver.info() != 0) {
        throw std::runtime_error("Eigensolver not converged");
    }
    const auto& eigenvals_vector = eigensolver.eigenvalues();
    std::vector<double> eigenvals(eigenvals_vector.size());
    std::copy_n(eigenvals_vector.data(), eigenvals_vector.size(), eigenvals.begin());
    std::sort(eigenvals.begin(), eigenvals.end());

    return eigenvals;
}

template<typename mat>
std::complex<double> closest_eigenval(mat m, std::complex<double> v) {
    std::complex<double> a, b, c, d;
    a = m(0,0);
    b = m(0,1);
    c = m(1,0);
    d = m(1,1);

    auto desc = std::sqrt(a*a + 4.0 * b * c - 2.0 * a * d + d * d);
    auto l_1 = (a + d + desc)/2.0;
    auto l_2 = (a + d - desc)/2.0;

    auto dist_1 = std::abs(l_1 - v);
    auto dist_2 = std::abs(l_2 - v);
    return dist_1 < dist_2 ? l_1 : l_2;
}

std::vector<double> QR_hamiltonian_eigenvals(const MatrixType& ham, MatrixType* Q_H,  double max_resid = 1e-7) {
    int max_iter = 1000 * ham.rows();
    assert(ham.rows() == ham.cols());
    assert(Q_H->rows() == ham.rows());
    assert(Q_H->cols() == ham.cols());

    // Eigen::HessenbergDecomposition<MatrixType> hd;
    // hd.compute(ham);
    // *Q_H = hd.matrixQ().adjoint();

    MatrixType temp(ham.rows(), ham.cols());
    MatrixType Q_curr(ham.rows(), ham.cols());

    MatrixType A = Q_H->adjoint() * ham * *Q_H;
    Eigen::HouseholderQR<MatrixType> qr(ham.rows(), ham.cols());

    // TODO: Use tridiagonal form
    int k;
    int deflate = ham.diagonal().size();
    for(k = 0; k < max_iter; ++k) {
        // // Rayleigh quotient
        // std::complex<double> shift = A(deflate-1, deflate-1);
        // Wilkinson
        auto shift = closest_eigenval(A.block(deflate-2, deflate-2, 2, 2), A(deflate-1, deflate-1));
        // QR
        qr.compute(A.topLeftCorner(deflate, deflate) - MatrixType::Identity(deflate, deflate) * shift);
        Q_curr = qr.householderQ();
        // Update Q_H with deflated matmul
        temp.topLeftCorner(deflate, deflate).noalias() = Q_H->topLeftCorner(deflate, deflate) * Q_curr;
        temp.bottomLeftCorner(Q_H->rows() - deflate, deflate).noalias() = Q_H->bottomLeftCorner(Q_H->rows() - deflate, deflate) * Q_curr;
        Q_H->topLeftCorner(Q_H->rows(), deflate) = temp.topLeftCorner(Q_H->rows(), deflate);
        // Update A = Q^* H Q
        A.noalias() = Q_H->adjoint() * ham * *Q_H; // TODO: Deflate this
        // Check convergence condition
        if(std::abs(A(deflate-1, deflate-2)) < max_resid && --deflate == 1) {
            break;
        }
    }

    double residual = 0;
    for(int j = 0; j < A.rows(); ++j) {
        for(int i = j+1; i < A.cols(); ++i) {
            residual += std::real(A(i,j) * std::conj(A(i,j)));
        }
    }
    residual = std::sqrt(residual) / A.diagonal().size();

    std::vector<double> eigenvals(A.diagonal().size());
    Eigen::Map<Eigen::VectorXd>(eigenvals.data(), eigenvals.size()) = A.diagonal().real();
    std::sort(eigenvals.begin(), eigenvals.end());

    return eigenvals;
}

std::vector<double> gpu_hamiltonian_eigenvals(const MatrixType& hamiltonian) {
    syk::GpuEigenValSolver solver;
    std::vector<double> eigenvals;
    #pragma omp critical
    eigenvals = solver.eigenvals(hamiltonian);
    return eigenvals;
}
}


#include "tasktorrent/tasktorrent.hpp"
#ifdef USE_MKL
#include <mkl_cblas.h>
#include <mkl_lapacke.h>
#else
#include <cblas.h>
#include <lapacke.h>
#endif
#include <Eigen/Core>
#include <Eigen/Cholesky>
#include <fstream>
#include <array>
#include <random>
#include <mutex>
#include <iostream>
#include <map>
#include <memory>
#include <set>
#include <mpi.h>
#include <string>

using namespace std;
using namespace Eigen;
using namespace ttor;

typedef array<int, 2> int2;
typedef array<int, 3> int3;

/*
 * Parametrized priorities for cholesky:
 * 0. No priority, only enforces potrf>trsm>gemm
 * 1. Row-based priority, prioritize tasks with smaller row number in addition to priority 0.
 * 2. Critical path priority, prioritize tasks with longest distance to the exit task. For references, check out the paper
    Beaumont, Olivier, et al. "A Makespan Lower Bound for the Scheduling of the Tiled Cholesky Factorization based on ALAP Schedule." (2020).
 * 3. Critical path and row priority, prioritize tasks with smaller row number in addition to priority 2. We also enforces potrf>trsm>gemm
 */

enum PrioKind { no = 0, row = 1, cp = 2, cp_row = 3};

/*
* 3D Cholesky based on the algorithm from the paper
    Kalluri Eswar, et al. "On Mapping Data and Computation for Parallel Sparse Cholesky Factorization." (1995)
* However, this implementation only uses 3D mapping for gemm tasks. 1D mapping is used for potrf tasks and 2D mapping is used for trsm tasks.
*/

void cholesky3d(int n_threads, int verb, int block_size, int num_blocks, int npcols, int nprows, PrioKind prio, int log, int debug)
{
    const int rank = comm_rank();
    const int n_ranks = comm_size();
    std::atomic<long long int> potrf_us_t(0);
    std::atomic<long long int> trsm_us_t(0);
    std::atomic<long long int> gemm_us_t(0);
    std::atomic<long long int> accu_us_t(0);
    assert(npcols * nprows == n_ranks);
    int q = static_cast<int>(cbrt(n_ranks));
    if (q * q * q != n_ranks)
    {
        if (rank == 0)
        {
            cerr << "Number of processes must be a perfect cube." << endl;
        }
        MPI_Finalize();
        exit(1);
    }

    int3 rank_3d;
    int2 rank_2d;
    rank_3d[0] = rank / (q * q);
    rank_3d[1] = (rank % (q * q)) / q;
    rank_3d[2] = (rank % (q * q)) % q;
    rank_2d[0] = rank % nprows;
    rank_2d[1] = rank / nprows;

    // Number of tasks
    int n_tasks_per_rank = 2;
    struct acc_data {
        vector<std::unique_ptr<MatrixXd>> to_accumulate; // to_accumulate[k] holds matrix result of gemm(k,i,j)
    };

    auto potf_block_2_prio = [&](int j) {
        if (prio == PrioKind::cp_row) {
            return (double)(9 * (num_blocks - j) - 1) + 18 * num_blocks * num_blocks;
        }
        else if(prio == PrioKind::cp) {
            return (double)(9 * (num_blocks - j) - 1);
        }
        else if(prio == PrioKind::row) {
            return 3.0 * (double)(num_blocks - j);
        }
        else {
            return 3.0;
        }
    };

    auto trsm_block_2_prio = [&](int2 ij) {
        if (prio == PrioKind::cp_row) {
            return (double)((num_blocks - ij[0]) + num_blocks * (9.0 * num_blocks - 9.0 * ij[1] - 2.0) + 9 * num_blocks * num_blocks);
        }
        else if(prio == PrioKind::cp) {
            return (double)(9 * (num_blocks - ij[1]) - 2);
        }
        else if(prio == PrioKind::row) {
            return 2.0 * (double)(num_blocks - ij[0]);
        }
        else {
            return 2.0;
        }
    };

    auto gemm_block_2_prio = [&](int3 kij) {
        if (prio == PrioKind::cp_row) {
            return (double)(num_blocks - kij[1]) + num_blocks * (9.0 * num_blocks - 3.0 * kij[2] - 6.0 * (kij[0] / q) - 2.0);
        }
        else if(prio == PrioKind::cp) {
            return (double)(9 * num_blocks - 9 * kij[2] - 2);
        }
        else if(prio == PrioKind::row) {
            return (double)(num_blocks - kij[1]);
        }
        else {
            return 1.0;
        }
    };

    std::vector<acc_data> gemm_results(num_blocks*num_blocks);
    auto val = [&](int i, int j) { return 1.0/(double)((i + j)*(i + j) + 1) + ((i == j) ? 1.0 * block_size * num_blocks : 0.0); };
    auto rank3d21 = [&](int i, int j, int k) { return ((j % q) * q + k % q) + (i % q) * q * q;};
    auto rank2d21 = [&](int i, int j) { return (j % npcols) * nprows + (i % nprows);};
    auto rank1d21 = [&](int j) { return j % n_ranks; };
    vector<unique_ptr<MatrixXd>> blocks(num_blocks*num_blocks);

    auto bloc_2_rank = [&](int i, int j) {
        int r = (j % npcols) * nprows + (i % nprows);
        assert(r >= 0 && r < n_ranks);
        return r;
    };

    auto block_2_thread = [&](int i, int j) {
        int ii = i / nprows;
        int jj = j / npcols;
        int num_blocksit = num_blocks / nprows;
        return (ii + jj * num_blocksit) % n_threads;
    };
    {
        Eigen::MatrixXd A = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd B = Eigen::MatrixXd::Identity(256,256);
        Eigen::MatrixXd C = Eigen::MatrixXd::Identity(256,256);
        for(int i = 0; i < 10; i++) {
            cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, 256, 256, 256, 1.0, A.data(), 256, B.data(), 256, 1.0, C.data(), 256);
        }
    }
    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<num_blocks; jj++) {
            auto val_loc = [&](int i, int j) { return val(ii*block_size+i,jj*block_size+j); };
            int dest = (ii == jj) ? rank1d21(ii) : rank2d21(ii,jj);
            if(dest == rank) {
                blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(block_size, block_size);
                *blocks[ii+jj*num_blocks]=MatrixXd::NullaryExpr(block_size, block_size, val_loc);
                gemm_results[ii+jj*num_blocks].to_accumulate= vector<std::unique_ptr<MatrixXd>>(q);
                for (int ll=0; ll<q; ll++) {
                    gemm_results[ii+jj*num_blocks].to_accumulate[ll]=make_unique<MatrixXd>(block_size, block_size);
                    *(gemm_results[ii+jj*num_blocks].to_accumulate[ll])=MatrixXd::Zero(block_size, block_size);
                }
            }
            else if (((ii % q) == rank_3d[0]) && ((jj % q) == rank_3d[1])) {
                blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(block_size, block_size);
                *blocks[ii+jj*num_blocks]=MatrixXd::Zero(block_size, block_size);
            }
            else {
                blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(block_size, block_size);
            }
        }
    }
    // Initialize the communicator structure
    Communicator comm(MPI_COMM_WORLD, verb);
    // Initialize the runtime structures
    Threadpool tp(n_threads, &comm, verb, "WkTuto_" + to_string(rank) + "_");
    Taskflow<int> potrf(&tp, verb);
    Taskflow<int2> trsm(&tp, verb);
    Taskflow<int3> gemm(&tp, verb);
    Taskflow<int3> accu(&tp, verb);
    DepsLogger dlog(1000000);
    Logger ttorlog(1000000);
    if (log)  {
        tp.set_logger(&ttorlog);
        comm.set_logger(&ttorlog);
    }
    // Create active message
    auto am_trsm = comm.make_large_active_msg(
        [&](int& j) {
                int offset = ((j + 1) / nprows + (((j + 1) % nprows) > rank_2d[0])) * nprows + rank_2d[0];
                for(int i = offset; i < num_blocks; i = i + nprows) {
                    if (debug) printf("Fulfilling trsm (%d, %d) on rank (%d, %d)\n", i, j, rank_2d[0], rank_2d[1]);
                    assert(rank2d21(i, j) == rank);
                    trsm.fulfill_promise({i,j});
                }
            },
            [&](int& j){
                return blocks[j+j*num_blocks]->data();
            },
            [&](int& j){
                return;
            });

    auto am_gemm = comm.make_large_active_msg(
        [&](int& i, int& k) {
            assert(k % q == rank_3d[2]);
            int offset_c = ((k + 1) / q + (((k + 1) % q) > rank_3d[1])) * q + rank_3d[1];
            if (i % q == rank_3d[0]) {
                for(int j = offset_c; j < i; j = j + q) {
                    if (debug) printf("TRSM (%d, %d) Fulfilling gemm (%d, %d, %d) on rank (%d, %d, %d)\n", i, k, k, i, j, rank_3d[2], rank_3d[0], rank_3d[1]);
                    assert(rank3d21(i,j,k) == rank);
                    gemm.fulfill_promise({k,i,j});
                }
            }
            int offset_r = (i / q + ((i % q) > rank_3d[0])) * q + rank_3d[0];
            if (i % q == rank_3d[1]) {
                for(int j = offset_r; j < num_blocks; j = j + q) {
                    if (debug) printf("TRSM (%d, %d) Fulfilling gemm (%d, %d, %d) on rank (%d, %d, %d)\n", i,k, k, j, i, rank_3d[2], rank_3d[0], rank_3d[1]);
                    assert(rank3d21(j,i,k) == rank);
                    gemm.fulfill_promise({k,j,i});
                }
            }
        },
        [&](int& i, int& k) {
            return blocks[i+k*num_blocks]->data();
        },
        [&](int& i, int& k) {
            return;
        });

    potrf.set_task([&](int j) {
            assert(rank1d21(j) == rank);
            timer t1 = wctime();
            int info = LAPACKE_dpotrf(LAPACK_COL_MAJOR, 'L', block_size, blocks[j+j*num_blocks]->data(), block_size);
            timer t2 = wctime();
            potrf_us_t += 1e6 * elapsed(t1, t2);
            assert(info == 0);
            if (debug) printf("Running POTRF %d on rank %d\n", j, rank);
        })
        .set_fulfill([&](int j) {
            for (int p = 0; p < nprows; p++)
            {
                int r = rank2d21(p,j);
                if (rank == r) {
                    int offset = ((j + 1) / nprows + ((j + 1) % nprows) / (rank_2d[0] + 1)) * nprows + rank_2d[0];
                    for(int i = offset; i < num_blocks; i = i + nprows) {
                        trsm.fulfill_promise({i,j});
                    }
                }
                else {
                    auto Ljjv = view<double>(blocks[j+j*num_blocks]->data(), block_size*block_size);
                    am_trsm->send_large(r, Ljjv, j);
                }
            }
        })
        .set_indegree([&](int j) {
            if (j==0) {
                return 1;
            }
            else if (j < q) {
                return j;
            }
            else {
                return q;
            }
        })
        .set_mapping([&](int j) {
            return block_2_thread(j,j);
        })
        .set_binding([&](int j) {
            return false;

        })
        .set_priority(potf_block_2_prio)
        .set_name([&](int j) {
            return "POTRF" + to_string(j) + "_" + to_string(rank);
        });

    trsm.set_task([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            assert(rank2d21(i,j) == rank);
            timer t1 = wctime();
            cblas_dtrsm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit, block_size, block_size, 1.0, blocks[j + j * num_blocks]->data(),block_size, blocks[i + j * num_blocks]->data(), block_size);
            timer t2 = wctime();
            trsm_us_t += 1e6 * elapsed(t1, t2);
            if (debug) printf("Running trsm (%d, %d) on rank %d, %d\n", i, j, rank_2d[0], rank_2d[1]);
        })
        .set_fulfill([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            for (int ri = 0; ri < q; ri++)  {
                for (int rj = 0; rj < q; rj++) {
                    int r = rank3d21(ri, rj, j);
                    if (r == rank) {
                        int offset_c = ((j + 1) / q + (((j + 1) % q) > rank_3d[1])) * q + rank_3d[1];
                        if (i % q == rank_3d[0]) {
                            for(int k = offset_c; k < i; k = k + q) {
                                gemm.fulfill_promise({j,i,k});
                            }
                        }
                        int offset_r = (i / q + ((i % q) > rank_3d[0])) * q + rank_3d[0];
                        if (i % q == rank_3d[1]) {
                            for(int k = offset_r; k < num_blocks; k = k + q) {
                                gemm.fulfill_promise({j,k,i});
                            }
                        }
                    }
                    else {
                        auto Lijv = view<double>(blocks[i + j * num_blocks]->data(), block_size*block_size);
                        am_gemm->send_large(r, Lijv, i, j);
                    }
                }
            }
        })
        .set_indegree([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            return (j < q) ? j + 1 : q + 1;
        })
        .set_mapping([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            return block_2_thread(i,j);
        })
        .set_binding([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            return false;

        })
        .set_priority(trsm_block_2_prio)
        .set_name([&](int2 ij) {
            int i=ij[0];
            int j=ij[1];
            return "TRSM" + to_string(j) + "_" + to_string(i) + "_" +to_string(rank);
        });

    auto am_accu = comm.make_large_active_msg(
        [&](int& i, int& j, int& from) {
            accu.fulfill_promise({from, i, j});
        },
        [&](int& i, int& j, int& from){
            return gemm_results[i+j*num_blocks].to_accumulate[from]->data();
        },
        [&](int& i, int& j, int& from){
            return;
        });

    gemm.set_task([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            assert(rank3d21(i,j,k) == rank);
            timer t1 = wctime();
            if (i==j) {
                cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, block_size, block_size, -1.0, blocks[i+k*num_blocks]->data(), block_size, 1.0, blocks[i+j*num_blocks]->data(), block_size);
            }
            else {
                cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans, block_size, block_size, block_size, -1.0,blocks[i+k*num_blocks]->data(), block_size, blocks[j+k*num_blocks]->data(), block_size, 1.0, blocks[i+j*num_blocks]->data(), block_size);
            }
            timer t2 = wctime();
            if (debug) printf("Running gemm (%d, %d, %d) on rank %d, %d, %d\n", k, i, j, rank_3d[2], rank_3d[0], rank_3d[1]);
            gemm_us_t += 1e6 * elapsed(t1, t2);
        })
        .set_fulfill([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            if (k+q<=j-1) {
                gemm.fulfill_promise({k+q, i, j});
            }
            else {
                int dest = (i == j) ? rank1d21(i) : rank2d21(i, j);
                if (dest == rank) {
                    if (debug) printf("gemm (%d, %d, %d) fulfilling accumu (%d, %d, %d) on rank %d, %d, %d\n", k, i, j, rank_3d[2], i, j, rank_3d[2], rank_3d[0], rank_3d[1]);
                    accu.fulfill_promise({rank_3d[2], i, j});
                }
                else {
                    int kk = rank_3d[2];
                    auto Lij = view<double>(blocks[i+j*num_blocks]->data(), block_size*block_size);
                    if (debug) printf("gemm (%d, %d, %d) Sending accumu (%d, %d, %d) to rank %d, %d\n", k, i, j, rank_3d[2], i, j, dest % nprows, dest / nprows);
                    am_accu->send_large(dest, Lij, i, j, kk);
                }
            }
        })
        .set_indegree([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            return 3 - (k/q == 0) - (i == j);
        })
        .set_mapping([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            return block_2_thread(i,j);
        })
        .set_binding([&](int3 kij) {
            return false;

        })
        .set_priority(gemm_block_2_prio)
        .set_name([&](int3 kij) { // This is just for debugging and profiling
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            return "gemm" + to_string(k) + "_" + to_string(i)+"_"+to_string(j)+"_"+to_string(comm_rank());
        });

    accu.set_task([&](int3 kij) {
            int k=kij[0]; // Step (gemm's pivot)
            int i=kij[1]; // Row
            int j=kij[2]; // Col
            int dest = (i == j) ? rank1d21(i) : rank2d21(i,j);
            assert(dest == rank);
            assert(j <= i);
            if (debug) printf("Running accumu (%d, %d, %d) on rank %d, %d\n", k, i, j, rank % nprows, rank / nprows);
            {
                timer t_ = wctime();
                *blocks[i+j*num_blocks] += (*gemm_results[i+j*num_blocks].to_accumulate[k]);
                timer t__ = wctime();
                accu_us_t += 1e6 * elapsed(t_, t__);
            }
        })
        .set_fulfill([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            assert(j <= i);
            if(i == j) {
                potrf.fulfill_promise(i);
            } else {
                trsm.fulfill_promise({i,j});
            }
        })
        .set_indegree([&](int3 kij) {
            return 1;
        })
        .set_mapping([&](int3 kij) {
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            return block_2_thread(i,j);
        })
        .set_priority(gemm_block_2_prio)
        .set_binding([&](int3 kij) {
            return true; // IMPORTANT
        })
        .set_name([&](int3 kij) { // This is just for debugging and profiling
            int k=kij[0];
            int i=kij[1];
            int j=kij[2];
            return "accumu" + to_string(k) + "_" + to_string(i)+"_"+to_string(j)+"_"+to_string(comm_rank());
        });

    MPI_Barrier(MPI_COMM_WORLD);
    timer t0 = wctime();
    if (rank == 0){
        potrf.fulfill_promise(0);
    }
    tp.join();
    MPI_Barrier(MPI_COMM_WORLD);
    timer t1 = wctime();
    MPI_Status status;
    if (rank==0) {
        cout<<"3D, Number of ranks "<<n_ranks<<", block_size "<<block_size<<", num_blocks "<<num_blocks<<", n_threads "<<n_threads<<", Priority "<<prio<<", Elapsed time: "<<elapsed(t0,t1)<<endl;
    }
    MatrixXd A;
    A = MatrixXd::NullaryExpr(block_size*num_blocks,block_size*num_blocks, val);
    MatrixXd L = A;
    for (int ii=0; ii<num_blocks; ii++) {
        for (int jj=0; jj<num_blocks; jj++) {
            if (jj<=ii)  {
                int dest = (ii == jj) ? rank1d21(ii) : rank2d21(ii,jj);
                if (rank==0 && rank!=dest) {
                    blocks[ii+jj*num_blocks]=make_unique<MatrixXd>(block_size,block_size);
                    MPI_Recv(blocks[ii+jj*num_blocks]->data(), block_size*block_size, MPI_DOUBLE, dest, 0, MPI_COMM_WORLD, &status);
                }
                else if (rank==dest && rank != 0) {
                    MPI_Send(blocks[ii+jj*num_blocks]->data(), block_size*block_size, MPI_DOUBLE, 0, 0, MPI_COMM_WORLD);
                }
            }
        }
    }
    if (rank == 0)  {
        for (int ii=0; ii<num_blocks; ii++) {
            for (int jj=0; jj<num_blocks; jj++) {
                if (jj<=ii)  {
                    L.block(ii*block_size,jj*block_size,block_size,block_size)=*blocks[ii+jj*num_blocks];
                }
            }
        }
        auto L1=L.triangularView<Lower>();
        VectorXd x = VectorXd::Random(block_size * num_blocks);
        VectorXd b = A*x;
        L1.solveInPlace(b);
        L1.transpose().solveInPlace(b);
        double error = (b - x).norm() / x.norm();
        cout << "Error solve: " << error << endl;
    }
    if (log)  {
        std::ofstream logfile;
        string filename = "ttor_3Dcholesky_Priority_"+to_string(block_size)+"_"+to_string(num_blocks)+"_"+ to_string(n_threads)+"_"+ to_string(n_ranks)+"_"+ to_string(prio)+".log."+to_string(rank);
        logfile.open(filename);
        logfile << ttorlog;
        logfile.close();
    }
}

int main(int argc, char **argv)
{
    int req = MPI_THREAD_FUNNELED;
    int prov = -1;
    MPI_Init_thread(NULL, NULL, req, &prov);
    assert(prov == req);
    int n_threads = 2;
    int verb = 0; // Can be changed to vary the verbosity of the messages
    int block_size = 5;
    int num_blocks = 10;
    int npcols = 1;
    int nprows = ttor::comm_size();
    PrioKind prio = PrioKind::no;
    int log = 0;
    int debug = 0;
    if (argc >= 2)
    {
        block_size = atoi(argv[1]);
        assert(block_size > 0);
    }
    if (argc >= 3)
    {
        num_blocks = atoi(argv[2]);
        assert(num_blocks > 0);
    }
    if (argc >= 5) {
        n_threads=atoi(argv[3]);
        assert(n_threads > 0);
        verb=atoi(argv[4]);
        assert(verb >= 0);
    }
    if (argc >= 7) {
        npcols=atoi(argv[5]);
        assert(npcols > 0);
        nprows=atoi(argv[6]);
        assert(nprows > 0);
    }
    if (argc >= 8) {
        prio=(PrioKind)atoi(argv[7]);
        assert(prio >= 0 && prio < 4);
    }
    if (argc >= 9) {
        log = atoi(argv[9]);
        assert(log == 0 || log == 1);
    }
    if (argc >= 10) {
        debug = atoi(argv[10]);
        assert(debug == 0 || debug == 1);
    }
    if(comm_rank() == 0) printf("Usage: ./3d_cholesky block_size num_blocks n_threads verb nprows npcols priority log debug\n");
    cholesky3d(n_threads, verb, block_size, num_blocks, npcols, nprows, prio, log, debug);
    MPI_Finalize();
}


#include <iostream>
#include <fstream>
#include <Eigen/Dense>
#include <json.hpp>

using namespace Eigen;
using namespace std;
using json = nlohmann::json;

// Converts an std::vector of std::vector to Eigen Matrix
template <typename T>
Matrix<T, Dynamic, Dynamic> vv_to_matrix(vector<vector<T>> vv){
    Matrix <T, Dynamic, Dynamic>  m (vv.size(), vv[0].size());
    for (int i = 0; i< vv.size(); i++){
        for (int j = 0; j < vv[i].size(); j++){
            m(i,j)= vv[i][j];
        }
    }
    return m;
}

// Converts a std::vector to Eigen Vector
template <typename T>
Vector<T,Dynamic> std_vector_to_eigen(vector <T> v){
    Vector<T,Dynamic> ev (v.size());
    for (int i = 0; i< v.size(); i++){
        ev(i)= v[i];
    }
    return ev;
}

VectorXd ReLu(VectorXd x){
    for (int i = 0; i < x.size(); i++)
        if (x(i) < 0) x(i) = 0;
    return x;
}

class DoubleLayerMLP{
    private:
        MatrixXd w1;
        VectorXd b1;
        MatrixXd w2;
        VectorXd b2;
    public:

        DoubleLayerMLP(MatrixXd w1, VectorXd b1, MatrixXd w2, VectorXd b2){
            this->w1 = w1;
            this->b1 = b1;
            this->w2 = w2;
            this->b2 = b2;
        }

        DoubleLayerMLP(const DoubleLayerMLP & nn){
            this->w1 = nn.w1;
            this->b1 = nn.b1;
            this->w2 = nn.w2;
            this->b2 = nn.b2;
        }

        DoubleLayerMLP(){
            this->w1 = MatrixXd();
            this->b1 = VectorXd();
            this->w2 = MatrixXd();
            this->b2 = VectorXd();
        }

        DoubleLayerMLP(json json_data, string s){
            vector<vector<double>> vv1 = json_data[s]["w1"].get<vector<vector<double>>>();
            MatrixXd w1 = vv_to_matrix(vv1);

            vector<vector<double>> vv2 = json_data[s]["w2"].get<vector<vector<double>>>();
            MatrixXd w2 = vv_to_matrix(vv2);

            vector<double> v1 = json_data[s]["b1"].get<vector<double>>();
            VectorXd b1 = std_vector_to_eigen(v1);

            vector<double> v2 = json_data[s]["b2"].get<vector<double>>();
            VectorXd b2 = std_vector_to_eigen(v2);

            this->w1 = w1;
            this->b1 = b1;
            this->w2 = w2;
            this->b2 = b2;
        }

        VectorXd forward(VectorXd x){
            //cout<< endl << x << endl;
            //cout<< endl << w1 << endl;
            //cout<< endl << w1 * x << endl;

            return w2*ReLu(w1 * x + b1) + b2;
        }
};

class MPNN{
    private:
        int V_attributes;
        int E_attributes;
        DoubleLayerMLP edge_update_nn;
        DoubleLayerMLP vertice_update_nn;
        DoubleLayerMLP output_update_nn;
        int V_hidden;
        int E_hidden;
    public:
        MPNN(int V_attributes, int E_attributes, DoubleLayerMLP edge_update_nn, DoubleLayerMLP vertice_update_nn, DoubleLayerMLP output_update_nn, int V_hidden = 0, int E_hidden = 0){
                this -> V_attributes = V_attributes;
                this -> E_attributes = E_attributes;
                this -> edge_update_nn = edge_update_nn;
                this -> vertice_update_nn = vertice_update_nn;
                this -> output_update_nn = output_update_nn;
                this -> V_hidden = V_hidden ? V_hidden : V_attributes;
                this -> E_hidden = E_hidden ? E_hidden : E_attributes;
            }
        tuple<MatrixXd, MatrixXd, VectorXd>  forward (MatrixXd E, MatrixXi E_V, MatrixXd V){
            int E_n = E.rows();
            int V_n = V.rows();

            MatrixXd E_new = MatrixXd::Zero (E_n, this->E_hidden);
            MatrixXd V_new = MatrixXd::Zero (V_n, this->V_hidden);

            for (int i = 0; i < E_n; i++){
                VectorXd E_concat(V.cols() + V.cols() + E.cols());
                E_concat << V.row(E_V(i,0)).transpose(), V.row(E_V(i,1)).transpose(), E.row(i).transpose();

                //cout<<V.row(E_V(i,0)) <<endl;
                E_new.row(i) = this->edge_update_nn.forward(E_concat);
            }

            VectorXd V_agregated = VectorXd::Zero(this->V_hidden);

            for (int i = 0; i < V_n; i++){
                VectorXd E_agregated = VectorXd::Zero(this->E_hidden);
                for (int j = 0; j < E_n; j++){
                    if (E(j,1) == i)
                        E_agregated += E_new.row(i);
                }
                VectorXd V_concat(E_agregated.size() + V.cols());
                V_concat << E_agregated, V.row(i).transpose();

                V_new.row(i) = this->vertice_update_nn.forward(V_concat);
                V_agregated += V_new.row(i);
            }

            tuple<MatrixXd, MatrixXd, VectorXd> res = tuple<MatrixXd, MatrixXd, VectorXd> (E_new, V_new, this->output_update_nn.forward(V_agregated));
            return res;
        }
};

int main(int argc, char* argv[]){

    std::ifstream ifs ("weights.json", std::ifstream::in);
    json json_data = json::parse(ifs);

    DoubleLayerMLP edge_update_nn = DoubleLayerMLP(json_data, "edge_update");
    DoubleLayerMLP vertice_update_nn = DoubleLayerMLP(json_data, "vertice_update");
    DoubleLayerMLP output_update_nn = DoubleLayerMLP(json_data, "output_update");
    MPNN model = MPNN(3, 3, edge_update_nn, vertice_update_nn, output_update_nn);
    std::ifstream graphifstream ("test_graph.json", std::ifstream::in);
    json json_graph = json::parse(graphifstream);

    vector<vector<double>> E_vv = json_graph["E"].get<vector<vector<double>>>();
    vector<vector<int>> V_E_vv = json_graph["V_E"].get<vector<vector<int>>>();
    vector<vector<double>> V_vv = json_graph["V"].get<vector<vector<double>>>();

    MatrixXd E = vv_to_matrix<double>(E_vv);
    MatrixXi V_E = vv_to_matrix<int> (V_E_vv);
    MatrixXd V = vv_to_matrix<double>(V_vv);

    tuple<MatrixXd, MatrixXd, VectorXd> res = model.forward(E,V_E,V);
    MatrixXd new_E = get<0> (res);
    MatrixXd new_V = get<1> (res);
    MatrixXd new_u = get<2> (res);
    cout << "Edge attributes: "<< endl << new_E  << endl;
    cout << "Vertice attributes: "<< endl << new_V << endl;
    cout << "Global attributes: "<< endl << new_u << endl;
}


/* p_integrand.c */
#include <math.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <boost/math/special_functions/gamma.hpp>
extern "C"
{
/* Survival function for R^phi*W */
int RW_marginal_C(double *xval, double phi, double gamma, int n_xval, double *result){
    double tmp2 = pow(gamma/2, phi)/boost::math::tgamma(0.5);
    double tmp1, tmp0, a;
    a = 0.5-phi;

    for(int i=0; i<n_xval; i++){
        tmp1 = gamma/(2*pow(xval[i],1/phi));
        tmp0 = tmp2/(a*xval[i]);
        result[i] = boost::math::gamma_p(0.5L,tmp1) + boost::math::tgamma((long double)(a+1),tmp1)*tmp0-pow(tmp1,a)*exp(-tmp1)*tmp0;
    }
    return 1;
}

/* Marginal distribution function for R^phi*W + epsilon */
int pRW_me_interp_C(double *xval, double *xp, double *surv_p, double tau_sqd, double phi, double gamma, int n_xval, int n_grid, double *result){
    bool tau_bool = (tau_sqd > 0.05);
    double tp[n_grid];
    double integrand_p[n_grid];
    double tmp, tmp_res; /* temporary constant */
    double tmp_sum = 0; /* temporary trapesoid sum */
    double sd = sqrt(tau_sqd);
    double sd_const = sqrt(2)*sd;
    double sd_const_pi =sqrt(2*M_PI)*sd;
    int i,j, tmp_int; /* iterative constants */

    for (i = 0; i < n_xval; i++) {
        if(tau_bool & (xval[i]<820)){
            /* Calculate integrand on a grid */
            for(j=0; j<n_grid;j++){
                tmp = xval[i]-xp[j];
                tp[j] = tmp;
                integrand_p[j] = exp(-tmp*tmp/(2*tau_sqd)) * surv_p[j];
            }

            /* Numerical integral using the trapesoid method */
            for(j=0; j<(n_grid-1);j++){
                tmp_sum+= (tp[j+1]-tp[j])*(integrand_p[j] + integrand_p[j+1])/2;
            }
            tmp_res = 0.5*erfc(-xval[i]/sd_const)-tmp_sum/sd_const_pi;
            tmp_sum = 0;

            /* CDF value must be greater than 0 */
            if(tmp_res < 0){
                tmp_res = 0;
            }
            result[i] = tmp_res;
        }
        else{
            tmp_int = RW_marginal_C(&xval[i], phi, gamma, 1, &tmp_res);
            result[i] = 1-tmp_res;
        }
    }

    return 1;
}

/* Get the quantile range for certain probability levels */
int find_xrange_pRW_me_C(double min_p, double max_p, double min_x, double max_x, double *xp, double *surv_p, double tau_sqd, double phi, double gamma, int n_grid, double *x_range){
    if (min_x >= max_x){
        printf("Initial value of mix_x must be smaller than max_x.\n");
        exit(EXIT_FAILURE);
    }

    /* First the min */
    double p_min_x;
    int tmp_int;
    tmp_int = pRW_me_interp_C(&min_x, xp, surv_p, tau_sqd, phi, gamma, 1, n_grid, &p_min_x);
    while (p_min_x > min_p){
        min_x = min_x-40/phi;
        tmp_int = pRW_me_interp_C(&min_x, xp, surv_p, tau_sqd, phi, gamma, 1, n_grid, &p_min_x);
    }

    x_range[0] = min_x;

    /* Now the max */
    double p_max_x;
    tmp_int = pRW_me_interp_C(&max_x, xp, surv_p, tau_sqd, phi, gamma, 1, n_grid, &p_max_x);
    while (p_max_x < max_p){
        max_x = max_x*2; /* Upper will set to 20 initially */
        tmp_int = pRW_me_interp_C(&max_x, xp, surv_p, tau_sqd, phi, gamma, 1, n_grid, &p_max_x);
    }

    x_range[1] = max_x;
    return 1;
}

/* Density function for R^phi*W */
int RW_density_C(double *xval, double phi, double gamma, int n_xval, double *result){
    double tmp2 = pow(gamma/2, phi)/boost::math::tgamma(0.5);
    double tmp1, tmp0, a;
    a = 0.5-phi;

    for(int i=0; i<n_xval; i++){
        tmp1 = gamma/(2*pow(xval[i],1/phi));
        tmp0 = tmp2/(a*pow(xval[i],2));
        result[i] = (boost::math::tgamma((long double)(a+1),tmp1)-pow(tmp1,a)*exp(-tmp1))*tmp0;
    }
    return 1;
}

/* Marginal density function for R^phi*W + epsilon */
int dRW_me_interp_C(double *xval, double *xp, double *den_p, double tau_sqd, double phi, double gamma, int n_xval, int n_grid, double *result){
    double thresh_large = 820;
    if(tau_sqd < 1) {
        thresh_large = 50;
    }
    bool tau_bool = (tau_sqd > 0.05);

    double tp[n_grid];
    double integrand_p[n_grid];
    double tmp, tmp_res; /* temporary constant */
    double tmp_sum = 0; /* temporary trapesoid sum */
    double sd = sqrt(tau_sqd);
    double sd_const_pi =sqrt(2*M_PI)*sd;
    int i,j, tmp_int; /* iterative constants */

    for (i = 0; i < n_xval; i++) {
        if(tau_bool & (xval[i]<thresh_large)){
            /* Calculate integrand on a grid */
            for(j=0; j<n_grid;j++){
                tmp = xval[i]-xp[j];
                tp[j] = tmp;
                integrand_p[j] = exp(-tmp*tmp/(2*tau_sqd)) * den_p[j];
            }

            /* Numerical integral using the trapesoid method */
            for(j=0; j<(n_grid-1);j++){
                tmp_sum+= (tp[j+1]-tp[j])*(integrand_p[j] + integrand_p[j+1])/2;
            }
            tmp_res = tmp_sum/sd_const_pi;
            tmp_sum = 0;
            result[i] = tmp_res;
        }else if((tau_bool & (xval[i]>=thresh_large))|(!tau_bool & (xval[i]>0))){
            tmp_int = RW_density_C(&xval[i], phi, gamma, 1, &tmp_res);
            result[i] = tmp_res;
        }else{
            result[i] = 0;
        }
    }

    return 1;
}

}


#ifndef INCLUDE_MIMKL_MODELS_EASY_MKL_HPP_
#define INCLUDE_MIMKL_MODELS_EASY_MKL_HPP_

#include <Eigen/StdVector>
#include <mimkl/data_structures.hpp>
#include <mimkl/definitions.hpp>
#include <mimkl/linear_algebra.hpp>
#include <mimkl/models/model.hpp>
#include <mimkl/solvers.hpp>
#include <spdlog/fmt/ostr.h>
#include <spdlog/spdlog.h>

using mimkl::definitions::Indexing;

namespace mimkl
{
namespace models
{

static std::shared_ptr<spdlog::logger> logger_easy_mkl =
spdlog::stdout_color_mt("EasyMKL");

template <typename Scalar, typename Kernel>
class EasyMKL : public Model<Scalar, Kernel>
{

    private:
    // inheritance of templatized base members
    using Model<Scalar, Kernel>::_kernels_handler;
    using Model<Scalar, Kernel>::_trained;
    using Model<Scalar, Kernel>::_precompute;
    using Model<Scalar, Kernel>::_trace_normalization;
    using Model<Scalar, Kernel>::_number_of_support_vectors;
    using Model<Scalar, Kernel>::_number_of_kernels;

    // logger instance
    std::shared_ptr<spdlog::logger> _logger = spdlog::get("EasyMKL");
    // typedefs
    typedef MATRIX(Scalar) Matrix;
    typedef COLUMN(Scalar) Column;
    typedef ROW(Scalar) Row;
    typedef Eigen::Map<Column> MapColumn;

    // parameters
    double _lambda = 0.8;
    double _epsilon = 0.0001;
    bool _regularization_factor = false;

    Matrix _gammas;
    Matrix _etas;
    Row _biases;
    Indexing _class_map;

    std::vector<std::string> _unique_labels;
    Index _number_of_classes;
    Index _number_of_dichotomies;
    bool _binary;
    Matrix _kernels_sum;

    Column get_dichotomy(const std::string &);
    std::pair<Column, Column>
    optimize_gamma(const bool &, const Column &, const Column & = Column());
    Column compute_weights(const Column &);
    Scalar compute_bias(const Column &, const Column &, const Column &);

    void setup(std::vector<std::string>);
    void train(const Index &);
    void fit(const std::vector<std::string> &);
    Matrix decision_function();
    Column
    distances(const Column &, const Column &, const Column &, const Scalar &);

    public:
    EasyMKL(const std::vector<Kernel> & = std::vector<Kernel>(),
            const bool precompute = true,
            const bool trace_normalization = true,
            const double lambda = 0.8,
            const double epsilon = 0.0001,
            const bool regularization_factor = false);
    EasyMKL(const std::vector<Matrix> & = std::vector<Matrix>(),
            const bool precompute = true,
            const bool trace_normalization = true,
            const double lambda = 0.8,
            const double epsilon = 0.0001,
            const bool regularization_factor = false);

    // Functional, call with data matrix
    void fit(const Matrix &, const std::vector<std::string> &);
    std::vector<std::string> predict(const Matrix &);
    Matrix predict_proba(const Matrix &);
    Matrix decision_function(const Matrix &);
    // Matricial, call with kernel matrices
    void fit(const std::vector<Matrix> &, const std::vector<std::string> &);
    std::vector<std::string> predict(const std::vector<Matrix> &);
    Matrix predict_proba(const std::vector<Matrix> &);
    Matrix decision_function(const std::vector<Matrix> &);

    double get_lambda() const;
    double get_epsilon() const;
    bool get_regularization_factor() const;
    Matrix get_gammas() const;
    Matrix get_etas() const;
    Row get_biases() const;
    Indexing get_class_map() const;
    std::vector<std::string> get_one_versus_rest_order() const;
    Matrix get_optimal_kernel();
    Matrix get_optimal_kernel_by_class_index(const Index i);
    std::vector<Matrix> get_optimal_kernels();

    void set_lambda(const double);
    void set_epsilon(const double);
    void set_regularization_factor(const bool);
    void set_parameters(const double,
                        const double epsilon = 0.0001,
                        const bool regularization_factor = false);
    void set_gammas(const Matrix &);
    void set_etas(const Matrix &);
    void set_biases(const Row &);
    void set_class_map(const Indexing &);
};

template <typename Scalar, typename Kernel>
EasyMKL<Scalar, Kernel>::EasyMKL(const std::vector<Kernel> &kernel_functions,
                                 const bool precompute,
                                 const bool trace_normalization,
                                 const double lambda,
                                 const double epsilon,
                                 const bool regularization_factor)
: Model<Scalar, Kernel>(kernel_functions, precompute, trace_normalization)
{
    set_parameters(lambda, epsilon, regularization_factor);
}

template <typename Scalar, typename Kernel>
EasyMKL<Scalar, Kernel>::EasyMKL(const std::vector<Matrix> &kernel_matrices,
                                 const bool precompute,
                                 const bool trace_normalization,
                                 const double lambda,
                                 const double epsilon,
                                 const bool regularization_factor)
: Model<Scalar, Kernel>(kernel_matrices, precompute, trace_normalization)
{
    set_parameters(lambda, epsilon, regularization_factor);
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Column
EasyMKL<Scalar, Kernel>::get_dichotomy(const std::string &label)
{
    auto range = _class_map.equal_range(label);
    Column y = Column::Constant(_number_of_support_vectors, 1, -1);
    for (auto i = range.first; i != range.second; ++i)
    {
        y(i->second) = 1;
    }
    return y;
}

template <typename Scalar, typename Kernel>
std::pair<typename EasyMKL<Scalar, Kernel>::Column,
          typename EasyMKL<Scalar, Kernel>::Column>
EasyMKL<Scalar, Kernel>::optimize_gamma(const bool &regularization_factor,
                                        const Column &y,
                                        const Column &eta)
{
    mimkl::solvers::KOMD<Scalar> optimizer((eta.rows() > 0) ?
                                           _kernels_handler.sum(eta) :
                                           _kernels_sum,
                                           _lambda, _epsilon,
                                           regularization_factor);
    optimizer.solve(y);
    Column gamma = optimizer.get_result();
    return std::make_pair(gamma, y.cwiseProduct(gamma));
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Column
EasyMKL<Scalar, Kernel>::compute_weights(const Column &directed_gamma)
{
    // TODO maybe after all gammas are known and then compute etas rowwise, so
    // we
    // compute each kernel once? ->cycle through directed_gamma instead of
    // through
    // kernels
    Column eta = Column::Constant(_number_of_kernels, 1, 0.0);
    for (Index eta_index = 0; eta_index < _number_of_kernels; ++eta_index)
    {
        // compute current kernel matrix K
        _logger->trace("base kernel :\n{}", _kernels_handler[eta_index]);
        eta(eta_index) = directed_gamma.transpose() *
                         _kernels_handler[eta_index] * directed_gamma;

        _logger->trace("d(y)_r :\n{}", eta(eta_index));
    }
    _logger->trace("pre norming eta :\n{}", eta);
    eta /=
    eta.sum(); // l1-norm, as eta_i is >=0 by construction ( kernels are SDP)
    _logger->trace("eta :\n{}", eta);
    _logger->debug("compute_weights() done");
    return eta;
}

template <typename Scalar, typename Kernel>
Scalar EasyMKL<Scalar, Kernel>::compute_bias(const Column &gamma,
                                             const Column &directed_gamma,
                                             const Column &eta)
{
    return 0.5 * gamma.transpose() * (_kernels_handler.sum(eta)) * directed_gamma;
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::setup(std::vector<std::string> labels)
{
    _logger->debug("setup() start");
    _number_of_support_vectors = _kernels_handler.get_lhs_size();
    _number_of_kernels = _kernels_handler.get_number_of_kernels();
    if (_number_of_support_vectors != labels.size())
        throw std::length_error("sample size and lables size do not match");

    // manage labels
    _class_map = mimkl::data_structures::indexing_from_vector_of_strings(labels);

    std::sort(labels.begin(), labels.end()); // pass labels as value
    _unique_labels = labels;
    auto last = std::unique(_unique_labels.begin(), _unique_labels.end());
    _unique_labels.erase(last, _unique_labels.end());
    _number_of_classes = _unique_labels.size();
    _binary = _number_of_classes < 3;

    // treat the binary classification case
    _number_of_dichotomies = _binary ? 1 : _number_of_classes;
    _logger->debug("is binary : {}", _binary);
    _logger->debug("numer of dichotomies: {}", _number_of_dichotomies);

    _gammas.resize(_number_of_support_vectors, _number_of_dichotomies);
    _etas.resize(_number_of_kernels, _number_of_dichotomies);
    _biases.resize(1, _number_of_dichotomies);

    _kernels_sum = _kernels_handler.sum();
    _logger->debug("setup() done");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::train(const Index &class_index)
{
    Column y = get_dichotomy(_unique_labels[class_index]);
    _logger->debug("get_dichotomy() done");
    // gamma on plain sum
    std::pair<Column, Column> gamma_pair =
    optimize_gamma(_regularization_factor, y);
    _logger->debug("optimize_gamma() done");
    Column eta = compute_weights(gamma_pair.second);
    _logger->debug("compute_weights() done");
    // gamma on weighted sum, like KOMD on given kernel
    gamma_pair = optimize_gamma(false, y, eta);
    _logger->debug("optimize_gamma() weighted done");
    Scalar bias = compute_bias(gamma_pair.first, gamma_pair.second, eta);
    _logger->debug("compute_bias() done");

    // update state for the given dichotomy
    _gammas.col(class_index) << gamma_pair.first;
    _etas.col(class_index) << eta;
    _biases.col(class_index) << bias;
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::fit(const std::vector<std::string> &labels)
{
    _logger->debug("fit() start");
    setup(labels);
    // or iterate over class_index?
    for (Index i = 0; i < _number_of_dichotomies; ++i)
        train(i);
    _trained = true;
    _logger->debug("fit() done");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::fit(const Matrix &X,
                                  const std::vector<std::string> &labels)
{
    _logger->debug("fit() data start");
    _kernels_handler.set_lhs(X);
    fit(labels);
    _logger->debug("fit() data done");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::fit(const std::vector<Matrix> &kernel_matrices,
                                  const std::vector<std::string> &labels)
{
    _logger->debug("fit() kernel matrices start");
    _kernels_handler.set_matrices(kernel_matrices, true);
    fit(labels);
    _logger->debug("fit() kernel matrices done");
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Column EasyMKL<Scalar, Kernel>::distances(
const Column &y, const Column &gamma, const Column &eta, const Scalar &bias)
{
    _logger->debug("about to break?");
    return (_kernels_handler.sum(eta).transpose() * y.cwiseProduct(gamma)).array() -
           bias;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::decision_function()
{
    Matrix D(_kernels_handler.get_rhs_size(), _number_of_classes);
    _logger->debug("decision_function() start");
    for (Index i = 0; i < _number_of_dichotomies; ++i)
    {
        D.col(i) << distances(get_dichotomy(_unique_labels[i]), _gammas.col(i),
                              _etas.col(i), _biases(i));
    }
    if (_binary)
    {
        D.col(1) << -D.col(0);
    }
    _logger->debug("decision_function() done");
    return D;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::decision_function(const Matrix &X)
{
    if (!_trained)
        throw std::logic_error("The model should be trained first (after "
                               "instantiation or change in parameters)");
    _kernels_handler.set_rhs(X);
    return decision_function();
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::decision_function(const std::vector<Matrix> &kernel_matrices)
{
    if (!_trained)
        throw std::logic_error("The model should be trained first (after "
                               "instantiation or change in parameters)");
    if (kernel_matrices[0].rows() != _number_of_support_vectors)
        throw std::length_error("Similarities must be provided for all support "
                                "vectors; matrices have wrong number of rows.");
    if (kernel_matrices.size() != _number_of_kernels)
        throw std::length_error(
        "Same number of kernels as on training is required");
    _kernels_handler.set_matrices(kernel_matrices);
    return decision_function();
}

template <typename Scalar, typename Kernel>
std::vector<std::string>
EasyMKL<Scalar, Kernel>::predict(const std::vector<Matrix> &kernel_matrices)
{
    _logger->debug("predict() matricial start");
    Matrix D = decision_function(kernel_matrices);
    std::vector<std::string> prediction;
    prediction.reserve(_kernels_handler.get_rhs_size());
    // argmax
    typename Matrix::Index max_index;
    for (Index i = 0; i < _kernels_handler.get_rhs_size(); ++i)
    {
        D.row(i).maxCoeff(&max_index);
        _logger->trace("max_index[{}]:\t {}", i, max_index);
        prediction.push_back(_unique_labels[max_index]);
    }
    _logger->debug("predict() matricial done");
    return prediction;
}

template <typename Scalar, typename Kernel>
std::vector<std::string> EasyMKL<Scalar, Kernel>::predict(const Matrix &X)
{
    _logger->debug("predict() functional start");
    Matrix D = decision_function(X);
    std::vector<std::string> prediction;
    prediction.reserve(_kernels_handler.get_rhs_size());
    // argmax
    typename Matrix::Index max_index;
    for (Index i = 0; i < _kernels_handler.get_rhs_size(); ++i)
    {
        D.row(i).maxCoeff(&max_index);
        _logger->trace("max_index[{}]:\t {}", i, max_index);
        prediction.push_back(_unique_labels[max_index]);
    }
    _logger->debug("predict() functional done");
    return prediction;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::predict_proba(const std::vector<Matrix> &kernel_matrices)
{
    _logger->debug("predict_proba() matricial done");
    Matrix D = decision_function(kernel_matrices);
    _logger->debug("predict_proba() matricial done");
    return mimkl::linear_algebra::rowwise_soft_max(D);
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::predict_proba(const Matrix &X)
{
    _logger->debug("predict_proba() functional done");
    Matrix D = decision_function(X);
    _logger->debug("predict_proba() functional done");
    return mimkl::linear_algebra::rowwise_soft_max(D);
}

template <typename Scalar, typename Kernel>
double EasyMKL<Scalar, Kernel>::get_lambda() const
{
    return _lambda;
}

template <typename Scalar, typename Kernel>
double EasyMKL<Scalar, Kernel>::get_epsilon() const
{
    return _epsilon;
}

template <typename Scalar, typename Kernel>
bool EasyMKL<Scalar, Kernel>::get_regularization_factor() const
{
    return _regularization_factor;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix EasyMKL<Scalar, Kernel>::get_gammas() const
{
    return _gammas;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix EasyMKL<Scalar, Kernel>::get_etas() const
{
    return _etas;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Row EasyMKL<Scalar, Kernel>::get_biases() const
{
    return _biases;
}

template <typename Scalar, typename Kernel>
Indexing EasyMKL<Scalar, Kernel>::get_class_map() const
{
    return _class_map;
}

template <typename Scalar, typename Kernel>
std::vector<std::string> EasyMKL<Scalar, Kernel>::get_one_versus_rest_order() const
{
    return _unique_labels;
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::get_optimal_kernel()
{
    if (!_trained)
        throw std::logic_error(
        "The model has not been fit, kernel weights are undetermined");
    return _kernels_handler.sum(
    _etas.rowwise().mean()); // TODO treat exception for binary problem
}

template <typename Scalar, typename Kernel>
typename EasyMKL<Scalar, Kernel>::Matrix
EasyMKL<Scalar, Kernel>::get_optimal_kernel_by_class_index(const Index i)
{
    if (!_trained)
        throw std::logic_error(
        "The model has not been fit, kernel weights are undetermined");
    return _kernels_handler.sum(_etas.col(i));
}

template <typename Scalar, typename Kernel>
std::vector<typename EasyMKL<Scalar, Kernel>::Matrix>
EasyMKL<Scalar, Kernel>::get_optimal_kernels()
{
    if (!_trained)
        throw std::logic_error(
        "The model has not been fit, kernel weights are undetermined");
    std::vector<Matrix> all_optimal_kernels;
    all_optimal_kernels.reserve(_number_of_dichotomies);
    for (Index i = 0; i < _number_of_dichotomies; ++i)
        all_optimal_kernels.push_back(_kernels_handler.sum(_etas.col(i)));
    return all_optimal_kernels;
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_lambda(const double lambda)
{
    _lambda = lambda;
    _trained = false;
    _logger->debug("changing parameters requires refitting before prediction");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_epsilon(const double epsilon)
{
    _epsilon = epsilon;
    _trained = false;
    _logger->debug("changing parameters requires refitting before prediction");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_regularization_factor(
const bool regularization_factor)
{
    _regularization_factor = regularization_factor;
    _trained = false;
    _logger->debug("changing parameters requires refitting before prediction");
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_parameters(const double lambda,
                                             const double epsilon,
                                             const bool regularization_factor)
{
    _lambda = lambda;
    _epsilon = epsilon;
    _regularization_factor = regularization_factor;
    _trained = false;
    _logger->debug("changing parameters requires refitting before prediction");
    _logger->debug("lambda:\n{}", _lambda);
    _logger->debug("epsilon:\n{}", _epsilon);
}

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_gammas(const Matrix &gammas)
{
    _gammas = gammas;
};

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_etas(const Matrix &etas)
{
    _etas = etas;
};

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_biases(const Row &biases)
{
    _biases = biases;
};

template <typename Scalar, typename Kernel>
void EasyMKL<Scalar, Kernel>::set_class_map(const Indexing &class_map)
{
    _class_map = class_map;
    // get companion objects
    std::vector<std::string> labels;
    labels.reserve(class_map.size());
    // keep linear complexity
    for (Indexing::const_iterator it = _class_map.begin();
         it != _class_map.end();)
    {
        const auto key = it->first;
        labels.push_back(key);
        do
        {
            ++it;
        } while (it != _class_map.end() && key == it->first);
    }
    std::sort(labels.begin(), labels.end()); // pass labels as value
    _unique_labels = labels;
    auto last = std::unique(_unique_labels.begin(), _unique_labels.end());
    _unique_labels.erase(last, _unique_labels.end());

    _number_of_classes = _unique_labels.size();
    _binary = _number_of_classes < 3;
    _number_of_dichotomies = _binary ? 1 : _number_of_classes;
};

} // namespace models
} // namespace mimkl

#endif /* INCLUDE_MIMKL_MODELS_EASY_MKL_HPP_ */


#include <g2o/core/base_vertex.h>
#include <g2o/core/base_binary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/auto_differentiation.h>
#include <g2o/solvers/csparse/linear_solver_csparse.h>
#include <g2o/core/robust_kernel_impl.h>
#include <g2o/types/slam3d/se3quat.h>
#include <g2o/types/slam3d/vertex_pointxyz.h>
#include <g2o/types/sba/types_six_dof_expmap.h>
#include <iostream>

#include "common.h"
#include <sophus/se3.hpp>
#include <sophus/so3.hpp>
#include <Eigen/Dense>

using namespace Sophus;
using namespace Eigen;
using namespace std;

struct Camera
{
    Camera() {}

    Camera(double* data)
    {
        Rt = g2o::SE3Quat::exp(Eigen::Map<Eigen::Matrix<double, 6, 1>>(data));

        f = data[6];
        k1 = data[7];
        k2 = data[8];
    }

    void set_to(double* data) const
    {
        Eigen::Matrix<double, 6, 1> rt = Rt.log();

        for (int i = 0; i < 6; ++i)
            data[i] = rt[i];

        data[6] = f;
        data[7] = k1;
        data[8] = k2;
    }

    g2o::SE3Quat Rt; // g2o::SE3Quat stores rotation first and then translation
    double f = 0.0, k1 = 0.0, k2 = 0.0;
};

class VertexCamera: public g2o::BaseVertex<9, Camera>
{
    public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    virtual void setToOriginImpl() override {

        _estimate = Camera();
    }

    virtual void oplusImpl(const double *update) override {
        _estimate.Rt = g2o::SE3Quat::exp(Eigen::Map<const Eigen::Matrix<double, 6, 1>>(update)) * _estimate.Rt;
        _estimate.f += update[6];
        _estimate.k1 += update[7];
        _estimate.k2 += update[8];
    }

    virtual bool read(std::istream&) override {}
    virtual bool write(std::ostream&) const override {}
};

class VertexLandmark: public g2o::BaseVertex<3, Eigen::Vector3d>
{
    public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    virtual void setToOriginImpl() override {
        _estimate = Eigen::Vector3d::Zero();
    }

    virtual void oplusImpl(const double *update) override {
        _estimate += Eigen::Map<const Eigen::Vector3d>(update);
    }

    virtual bool read(std::istream&) override {}
    virtual bool write(std::ostream&) const override {}
};

class EdgeReprojection: public g2o::BaseBinaryEdge<2, Eigen::Vector2d, VertexLandmark, VertexCamera>
{
  public:
    EIGEN_MAKE_ALIGNED_OPERATOR_NEW;

    EdgeReprojection(double f, double k1, double k2) : f_(f), k1_(k1), k2_(k2) { }

    virtual void computeError() override {
        const VertexCamera* v_cam = static_cast<VertexCamera*>(_vertices[1]);
        const VertexLandmark* v_point = static_cast<VertexLandmark*>(_vertices[0]);
        auto cam = v_cam->estimate();
        Eigen::Vector3d X = v_point->estimate();
        Eigen::Vector3d X_cam = cam.Rt * X;
        X_cam /= X_cam.z();

        auto p2 = X_cam.x() * X_cam.x() + X_cam.y() * X_cam.y();
        auto r = 1.0 + p2 * (cam.k1 + (p2 * cam.k2));

        Eigen::Vector2d uv = -X_cam.head<2>() *  cam.f * r; // minus because of the dataset projection
        _error = _measurement - uv;
    }

    virtual void linearizeOplus()  override  {
        const VertexCamera *v_cam = static_cast<VertexCamera*>(_vertices[1]);
        const VertexLandmark *v_landmark = static_cast<VertexLandmark*>(_vertices[0]);
        Eigen::Vector3d X3d = v_landmark->estimate();
        auto cam = v_cam->estimate();
        Eigen::Vector3d Xc = cam.Rt * X3d;

        double X = Xc.x();
        double Y = Xc.y();
        double Z = Xc.z();
        double Z_2 = Z * Z;
        double f = cam.f;
        double k1 = cam.k1;
        double k2 = cam.k2;

        double x = X / Z;
        double y = Y / Z;
        double x2 = x*x;
        double y2 = y*y;
        double n2 = x*x + y*y;
        double n4 = n2 * n2;
        double r = 1.0 + n2 * k1 + n4 * k2;

        Eigen::Matrix2d dedxd = -Eigen::Matrix2d::Identity() * f; // minus because of the dataset projection

        Eigen::Matrix2d dxddxp = Eigen::Matrix2d::Identity();

        dxddxp(0, 0) = (2*k1*x+4*k2*x2*x+4*y2*x*k2)*x + r;
        dxddxp(0, 1) = x*(2*k1*y+4*x2*k2*y+4*k2*y2*y);
        dxddxp(1, 0) = y*(2*k1*x+4*k2*x2*x+4*y2*k2*x);
        dxddxp(1, 1) = (2*k1*y+4*x2*y*k2+4*k2*y2*y)*y + r;

        Eigen::Matrix<double, 2, 3> dxpdXc;
        dxpdXc << 1.0/Z, 0.0, -X / (Z_2),
                  0.0, 1.0/Z, -Y / (Z_2);

        Eigen::Matrix<double, 2, 3> dedXc = dedxd * dxddxp * dxpdXc;

        Eigen::Matrix3d dXcdR;
        dXcdR << 0.0, -Z, Y,
                 Z, 0.0, -X,
                 -Y, X, 0.0;

        Eigen::Matrix3d dXcdt = Eigen::Matrix3d::Identity();

        Eigen::Vector2d dedf(-r*x, -r*y);   // minus because of the dataset projection
        Eigen::Vector2d dxddk1(x * n2, y * n2);
        Eigen::Vector2d dxddk2(x * n4, y * n4);
        Eigen::Vector2d dedk1 = dedxd * dxddk1;
        Eigen::Vector2d dedk2 = dedxd * dxddk2;

        _jacobianOplusXj.block<2, 3>(0, 0) = dedXc * -dXcdR; // -dXcdR <=> -Xc^
        _jacobianOplusXj.block<2, 3>(0, 3) = dedXc * dXcdt;
        _jacobianOplusXj.col(6) = dedf;
        _jacobianOplusXj.col(7) = dedk1;
        _jacobianOplusXj.col(8) = dedk2;

        _jacobianOplusXi = dedXc * cam.Rt.rotation().matrix();

         _jacobianOplusXi *= -1; // take the negative jacobian because the error is defined as (measurement - projection)
         _jacobianOplusXj *= -1;

    }

    virtual bool read(std::istream&) override {}
    virtual bool write(std::ostream&) const override {}
    private:
        double f_, k1_, k2_;

};

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

    if (argc != 2) {
        cout << "usage: bundle_adjustment_g2o bal_data.txt" << endl;
        return 1;
    }

    BALProblem dataset(argv[1]);
    dataset.Normalize();
    dataset.Perturb(0.1, 0.5, 0.5);
    dataset.WriteToPLYFile("initial_pc.ply");

    std::cout << "\n";
    std::cout << "nb cameras: " << dataset.num_cameras() << std::endl;
    std::cout << "nb landmarks: " << dataset.num_points() << std::endl;
    std::cout << "nb observations: " << dataset.num_observations() << std::endl;
    std::cout << "nb parameters: " << dataset.num_parameters() << std::endl;
    std::cout << "check: " << dataset.num_cameras() * 9 + dataset.num_points()*3 << std::endl;

    // pose dimension 9, landmark is 3
    typedef g2o::BlockSolver<g2o::BlockSolverTraits<9, 3>> BlockSolverType;
    typedef g2o::LinearSolverCSparse<BlockSolverType::PoseMatrixType> LinearSolverType;

    auto solver = new g2o::OptimizationAlgorithmLevenberg(
        g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>())
    );
    g2o::SparseOptimizer optimizer;
    optimizer.setAlgorithm(solver);
    optimizer.setVerbose(true);

    auto* cameras = dataset.mutable_cameras();
    std::vector<VertexCamera*> camera_vertices;
    for (int i = 0; i < dataset.num_cameras(); ++i)
    {
        auto *c = new VertexCamera();
        c->setId(i);
        c->setEstimate(Camera(cameras + (i*dataset.camera_block_size())));
        optimizer.addVertex(c);
        camera_vertices.push_back(c);
    }

    auto* landmarks = dataset.mutable_points();
    std::vector<VertexLandmark*> landmark_vertices;
    for (int i = 0; i < dataset.num_points(); ++i)
    {
        auto* l = new VertexLandmark();
        l->setId(dataset.num_cameras() + i);
        l->setEstimate(Eigen::Map<Eigen::Vector3d>(landmarks + i*dataset.point_block_size()));
        l->setMarginalized(true);
        optimizer.addVertex(l);
        landmark_vertices.push_back(l);
    }

    auto* observations = dataset.observations();
    auto* cam_indices = dataset.camera_index();
    auto* landmark_indices = dataset.point_index();
    for (int i = 0; i < dataset.num_observations(); ++i)
    {
        auto c = Camera(cameras + (cam_indices[i]*dataset.camera_block_size()));
        auto* e = new EdgeReprojection(c.f, c.k1, c.k2);
        e->setVertex(1, camera_vertices[cam_indices[i]]);
        e->setVertex(0, landmark_vertices[landmark_indices[i]]);
        e->setMeasurement(Eigen::Map<const Eigen::Vector2d>(observations + i*2));
        e->setInformation(Eigen::Matrix2d::Identity());
        optimizer.addEdge(e);
    }

    optimizer.initializeOptimization();
    optimizer.optimize(40);

    for (int i = 0; i < dataset.num_cameras(); ++i)
    {
        camera_vertices[i]->estimate().set_to(cameras + (i * dataset.camera_block_size()));
    }
    for (int i = 0; i < dataset.num_points(); ++i)
    {
        Eigen::Vector3d X = landmark_vertices[i]->estimate();
        landmarks[i*3] = X.x();
        landmarks[i*3+1] = X.y();
        landmarks[i*3+2] = X.z();
    }

    dataset.WriteToPLYFile("after_ba_g2o.ply");

    return 0;
}


#include <gmp_lib/CanonicalSystem/CanonicalSystem.h>

#include <cstdlib>
#include <cmath>
#include <array>
#include <exception>

#include <boost/numeric/odeint.hpp>

namespace as64_
{

namespace gmp_
{

typedef std::array<double,2> can_sys_state;
typedef boost::numeric::odeint::runge_kutta_cash_karp54< can_sys_state > can_sys_error_stepper_type;
typedef boost::numeric::odeint::controlled_runge_kutta< can_sys_error_stepper_type > can_sys_controlled_stepper_type;

CanonicalSystem::CanonicalSystem(double T, double Ds)
{
  this->s0 = 0;
  this->sf = 1;
  this->Ds = Ds;

  this->sd_dot = 1/T;

  this->s = 0;
  this->s_dot = this->sd_dot;
}

void CanonicalSystem::integrate(double t0, double tf)
{
  static can_sys_controlled_stepper_type ctrl_stepper;

  if (tf < t0) throw std::runtime_error("Cannot integrate backwards in time: t0 > tf");

  if (std::isinf(this->Ds))
  {
    this->s_dot = this->sd_dot;
    this->s = this->s + this->s_dot*(tf - t0);

    if (this->s > this->sf) this->s = this->sf;
    else if (this->s < this->s0) this->s = this->s0;

    return;
  }

  auto ode_fun = [this](const can_sys_state &state, can_sys_state &state_dot, double t)
  {
    state_dot[0] = state[1];
    state_dot[1] = this->getPhaseDDot(state[0], state[1]);
  };

  can_sys_state state = {s, s_dot};
  boost::numeric::odeint::integrate_adaptive(ctrl_stepper, ode_fun, state, t0, tf, tf-t0);
  s = state[0];
  s_dot = state[1];
}

void CanonicalSystem::setDuration(double T, double t)
{
  if (T < t) throw std::runtime_error("The current time has already exceeded the duration");

  this->sd_dot = (this->sf - this->s) / (T - t);
}

void CanonicalSystem::setRemainingDuration(double T)
{
  this->sd_dot = (this->sf - this->s) / T;
}

double CanonicalSystem::getPhaseDDot(double s, double s_dot)
{
  double s_ddot = 0;
  if (s>=this->s0 && s < this->sf) s_ddot = -this->Ds*(s_dot - this->sd_dot);
  else if (s>=1)                   s_ddot = -400*s_dot - 1000*(s-this->sf);
  else                             s_ddot = -400*s_dot - 1000*(s-this->s0);

  return s_ddot;
}

void CanonicalSystem::reset()
{
  this->s = this->s0;
  this->s_dot = this->sd_dot;
}

} // namespace gmp_

} // namespace as64_