/*********************************************************************
* Software License Agreement (BSD License)
*
*  Copyright (c) 2014, Rice University
*  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 Rice University 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 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.
*********************************************************************/

/* Author: Caleb Voss */

#include "ompl/base/spaces/constraint/AtlasChart.h"
#include <boost/math/constants/constants.hpp>
#include <eigen3/Eigen/Dense>

/// AtlasChart::Halfspace

/// Public

ompl::base::AtlasChart::Halfspace::Halfspace(const AtlasChart *owner, const AtlasChart *neighbor) : owner_(owner)
{
    // Project neighbor's chart center onto our chart.
    Eigen::VectorXd u(owner_->k_);
    owner_->psiInverse(*neighbor->getOrigin(), u);

    // Compute the halfspace equation, which is the perpendicular bisector
    // between 0 and u (plus 5% to reduce cracks, see Jaillet et al.).
    setU(1.05 * u);
}

bool ompl::base::AtlasChart::Halfspace::contains(const Eigen::Ref<const Eigen::VectorXd> &v) const
{
    return v.dot(u_) <= rhs_;
}

void ompl::base::AtlasChart::Halfspace::checkNear(const Eigen::Ref<const Eigen::VectorXd> &v) const
{
    // Threshold is 10% of the distance from the boundary to the origin.
    if (distanceToPoint(v) < 1.0 / 20)
    {
        Eigen::VectorXd x(owner_->n_);
        owner_->psi(v, x);
        complement_->expandToInclude(x);
    }
}

bool ompl::base::AtlasChart::Halfspace::circleIntersect(const double r, Eigen::Ref<Eigen::VectorXd> v1,
                                                        Eigen::Ref<Eigen::VectorXd> v2) const
{
    if (owner_->getManifoldDimension() != 2)
        throw ompl::Exception("ompl::base::AtlasChart::Halfspace::circleIntersect() "
                              "Only works on 2D manifolds.");

    // Check if there will be no solutions.
    double discr = 4 * r * r - usqnorm_;
    if (discr < 0)
        return false;
    discr = std::sqrt(discr);

    // Compute the 2 solutions (possibly 1 repeated solution).
    double unorm = std::sqrt(usqnorm_);
    v1[0] = -u_[1] * discr;
    v1[1] = u_[0] * discr;
    v2 = -v1;
    v1 += u_ * unorm;
    v2 += u_ * unorm;
    v1 /= 2 * unorm;
    v2 /= 2 * unorm;

    return true;
}

/// Public static

void ompl::base::AtlasChart::Halfspace::intersect(const Halfspace &l1, const Halfspace &l2,
                                                  Eigen::Ref<Eigen::VectorXd> out)
{
    if (l1.owner_ != l2.owner_)
        throw ompl::Exception("Cannot intersect linear inequalities on different charts.");
    if (l1.owner_->getManifoldDimension() != 2)
        throw ompl::Exception("AtlasChart::Halfspace::intersect() only works on 2D manifolds.");

    // Computer the intersection point of these lines.
    Eigen::MatrixXd A(2, 2);
    A.row(0) = l1.u_.transpose();
    A.row(1) = l2.u_.transpose();
    out[0] = l1.u_.squaredNorm();
    out[1] = l2.u_.squaredNorm();
    out = 0.5 * A.inverse() * out;
}

/// Private

void ompl::base::AtlasChart::Halfspace::setU(const Eigen::Ref<const Eigen::VectorXd> &u)
{
    u_ = u;

    // Precompute the squared norm of u.
    usqnorm_ = u_.squaredNorm();

    // Precompute the right-hand side of the linear inequality.
    rhs_ = usqnorm_ / 2;
}

double ompl::base::AtlasChart::Halfspace::distanceToPoint(const Eigen::Ref<const Eigen::VectorXd> &v) const
{
    // Result is a scalar factor of u_.
    return (0.5 - v.dot(u_)) / usqnorm_;
}

void ompl::base::AtlasChart::Halfspace::expandToInclude(const Eigen::Ref<const Eigen::VectorXd> &x)
{
    // Compute how far v = psiInverse(x) lies past the boundary, if at all.
    Eigen::VectorXd v(owner_->k_);
    owner_->psiInverse(x, v);
    const double t = -distanceToPoint(v);

    // Move u_ further out by twice that much.
    if (t > 0)
        setU((1 + 2 * t) * u_);
}

/// AtlasChart

/// Public

ompl::base::AtlasChart::AtlasChart(const AtlasStateSpace *atlas, const AtlasStateSpace::StateType *state)
  : constraint_(atlas->getConstraint().get())
  , n_(atlas->getAmbientDimension())
  , k_(atlas->getManifoldDimension())
  , state_(state)
  , bigPhi_([&]() -> const Eigen::MatrixXd {
      Eigen::MatrixXd j(n_ - k_, n_);
      constraint_->jacobian(*state_, j);

      Eigen::FullPivLU<Eigen::MatrixXd> decomp = j.fullPivLu();
      if (!decomp.isSurjective())
          throw ompl::Exception("Cannot compute full-rank tangent space.");

      // Compute the null space and orthonormalize, which is a basis for the tangent space.
      return decomp.kernel().householderQr().householderQ() * Eigen::MatrixXd::Identity(n_, k_);
  }())
  , radius_(atlas->getRho_s())
{
}

ompl::base::AtlasChart::~AtlasChart()
{
    clear();
}

void ompl::base::AtlasChart::clear()
{
    for (auto h : polytope_)
        delete h;

    polytope_.clear();
}

void ompl::base::AtlasChart::phi(const Eigen::Ref<const Eigen::VectorXd> &u, Eigen::Ref<Eigen::VectorXd> out) const
{
    out = *state_ + bigPhi_ * u;
}

bool ompl::base::AtlasChart::psi(const Eigen::Ref<const Eigen::VectorXd> &u, Eigen::Ref<Eigen::VectorXd> out) const
{
    // Initial guess for Newton's method
    Eigen::VectorXd x0(n_);
    phi(u, x0);

    // Newton-Raphson to solve Ax = b
    unsigned int iter = 0;
    double norm = 0;
    Eigen::MatrixXd A(n_, n_);
    Eigen::VectorXd b(n_);

    const double tolerance = constraint_->getTolerance();
    const double squaredTolerance = tolerance * tolerance;

    // Initialize output to initial guess
    out = x0;

    // Initialize A with orthonormal basis (constant)
    A.block(n_ - k_, 0, k_, n_) = bigPhi_.transpose();

    // Initialize b with initial f(out) = b
    constraint_->function(out, b.head(n_ - k_));
    b.tail(k_).setZero();

    while ((norm = b.squaredNorm()) > squaredTolerance && iter++ < constraint_->getMaxIterations())
    {
        // Recompute the Jacobian at the new guess.
        constraint_->jacobian(out, A.block(0, 0, n_ - k_, n_));

        // Move in the direction that decreases F(out) and is perpendicular to
        // the chart.
        out -= A.partialPivLu().solve(b);

        // Recompute b with new guess.
        constraint_->function(out, b.head(n_ - k_));
        b.tail(k_) = bigPhi_.transpose() * (out - x0);
    }

    return norm < squaredTolerance;
}

void ompl::base::AtlasChart::psiInverse(const Eigen::Ref<const Eigen::VectorXd> &x,
                                        Eigen::Ref<Eigen::VectorXd> out) const
{
    out = bigPhi_.transpose() * (x - *state_);
}

bool ompl::base::AtlasChart::inPolytope(const Eigen::Ref<const Eigen::VectorXd> &u, const Halfspace *const ignore1,
                                        const Halfspace *const ignore2) const
{
    if (u.norm() > radius_)
        return false;

    for (Halfspace *h : polytope_)
    {
        if (h == ignore1 || h == ignore2)
            continue;

        if (!h->contains(u))
            return false;
    }

    return true;
}

void ompl::base::AtlasChart::borderCheck(const Eigen::Ref<const Eigen::VectorXd> &v) const
{
    for (Halfspace *h : polytope_)
        h->checkNear(v);
}

const ompl::base::AtlasChart *ompl::base::AtlasChart::owningNeighbor(const Eigen::Ref<const Eigen::VectorXd> &x) const
{
    Eigen::VectorXd projx(n_), proju(k_);
    for (Halfspace *h : polytope_)
    {
        // Project onto the neighboring chart.
        const AtlasChart *c = h->getComplement()->getOwner();
        c->psiInverse(x, proju);
        c->phi(proju, projx);

        // Check if it's within the validity region and polytope boundary.
        const bool withinTolerance = (projx - x).norm();
        const bool inPolytope = c->inPolytope(proju);

        if (withinTolerance && inPolytope)
            return c;
    }

    return nullptr;
}

bool ompl::base::AtlasChart::toPolygon(std::vector<Eigen::VectorXd> &vertices) const
{
    if (k_ != 2)
        throw ompl::Exception("AtlasChart::toPolygon() only works on 2D manifold/charts.");

    // Compile a list of all the vertices in P and all the times the border
    // intersects the circle.
    Eigen::VectorXd v(2);
    Eigen::VectorXd intersection(n_);
    vertices.clear();
    for (std::size_t i = 0; i < polytope_.size(); i++)
    {
        for (std::size_t j = i + 1; j < polytope_.size(); j++)
        {
            // Check if intersection of the lines is a part of the boundary and
            // within the circle.
            Halfspace::intersect(*polytope_[i], *polytope_[j], v);
            phi(v, intersection);
            if (inPolytope(v, polytope_[i], polytope_[j]))
                vertices.push_back(intersection);
        }

        // Check if intersection with circle is part of the boundary.
        Eigen::VectorXd v1(2), v2(2);
        if ((polytope_[i])->circleIntersect(radius_, v1, v2))
        {
            if (inPolytope(v1, polytope_[i]))
            {
                phi(v1, intersection);
                vertices.push_back(intersection);
            }
            if (inPolytope(v2, polytope_[i]))
            {
                phi(v2, intersection);
                vertices.push_back(intersection);
            }
        }
    }

    // Include points approximating the circle, if they're inside the polytope.
    bool is_frontier = false;
    Eigen::VectorXd v0(2);
    v0 << radius_, 0;
    const double step = boost::math::constants::pi<double>() / 32.;
    for (double a = 0.; a < 2. * boost::math::constants::pi<double>(); a += step)
    {
        const Eigen::VectorXd vn = Eigen::Rotation2Dd(a) * v0;

        if (inPolytope(vn))
        {
            is_frontier = true;
            phi(vn, intersection);
            vertices.push_back(intersection);
        }
    }

    // Put all the points in order.
    std::sort(vertices.begin(), vertices.end(),
              [&](const Eigen::Ref<const Eigen::VectorXd> &x1, const Eigen::Ref<const Eigen::VectorXd> &x2) -> bool {
                  // Check the angles to see who should come first.
                  Eigen::VectorXd v1(2), v2(2);
                  psiInverse(x1, v1);
                  psiInverse(x2, v2);
                  return std::atan2(v1[1], v1[0]) < std::atan2(v2[1], v2[0]);
              });

    return is_frontier;
}

bool ompl::base::AtlasChart::estimateIsFrontier() const
{
    RNG rng;
    Eigen::VectorXd ru(k_);
    for (int k = 0; k < 1000; k++)
    {
        for (int i = 0; i < ru.size(); i++)
            ru[i] = rng.gaussian01();
        ru *= radius_ / ru.norm();
        if (inPolytope(ru))
            return true;
    }
    return false;
}

/// Public Static

void ompl::base::AtlasChart::generateHalfspace(AtlasChart *c1, AtlasChart *c2)
{
    if (c1 == c2)
        throw ompl::Exception("ompl::base::AtlasChart::generateHalfspace(): "
                              "Must use two different charts.");

    // c1, c2 will delete l1, l2, respectively, upon destruction.
    Halfspace *l1, *l2;
    l1 = new Halfspace(c1, c2);
    l2 = new Halfspace(c2, c1);
    l1->setComplement(l2);
    l2->setComplement(l1);
    c1->addBoundary(l1);
    c2->addBoundary(l2);
}

/// Protected

void ompl::base::AtlasChart::addBoundary(Halfspace *halfspace)
{
    polytope_.push_back(halfspace);
}


#include <iostream>
#include <iterator>
#include <iomanip>
#include <fstream>
#include <vector>
#include <numeric>
#include <cmath>
#ifdef _OPENMP
#include <omp.h>
#endif

#include <boost/program_options.hpp>
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/json_parser.hpp>
#include <boost/range/iterator_range.hpp>
#include <boost/scope_exit.hpp>

#if defined(SOLVER_BACKEND_VEXCL)
#  include <amgcl/backend/vexcl.hpp>
   typedef amgcl::backend::vexcl<double> Backend;
#elif defined(SOLVER_BACKEND_CUDA)
#  include <amgcl/backend/cuda.hpp>
#  include <amgcl/relaxation/cusparse_ilu0.hpp>
   typedef amgcl::backend::cuda<double> Backend;
#else
#  ifndef SOLVER_BACKEND_BUILTIN
#    define SOLVER_BACKEND_BUILTIN
#  endif
#  include <amgcl/backend/builtin.hpp>
   typedef amgcl::backend::builtin<double> Backend;
#endif

#include <amgcl/io/binary.hpp>
#include <amgcl/io/mm.hpp>
#include <amgcl/adapter/crs_tuple.hpp>
#include <amgcl/amg.hpp>
#include <amgcl/solver/runtime.hpp>
#include <amgcl/coarsening/runtime.hpp>
#include <amgcl/relaxation/runtime.hpp>
#include <amgcl/relaxation/as_preconditioner.hpp>
#include <amgcl/mpi/make_solver.hpp>
#include <amgcl/mpi/schur_pressure_correction.hpp>
#include <amgcl/mpi/block_preconditioner.hpp>
#include <amgcl/mpi/subdomain_deflation.hpp>
#include <amgcl/mpi/direct_solver/runtime.hpp>
#include <amgcl/profiler.hpp>

namespace amgcl {
    profiler<> prof;
}

using amgcl::prof;
using amgcl::precondition;

//---------------------------------------------------------------------------
std::vector<ptrdiff_t> read_problem(
        const amgcl::mpi::communicator &world,
        const std::string &A_file,
        const std::string &rhs_file,
        const std::string &part_file,
        std::vector<ptrdiff_t> &ptr,
        std::vector<ptrdiff_t> &col,
        std::vector<double>    &val,
        std::vector<double>    &rhs
        )
{
    // Read partition
    ptrdiff_t n, m;
    std::vector<ptrdiff_t> domain(world.size + 1, 0);
    std::vector<int> part;

    std::tie(n, m) = amgcl::io::mm_reader(part_file)(part);
    for(int p : part) {
        ++domain[p+1];
        precondition(p < world.size, "MPI world does not correspond to partition");
    }
    std::partial_sum(domain.begin(), domain.end(), domain.begin());

    ptrdiff_t chunk_beg = domain[world.rank];
    ptrdiff_t chunk_end = domain[world.rank + 1];
    ptrdiff_t chunk     = chunk_end - chunk_beg;

    // Reorder unknowns
    std::vector<ptrdiff_t> order(n);
    for(ptrdiff_t i = 0; i < n; ++i)
        order[i] = domain[part[i]]++;

    std::rotate(domain.begin(), domain.end()-1, domain.end());
    domain[0] = 0;

    // Read matrix chunk
    {
        using namespace amgcl::io;

        std::ifstream A(A_file.c_str(), std::ios::binary);
        precondition(A, "Failed to open matrix file (" + A_file + ")");

        std::ifstream b(rhs_file.c_str(), std::ios::binary);
        precondition(b, "Failed to open rhs file (" + rhs_file + ")");

        ptrdiff_t rows;
        precondition(read(A, rows), "File I/O error");
        precondition(rows == n, "Matrix and partition have incompatible sizes");

        ptr.clear(); ptr.reserve(chunk + 1); ptr.push_back(0);

        std::vector<ptrdiff_t> gptr(n + 1);
        precondition(read(A, gptr), "File I/O error");

        size_t col_beg = sizeof(rows) + sizeof(gptr[0]) * (n + 1);
        size_t val_beg = col_beg + sizeof(col[0]) * gptr.back();
        size_t rhs_beg = 2 * sizeof(ptrdiff_t);

        // Count local nonzeros
        for(ptrdiff_t i = 0; i < n; ++i)
            if (part[i] == world.rank)
                ptr.push_back(gptr[i+1] - gptr[i]);

        std::partial_sum(ptr.begin(), ptr.end(), ptr.begin());

        col.clear(); col.reserve(ptr.back());
        val.clear(); val.reserve(ptr.back());
        rhs.clear(); rhs.reserve(chunk);

        // Read local matrix and rhs stripes
        for(ptrdiff_t i = 0; i < n; ++i) {
            if (part[i] != world.rank) continue;

            ptrdiff_t c;
            A.seekg(col_beg + gptr[i] * sizeof(c));
            for(ptrdiff_t j = gptr[i], e = gptr[i+1]; j < e; ++j) {
                precondition(read(A, c), "File I/O error (1)");
                col.push_back(order[c]);
            }
        }

        for(ptrdiff_t i = 0; i < n; ++i) {
            if (part[i] != world.rank) continue;

            double v;
            A.seekg(val_beg + gptr[i] * sizeof(v));
            for(ptrdiff_t j = gptr[i], e = gptr[i+1]; j < e; ++j) {
                precondition(read(A, v), "File I/O error (2)");
                val.push_back(v);
            }
        }

        for(ptrdiff_t i = 0; i < n; ++i) {
            if (part[i] != world.rank) continue;

            double f;
            b.seekg(rhs_beg + i * sizeof(f));
            precondition(read(b, f), "File I/O error (3)");
            rhs.push_back(f);
        }
    }

    return domain;
}

//---------------------------------------------------------------------------
int main(int argc, char *argv[]) {
    int provided;
    MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided);
    BOOST_SCOPE_EXIT(void) {
        MPI_Finalize();
    } BOOST_SCOPE_EXIT_END

    amgcl::mpi::communicator world(MPI_COMM_WORLD);

    if (world.rank == 0)
        std::cout << "World size: " << world.size << std::endl;

    // Read configuration from command line
    namespace po = boost::program_options;
    using std::string;
    po::options_description desc("Options");

    desc.add_options()
        ("help,h", "show help")
        (
         "matrix,A",
         po::value<string>()->required(),
         "The system matrix in binary format"
        )
        (
         "rhs,f",
         po::value<string>(),
         "The right-hand side in binary format"
        )
        (
         "part,s",
         po::value<string>()->required(),
         "Partitioning of the problem in MatrixMarket format"
        )
        (
         "pmask,m",
         po::value<string>(),
         "The pressure mask in binary format. Or, if the parameter has "
         "the form '%n:m', then each (n+i*m)-th variable is treated as pressure."
        )
        (
         "params,P",
         po::value<string>(),
         "parameter file in json format"
        )
        (
         "prm,p",
         po::value< std::vector<string> >()->multitoken(),
         "Parameters specified as name=value pairs. "
         "May be provided multiple times. Examples:\n"
         "  -p solver.tol=1e-3\n"
         "  -p precond.coarse_enough=300"
        )
        ;

    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);

    if (vm.count("help")) {
        if (world.rank == 0)
            std::cout << desc << std::endl;
        return 0;
    }

    po::notify(vm);

    boost::property_tree::ptree prm;
    if (vm.count("params")) read_json(vm["params"].as<string>(), prm);

    if (vm.count("prm")) {
        for(const string &v : vm["prm"].as<std::vector<string> >()) {
            amgcl::put(prm, v);
        }
    }

    prof.tic("read problem");
    std::vector<ptrdiff_t> ptr;
    std::vector<ptrdiff_t> col;
    std::vector<double>    val;
    std::vector<double>    rhs;

    std::vector<ptrdiff_t> domain = read_problem(
            world,
            vm["matrix"].as<string>(), vm["rhs"].as<string>(), vm["part"].as<string>(),
            ptr, col, val, rhs
            );

    ptrdiff_t chunk = domain[world.rank + 1] - domain[world.rank];
    prof.toc("read problem");

    std::vector<char> pm;
    if(vm.count("pmask")) {
        std::string pmask = vm["pmask"].as<string>();
        prm.put("precond.pmask_size", chunk);

        switch (pmask[0]) {
            case '%':
            case '<':
            case '>':
                prm.put("precond.pmask_pattern", pmask);
                break;
            default:
                precondition(false, "Pressure mask may only be set with a pattern");
        }
    }

    std::function<double(ptrdiff_t,unsigned)> dv = amgcl::mpi::constant_deflation(1);
    prm.put("precond.psolver.num_def_vec", 1);
    prm.put("precond.psolver.def_vec", &dv);

    Backend::params bprm;

#if defined(SOLVER_BACKEND_VEXCL)
    vex::Context ctx(vex::Filter::Env);
    std::cout << ctx << std::endl;
    bprm.q = ctx;
#elif defined(SOLVER_BACKEND_CUDA)
    cusparseCreate(&bprm.cusparse_handle);
#endif

    auto f = Backend::copy_vector(rhs, bprm);
    auto x = Backend::create_vector(chunk, bprm);

    amgcl::backend::clear(*x);

    MPI_Barrier(world);

    prof.tic("setup");
    typedef
        amgcl::mpi::make_solver<
            amgcl::mpi::schur_pressure_correction<
                amgcl::mpi::make_solver<
                    amgcl::mpi::block_preconditioner<
                        amgcl::relaxation::as_preconditioner<Backend, amgcl::runtime::relaxation::wrapper>
                        >,
                    amgcl::runtime::solver::wrapper
                    >,
                amgcl::mpi::subdomain_deflation<
                    amgcl::amg<Backend, amgcl::runtime::coarsening::wrapper, amgcl::runtime::relaxation::wrapper>,
                    amgcl::runtime::solver::wrapper,
                    amgcl::runtime::mpi::direct::solver<double>
                    >
                >,
            amgcl::runtime::solver::wrapper
            > Solver;

    Solver solve(world, std::tie(chunk, ptr, col, val), prm, bprm);
    double tm_setup = prof.toc("setup");

    prof.tic("solve");
    size_t iters;
    double resid;
    std::tie(iters, resid) = solve(*f, *x);
    double tm_solve = prof.toc("solve");

    if (world.rank == 0) {
        std::cout
            << "Iterations: " << iters << std::endl
            << "Error:      " << resid << std::endl
            << std::endl
            << prof << std::endl;

#ifdef _OPENMP
        int nt = omp_get_max_threads();
#else
        int nt = 1;
#endif
        std::ostringstream log_name;
        log_name << "schur_" << domain.back() << "_" << nt << "_" << world.size << ".txt";
        std::ofstream log(log_name.str().c_str(), std::ios::app);
        log << domain.back() << "\t" << nt << "\t" << world.size
            << "\t" << tm_setup << "\t" << tm_solve
            << "\t" << iters << "\t" << std::endl;
    }
}


#include <armadillo>
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <algorithm>
//~ #include <thread>

#include <boost/filesystem.hpp>
#include <boost/filesystem/fstream.hpp>

#include <iomanip>
#include <sstream>

using namespace std;
using namespace arma;
using namespace boost::filesystem;

class Pdb_data {

    int anchor_len;
    int loop_len_1;
    int loop_len_2;
    int n_atoms;

    mat anchor_points_1;
    mat aligned_anchor_points_1;
    mat anchor_points_2;

    mat loop_points_1;
    mat aligned_loop_points_1;
    mat loop_points_2;

    rowvec anchor_atom_found_1;
    rowvec anchor_atom_found_2;

    rowvec residue_atom_found_1;
    rowvec residue_atom_found_2;

    void get_atom(char[], char[]);
    bool correct_atom(char[]);
    double get_double(string, int, int);
    int atom_index(char[]);
    int get_index(char*, char[][5], int);
    void read_number(char[], int, int, double*);
    void read_file(string, mat&, mat&, rowvec&, rowvec&, char[][5], int, int);
    void Kabsch(mat, mat, mat&, mat&, mat&);
    void Align_anchors();
    template <typename T> int sgn(T);
    void ResidueRM_TSD(mat, mat, rowvec, rowvec, double&, double&, int&);
    string format_number(double);

    public:

    Pdb_data(string, string, char[][5], int, int);
    void Retrieve_Residue(int, int, mat&, rowvec&);
    void calculate_DTW(double*);
    void write_modified_file(const char*, const char*, int);

};

void Pdb_data::get_atom(char line[], char atom[]){
    int atom_index = 0;
    for(int i=12; i<16; i++){
        if(isalnum(line[i])){
            atom[atom_index++] = line[i];
        }
    }
    atom[atom_index] = '\0';
}

bool Pdb_data::correct_atom(char atom[]){
    return strcmp(atom,"C")==0 || strcmp(atom,"CA")==0 || strcmp(atom,"N")==0 || strcmp(atom,"O")==0;
}

double Pdb_data::get_double(string line, int start, int length){
    double x = atof(line.substr(start, length).c_str());
    return x;
}

int Pdb_data::atom_index(char atom[]){
    if(strcmp(atom,"N")==0){
        return 0;
    } else if(strcmp(atom,"CA")==0){
        return 1;
    } else if(strcmp(atom,"C")==0){
        return 2;
    } else if(strcmp(atom,"O")==0){
        return 3;
    }

    cout << "Unrecognized atom " << atom << endl;
}

int Pdb_data::get_index(char* atom, char atoms[][5], int n_atoms){
    for(int i=0; i<n_atoms; i++){
        if(strcmp(atom, atoms[i]) == 0){
            return i;
        }
    }
    return -1;
}

void Pdb_data::read_number(char line[], int from, int length, double *result){
    char num[length+1];
    memcpy(num, &line[from], length);
    num[length] = '\0';
    *result = atof(num);
}

void Pdb_data::read_file(string file, mat &anchor_points, mat &loop_points, rowvec &anchor_atom_found, rowvec &residue_atom_found, char atoms[][5], int n_atoms, int anchor_len){
    boost::filesystem::ifstream input(file);
    char buf[81];
    char c_aa[6], l_aa[6], atom[5];
    char beginning[7];
    const char *ATOM = "ATOM  ";
    const char *HETATM = "HETATM";
    int aa_index = -1;
    rowvec g_atom_found;
    mat points;

    c_aa[5] = '\0';
    l_aa[5] = '\0';
    atom[4] = '\0';
    beginning[4] = '\0';

    mat aa_atoms = zeros<mat>(3, n_atoms);
    //mat empty = zeros<mat>(n_atoms,3);
    rowvec l_atom_found = zeros<rowvec>(n_atoms);
    //rowvec empty_vec = zeros<rowvec>(n_atoms);

    // cout << "Reading file " << file << endl;

    while(input.good()){

        input.getline(buf, 81);

        memcpy(beginning, &buf[0], 6);

        beginning[6] = '\0';

        if (strcmp(beginning, ATOM) == 0 || strcmp(beginning, HETATM) == 0)
        {

            // Change Amino Acid
            memcpy(l_aa, &buf[22], 5);

            //cout << c_aa << endl;

            if(strcmp(c_aa, l_aa) != 0){
                //cout << "Change from" << c_aa << "to" << l_aa << " " << strcmp(c_aa, l_aa) << endl;
                memcpy(c_aa, &l_aa[0], 5);
                aa_index++;
                //cout << aa_index << '\n';
                if(aa_index != 0){
                    points = join_horiz(points, aa_atoms);

                    //cout << points << '\n';

                    g_atom_found = join_horiz(g_atom_found, l_atom_found);
                }

                l_atom_found = zeros<rowvec>(n_atoms);
                aa_atoms = zeros<mat>(3, n_atoms);

            }

            get_atom(buf, atom);
            int atom_index = get_index(atom, atoms, n_atoms);

            if(atom_index != -1){
                double x, y, z;
                read_number(buf,30, 8,&x);
                read_number(buf,38, 8,&y);
                read_number(buf,46, 8,&z);

                colvec point = zeros<colvec>(3);
                point(0) = x;
                point(1) = y;
                point(2) = z;

                aa_atoms.col(atom_index) = point;

                l_atom_found(atom_index) = 1.0;
            }

        }
    }

    points = join_horiz(points, aa_atoms);

    g_atom_found = join_horiz(g_atom_found, l_atom_found);

    int points_len = points.n_cols;

    anchor_points = join_horiz(points.cols(0, anchor_len*n_atoms - 1), points.cols(points_len - anchor_len*n_atoms, points_len - 1));

    loop_points = points.cols(anchor_len*n_atoms, points_len - anchor_len*n_atoms - 1);

    anchor_atom_found = join_horiz(g_atom_found.subvec(0, anchor_len*n_atoms - 1), g_atom_found.subvec(points_len - anchor_len*n_atoms, points_len - 1));

    residue_atom_found = g_atom_found.subvec(anchor_len*n_atoms, points_len - anchor_len*n_atoms - 1);

    //~ cout << loop_points << endl;

}

void Pdb_data::Kabsch(mat P, mat Q, mat &anchor_mass_centre_1, mat &anchor_mass_centre_2, mat &Rotation_mat)
{
    mat Middle_mat;

    mat P_translated;
    mat Q_translated;
    mat P_rotated_translated;
    mat P_finish;

    mat U;
    vec s;
    mat V;

    mat R;

    colvec P_centroid;
    colvec Q_centroid;
    rowvec v_ones;

    mat Covariance_matrix;

    int d;
    int n_points;

    // Calculate the number of datapoints
    n_points = P.n_cols;

    //Calculate the coordinates of the centres of mass
    P_centroid=sum(P,1)/n_points;

    Q_centroid=sum(Q,1)/n_points;

    v_ones = ones<rowvec>(n_points);

    //Translate the points so that the cetre of mass coincides with the origin
    P_translated = P - P_centroid*v_ones;

    Q_translated = Q - Q_centroid*v_ones;

    //Calculate the covariance matrix
    Covariance_matrix=P_translated*Q_translated.t()/n_points;

    //Calculate svd
    svd(U,s,V,Covariance_matrix);

    //Ensure a right-handed coordinate system
    d = sgn(det( V*U.t() ));

    Middle_mat = eye<mat>(3,3);

    Middle_mat(2,2) =  d;

    //Calculate the rotation matrix
    R=V*Middle_mat*U.t();

    //~ cout << Q << endl;
    //~ cout << R*(P - P_centroid*v_ones) + Q_centroid*v_ones << endl;

    aligned_anchor_points_1 = R*(P - P_centroid*v_ones) + Q_centroid*v_ones;

    anchor_mass_centre_1 = P_centroid;
    anchor_mass_centre_2 = Q_centroid;
    Rotation_mat = R;

}

void Pdb_data::Align_anchors(){

    mat anchor_mass_centre_1;
    mat anchor_mass_centre_2;
    mat Rotation_mat;
    int len_1;
    uvec logic_vec_1;

    Kabsch(anchor_points_1, anchor_points_2, anchor_mass_centre_1, anchor_mass_centre_2, Rotation_mat);

    rowvec v_ones = ones<rowvec>(sum(residue_atom_found_1));

    logic_vec_1 = find(residue_atom_found_1 == 1);

    aligned_loop_points_1 = loop_points_1;

    aligned_loop_points_1.cols(logic_vec_1) = aligned_loop_points_1.cols(logic_vec_1) - anchor_mass_centre_1*v_ones;

    aligned_loop_points_1.cols(logic_vec_1) = Rotation_mat*aligned_loop_points_1.cols(logic_vec_1);

    aligned_loop_points_1.cols(logic_vec_1) = aligned_loop_points_1.cols(logic_vec_1) + anchor_mass_centre_2*v_ones;

}

template <typename T> int Pdb_data::sgn(T val) {
    return (T(0) < val) - (val < T(0));
}

void Pdb_data::Retrieve_Residue(int data_index, int res_index, mat &Residue_points, rowvec &individual_atoms_found){

    if(data_index == 1)
    {
        Residue_points = aligned_loop_points_1.cols(res_index*n_atoms, (res_index + 1)*n_atoms - 1);

        individual_atoms_found = residue_atom_found_1.subvec(res_index*n_atoms, (res_index + 1)*n_atoms - 1);

    }
    else if(data_index == 2)
    {
        Residue_points = loop_points_2.cols(res_index*n_atoms, (res_index + 1)*n_atoms - 1);

        individual_atoms_found = residue_atom_found_2.subvec(res_index*n_atoms, (res_index + 1)*n_atoms - 1);

    }
    else
    {
        cout << "There are only two datasets" << endl;

    }

}

void Pdb_data::ResidueRM_TSD(mat Res1, mat Res2, rowvec individual_atoms_found_1, rowvec individual_atoms_found_2, double &RMSD, double &TSD, int &N){

    uvec logic_vec_both;

    mat Diff;

    rowvec total_individual_atoms_found;

    mat Res_inboth_1;

    mat Res_inboth_2;

    total_individual_atoms_found = individual_atoms_found_1 % individual_atoms_found_2;

    N = sum(total_individual_atoms_found);

    logic_vec_both = find(total_individual_atoms_found == 1);

    Res_inboth_1 = Res1.cols(logic_vec_both);

    Res_inboth_2 = Res2.cols(logic_vec_both);

    Diff = Res_inboth_1 - Res_inboth_2;

    RMSD = sqrt(1.0/N * sum(sum(Diff % Diff)));

    TSD = sum(sum(Diff % Diff));

}

void Pdb_data::calculate_DTW(double *result)
{

    double RMSD_val;

    double SD_val;

    int N;

    mat Residue_points_1;
    rowvec individual_atoms_found_1;

    mat Residue_points_2;
    rowvec individual_atoms_found_2;

    mat cost_mat = zeros<mat>(loop_len_1, loop_len_2);

    mat SD_mat = zeros<mat>(loop_len_1, loop_len_2);

    mat N_mat = zeros<mat>(loop_len_1, loop_len_2);

    uword index;

    double min_out_of_three;

    rowvec three_values = zeros<rowvec>(3);
    rowvec SD_three_values = zeros<rowvec>(3);
    rowvec N_three_values = zeros<rowvec>(3);

    Retrieve_Residue(1, 0 ,Residue_points_1, individual_atoms_found_1);
    Retrieve_Residue(2, 0 ,Residue_points_2, individual_atoms_found_2);

    ResidueRM_TSD(Residue_points_1, Residue_points_2, individual_atoms_found_1, individual_atoms_found_2, RMSD_val, SD_val, N);

    cost_mat(0,0) = RMSD_val;
    SD_mat(0,0) = SD_val;
    N_mat(0,0) = N;

    for(int i = 1; i < loop_len_1; i++)
    {
        Retrieve_Residue(1, i ,Residue_points_1, individual_atoms_found_1);

        ResidueRM_TSD(Residue_points_1, Residue_points_2, individual_atoms_found_1, individual_atoms_found_2, RMSD_val, SD_val, N);

        cost_mat(i,0) = cost_mat(i-1,0) + RMSD_val;

        SD_mat(i,0) = SD_mat(i-1,0) + SD_val;

        N_mat(i,0) = N_mat(i-1,0) + N;

    }

    Retrieve_Residue(1, 0 ,Residue_points_1, individual_atoms_found_1);

    for(int j = 1; j < loop_len_2; j++)
    {
        Retrieve_Residue(2, j ,Residue_points_2, individual_atoms_found_2);

        ResidueRM_TSD(Residue_points_1, Residue_points_2, individual_atoms_found_1, individual_atoms_found_2, RMSD_val, SD_val, N);

        cost_mat(0,j) = cost_mat(0, j-1) + RMSD_val;

        SD_mat(0,j) = SD_mat(0, j-1) + SD_val;

        N_mat(0,j) = N_mat(0, j-1) + N;

    }

    for(int i = 1; i < loop_len_1; i++)
    {
        // cout << "Here" << endl;
        for(int j = 1; j< loop_len_2; j++)
        {
            Retrieve_Residue(1, i ,Residue_points_1, individual_atoms_found_1);
            Retrieve_Residue(2, j ,Residue_points_2, individual_atoms_found_2);

            ResidueRM_TSD(Residue_points_1, Residue_points_2, individual_atoms_found_1, individual_atoms_found_2, RMSD_val, SD_val, N);

            three_values(0) = cost_mat(i-1, j);
            three_values(1) = cost_mat(i, j-1);
            three_values(2) = cost_mat(i-1, j-1);

            SD_three_values(0) = SD_mat(i-1, j);
            SD_three_values(1) = SD_mat(i, j-1);
            SD_three_values(2) = SD_mat(i-1, j-1);

            N_three_values(0) = N_mat(i-1, j);
            N_three_values(1) = N_mat(i, j-1);
            N_three_values(2) = N_mat(i-1, j-1);

            min_out_of_three = three_values.min(index);

            //cout << i << ' ' << j <<' '  << SD_val << three_values << endl;

            cost_mat(i,j) = min_out_of_three + RMSD_val;

            SD_mat(i, j) = SD_three_values(index) + SD_val;

            N_mat(i, j) = N_three_values(index) + N;

        }
    }

    //~ cout << cost_mat << endl;
    //~
    //~ cout << SD_mat << endl;
    //~
    //~ cout << N_mat << endl;

    N = N_mat(loop_len_1 - 1, loop_len_2 - 1);

    //~ cout << N << endl;
    //~
    //~ cout << sqrt(1.0/N * SD_mat(loop_len_1 - 1, loop_len_2 - 1)) << endl;

    *result = sqrt(1.0/N * SD_mat(loop_len_1 - 1, loop_len_2 - 1));

}

string Pdb_data::format_number(double Number_to_format)
{
    int wspc_missing;

    string Result;

    string complete_string;

    ostringstream Convert;

    Convert << fixed << setprecision(3) << Number_to_format;

    Result = Convert.str();

    wspc_missing = 7 - Result.length();

    complete_string = string(wspc_missing, ' ') + Result;

    return complete_string;

}

void Pdb_data::write_modified_file(const char *original_file, const char *result_file, int data_index)
{
    mat P_mat;

    if(data_index == 1)
    {
        P_mat = join_horiz(aligned_anchor_points_1.cols(0, anchor_len*n_atoms - 1), aligned_loop_points_1);
        P_mat = join_horiz(P_mat, aligned_anchor_points_1.cols(anchor_len*n_atoms, 2 * anchor_len*n_atoms - 1));

    }
    else if(data_index == 2)
    {
        P_mat = join_horiz(anchor_points_2.cols(0, anchor_len*n_atoms - 1), loop_points_2);
        P_mat = join_horiz(P_mat, anchor_points_2.cols(anchor_len*n_atoms, 2 * anchor_len*n_atoms - 1));

    }
    else
    {
        cout << "There are only two datasets" << endl;

    }

    boost::filesystem::ifstream infile(original_file);
    boost::filesystem::ofstream pdb_modified;

    double x,y,z;

    int i;

    pdb_modified.open(result_file);

    string line;

    string x_string;

    string y_string;

    string z_string;

    i=0;

    while(getline(infile, line))
    {

        if (line.compare(13, 3, "CA ")==0 or line.compare(13, 3, "N  ")==0 or line.compare(13, 3, "C  ")==0 or line.compare(13, 3, "O  ")==0)

        {
            x = P_mat(0,i);

            y = P_mat(1,i);

            z = P_mat(2,i);

            x_string = format_number(x);

            y_string = format_number(y);

            z_string = format_number(z);

            line.replace(31,7,x_string);

            line.replace(39,7,y_string);

            line.replace(47,7,z_string);

            //~ cout << line << line.length() << '\n';

            pdb_modified << line << '\n';

            i++;
        }
    }

    pdb_modified.close();
}

Pdb_data::Pdb_data(string file1, string file2, char atoms[][5], int n_atoms, int anchor_len){

    this->anchor_len = anchor_len;
    this->n_atoms = n_atoms;

    uvec logic_vec;
    rowvec total_anchor_atoms_found;

    //~ thread thread1(read_file, file1, ref(pos1), ref(atom_found1), atoms, n_atoms);
    //~ thread thread2(read_file, file2, ref(pos2), ref(atom_found2), atoms, n_atoms);
    //~
    //~ thread1.join();
    //~ thread2.join();

    // cout << file1 << endl;

    // cout << file2 << endl;

    read_file( file1, anchor_points_1, loop_points_1, anchor_atom_found_1, residue_atom_found_1, atoms, n_atoms, anchor_len);
    read_file( file2, anchor_points_2, loop_points_2, anchor_atom_found_2, residue_atom_found_2, atoms, n_atoms, anchor_len);

    total_anchor_atoms_found = anchor_atom_found_1 % anchor_atom_found_2;

    // cout << anchor_points_1 << endl;

    // cout << anchor_points_2 << endl;

    logic_vec = find(total_anchor_atoms_found == 1);

    anchor_points_1 = anchor_points_1.cols(logic_vec);

    anchor_points_2 = anchor_points_2.cols(logic_vec);

    //~ cout << anchor_points_1 << endl;

    if(anchor_points_1.n_elem != anchor_points_2.n_elem)
    {
        cout << "Anchor Lengths are not the same " << file1 << " " << file2 << endl;
    }

    //~ cout << anchor_points_2 << endl;
    //~
    //~ cout << loop_points_1 << endl;
    //~
    //~ cout << loop_points_2 << endl;

    if(loop_points_1.n_cols % 4 != 0 || loop_points_2.n_cols % 4 != 0)
    {
        cout << "Something's wrong with the number of residues " << file1 << " " << file2 << endl;
    }

    loop_len_1 = loop_points_1.n_cols / 4.0;

    loop_len_2 = loop_points_2.n_cols / 4.0;

    //~ cout << loop_len_1 << " " << loop_len_2 << endl;

    Align_anchors();

    //~ cout << aligned_loop_points_1 << endl;
    //~
    //~ cout << loop_points_2 << endl;

}

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

    boost::filesystem::ofstream distmat_file;
    boost::filesystem::ofstream file_list;

    mat P1_kabsch;

    int anchor_len = 5;

    double result;

    char atoms[4][5];

    strcpy(atoms[0],"N");
    strcpy(atoms[1],"CA");
    strcpy(atoms[2],"C");
    strcpy(atoms[3],"O");

    string filename;

    string extension = "pdb";

    string current_extension;

    int filelen;

    int start;

    vector<string> filenames;

    mat distmat;

    int n_files;

    int row, col;

    int res_anchor_len = 5;
    int start_atom_anchor_len, end_atom_anchor_len;

    const char* directory = argv[1];

    const char* output_file_list_name = argv[2];

    const char* output_distmat_name = argv[3];

    path p(directory);

    // Read old file list
    vector<string> filenames_old;
    int n_oldfiles = 0;
    bool found_old = false;
    if (boost::filesystem::exists(output_file_list_name))
    {
        found_old = true;
	cout << "Found output file list\n";
        boost::filesystem::ifstream file_list_old(output_file_list_name);
        while(getline(file_list_old,filename))
        {
            filenames.push_back(filename);
        }
        n_oldfiles = filenames.size();
        cout << "Number of files: "<<n_oldfiles<<endl;
    }
    // Read old distmat
    mat distmat_old;
    if (found_old & boost::filesystem::exists(output_distmat_name))
    {
        boost::filesystem::ifstream distmat_file_old(output_distmat_name);
        distmat_old.load(distmat_file_old);
    }

    for (auto i = directory_iterator(p); i != directory_iterator(); i++)
    {
        if (!is_directory(i->path())) //we eliminate directories
        {

            filename = i->path().string();

            filelen = filename.length();

            start = filelen - 3;

            current_extension = filename.substr(start,filelen);

            if(!current_extension.compare(extension) & std::find(filenames.begin(),filenames.end(),filename) == filenames.end() ){
                filenames.push_back(filename);
            }

        }
        else
            continue;
    }

    cout << "Reading files and calculating distmat..." << endl;

    file_list.open(output_file_list_name);

    n_files = filenames.size();

    distmat = zeros<mat>(n_files, n_files);

    row = 0;

    for(auto i = filenames.begin(); i != filenames.end(); ++i)
    {

        //cout << row << " " << *i << endl;

        col = row + 1;
        file_list << *i << endl;

        for(auto j = i + 1; j != filenames.end(); ++j)
        {
            //cout << *i << " " << *j << " " << row << " " << col << " ";
            if (found_old & col<n_oldfiles)
            {
		//cout << "skipped with "<<distmat_old(row,col)<<endl;
                distmat(row,col) = distmat_old(row,col);
                distmat(col,row) = distmat_old(col,row);
                col++;
                continue;
            }
            //cout << "calculating"<<endl;
            Pdb_data x(*i, *j, atoms, 4, anchor_len);

            x.calculate_DTW(&result);

            distmat(row,col) = result;

            distmat(col,row) = result;

            col++;

        }

        row++;
    }

    distmat_file.open(output_distmat_name);

    distmat_file << distmat;

    distmat_file.close();

    // cout << distmat << endl;

    // cout << "Done !" << endl;

    // Kabsch(all_Ps[0], all_Ps[2], &rmsd);

    // cout << filenames[0] << " " << filenames[1] << endl;

    // cout << rmsd << endl;
}


#include <fstream>
#include <iostream>
#include <filesystem>
#include <string>
#include <random>
#include <unordered_map>
#include <algorithm>
#include <chrono>
#include <future>
#include <utility>
#include <cmath>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/filter/gzip.hpp>
#include <xtensor/xarray.hpp>
#include <xtensor/xmath.hpp>
#include "csv.hpp"
#include "../assembly/assembly.cpp"
#include "../distance/distance.cpp"
#include "../distribution/distribution.cpp"
#include "../probability/probability.cpp"
#include "../probability/probability_util.cpp"
#include "../domain/domain.cpp"
#include "../fasta/fasta.cpp"
#include "../count/tri_count.cpp"
#include "../codons/swap_codons.cpp"

using namespace std;
using namespace std::chrono;
using namespace csv;

unordered_map<string, vector<double>> MakeInterventions(
	FastaParser parser,
	const xt::xarray<double>& genome_wide_dist,
	const string& intervention,
	const int n_samples,
	mt19937 rng
) {
	unordered_map<string, vector<double>> protein_id_to_distances;
	FastaRecord cds_record{};
	while (parser.Get(cds_record)) {
		auto& sequence = cds_record.content;
		string protein_id = cds_record.id;

		// Perform a number of random interventions on the CDS
		vector<double> distances;
		for (int s = 0; s < n_samples; ++s) {
			string seq = "";
			if (intervention == "swap") {
				seq = SwapSynonymousCodons(sequence, rng);
			} else if (intervention == "shuffle") {
				seq = ShuffleCodons(sequence, rng);
			} else {
				throw runtime_error("Unknown intervention type: " + intervention);
			}

			// Count tri-nucleotides
			bool overlap = true;
			bool use_async = false;
			vector<int> counts = CountTriNucleotides(seq, overlap, use_async);

			// Compute distribution
			xt::xarray<double> cds_distribution = xt::zeros<double>({counts.size()});
			int sum = accumulate(counts.begin(), counts.end(), 0);
			for (int cix = 0; cix < counts.size(); ++cix) {
				cds_distribution[cix] = (double) counts[cix] / (double) sum;
			}

			// Compute distance
			double distance = jensen_shannon_distance(genome_wide_dist, cds_distribution);
			if (isinf(distance) || isnan(distance)) {
				cerr << "MakeInterventions: Nan or Inf encountered for protein id: " << protein_id << endl;
				continue;
			} else {
				distances.push_back(distance);
			}
		}
		if (distances.size() > 0) {
			protein_id_to_distances[protein_id] = distances;
		}
	}
	return protein_id_to_distances;
}

bool task_intervention_per_assembly(
	const int task_nb,
	const string& query,
	const string& tail,
	const string intervention,
	const int n_samples,
	const vector<string>& assembly_ids
) {
	auto start_clock = system_clock::now();

	cerr << "Thread " << task_nb << " started." << endl;

	string dataFolder = "../data/";
	string sequencesFolder = dataFolder + "sequences/";
	auto n_assemblies = assembly_ids.size();

	// Load genome-wide distributions
	string assembly_dist_path = (
		dataFolder + "tri_nucleotide_dist_genome_wide_with_overlap.csv"
	);
	ifstream assembly_dist_f(assembly_dist_path);
	Distributions assembly_distributions(assembly_dist_f);

	// Load Protein domains metadata
	string metadata_path = dataFolder + query + "_master.csv";
	auto metadata = LoadDomainMetadata(metadata_path);

	// Random number generator
	random_device rd;
	mt19937 rng(rd());

	// Compute interventions on every assemblies
	int i = 0;
	for (auto& accession : assembly_ids) {
		auto elapsed = duration_cast<seconds>(system_clock::now() - start_clock).count();
		cerr << "Thread " << task_nb << ": ";
		cerr << "Processing assembly " << i + 1 << " / " << n_assemblies;
		cerr << " (elapsed: " << elapsed << " seconds)" << endl;
		++i;

		// Get genome-wide distribution
		xt::xarray<double> genome_wide_dist = assembly_distributions[accession];

		// Load protein domains to protein ids map
		string protein_domains_path = (
			sequencesFolder + accession + "/" +
			accession + "_" + query + ".csv.gz"
		);
		ifstream protein_domains_file(protein_domains_path);
		boost::iostreams::filtering_streambuf<boost::iostreams::input> inbuf;
		inbuf.push(boost::iostreams::gzip_decompressor());
		inbuf.push(protein_domains_file);
		istream instream(&inbuf);
		ProteinDomains domains(instream);

		// Instanciate CDS parser
		string cds_path = (
			sequencesFolder +
			accession + "/" +
			accession + "_cds_from_genomic.fna.gz"
		);
		ifstream input_file(cds_path, ios_base::in | ios_base::binary);
		boost::iostreams::filtering_streambuf<boost::iostreams::input> inbuf_cds;
		inbuf_cds.push(boost::iostreams::gzip_decompressor());
		inbuf_cds.push(input_file);
		istream instream_cds(&inbuf_cds);
		FastaParser parser(instream_cds);

		// Make interventions: randomly modify sequences a bunch of times and return results as a
		// map of protein ids to the list of distances from the modified sequence distributions
		// the genome-wide distribution.
		unordered_map<string, vector<double>> protein_id_to_distances = MakeInterventions(
			parser,
			genome_wide_dist,
			intervention,
			n_samples,
			rng
		);

		// Compute mean probability per sample
		vector<double> means(n_samples);
		for (int s = 0; s < n_samples; ++s) {
			xt::xarray<double> probs = xt::zeros<double>({protein_id_to_distances.size()});
			int ix = 0;
			for (auto& it : protein_id_to_distances) {
				double distance = it.second[s];
				probs[ix] = compute_probability_from_distance(distance, tail == "left");
				++ix;
			}
			double mean = xt::mean(probs)();
			means[s] = mean;
		}

		// Compute domain probability
		vector<DomainProbability> records;
		for (ProteinDomain& domain : domains.Keys()) {
			// Set description from metadata rather than from annotations
			if (metadata.find(domain.id) != metadata.end()) {
				auto& [domain_query, domain_description] = metadata[domain.id];
				domain.query = domain_query;
				domain.description = domain_description;
			}

			// Get protein ids containing this domain
			auto protein_ids = domains.ProteinIds(domain);

			// Compute gene probabilities for each sample
			int n_records = -1;
			xt::xarray<double> log_probs = xt::zeros<double>({n_samples});
			xt::xarray<double> log_probs_baseline = xt::zeros<double>({n_samples});
			for (int s = 0; s < n_samples; ++s) {
				double mean = means[s];

				xt::xarray<double> probabilities = xt::zeros<double>({protein_ids.size()});
				xt::xarray<double> probabilities_baseline = xt::zeros<double>({protein_ids.size()});
				for (int k = 0; k < protein_ids.size(); ++k) {
					auto& protein_id = protein_ids[k];
					double distance = mean; // Default in case there is no data for protein id
					if (protein_id_to_distances.find(protein_id) != protein_id_to_distances.end()) {
						distance = protein_id_to_distances[protein_id][s];
					}
					probabilities[k] = compute_probability_from_distance(distance, tail == "left");
					probabilities_baseline[k] = mean;
				}

				xt::xarray<double> log_probabilities = xt::eval(xt::log(probabilities));
				xt::xarray<double> log_probabilities_baseline = xt::eval(xt::log(probabilities_baseline));

				double log_prob = product_rule_log(log_probabilities);
				double log_prob_baseline = product_rule_log(log_probabilities_baseline);

				if (n_records == -1) {
					n_records = log_probabilities.size();
				}
				log_probs[s] = log_prob;
				log_probs_baseline[s] = log_prob_baseline;
			}

			// Take the average of all samples
			DomainProbability record(
				domain,
				xt::mean(log_probs)(),
				xt::mean(log_probs_baseline)(),
				n_records
			);
			records.push_back(record);
		}

		// Sort records from best to worse evidence
		sort(records.begin(), records.end(), greater<DomainProbability>());

		// Create assembly directory if it does not exist
		string assembly_domain_prob_out_folder = (
			dataFolder + "intervention/" +  intervention + "/sequences/" + accession + "/"
		);
		filesystem::create_directory(assembly_domain_prob_out_folder);

		// Prepare output writer
		string assembly_domain_prob_out_path = (
			assembly_domain_prob_out_folder +
			accession + "_" + query + "_probability_" + tail + ".csv"
		);
		ofstream of(assembly_domain_prob_out_path);
		auto writer = make_csv_writer(of);
		writer << DomainProbability::RecordHeader();

		// Write assembly outputs
		for (auto& record : records) {
			writer << record.Record();
		}
	}
	return true;
}

bool task_compute_domain_probabilities_per_phylum(
	const int task_nb,
	const string& query,
	const string& tail,
	const string intervention,
	const vector<string>& phyla,
	const unordered_map<string, vector<string>>& assemblies_per_phylum
) {
	cerr << "Thread " << task_nb << " started." << endl;

	string dataFolder = "../data/intervention/" + intervention + "/";
	string sequencesFolder = dataFolder + "sequences/";
	string phylumFolder = dataFolder + "phylum/";
	auto n_phyla = phyla.size();

	// Create phylum directory if it does not exist.
	filesystem::create_directory(phylumFolder);

	int i = 0;
	for (auto& phylum : phyla) {
		auto& assembly_ids = assemblies_per_phylum.at(phylum);
		auto n_assemblies = assembly_ids.size();

		cerr << "Thread " << task_nb << ": ";
		cerr << "Processing phylum " << i + 1 << " / " << n_phyla;
		cerr << ": " << phylum << " (" << n_assemblies << " assemblies)";
		cerr << endl;
		++i;

		set<ProteinDomain> protein_domains;
		unordered_map<ProteinDomain, vector<DomainProbability>> protein_domain_probs;
		for (auto& accession : assembly_ids) {
			string path = (
				sequencesFolder + accession + "/" +
				accession + "_" + query + "_probability_" + tail + ".csv"
			);
			vector<DomainProbability> domains = LoadDomainProbabilities(path);
			for (auto& domain_prob : domains) {
				auto& domain = domain_prob.domain;
				protein_domains.insert(domain);

				if (protein_domain_probs.find(domain) == protein_domain_probs.end()) {
					protein_domain_probs[domain] = vector<DomainProbability>{domain_prob};
				} else {
					protein_domain_probs[domain].push_back(domain_prob);
				}
			}
		}

		string phylum_lower = phylum;
		transform(phylum_lower.begin(), phylum_lower.end(), phylum_lower.begin(), ::tolower);
		transform(phylum_lower.begin(), phylum_lower.end(), phylum_lower.begin(), [](char ch) {
		    return ch == ' ' ? '_' : ch;
		});

		string phylumDir = phylumFolder + phylum_lower + "/";

		filesystem::create_directory(phylumDir);

		string phylum_domain_prob_out_path = (
			phylumDir +
			phylum_lower + "_" + query + "_probability_" + tail + ".csv"
		);
		ofstream of(phylum_domain_prob_out_path);
		auto writer = make_csv_writer(of);
		writer << DomainProbability::RecordHeader();

		vector<DomainProbability> records;
		for (auto& domain : protein_domains) {
			auto& domain_probs = protein_domain_probs[domain];
			auto n_probs = domain_probs.size();
			xt::xarray<double> log_probs = xt::zeros<double>({n_probs});
			xt::xarray<double> log_probs_random = xt::zeros<double>({n_probs});
			for (int ix = 0; ix < n_probs; ++ix) {
				log_probs[ix] = domain_probs[ix].log_probability;
				log_probs_random[ix] = domain_probs[ix].log_probability_random;
			}

			double log_prob = product_rule_log(log_probs);
			double log_prob_random = product_rule_log(log_probs_random);

			try {
				DomainProbability record(
					domain,
					log_prob,
					log_prob_random,
					n_probs
				);
				records.push_back(record);
			}
			catch (exception& e) {
				cerr << "Thread " << task_nb << " | Phylum: " << phylum << " | ";
				cerr << "Exception: " << e.what() << endl;
				throw;
			}
		}

		sort(records.begin(), records.end(), greater<DomainProbability>());

		for (auto& record : records) {
			writer << record.Record();
		}
	}
	cerr << "Thread " << task_nb << ": DONE" << endl;
	return true;
}

void compute_intervention(
	const string query,
	const string tail,
	const string intervention,
	const int n_samples,
	const int n_threads
) {
	auto start_clock = system_clock::now();

	string dataFolder = "../data/";

	bool complete_genome_only = true;
	Assemblies assemblies(dataFolder + "intervention/" + "sample_assemblies.csv");
	auto assembly_ids = assemblies.GetIds();
	auto n_per_thread = ceil((double) assembly_ids.size() / (double) n_threads);

	cerr << "Processing " << assemblies.Size() << " assemblies" << endl;

	//
	// 1) Proceed with intervention and compute domain probability per assembly.
	//
	vector<future<bool>> futures;
	for (int i = 0; i < n_threads; ++i) {
		auto start = assembly_ids.begin() + i * n_per_thread;
		auto end = assembly_ids.end();
		int endInt = i * n_per_thread + n_per_thread;
		if (endInt < assembly_ids.size()) {
			end = assembly_ids.begin() + endInt;
		}
		auto ids = vector<string>(start, end);
		futures.push_back(async(
			task_intervention_per_assembly,
			i+1,
			query,
			tail,
			intervention,
			n_samples,
			ids
		));
	}
	for (auto& f : futures) {
		if(!f.get()) {
			throw runtime_error("Unexpected error while processing assembly output");
		}
	}

	//
	// 2) Compute probability of domains for each phylum.
	//
	cerr << "Processing of domain probabilities per phylum" << endl;

	unordered_map<string, vector<string>> assemblies_per_phylum;
	for (auto& assembly_id : assembly_ids) {
		Assembly& assembly = assemblies.Get(assembly_id);
		string phylum = assembly.phylum;
		if (phylum.empty()) {
			continue;
		}
		if (assemblies_per_phylum.find(phylum) == assemblies_per_phylum.end()) {
			assemblies_per_phylum[phylum] = vector<string>{assembly_id};
		} else {
			assemblies_per_phylum[phylum].push_back(assembly_id);
		}
	}

	set<string> phyla_set;
	for (auto& assembly_id : assembly_ids) {
		Assembly& assembly = assemblies.Get(assembly_id);
		string phylum = assembly.phylum;
		if (!phylum.empty()) {
			phyla_set.insert(phylum);
		}
	}

	vector<string> phyla;
	phyla.assign(phyla_set.begin(), phyla_set.end());
	auto n_phyla = phyla.size();

	cerr << "Processing " << n_phyla << " phyla" << endl;
	cerr << "Starting " << n_threads << " threads" << endl;

	n_per_thread = ceil((double) n_phyla / (double) n_threads);

	vector<future<bool>> futuresP;
	for (int i = 0; i < n_threads; ++i) {
		auto start = phyla.begin() + i * n_per_thread;
		auto end = phyla.end();
		int endInt = i * n_per_thread + n_per_thread;
		if (endInt < n_phyla) {
			end = phyla.begin() + endInt;
		}
		futuresP.push_back(async(
			task_compute_domain_probabilities_per_phylum,
			i+1,
			query,
			tail,
			intervention,
			vector<string>(start, end),
			assemblies_per_phylum
		));
	}
	for (auto& f : futuresP) {
		if(!f.get()) {
			throw runtime_error("Unexpected error while processing phylum output");
		}
	}
	auto elapsed = duration_cast<seconds>(system_clock::now() - start_clock).count();
	cerr << "Processing of domain probabilities per phylum is complete" << endl;
	cerr << "Elapsed: " << elapsed << " seconds" << endl;

	//
	// 3) Compute global probability of domains.
	//
	cerr << "Processing of domain probabilities globally" << endl;
	set<ProteinDomain> protein_domains;
	unordered_map<ProteinDomain, vector<DomainProbability>> protein_domain_probs;
	for (auto& phylum : phyla) {
		string phylum_lower = phylum;
		transform(phylum_lower.begin(), phylum_lower.end(), phylum_lower.begin(), ::tolower);
		transform(phylum_lower.begin(), phylum_lower.end(), phylum_lower.begin(), [](char ch) {
		    return ch == ' ' ? '_' : ch;
		});
		string phylumDir = (
			dataFolder + "intervention/" + intervention + "/phylum/" + phylum_lower + "/"
		);
		string phylum_domain_prob_path = (
			phylumDir +
			phylum_lower + "_" + query + "_probability_" + tail + ".csv"
		);

		vector<DomainProbability> domains = LoadDomainProbabilities(phylum_domain_prob_path);
		for (auto& domain_prob : domains) {
			auto& domain = domain_prob.domain;
			protein_domains.insert(domain);

			if (protein_domain_probs.find(domain) == protein_domain_probs.end()) {
				protein_domain_probs[domain] = vector<DomainProbability>{domain_prob};
			} else {
				protein_domain_probs[domain].push_back(domain_prob);
			}
		}
	}

	const string protein_out_path = (
		dataFolder + "intervention/" + intervention + "/" +
		query + "_probability_" + tail + ".csv"
	);
	ofstream output_file(protein_out_path);
	auto writer = make_csv_writer(output_file);
	writer << DomainProbability::RecordHeader();

	xt::xarray<double> uniform_log_prior = xt::eval(
		xt::log(make_uniform_prior(n_phyla))
	);

	vector<DomainProbability> records;
	for (auto& domain : protein_domains) {
		auto& domain_probs = protein_domain_probs[domain];
		auto n_probs = domain_probs.size();
		xt::xarray<double> log_probs = xt::zeros<double>({n_phyla});
		xt::xarray<double> log_probs_random = xt::zeros<double>({n_phyla});
		for (int ix = 0; ix < n_probs; ++ix) {
			log_probs[ix] = domain_probs[ix].log_probability;
			log_probs_random[ix] = domain_probs[ix].log_probability_random;
		}

		double log_prob = marginalization_log(
			uniform_log_prior,
			log_probs
		);
		double log_prob_random = marginalization_log(
			uniform_log_prior,
			log_probs_random
		);

		try {
			DomainProbability record(
				domain,
				log_prob,
				log_prob_random,
				n_probs
			);
			records.push_back(record);
		}
		catch (exception& e) {
			cerr << "Global computation | ";
			cerr << "Exception: " << e.what() << endl;
			throw;
		}
	}

	sort(records.begin(), records.end(), greater<DomainProbability>());

	for (auto& record : records) {
		writer << record.Record();
	}

	auto elapsed_final = duration_cast<seconds>(system_clock::now() - start_clock).count();
	cerr << "Elapsed: " << elapsed_final << " seconds" << endl;
	cerr << "DONE" << endl;
}


#include "dnn.hpp"
#include "fixedwavefunction.hpp"
#include "hardspherewavefunction.hpp"
#include "inputsorter.hpp"
#include "jastrowmcmillian.hpp"
#include "jastroworion.hpp"
#include "jastrowpade.hpp"
#include "rbmsymmetricwavefunction.hpp"
#include "rbmwavefunction.hpp"
#include "simplegaussian.hpp"
#include "wavefunction.hpp"
#include "wavefunctionpooling.hpp"
#include "wavefunctionproduct.hpp"

#include <Eigen/Dense>
#include <pybind11/eigen.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

namespace py = pybind11;

void init_nn(py::module&);

void init_wavefunction(py::module& main)
{
    auto m  = main.def_submodule("wavefunctions");
    m.doc() = R"doc(
Wavefunctions
-------------

The wavefunction module defines a set of pre-built components that can be used to
set up a variety of different trial wavefunctions.

)doc";

    init_nn(m);

    py::class_<Wavefunction>(m, "Wavefunction", R"doc(
The :class:`Wavefunction` class provides a unified abstraction for all wavefunctions.

All other code expecting a wavefunction instance will take a :class:`Wavefunction` reference.

In this context, a trial wavefunction is defied to be a class that can:

- Have zero or more parameters, :math:`\vec\alpha`
- Evaluate given some system configuraiton :math:`\mathbf{X}`:

.. math::
    \text{Evaluation}_\Psi(\mathbf{X} = \Psi(\mathbf{X})

- Compute the gradient w.r.t. each variational parameter, :math:`\vec{\alpha}`:

.. math::

    \text{Gradient}_\Psi(\mathbf{X}) = \frac{1}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial\vec\alpha}

- Compute the drift force w.r.t. particle :math:`k`'s dimensions ional coordinate :math:`l`:

.. math::

    \text{Drift}_{\Psi(\mathbf{X}), k,l} = \frac{2}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial X_{kl}}

- Compute the laplacian of a system of :math:`N` particles in :math:`D` dimensions:

.. math::

    \text{Laplacian}_\Psi(\mathbf{X}) = \sum_{k=1}^N\sum_{l=1}^D
        \frac{1}{\Psi(\mathbf{X})}\frac{\partial^2 \Psi(\mathbf{X})}{\partial X_{kl}^2}

.. note::
    Wavefunctions are limited to real-valued functions only - there is no
    support for complex valued wavefunctions at this time.

)doc")

        .def("__call__", &Wavefunction::operator(), py::arg("system"), R"doc(
Return the evaluation of the wavefunction for the given system

.. math::

    \text{Evaluation}_\Psi(\mathbf{X} = \Psi(\mathbf{X})
)doc")

        .def("gradient", &Wavefunction::gradient, py::arg("system"), R"doc(
Return the gradient w.r.t. each variational parameter, divided by the evaluation:

.. math::
    \text{Gradient}_\Psi(\mathbf{X}) = \frac{1}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial\vec\alpha}
)doc")

        .def("laplacian", &Wavefunction::laplacian, py::arg("system"), R"doc(

.. math::

    \text{Laplacian}_\Psi(\mathbf{X}) = \sum_{k=1}^N\sum_{l=1}^D
        \frac{1}{\Psi(\mathbf{X})}\frac{\partial^2 \Psi(\mathbf{X})}{\partial X_{kl}^2}
)doc")
        .def("drift_force",
             py::overload_cast<const System&, int, int>(&Wavefunction::drift_force),
             py::arg("system"),
             py::arg("k"),
             py::arg("l"),
             R"doc(

Return the drift force w.r.t. particle :math:`k`'s dimensions ional coordinate :math:`l`:

.. math::

    \text{Drift}_{\Psi(\mathbf{X}), k, l} = \frac{2}{\Psi(\mathbf{X})}
        \frac{\partial \Psi(\mathbf{X})}{\partial X_{kl}}
 )doc")

        .def("drift_force",
             py::overload_cast<const System&>(&Wavefunction::drift_force),
             py::arg("system"),
             R"doc(
Return a list of drift forces for all particles and all dimensions.

The list will be one-dimensional, such that
:math:`\text{Drift}_{\Psi(\mathbf{X}), k,l}` is at index ``k * D + l``.
 )doc")
        .def("symmetry_metric",
             &Wavefunction::symmetry_metric,
             py::arg("sampler"),
             py::arg("samples"),
             py::arg("max_permutations") = 100,
             R"doc(
Return an estimate of the symmetry of the wavefunction.

This is defined as follows:

.. math::
    S(\Psi) =  \frac{\int_{-\infty}^\infty \text{d}\mathbf X\left|\frac{1}{n!}
    \sum_{\mathbf\alpha\in \mathcal{P}_n} P_{\mathbf\alpha}\Psi\right|^2}{\int_{-\infty}^\infty\text{d}\mathbf X
    \max_{\mathbf\alpha\in \mathcal{P}} \left|P_{\mathbf\alpha}\Psi\right|^2}

Here :math:`P_\mathbf{\alpha}\Psi` denotes applying the permutation
:math:`\mathbf{\alpha}` to the input system before evaluating the wavefunction,
and :math:`\mathcal{P}_n` is the set of all permutations of :math:`n` particles.

Properties of this metric:

    - It takes values in :math:`[0, 1]`
    - For symmetric wavefunctions, it equals exactly 1
    - For anti-symmetric wavefunctions, it equals exactly 0
    - Any other function evaluates in :math:`(0, 1)`

This integral will be approximated with Monte Carlo integration using samples from the provided
sampling strategy, and using the specified number of samples. While the integral is defied over all
permutations, a maximum of ``max_permutations`` will be used to make it tractable for large ``N``.
 )doc")
        .def_property(
            "parameters",
            py::overload_cast<>(&Wavefunction::get_parameters, py::const_),
            py::overload_cast<const RowVector&>(&Wavefunction::set_parameters),
            R"doc(
List of all variational parameters. Supports both read and write.
)doc");

    py::class_<FixedWavefunction, Wavefunction>(m, "FixedWavefunction", R"doc(
Wrapper for any wavefunction that fixes the variational parameters.

This means that if this wavefunction is used in a optimization call, its
parameters will not change. This is useful if parts of a wavefunction should be
held constant, while another is allowed to be optimized.
)doc")
        .def(py::init<Wavefunction&>(), py::arg("wavefunction"));

    py::class_<WavefunctionProduct, Wavefunction>(m, "WavefunctionProduct", R"doc(
Defines a wavefunction that acts like the product of two others.

All derivatives will be suitably derived, so this is a simple way to produce
compound expressions.
)doc")
        .def(py::init<Wavefunction&, Wavefunction&>(),
             py::arg("Psi_1"),
             py::arg("Psi_2"));

    py::class_<SumPooling, Wavefunction>(m, "SumPooling", R"doc(
This wavefunction expects a wavefunction of two particles,
:math:`f(\mathbf{x}_1, \mathbf{x}_2)`, and represents the following compound
expression:

.. math::
    \Psi(\mathbf{X}) = \sum_{i \neq j}^N f(\mathbf{X}_i, \mathbf{X}_j)

This is guaranteed to produce a permutation symmetric wavefunction, given any suitable
inner wavefunction :math:`f`.
)doc")
        .def(py::init<Wavefunction&>());

    py::class_<InputSorter, Wavefunction>(m, "InputSorter")
        .def(py::init<Wavefunction&>());

    py::class_<SimpleGaussian, Wavefunction>(m, "SimpleGaussian", R"doc(
A product of gaussians:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha ||\mathbf{X}_i||^2}

The only variational parameter is :math:`\alpha`.

For 3D systems, an optional *fixed* parameter :math:`\beta` can be specified, which changes the above definition to:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha\left(X_{i,1}^2 + X_{i,2}^2 + \beta X_{i,3}^2\right)}
)doc")
        .def(py::init<Real, Real>(), py::arg("alpha") = 0.5, py::arg("beta") = 1);

    py::class_<HardSphereWavefunction, SimpleGaussian>(
        m, "HardSphereWavefunction", R"doc(
A product of gaussians and pairwise correlation factors:

.. math::

    \Psi(\mathbf{X}) = \prod_{i=1}^N e^{-\alpha\left(X_{i,1}^2 + X_{i,2}^2 + \beta X_{i,3}^2\right)}
        \prod_{j = i + 1}^{N} \begin{cases} 0 &\text{if}\ \  ||\mathbf{X}_i - \mathbf{X}_j|| \leq a\\
                                            1 - \frac{a}{||\mathbf{X}_i - \mathbf{X}_j||} &\text{otherwise}
                              \end{cases}

Similarly to :class:`SimpleGaussian`, the only variational parameter is
:math:`\alpha`, while the other two parameters :math:`\beta` and :math:`a` are
assumed constant.
)doc")
        .def(py::init<Real, Real, Real>(),
             py::arg("alpha") = 0.5,
             py::arg("beta")  = 1,
             py::arg("a")     = 0);

    py::class_<JastrowPade, Wavefunction>(m, "JastrowPade", R"doc(
A correlation term meant to be suitable for particles in a harmonic oscillator potential with
a repulsive Coulomb force causing interactions:

.. math::

    \Psi(\mathbf{X}) = \prod_{i < j}^N e^{\frac{\alpha r_ij}{1 + \beta r_ij}}

)doc")
        .def(py::init<Real, Real, bool>(),
             py::arg("alpha")             = 0.5,
             py::arg("beta")              = 1,
             py::arg("alpha_is_constant") = true);

    py::class_<JastrowOrion, Wavefunction>(m, "JastrowOrion", R"doc(
A correlation term meant to be suitable for particles in a harmonic oscillator potential with
a repulsive Coulomb force causing interactions:

.. math::

    \Psi(\mathbf{X}) = \prod_{i < j}^N
        e^{-\frac{\beta^2}{2}||\mathbf{X_i}-\mathbf{X_j}||^2 +
            |\beta\gamma|||\mathbf{X_i}-\mathbf{X_j}||}

)doc")
        .def(py::init<Real, Real>(), py::arg("beta") = 1.0, py::arg("gamma") = 0);

    py::class_<JastrowMcMillian, Wavefunction>(m, "JastrowMcMillian")
        .def(py::init<int, Real, Real>(), py::arg("n"), py::arg("beta"), py::arg("L"));

    py::class_<RBMWavefunction, Wavefunction>(m, "RBMWavefunction")
        .def(py::init<int, int, Real, Real>(),
             py::arg("M"),
             py::arg("N"),
             py::arg("sigma2")      = 1,
             py::arg("root_factor") = 1);

    py::class_<RBMSymmetricWavefunction, Wavefunction>(m, "RBMSymmetricWavefunction")
        .def(py::init<int, int, int, Real, Real>(),
             py::arg("M"),
             py::arg("N"),
             py::arg("f"),
             py::arg("sigma2")      = 1,
             py::arg("root_factor") = 1);

    py::class_<Dnn, Wavefunction>(m, "Dnn")
        .def(py::init<>())
        .def("add_layer", &Dnn::addLayer)
        .def_property_readonly("layers", &Dnn::getLayers);
}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

#include <iostream>
#include <cmath>
#include <string>
#include <vector>
#include <boost/timer.hpp>

#include <boost/numeric/mtl/matrix/parameter.hpp>
#include <boost/numeric/mtl/matrix/dense2D.hpp>
#include <boost/numeric/mtl/matrix/morton_dense.hpp>
#include <boost/numeric/mtl/operation/print_matrix.hpp>
#include <boost/numeric/mtl/operation/matrix_mult.hpp>
#include <boost/numeric/mtl/matrix/hessian_setup.hpp>
#include <boost/numeric/mtl/operation/assign_mode.hpp>
// #include <boost/numeric/mtl/recursion/recursive_matrix_mult.hpp>

#ifdef __INTEL_COMPILER
#  define MTL_DECL16  __declspec(align(16))
#else
#  define MTL_DECL16
#endif

using namespace mtl;
using namespace mtl::recursion;
using namespace std;

// Maximum time for a single measurement
// is 20 min
const double max_time= 900;

    // Bitmasks:
    const unsigned long morton_mask= generate_mask<true, 0, row_major, 0>::value,
	morton_z_mask= generate_mask<false, 0, row_major, 0>::value,
	doppled_16_row_mask= generate_mask<true, 4, row_major, 0>::value,
	doppled_16_col_mask= generate_mask<true, 4, col_major, 0>::value,
	doppled_32_row_mask= generate_mask<true, 5, row_major, 0>::value,
	doppled_32_col_mask= generate_mask<true, 5, col_major, 0>::value,
	doppled_z_32_row_mask= generate_mask<false, 5, row_major, 0>::value,
	doppled_z_32_col_mask= generate_mask<false, 5, col_major, 0>::value,
	doppled_64_row_mask= generate_mask<true, 6, row_major, 0>::value,
	doppled_64_col_mask= generate_mask<true, 6, col_major, 0>::value,
	doppled_z_64_row_mask= generate_mask<false, 6, row_major, 0>::value,
	doppled_z_64_col_mask= generate_mask<false, 6, col_major, 0>::value,
	doppled_128_row_mask= generate_mask<true, 7, row_major, 0>::value,
	doppled_128_col_mask= generate_mask<true, 7, col_major, 0>::value,
	shark_32_row_mask= generate_mask<true, 5, row_major, 1>::value,
	shark_32_col_mask= generate_mask<true, 5, col_major, 1>::value,
	shark_z_32_row_mask= generate_mask<false, 5, row_major, 1>::value,
	shark_z_32_col_mask= generate_mask<false, 5, col_major, 1>::value,
	shark_64_row_mask= generate_mask<true, 6, row_major, 1>::value,
	shark_64_col_mask= generate_mask<true, 6, col_major, 1>::value,
	shark_z_64_row_mask= generate_mask<false, 6, row_major, 1>::value,
	shark_z_64_col_mask= generate_mask<false, 6, col_major, 1>::value;

typedef assign::plus_sum                            ama_t;

typedef recursion::bound_test_static<32>                    test32_t;
typedef recursion::bound_test_static<64>                    test64_t;

typedef gen_dense_mat_mat_mult_t<assign::plus_sum>  base_mult_t;
typedef gen_recursive_dense_mat_mat_mult_t<base_mult_t>     rec_mult_t;

typedef gen_tiling_22_dense_mat_mat_mult_t<assign::plus_sum>  tiling_22_base_mult_t;
typedef gen_tiling_44_dense_mat_mat_mult_t<assign::plus_sum>  tiling_44_base_mult_t;
typedef gen_tiling_dense_mat_mat_mult_t<2, 4, assign::plus_sum>  tiling_24_base_mult_t;

void print_time_and_mflops(double time, double size)
{
    // time and MFlops of single measure
    std::cout << time << ", " << 2.0 * size * size * size / time / 1e6f << ", ";
}

// Matrices are only placeholder to provide the type
template <typename MatrixA, typename MatrixB, typename MatrixC, typename Mult>
void single_measure(MatrixA&, MatrixB&, MatrixC&, Mult mult, unsigned size, std::vector<int>& enabled, int i)
{
    MatrixA a(size, size);
    MatrixB b(size, size);
    MatrixC c(size, size);
    hessian_setup(a, 1.0);
    hessian_setup(b, 1.0);

    if (enabled[i]) {
	int reps= 0;
	boost::timer start;
	for (; start.elapsed() < 1; reps++)
	    mult(a, b, c);
	double time= start.elapsed() / double(reps);
	print_time_and_mflops(time, a.num_rows());
	if (time > max_time)
	    enabled[i]= 0;
    } else
	std::cout << ", , ";
}

template <typename Functor, typename Result, typename Arg1, typename Arg2, typename Arg3>
struct no_inline3
{
    Result operator()(Arg1& arg1, Arg2& arg2, Arg3& arg3)
    {
	static Functor* f= new(Functor);
	return apply(f, arg1, arg2, arg3);
    }

    Result apply(Functor* f, Arg1& arg1, Arg2& arg2, Arg3& arg3)
    {
	return (*f)(arg1, arg2, arg3);
    }
};

#if 0
// Specialization not needed, at least not with g++
template <typename Functor, typename Arg1, typename Arg2, typename Arg3>
struct no_inline3<Functor, void, Arg1, Arg2, Arg3>
{

    void operator()(Arg1& arg1, Arg2& arg2, Arg3& arg3)
    {
	Functor* f= new(Functor);
	apply(f, arg1, arg2, arg3);
	free(f);
    }

    void apply(Functor* f, Arg1& arg1, Arg2& arg2, Arg3& arg3)
    {
	(*f)(arg1, arg2, arg3);
    }
};
#endif

#if 0
void hybrid_ext_mult_44(const morton_dense<double,  doppled_64_row_mask>& a,
			const morton_dense<double,  doppled_64_col_mask>& b,
			morton_dense<double,  doppled_64_row_mask>& c);
#endif

void dense_ext_mult_44(const dense2D<double>& a,
		       const dense2D<double, mat::parameters<col_major> >& b,
		       dense2D<double>& c);

struct ext_mult_44
{
#if 0
    void operator()(const morton_dense<double,  doppled_64_row_mask>& a,
		    const morton_dense<double,  doppled_64_col_mask>& b,
		    morton_dense<double,  doppled_64_row_mask>& c)
    {
	hybrid_ext_mult_44(a, b, c);
    }
#endif

    void operator()(const dense2D<double>& a,
		    const dense2D<double, mat::parameters<col_major> >& b,
		    dense2D<double>& c)
    {
	dense_ext_mult_44(a, b, c);
    }
};

typedef dense2D<double> rmt;
typedef dense2D<double, mat::parameters<col_major> > cmt;

// C must have even dimensions
template <typename MatrixA, typename MatrixB, typename MatrixC>
void mult_simple_ptu22t(const MatrixA& a, const MatrixB& b, MatrixC& c)
{
    typedef typename MatrixC::value_type  value_type;
    const value_type z= math::zero(c[0][0]);    // if this are matrices we need their size

    set_to_zero(c);

    // Temporary solution; dense matrices need to return const referencens
    MatrixA& aref= const_cast<MatrixA&>(a);
    MatrixB& bref= const_cast<MatrixB&>(b);

    size_t ari= &aref(1, 0) - &aref(0, 0), // how much is the offset of A's entry increased by incrementing row
	aci= &aref(0, 1) - &aref(0, 0), bri= &bref(1, 0) - &bref(0, 0), bci= &bref(0, 1) - &bref(0, 0);

    for (unsigned i= 0; i < c.num_rows(); i+=2)
	for (unsigned k= 0; k < c.num_cols(); k+=2) {
	    int ld= b.num_rows();
	    value_type tmp00= z, tmp01= z, tmp10= z, tmp11= z;

	    MTL_DECL16 const value_type *begin_a= &aref[i][0], *end_a= &aref[i][a.num_cols()];
	    MTL_DECL16 const value_type *begin_b= &bref[0][k];
	    for (; begin_a != end_a; begin_a+= aci, begin_b+= bri) {
		tmp00+= *begin_a * *begin_b;
		tmp01+= *begin_a * *(begin_b+bci);
		tmp10+= *(begin_a+ari) * *begin_b;
		tmp11+= *(begin_a+ari) * *(begin_b+bci);
	    }
	    assign::assign_sum::update(c[i][k], tmp00);
	    assign::assign_sum::update(c[i][k+1], tmp01);
	    assign::assign_sum::update(c[i+1][k], tmp10);
	    assign::assign_sum::update(c[i+1][k+1], tmp11);

#if 0
	    c[i][k]= tmp00; c[i][k+1]= tmp01;
	    c[i+1][k]= tmp10; c[i+1][k+1]= tmp11;
#endif
	}
}

template <typename Matrix, typename MatrixB>
void measure_unrolling(unsigned size, std::vector<int>& enabled, Matrix& matrix, MatrixB& matrixb)
{
    std::cout << size << ", ";

    dense2D<double> dense(4, 4);
    dense2D<double, mat::parameters<col_major> >    denseb(4, 4);

    gen_recursive_dense_mat_mat_mult_t<base_mult_t>           mult;
    gen_recursive_dense_mat_mat_mult_t<tiling_22_base_mult_t> mult_22;
    gen_recursive_dense_mat_mat_mult_t<tiling_44_base_mult_t> mult_44;
    gen_recursive_dense_mat_mat_mult_t<tiling_24_base_mult_t> mult_24;

    typedef gen_tiling_22_dense_mat_mat_mult_ft<Matrix, MatrixB, Matrix, assign::plus_sum>
      tiling_22_t;
    typedef no_inline3<tiling_22_t, void, const Matrix, const MatrixB, Matrix> tiling_22_no_inline_t;

    typedef gen_tiling_44_dense_mat_mat_mult_ft<Matrix, MatrixB, Matrix, assign::plus_sum>
      tiling_44_t;
    typedef no_inline3<tiling_44_t, void, const Matrix, const MatrixB, Matrix> tiling_44_no_inline_t;

    gen_recursive_dense_mat_mat_mult_t<ext_mult_44>           rec_ext_mult_44;

    typedef typename base_case_matrix<Matrix, test64_t>::type    BaseMatrix;
    typedef typename base_case_matrix<MatrixB, test64_t>::type   BaseMatrixB;
    typedef no_inline3<tiling_44_t, void, const BaseMatrix, const BaseMatrixB, BaseMatrix> tiling_44_base_t;
    gen_recursive_dense_mat_mat_mult_t<tiling_44_base_mult_t>    rec_no_inline_mult_44;

    single_measure(matrix, matrixb, matrix, mult, size, enabled, 0);
    single_measure(matrix, matrixb, matrix, mult_22, size, enabled, 1);
    single_measure(matrix, matrixb, matrix, mult_44, size, enabled, 2);
    single_measure(matrix, matrixb, matrix, tiling_22_no_inline_t(), size, enabled, 3);
    single_measure(matrix, matrixb, matrix, tiling_44_no_inline_t(), size, enabled, 4);
    single_measure(matrix, matrixb, matrix, rec_ext_mult_44, size, enabled, 5);
    single_measure(matrix, matrixb, matrix, rec_no_inline_mult_44, size, enabled, 6);
    single_measure(dense, denseb, dense, tiling_22_base_mult_t(), size, enabled, 7);
    single_measure(dense, denseb, dense, tiling_44_base_mult_t(), size, enabled, 8);
    single_measure(dense, denseb, dense, mult_simple_ptu22t<rmt, cmt, rmt>, size, enabled, 9);
    //single_measure(matrix, matrixb, matrix, mult_24, size, enabled, 10);
    std::cout << "0\n";  std::cout.flush();
}

void measure_unrolling_hybrid(unsigned size, std::vector<int>& enabled)
{
    morton_dense<double,  doppled_64_row_mask>     d64r(4, 4);
    morton_dense<double,  doppled_64_col_mask>     d64c(4, 4);
    measure_unrolling(size, enabled, d64r, d64c);
}

void measure_unrolling_dense(unsigned size, std::vector<int>& enabled)
{
    dense2D<double> dense(4, 4);
    dense2D<double, mat::parameters<col_major> >    b(4, 4);
    measure_unrolling(size, enabled, dense, b);
}

template <typename Measure>
void series(unsigned steps, unsigned max_size, Measure measure, const string& comment)
{
    std::cout << "# " << comment << '\n';
    std::cout << "# Gnu-Format size, time, MFlops\n"; std::cout.flush();

    std::vector<int> enabled(16, 1);
    for (unsigned i= steps; i <= max_size; i+= steps)
	measure(i, enabled);
}

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

    std::vector<std::string> scenarii;
    scenarii.push_back(string("Comparing different unrolling for hybrid row-major matrices"));

    using std::cout;
    if (argc < 3) {
	cerr << "usage: recursive_mult_timing  <steps> <max_size>\n";
	exit(1);
    }
    unsigned int steps= atoi(argv[1]), max_size= atoi(argv[2]), size= 32;
    series(steps, max_size, measure_unrolling_hybrid, scenarii[0]);

    return 0;

}