/**
 * Copyright (c) 2011, 2012
 * Claudio Kopper <claudio.kopper@icecube.wisc.edu>
 * and the IceCube Collaboration <http://www.icecube.wisc.edu>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
 * OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
 * CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 *
 *
 * $Id: I3ExtraGeometryItemCylinder.cxx 141406 2016-02-01 21:46:00Z david.schultz $
 *
 * @file I3ExtraGeometryItemCylinder.cxx
 * @version $Revision: 141406 $
 * @date $Date: 2016-02-01 15:46:00 -0600 (Mon, 01 Feb 2016) $
 * @author Claudio Kopper
 */

#include <limits>

#include <icetray/serialization.h>
#include <icetray/I3Units.h>
#include <simclasses/I3ExtraGeometryItemCylinder.h>

#include <boost/foreach.hpp>

I3ExtraGeometryItemCylinder::I3ExtraGeometryItemCylinder()
:
radius_(NAN),
boundingBoxCalculated_(false)
{;}

I3ExtraGeometryItemCylinder::
I3ExtraGeometryItemCylinder(const I3Position &from, const I3Position &to, double radius)
:
from_(from),
to_(to),
radius_(radius),
boundingBoxCalculated_(false)
{
}

I3ExtraGeometryItemCylinder::~I3ExtraGeometryItemCylinder() { }

void I3ExtraGeometryItemCylinder::CalculateBoundingBox() const
{
    if (boundingBoxCalculated_) return;

    boundingBoxLower_=I3Position();
    boundingBoxUpper_=I3Position();

    boundingBoxCalculated_=true;
}

unsigned int I3ExtraGeometryItemCylinder::FindIntersections(const I3Position &lineStart,
                                                            const I3Position &lineEnd,
                                                            double &p0, double &p1,
                                                            double &line_segment_length
                                                            ) const
{
    // This function is based on code from http://www.geometrictools.com/
    // Geometric Tools, LLC
    // Copyright (c) 1998-2011
    // Distributed under the Boost Software License, Version 1.0.
    // http://www.boost.org/LICENSE_1_0.txt

    double line_dir_x = lineEnd.GetX()-lineStart.GetX();
    double line_dir_y = lineEnd.GetY()-lineStart.GetY();
    double line_dir_z = lineEnd.GetZ()-lineStart.GetZ();

    line_segment_length = std::sqrt(line_dir_x*line_dir_x + line_dir_y*line_dir_y + line_dir_z*line_dir_z);
    line_dir_x /= line_segment_length;
    line_dir_y /= line_segment_length;
    line_dir_z /= line_segment_length;

    // the cylinder direction
    double Wx = to_.GetX() - from_.GetX();
    double Wy = to_.GetY() - from_.GetY();
    double Wz = to_.GetZ() - from_.GetZ();
    double halfHeight;
    {
        const double W_len = std::sqrt(Wx*Wx + Wy*Wy + Wz*Wz);
        Wx /= W_len; Wy /= W_len; Wz /= W_len;
        halfHeight = W_len/2.; // the cylinder height
    }
    const double rSqr = radius_*radius_;

    // two orthogonal directions
    double Ux, Uy, Uz, Vx, Vy, Vz;
    {
        const double x = Wx < 0.0 ? -Wx : Wx;
        const double y = Wy < 0.0 ? -Wy : Wy;
        const double z = Wz < 0.0 ? -Wz : Wz;
        if (x < y) {
            if (x < z) {
                Ux = 0.; Uy = Wz; Uz=-Wy;
            } else {
                Ux = Wy; Uy = -Wx; Uz=0.;
            }
        } else {
            if (y < z) {
                Ux = -Wz; Uy = 0.; Uz=Wx;
            } else {
                Ux = Wy; Uy = -Wx; Uz=0.;
            }
        }

        const double U_len = std::sqrt(Ux*Ux + Uy*Uy + Uz*Uz);
        Ux /= U_len; Uy /= U_len; Uz /= U_len;

        // V = W x U  =>  U x V = W
        Vx = Wy*Uz - Wz*Uy;
        Vy = Wz*Ux - Wx*Uz;
        Vz = Wx*Uy - Wy*Ux;
    }

    // diff from line origin to cylinder center
    const double cylOrigin_x = from_.GetX() + Wx;
    const double cylOrigin_y = from_.GetY() + Wy;
    const double cylOrigin_z = from_.GetZ() + Wz;

    const double diff_x = lineStart.GetX() - cylOrigin_x;
    const double diff_y = lineStart.GetY() - cylOrigin_y;
    const double diff_z = lineStart.GetZ() - cylOrigin_z;

    // convert to local cylinder coordinates
    const double Px = Ux*diff_x + Uy*diff_y + Uz*diff_z;
    const double Py = Vx*diff_x + Vy*diff_y + Vz*diff_z;
    const double Pz = Wx*diff_x + Wy*diff_y + Wz*diff_z;

    // Get the z-value, in cylinder coordinates, of the incoming line's
    // unit-length direction.
    const double dz = Wx*line_dir_x + Wy*line_dir_y + Wz*line_dir_z;

    const double EPSILON=1e-8;

    if (std::fabs(dz) >= 1. - EPSILON)
    {
        // The line is parallel to the cylinder axis.  Determine if the line
        // intersects the cylinder end disks.
        double radialSqrDist = rSqr - Px*Px - Py*Py;
        if (radialSqrDist < 0.)
        {
            // Line outside the cylinder, no intersection.
            return 0;
        }

        // Line intersects the cylinder end disks.
        if (dz > 0.)
        {
            p0 = -Pz - halfHeight;
            p1 = -Pz + halfHeight;
        }
        else
        {
            p0 = Pz - halfHeight;
            p1 = Pz + halfHeight;
        }
        return 2;
    }

    // convert incoming line unit-length direction to cylinder coordinates
    const double Dx = Ux*line_dir_x + Uy*line_dir_y + Uz*line_dir_z;
    const double Dy = Vx*line_dir_x + Vy*line_dir_y + Vz*line_dir_z;
    const double Dz = dz;

    double a0, a1, a2, discr, root, inv, tValue;

    if (std::fabs(Dz) <= EPSILON)
    {
        // The line is perpendicular to the cylinder axis.
        if (std::fabs(Pz) > halfHeight)
        {
            // Line is outside the planes of the cylinder end disks.
            return 0;
        }

        // Test intersection of line P+t*D with infinite cylinder
        // x^2+y^2 = r^2.  This reduces to computing the roots of a
        // quadratic equation.  If P = (px,py,pz) and D = (dx,dy,dz),
        // then the quadratic equation is
        //   (dx^2+dy^2)*t^2 + 2*(px*dx+py*dy)*t + (px^2+py^2-r^2) = 0
        a0 = Px*Px + Py*Py - rSqr;
        a1 = Px*Dx + Py*Dy;
        a2 = Dx*Dx + Dy*Dy;
        discr = a1*a1 - a0*a2;
        if (discr < 0.)
        {
            // Line does not intersect cylinder.
            return 0;
        }
        else if (discr > EPSILON)
        {
            // Line intersects cylinder in two places.
            root = std::sqrt(discr);
            inv = 1./a2;
            p0 = (-a1 - root)*inv;
            p1 = (-a1 + root)*inv;
            return 2;
        }
        else
        {
            // Line is tangent to the cylinder.
            p0 = -a1/a2;
            p1 = NAN;
            return 1;
        }
    }

    // Test plane intersections first.
    int quantity = 0;
    inv = 1./Dz;

    double t0 = (-halfHeight - Pz)*inv;
    double xTmp = Px + t0*Dx;
    double yTmp = Py + t0*Dy;
    if (xTmp*xTmp + yTmp*yTmp <= rSqr)
    {
        // Planar intersection inside the top cylinder end disk.
        p0 = t0;
        quantity++;
    }

    double t1 = (+halfHeight - Pz)*inv;
    xTmp = Px + t1*Dx;
    yTmp = Py + t1*Dy;
    if (xTmp*xTmp + yTmp*yTmp <= rSqr)
    {
        // Planar intersection inside the bottom cylinder end disk.
        if (quantity==0) {
            p0 = t1;
        } else {
            p1 = t1;
        }
        quantity++;
    }

    if (quantity == 2)
    {
        // Line intersects both top and bottom cylinder end disks.
        if (p0 > p1)
        {
            double save = p0;
            p0 = p1;
            p1 = save;
        }
        return 2;
    }

    // If quantity == 1, then the line must intersect cylinder wall in a
    // single point somewhere between the end disks.  This case is detected
    // in the following code that tests for intersection between line and
    // cylinder wall.
    a0 = Px*Px + Py*Py - rSqr;
    a1 = Px*Dx + Py*Dy;
    a2 = Dx*Dx + Dy*Dy;
    discr = a1*a1 - a0*a2;
    if (discr < 0.)
    {
        // Line does not intersect cylinder wall.
        if (quantity>0) {
            log_fatal("Internal logic error: found wall intersection when there should have been none.");
        }
        return 0;
    }
    else if (discr > EPSILON)
    {
        root = std::sqrt(discr);
        inv = 1./a2;
        tValue = (-a1 - root)*inv;
        if (t0 <= t1)
        {
            if ((t0 <= tValue) && (tValue <= t1))
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }
        else
        {
            if ((t1 <= tValue) && (tValue <= t0))
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }

        if (quantity == 2)
        {
            // Line intersects one of the cylinder end disks and once on the
            // cylinder wall.
            if (p0 > p1)
            {
                double save = p0;
                p0 = p1;
                p1 = save;
            }
            return 2;
        }

        tValue = (-a1 + root)*inv;
        if (t0 <= t1)
        {
            if (t0 <= tValue && tValue <= t1)
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }
        else
        {
            if (t1 <= tValue && tValue <= t0)
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }
    }
    else
    {
        tValue = -a1/a2;
        if (t0 <= t1)
        {
            if ((t0 <= tValue) && (tValue <= t1))
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }
        else
        {
            if ((t1 <= tValue) && (tValue <= t0))
            {
                if (quantity==0) {
                    p0 = tValue;
                } else {
                    p1 = tValue;
                }
                quantity++;
            }
        }
    }

    if (quantity == 2)
    {
        if (p0 > p1)
        {
            double save = p0;
            p0 = p1;
            p1 = save;
        }
    }

    return quantity;

}

bool
I3ExtraGeometryItemCylinder::DoesLineIntersect
(const I3Position &lineStart,
 const I3Position &lineEnd) const
{
    if (std::isnan(radius_)) return false;

    double result[2];
    double line_segment_length;
    const unsigned int numResults = FindIntersections(lineStart,
                                                      lineEnd,
                                                      result[0],
                                                      result[1],
                                                      line_segment_length);

    for (unsigned int i=0; i<numResults; ++i)
    {
        if ((result[i] >= 0.) && (result[i] <= line_segment_length)) return true;
    }

    // no intersection found
    return false;
}

std::pair<I3Position, I3Position>
I3ExtraGeometryItemCylinder::GetBoundingBox
() const
{
    CalculateBoundingBox();
    return std::make_pair(boundingBoxLower_, boundingBoxUpper_);
}

template <class Archive>
void I3ExtraGeometryItemCylinder::load(Archive &ar, unsigned version)
{
    if (version > i3extrageometryitemcylinder_version_)
        log_fatal("Attempting to read version %u from file but running version %u of I3ExtraGeometryItem class.",version,i3extrageometryitemcylinder_version_);

    ar >> make_nvp("I3ExtraGeometryItem", base_object<I3ExtraGeometryItem>(*this));

    ar >> make_nvp("from", from_);
    ar >> make_nvp("to", to_);
    ar >> make_nvp("radius", radius_);
}

template <class Archive>
void I3ExtraGeometryItemCylinder::save(Archive &ar, unsigned version) const
{
    ar << make_nvp("I3ExtraGeometryItem", base_object<I3ExtraGeometryItem>(*this));

    ar << make_nvp("from", from_);
    ar << make_nvp("to", to_);
    ar << make_nvp("radius", radius_);
}

std::ostream& I3ExtraGeometryItemCylinder::operator<<(std::ostream& oss) const
{
    oss << "[ I3ExtraGeometryItemCylinder :" << std::endl;
    oss << "                         from : (" << from_.GetX()/I3Units::m << "," << from_.GetY()/I3Units::m << "," << from_.GetZ()/I3Units::m << ")m" << std::endl;
    oss << "                           to : (" << to_.GetX()/I3Units::m << "," << to_.GetY()/I3Units::m << "," << to_.GetZ()/I3Units::m << ")m" << std::endl;
    oss << "                       radius : " << radius_/I3Units::m << "m" << std::endl;
    oss << "]" << std::endl;
    return oss;
}

// the following line needs to go away once a more recent version of icetray makes it into a release:
#include "split_serializable_backport.h"

I3_SPLIT_SERIALIZABLE(I3ExtraGeometryItemCylinder);


#include <cctbx/boost_python/flex_fwd.h>

#include <boost/python/module.hpp>
#include <boost/python/class.hpp>
#include <boost/python/args.hpp>
#include <boost/python/return_value_policy.hpp>
#include <boost/python/return_by_value.hpp>
#include <boost/python/return_internal_reference.hpp>
#include <boost/python/def.hpp>
#include <cctbx/masks/around_atoms.h>
#include <cctbx/masks/flood_fill.h>

namespace cctbx { namespace masks {
namespace {

  template <typename DataType, typename FloatType>
  struct around_atoms_wrappers
  {
    typedef around_atoms<DataType, FloatType> w_t;

    static
    void
    wrap()
    {
      using namespace boost::python;
      typedef return_value_policy<return_by_value> rbv;
      class_<w_t>("around_atoms", no_init)
        .def(init<
          cctbx::uctbx::unit_cell const&,
          std::size_t,
          af::shared<scitbx::vec3<double> > const&,
          af::shared<double> const&,
          af::c_grid<3>::index_type const&,
          FloatType const&,
          FloatType const&,
          optional< bool, bool> >((
            arg("unit_cell"),
            arg("space_group_order_z"),
            arg("sites_frac"),
            arg("atom_radii"),
            arg("gridding_n_real"),
            arg("solvent_radius"),
            arg("shrink_truncation_radius"),
            arg("explicit_distance"),
            arg("debug")     )))
        .def_readonly("solvent_radius",
                 &w_t::solvent_radius)
        .def_readonly("n_atom_points",
                 &w_t::n_atom_points)
        .def_readonly("shrink_truncation_radius",
                 &w_t::shrink_truncation_radius)
        .add_property("data", make_getter(&w_t::data, rbv()))
        .def_readonly("contact_surface_fraction",
                 &w_t::contact_surface_fraction)
        .def_readonly("accessible_surface_fraction",
                 &w_t::accessible_surface_fraction)
      ;
    }
  };

  template <typename DataType, typename FloatType>
  struct flood_fill_wrappers
  {
    typedef flood_fill<DataType, FloatType> w_t;

    static
    void
    wrap()
    {
      using namespace boost::python;
      typedef return_value_policy<return_by_value> rbv;
      typedef return_internal_reference<> rir;
      class_<w_t>("flood_fill", no_init)
        .def(init<
          af::ref<DataType, af::c_grid_periodic<3> > const &,
          cctbx::uctbx::unit_cell const &>((
            arg("data"),
            arg("unit_cell"))))
        .def("n_voids", &w_t::n_voids)
        .def("centres_of_mass", &w_t::centres_of_mass)
        .def("centres_of_mass_frac", &w_t::centres_of_mass_frac)
        .def("centres_of_mass_cart", &w_t::centres_of_mass_cart)
        .def("covariance_matrices_frac", &w_t::covariance_matrices_frac)
        .def("covariance_matrices_cart", &w_t::covariance_matrices_cart)
        .def("inertia_tensors_frac", &w_t::inertia_tensors_frac)
        .def("inertia_tensors_cart", &w_t::inertia_tensors_cart)
        .def("grid_points_per_void", make_getter(&w_t::grid_points_per_void, rbv()))
        .def("gridding_n_real", make_getter(&w_t::gridding_n_real, rbv()))
        .def("unit_cell", &w_t::unit_cell, rir())
      ;
    }
  };

  void init_module()
  {
    using namespace boost::python;

    around_atoms_wrappers<int, double>::wrap();
    flood_fill_wrappers<int, double>::wrap();
  }

} // namespace <anonymous>
}} // namespace cctbx::masks

BOOST_PYTHON_MODULE(cctbx_masks_ext)
{
  cctbx::masks::init_module();
}


//    Copyright 2017 Rainer Gemulla
//
//    Licensed under the Apache License, Version 2.0 (the "License");
//    you may not use this file except in compliance with the License.
//    You may obtain a copy of the License at
//
//        http://www.apache.org/licenses/LICENSE-2.0
//
//    Unless required by applicable law or agreed to in writing, software
//    distributed under the License is distributed on an "AS IS" BASIS,
//    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//    See the License for the specific language governing permissions and
//    limitations under the License.
#include <iostream>

#include <boost/filesystem.hpp>
#include <boost/program_options.hpp>
#include <boost/random/uniform_real.hpp>

#include <util/evaluation.h>

#include <mpi2/mpi2.h>
#include <mf/mf.h>
#include <mf/matrix/io/generateDistributedMatrix.h>

// added by me ******************
#include "parse.h"
using namespace std;
using namespace mf;
using namespace mpi2;
using namespace rg;
//*******************************
log4cxx::LoggerPtr logger(log4cxx::Logger::getLogger("main"));

struct Args {
	std::string inputMatrixFile, inputTestMatrixFile, inputRowFacFile, inputColFacFile, outputRowFacFile,
		   outputColFacFile, traceFile, traceVar, lossString, balanceString, balanceMethodString;

	std::string lossName;
	std::vector<double> lossArgs;

	mf::mf_size_type epochs, blocks;
	unsigned seed;
	rg::Random32 random;
	mf::AlsRescale alsRescale;
	double lambda;
	int tasksPerRank;
	int worldSize;
	int worldRank;
	boost::mpi::communicator world;
	mf::BalanceType balanceType;
	mf::BalanceMethod balanceMethod;

	void createTraceFields(mf::Trace& trace) {
		trace.addField("loss", lossString);
	}
};

bool run(Args& args) {

	std::vector<DistributedSparseMatrix> dataVector;
	if (args.inputTestMatrixFile.length() == 0) {
		dataVector=getDataMatrices<SparseMatrix>(args.inputMatrixFile,"V",true,args.tasksPerRank, args.worldSize, args.blocks, 1,false,true);
	}else{
		dataVector=getDataMatrices<SparseMatrix>(args.inputMatrixFile,"V",true,args.tasksPerRank, args.worldSize, args.blocks, 1,
				false,true, &args.inputTestMatrixFile);
	}

	std::vector<DistributedSparseMatrixCM> dataVectorVC=getDataMatrices<SparseMatrixCM>(args.inputMatrixFile,"Vcm",false,
			args.tasksPerRank, args.worldSize, 1, args.blocks,false,true);

	std::pair<DistributedDenseMatrix, DistributedDenseMatrixCM> factorsPair= getFactors(args.inputRowFacFile,
			args.inputColFacFile,  args.tasksPerRank, args.worldSize,args.blocks,args.blocks,false);

//	// load the input matrices
//	DistributedSparseMatrix dv = loadMatrix<SparseMatrix>(args.inputMatrixFile,
//							"V", true, args.tasksPerRank, args.worldSize, args.blocks, 1);
//	LOG4CXX_INFO(logger, "Data matrix: "
//		<< dv.size1() << " x " << dv.size2() << ", " << nnz(dv) << " nonzeros, "
//		<< dv.blocks1() << " x " << dv.blocks2() << " blocks");
//
//	DistributedSparseMatrixCM dvc = loadMatrix<SparseMatrixCM>(args.inputMatrixFile,
//							"VC", false, args.tasksPerRank, args.worldSize, 1, args.blocks);
//	LOG4CXX_INFO(logger, "Data matrix (CM): "
//		<< dvc.size1() << " x " << dvc.size2() << ", " << nnz(dvc) << " nonzeros, "
//		<< dvc.blocks1() << " x " << dvc.blocks2() << " blocks");
//
//	DistributedDenseMatrix dw = loadMatrix<DenseMatrix>(args.inputRowFacFile,
//											"W", true, args.tasksPerRank, args.worldSize, args.blocks, 1);
//	LOG4CXX_INFO(logger, "Row factor matrix: "
//		<< dw.size1() << " x " << dw.size2() << ", " << dw.blocks1() << " x " << dw.blocks2() << " blocks");
//
//	DistributedDenseMatrixCM dh = loadMatrix<DenseMatrixCM>(args.inputColFacFile,
//											"H", false, args.tasksPerRank, args.worldSize, 1, args.blocks);
//	LOG4CXX_INFO(logger, "Column factor matrix: "
//		<< dh.size1() << " x " << dh.size2() << ", " << dh.blocks1() << " x " << dh.blocks2() << " blocks");

//	if (!checkConformity(dv, dw, dh)) {
	if (!checkConformity(dataVector[0], factorsPair.first, factorsPair.second)) {
		std::cerr << "Input matrices are not conforming." << std::endl;
		return false;
	}

	// create factorization data
//	DapFactorizationData<> data(dv, dw, dh, args.tasksPerRank, &dvc);
	DapFactorizationData<> data(dataVector[0], factorsPair.first, factorsPair.second, args.tasksPerRank, &dataVectorVC[0]);

	DsgdFactorizationData<>* testJob = NULL;
	DistributedSparseMatrix *dvTest = NULL;
	if (args.inputTestMatrixFile.compare("") != 0) {
//		dvTest = new DistributedSparseMatrix(
//				loadMatrix<SparseMatrix>(args.inputTestMatrixFile,
//						"Vtest", true, args.tasksPerRank, args.worldSize, args.blocks, args.blocks)
//		);

		dvTest=&dataVector[1];

		LOG4CXX_INFO(logger, "Test matrix: "
						<< dvTest->size1() << " x " << dvTest->size2() << ", " << nnz(*dvTest) << " nonzeros, "
						<< dvTest->blocks1() << " x " << dvTest->blocks2() << " blocks");

//		testJob = new DsgdFactorizationData<>(*dvTest, dw, dh, args.tasksPerRank);
		testJob = new DsgdFactorizationData<>(*dvTest, factorsPair.first, factorsPair.second, args.tasksPerRank);

	}

	// tracing
	Trace trace;
	args.createTraceFields(trace);
	Timer t;
	t.start();

	// go
	if (args.lossName.compare("Nzsl") == 0) {
		args.lambda = 0;
		AlsRegularizer regularizer = ALS_L2;
		dalsNzsl(data, args.epochs, trace, args.lambda, regularizer, args.balanceType, args.balanceMethod, testJob);
	} else if (args.lossName.compare("Nzsl_L2") == 0) {
		if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
			std::cout << "Invalid number of arguments in " << args.lossString << std::endl;
			return false;
		}
		args.lambda = args.lossArgs[0];
		AlsRegularizer regularizer = ALS_L2;
		dalsNzsl(data, args.epochs, trace, args.lambda, regularizer, args.balanceType, args.balanceMethod, testJob);
	} else if (args.lossName.compare("Nzsl_Nzl2") == 0) {
		// als with Nzsl_Nzl2
		if (args.lossArgs.size()<1 || args.lossArgs.size()>1) {
			std::cout << "Invalid number of arguments in " << args.lossString << std::endl;
			return false;
		}
		args.lambda = args.lossArgs[0];
		AlsRegularizer regularizer = ALS_NZL2;
		dalsNzsl(data, args.epochs, trace, args.lambda, regularizer, args.balanceType, args.balanceMethod, testJob);
	} else if (args.lossName.compare("Sl") == 0) {
		// gnmf with Sl
		if (args.lossArgs.size() != 0) {
			std::cout << "Invalid number of arguments in " << args.lossString << std::endl;
			return false;
		}
		dgnmf(data, args.epochs, trace, args.balanceType, args.balanceMethod, testJob);
	} else if (args.lossName.compare("Gkl") == 0) {
		// lee01 with Gkl
		if (args.lossArgs.size() != 0) {
			std::cout << "Invalid number of arguments in " << args.lossString << std::endl;
			return false;
		}
		dlee01Gkl(data, args.epochs, trace);
	} else {
		std::cout << "Invalid loss " << args.lossString << std::endl;
		return false;
	}
	t.stop();
	LOG4CXX_INFO(logger, "Total time: " << t);

	// write trace to an R file
	if (args.traceFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing trace to " << args.traceFile);
		trace.toRfile(args.traceFile, args.traceVar);
	}

	// write computed factors to file
	if (args.outputRowFacFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing row factors to " << args.outputRowFacFile);
		DenseMatrix w0;
//		unblock(dw, w0);
		unblock(factorsPair.first, w0);
		writeMatrix(args.outputRowFacFile, w0);
	}
	if (args.outputColFacFile.length() > 0) {
		LOG4CXX_INFO(logger, "Writing column factors to " << args.outputColFacFile);
		DenseMatrixCM h0;
//		unblock(dh, h0);
		unblock(factorsPair.second, h0);
		writeMatrix(args.outputColFacFile, h0);
	}

	return true;
}

int main(int argc, char *argv[]) {
	using namespace std;
	using namespace mf;
	using namespace mpi2;
	using namespace rg;
	using namespace boost::program_options;

	// initialize mf library and mpi2
	boost::mpi::communicator& world = mfInit(argc, argv);

	mfStart();
	bool result = true;

	if (world.rank() == 0) {
		boost::this_thread::sleep(boost::posix_time::milliseconds(100)); // so that arguments are logged nicely
		Args args;

		// parse command line
		options_description desc("Options");
		desc.add_options()
			("help", "produce help message")
			("input-file", value<string>(&args.inputMatrixFile), "filename of data matrix")
			("input-test-file", value<string>(&args.inputTestMatrixFile), "filename of test matrix")
			("input-row-file", value<string>(&args.inputRowFacFile), "filename of initial row factors")
			("input-col-file", value<string>(&args.inputColFacFile), "filename of initial column factors")
			("output-row-file", value<string>(&args.outputRowFacFile), "filename of final row factors")
			("output-col-file", value<string>(&args.outputColFacFile), "filename of final column factors")
			("trace", value<string>(&args.traceFile), "filename of trace [trace.R]")
			("trace-var", value<string>(&args.traceVar), "variable name for trace [traceVar]")
			("epochs", value<mf_size_type>(&args.epochs), "number of epochs to run [10]")
			("tasks-per-rank", value<int>(&args.tasksPerRank), "number of concurrent tasks per rank [1]")
			("seed", value<unsigned>(&args.seed), "seed for random number generator (system time if not set)")
			("loss", value<string>(&args.lossString), "loss function (e.g., \"Nzsl\", \"Nzsl_L2(0.5)\"))")
			("balance", value<string>(&args.balanceString), "Type of balancing (None, L2, Nzl2) [None]")
			("balance-method", value<string>(&args.balanceMethodString), "Balancing method (e.g., \"Simple\", \"Optimal\") [Simple]")
		;

		positional_options_description pdesc;
		pdesc.add("input-file", 1);
		pdesc.add("input-test-file", 2);
		pdesc.add("input-row-file", 3);
		pdesc.add("input-col-file", 4);
		pdesc.add("output-row-file", 5);
		pdesc.add("output-col-file", 6);

		variables_map vm;
		store(command_line_parser(argc, argv).options(desc).positional(pdesc).run(), vm);
		notify(vm);

		// check required arguments
		if (vm.count("help") || vm.count("input-file") == 0 || vm.count("input-row-file") == 0 ||
			vm.count("input-col-file") == 0 || vm.count("loss") == 0)
		{
			cerr << "Error: Options input-file, input-row-file, input-col-file, loss have to be present" << endl;
			cout << "alternating projection [options]" << endl << endl;
			cout << desc << endl;
			exit(1);
		}

		// set default arguments
		if (vm.count("seed") == 0) args.seed = time(NULL);
		if (vm.count("trace") == 0) { args.traceFile = ""; }
		if (vm.count("trace-var") == 0) args.traceVar = "trace";
		if (vm.count("input-test-file") == 0) args.inputTestMatrixFile = "";
		if (vm.count("seed") == 0) args.seed = time(NULL);
		if (vm.count("epochs") == 0) args.epochs = 10;
		if (vm.count("tasks-per-rank") == 0) args.tasksPerRank = 1;
		if (vm.count("output-row-file") == 0) { args.outputRowFacFile = ""; }
		if (vm.count("output-col-file") == 0) { args.outputColFacFile = ""; }
		if (vm.count("balance") == 0) { args.balanceString = "None"; }
		if (vm.count("balance-method") == 0) { args.balanceMethodString = "Simple"; }

		// print some information
		LOG4CXX_INFO(logger, "Input");
		LOG4CXX_INFO(logger, "    Input file: " << args.inputMatrixFile);
		LOG4CXX_INFO(logger, "    Input test file: " << (args.inputTestMatrixFile.length() == 0 ? "Disabled" :args.inputTestMatrixFile));
		LOG4CXX_INFO(logger, "    Input row factors: " << args.inputRowFacFile);
		LOG4CXX_INFO(logger, "    Input column factors: " << args.inputColFacFile);
		LOG4CXX_INFO(logger, "Output");
		LOG4CXX_INFO(logger, "    Output row factors: " << (args.outputRowFacFile.length() == 0 ? "Disabled" :args.outputRowFacFile));
		LOG4CXX_INFO(logger, "    Output column factors: " << (args.outputColFacFile.length() == 0 ? "Disabled" :args.outputColFacFile));
		if (args.traceFile.length() == 0) {
			LOG4CXX_INFO(logger, "    Trace: Disabled");
		} else {
			LOG4CXX_INFO(logger, "    Trace: " << args.traceFile << " (" << args.traceVar << ")");
		}
		LOG4CXX_INFO(logger, "Parallelization");
		LOG4CXX_INFO(logger, "    MPI ranks: " << world.size());
		LOG4CXX_INFO(logger, "    Tasks per rank: " << args.tasksPerRank);
		LOG4CXX_INFO(logger, "Alternating Projection options");
		LOG4CXX_INFO(logger, "    Seed: " << args.seed);
		LOG4CXX_INFO(logger, "    Epochs: " << args.epochs);
//		LOG4CXX_INFO(logger, "    Factorization rank: " << args.rank);
		LOG4CXX_INFO(logger, "    Loss function: " << args.lossString);

		// fill fields
		args.random = Random32(args.seed);
		args.world = world;
		args.worldSize = world.size();
		args.worldRank = world.rank();
		args.blocks = args.worldSize * args.tasksPerRank;

		// parsing
		parse::parseArg("loss", args.lossString, args.lossName, args.lossArgs);
		if (args.balanceString.compare("None") == 0) {
			args.balanceType = BALANCE_NONE;
		} else if (args.balanceString.compare("L2") == 0) {
			args.balanceType = BALANCE_L2;
		} else if (args.balanceString.compare("Nzl2") == 0) {
			args.balanceType = BALANCE_NZL2;
		} else {
			cerr << "Error: Invalid value for --balance: " << args.balanceString << endl;
			exit(1);
		}
		if (args.balanceMethodString.compare("Simple") == 0) {
			args.balanceMethod = BALANCE_SIMPLE;
		} else if (args.balanceMethodString.compare("Optimal") == 0) {
			args.balanceMethod = BALANCE_OPTIMAL;
		} else {
			cerr << "Error: Invalid value for --balance-method: " << args.balanceMethodString << endl;
			exit(1);
		}

		switch (args.balanceType){
		case BALANCE_NONE:
			LOG4CXX_INFO(logger, "    Balancing: Disabled");
			break;
		case BALANCE_L2:
			LOG4CXX_INFO(logger, "    Balancing: L2, " << (args.balanceMethod == BALANCE_SIMPLE ? "Simple" : "Optimal"));
			break;
		case BALANCE_NZL2:
			LOG4CXX_INFO(logger, "    Balancing: Nzl2, " << (args.balanceMethod == BALANCE_SIMPLE ? "Simple" : "Optimal"));
			break;
		}

		// let's go
		result = run(args);

	}

	mfStop();
	mfFinalize();

	// everything OK
	return result ? 0 : 1;

}


#include "atomistic.hpp"
#include "formats.hpp"
#include "types.hpp"
#include "io.hpp"

#include <iostream>
#include <fstream>
#include <iterator>
#include <vector>
#include <string>

#include <boost/program_options.hpp>
#include <blitz/array.h>
#include <fftw3.h>

#include <ctime>
time_t t = clock();
//#include <gsl/gsl_const_mksa.h>

namespace po = boost::program_options;
namespace at = atomistic;
namespace cp2k = formats::cp2k;

// Returns true, if parsing went ok
bool parse(int ac, char* av[], po::variables_map& vm);

// Prepares extrapolation
bool prepare(types::String levelFileName,
             types::String cubeListFileName,
             types::String hartreeFileName,
             double start,
             double width);

// Handles extrapolation
/**
 * Do a 2d Fourier transform at a given z-plane,
 * then propagate the Fourier coefficients \f$a_G\f$ according to
 * \f$ a_G(z) = a_G(z_0) e^{-\lambda (z-z_0)}\f$ with
 * \f$\lambda = \sqrt{G^2-\frac{2m}{\hbar^2}(E_n-V_{vac})}\f$
 */
bool extrapolate(types::String cubeFile,
                 types::Uint zStartIndex,
                 types::Uint zEndIndex,
                 const cp2k::Spectrum &spectrum);

int main(int ac, char* av[]) {

    po::variables_map args;
    if (parse(ac, av, args)) {
        prepare(args["levels"].as< types::String >(),
                args["cubelist"].as< types::String >(),
                args["hartree"].as< types::String >(),
                args["start"].as< double >(),
                args["width"].as< double >()
               );

    }

    return 0;
}

bool prepare(types::String levelFileName,
             types::String cubeListFileName,
             types::String hartreeFileName,
             types::Real start,
             types::Real width) {

    // Read energy levels
    cp2k::Spectrum spectrum = cp2k::Spectrum();
    spectrum.readFromCp2k(levelFileName.c_str());

    // Read hartree cube file
    formats::Cube hartree = formats::Cube();
    hartree.readCubeFile(hartreeFileName.c_str());

    // Find highest z-coordinate
    // Note: The slowest index in cube file format is x, so in terms of storage
    //       modifying the number of x-coordinates would be the easiest.
    std::vector< at::Atom >::const_iterator it = hartree.atoms.begin(),
        end = hartree.atoms.end();
    types::Real zTop = it->coordinates[2];
    ++it;
    while(it != end) {
        if( it->coordinates[2] > zTop ) {
            zTop = it->coordinates[2];
        }
        ++it;
    }

    // Define region of interpolation
    std::vector<types::Real> tempvector;
    tempvector.push_back(0);
    tempvector.push_back(0);
    tempvector.push_back(zTop + start);
    std::vector<types::Uint> indices;
    hartree.grid.getNearestIndices(tempvector, indices);
    types::Uint zStartIndex = indices[2];
    std::cout << "Fourier transform will be performed at z-index "
        << zStartIndex << "\n";
    tempvector[2] += width;
    hartree.grid.getNearestIndices(tempvector, indices);
    types::Uint zEndIndex = indices[2];

    // Shift levels by "vacuum" Hartree potential
    // (aveaged value at zStartIndex)
    std::vector<types::Real> hartreeZ;
    hartree.averageXY(hartreeZ);
    types::Real hartreeZStart = hartreeZ[zStartIndex];
    spectrum.shift(-hartreeZStart);
    std::cout << "Vacuum hartree potential is " << hartreeZStart << " Ha\n";
    types::Real hartreeZEnd = hartreeZ[zEndIndex];
    std::cout << "Hartree potential changes by " <<
        hartreeZEnd - hartreeZStart << " Ha over extrapolation region.\n";

    // Write Zprofile of hartree potential
    std::string hartreeZProfile = io::getFileName(hartreeFileName);
    hartreeZProfile += ".zprofile";
    std::cout << "Writing Z profile of Hartree potential to "
        << hartreeZProfile << std::endl;
    using namespace blitz;
    hartree.writeZProfile(hartreeZProfile);

    // Write z plane of hartree potential for gnuplot
    std::vector<types::Real> hartreePlane;
    hartree.getZPlane(zStartIndex, hartreePlane);
    types::String s = formats::gnuplot::writeMatrix(
            hartreePlane,
            hartree.grid.directions[0].getNElements(),
            hartree.grid.directions[1].getNElements());
    io::writeStream("hartree.zplane",s);

    // Read file with list of cubes
    std::ifstream cubeListFile;
    cubeListFile.open(cubeListFileName.c_str());
    if (!cubeListFile.is_open() )
            throw types::fileAccessError()
                << boost::errinfo_file_name(cubeListFileName);

    std::vector<types::String> cubeList;
    types::String fileName;
    while (getline(cubeListFile, fileName)) {
        cubeList.push_back(fileName);
    }
    cubeListFile.close();

    std::cout << "Time to prepare : " << (clock() -t)/1000.0 << " ms\n";
    t = clock();

    // Iterate over cubes
    for(std::vector<types::String>::const_iterator it = cubeList.begin();
            it != cubeList.end(); ++it) {
        extrapolate(*it,zStartIndex,zEndIndex,spectrum);
    }

    return true;
}

bool extrapolate(
    types::String cubeFile,
    types::Uint startIndex,
    types::Uint endIndex,
    const cp2k::Spectrum &spectrum) {

    using namespace blitz;

    std::cout << "--------\n Extrapolating " << cubeFile << "\n";

    t = clock();
    formats::WfnCube wfn = formats::WfnCube();
    wfn.readCubeFile(cubeFile.c_str());
    std::cout << "Time to read cube : " << (clock() -t)/1000.0 << " ms\n";

    wfn.energy = spectrum.spins[wfn.spin -1].getLevel(wfn.wfn);
    types::Uint nX = wfn.grid.directions[0].getNElements(),
                nY = wfn.grid.directions[1].getNElements(),
                nZ = endIndex;
    wfn.grid.resize(nX, nY, nZ);
    Array<types::Real,3> dataArray(
        &wfn.grid.data[0],
        shape(nX, nY, nZ),
        neverDeleteData);

    // Produce z profile
    std::string zProfile = io::getFileName(cubeFile);
    formats::WfnCube wfnSq = wfn;
    wfnSq.grid.squareValues();
    zProfile += ".zprofile";
    std::cout << "Writing Z profile to " << zProfile << std::endl;
    wfnSq.writeZProfile(zProfile, "Z profile of sqared wave function\n");

    // Get z-plane for interpolation
    Array<types::Real,2> planeDirect(nX, nY);
    planeDirect = dataArray(Range::all(), Range::all(), startIndex);

    // Do a real 2 complex fft
    // Since a(-k)=a(k)^* (or in terms of indices: a[n-k]=a[k]^*)
    // only a[k], k=0...n/2+1 (division rounded down) are retained
    types::Uint nXF = nX, nYF = nY/2 + 1;
    fftw_complex *out = (fftw_complex*) fftw_malloc(sizeof(fftw_complex) * nXF * nYF);
    fftw_plan plan_forward = fftw_plan_dft_r2c_2d(nX, nY, planeDirect.data(), out, FFTW_ESTIMATE);
    fftw_execute(plan_forward);

    Array<types::Complex,2> planeFourier(
        reinterpret_cast<types::Complex *>(out),
        shape(nXF, nYF),
        neverDeleteData);

    // Calculate the exponential prefactors.
    // Multiplication of Fourier coefficients with these prefactors
    // propagates them to next z plane.
    // The prefectors would only need dimensions (nX/2 +1, nY/2 +1)
    // since they are the same for G and -G. However for the sake of
    // simplicity of calculation, we prepare them here for (nX, nY/2+1) = (nXF,nYF).
    t = clock();
    Array<types::Real,2> prefactors(nXF, nYF);
    // SI units energy: 2 m E/hbar^2 a0 = 2 E/Ha
    types::Real energyTerm = 2 * wfn.energy;
    std::cout << "Energy level " << wfn.energy << " Ha\n";
    types::Real deltaZ = wfn.grid.directions[2].getIncrementVector()[2];
    la::Cell cell = la::Cell(wfn.grid.directions);
    vector<types::Real> X = cell.vector(0);
    vector<types::Real> Y = cell.vector(1);
    types::Real dKX = 2 * M_PI / X[0];
    types::Real dKY = 2 * M_PI / Y[1];

    // Notice that the order of storage is 0...G -G...-1 for uneven nX
    // and 0...G-1 G -(G-1)...-1 for even NX
    prefactors( Range(0, nXF/2), Range::all())= exp(- sqrt(tensor::i * dKX * tensor::i * dKX + tensor::j *dKY * tensor::j * dKY - energyTerm) * deltaZ);
    // tensor::i always starts from 0, i.e. it ranges from 0...nXF/2-1 (nX even) or 0...nXF/2 (nX uneven)
    if(nX % 2 == 1)
            prefactors( Range(nXF/2 + 1, nXF - 1), Range::all())= exp(- sqrt( (nXF/2 - tensor::i) * dKX * (nXF/2 - tensor::i) * dKX + tensor::j *dKY * tensor::j * dKY - energyTerm) * deltaZ);
    else
            prefactors( Range(nXF/2 + 1, nXF - 1), Range::all())= exp(- sqrt( (nXF/2 -1 - tensor::i) * dKX * (nXF/2 - 1 - tensor::i) * dKX + tensor::j *dKY * tensor::j * dKY - energyTerm) * deltaZ);

//    // Produce Fourier coefficient for gnuplot
//    types::Complex *pt =  reinterpret_cast<types::Complex *>(out);
//    for(types::Uint i = 0; i < nXF; ++i){
//            for(types::Uint j = 0;j< nYF; ++j){
//                    std::cout << std::abs(*pt) << " ";
//                    ++pt;
//            }
//            std::cout << std::endl;
//    }

    // Sequentially update the cube file
    Array<types::Real,2> tempDirect(nX, nY);
    Array<types::Complex,2> tempFourier(nXF, nYF);
    fftw_plan plan_backward;
    for(int zIndex = startIndex + 1; zIndex <= endIndex; ++zIndex) {
        // Propagate fourier coefficients
        planeFourier = prefactors(tensor::i, tensor::j) * planeFourier(tensor::i, tensor::j);

        // The c2r transform destroys its input
        tempFourier = planeFourier;
        plan_backward = fftw_plan_dft_c2r_2d(nX, nY, (fftw_complex*) tempFourier.data(), tempDirect.data(), FFTW_ESTIMATE);

        // Do Fourier-back transform
        fftw_execute(plan_backward); /* repeat as needed */
        tempDirect /= nX * nY;

        // Copy data
        dataArray(Range::all(), Range::all(), zIndex) = tempDirect(Range::all(), Range::all());
    }

    fftw_destroy_plan(plan_forward);
    fftw_destroy_plan(plan_backward);
    fftw_free(out);

    std::cout << "Time to extrapolate : " << (clock() -t)/1000.0 << " ms\n";

    t = clock();
    types::String outCubeFile = "extrapolated.";
    outCubeFile += io::getFileName(cubeFile);
    std::cout << "Writing extrapolated cube file to "
        << outCubeFile << std::endl;
    wfn.writeCubeFile(outCubeFile);
    std::cout << "Time to write cube : " << (clock() -t)/1000.0 << " ms\n";

    // Produce z profile
    wfn.grid.squareValues();
    std::string outZProfile = outCubeFile;
    outZProfile += ".zprofile";
    std::cout << "Writing Z profile of extrapolated cube file to "
        << outZProfile << std::endl;
    wfn.writeZProfile(outZProfile, "Z profile of squared extrpolated wave function\n");

    return true;
}

bool parse(int ac, char* av[], po::variables_map& vm) {
    types::String input_file;
    // Declare regular options
    po::options_description desc("Allowed options");
    desc.add_options()
    ("help,h", "produce help message")
    ("version,v", "print version information")
    ("input-file,i", po::value<types::String>(&input_file), "Input file specifying all or several of the following options")
     ("levels", po::value<types::String>(), "CP2K output file containing energy levels")
    ("cubelist", po::value<types::String>(), "file with list of wavefunction cubes you wish to extrapolate")
    ("hartree", po::value<types::String>(), "cube file of hartree potential from CP2K")
    ("start", po::value<double>(), "distance between extrapolation plane and outermost atom in a.u.")
    ("width", po::value<double>(), "length of extrapolation in a.u.")
    ;

    // Register positional options
    po::positional_options_description p;
    p	.add("input-file", -1);

    // Parse
    po::store(po::command_line_parser(ac,av).
              options(desc).positional(p).run(), vm);
    po::notify(vm);

    // If specified, try to parse conig file
    if (vm.count("input-file")){
        std::ifstream ifs(input_file.c_str());
        if(!ifs) throw types::fileAccessError() << boost::errinfo_file_name(input_file);
        store(parse_config_file(ifs, desc), vm);
        notify(vm);
    }

    // Display help message
    if (	vm.count("help") ||
            !vm.count("levels") ||
            !vm.count("hartree") ||
            !vm.count("start") ||
            !vm.count("width")	) {
        std::cout << "Usage: extrapolate [options]\n";
        std::cout << desc << "\n";
    } else if (vm.count("version")) {
        std::cout << "Mar 1st 2012\n";
    } else {
        return true;
    }

    return false;
}


//---------------------------------------------------------------------------//
// Copyright (c) 2018-2021 Mikhail Komarov <nemo@nil.foundation>
// Copyright (c) 2020-2021 Nikita Kaskov <nbering@nil.foundation>
//
// MIT License
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.
//---------------------------------------------------------------------------//

#define BOOST_TEST_MODULE basic_components_verification_test

#include <boost/test/unit_test.hpp>

#include <nil/crypto3/algebra/curves/bls12.hpp>
#include <nil/crypto3/algebra/fields/bls12/base_field.hpp>
#include <nil/crypto3/algebra/fields/bls12/scalar_field.hpp>
#include <nil/crypto3/algebra/fields/arithmetic_params/bls12.hpp>
#include <nil/crypto3/algebra/curves/params/multiexp/bls12.hpp>
#include <nil/crypto3/algebra/curves/params/wnaf/bls12.hpp>

#include <nil/crypto3/algebra/curves/mnt4.hpp>
#include <nil/crypto3/algebra/fields/mnt4/base_field.hpp>
#include <nil/crypto3/algebra/fields/mnt4/scalar_field.hpp>
#include <nil/crypto3/algebra/fields/arithmetic_params/mnt4.hpp>
#include <nil/crypto3/algebra/curves/params/multiexp/mnt4.hpp>
#include <nil/crypto3/algebra/curves/params/wnaf/mnt4.hpp>
#include <nil/crypto3/algebra/curves/mnt6.hpp>
#include <nil/crypto3/algebra/fields/mnt6/base_field.hpp>
#include <nil/crypto3/algebra/fields/mnt6/scalar_field.hpp>
#include <nil/crypto3/algebra/fields/arithmetic_params/mnt6.hpp>
#include <nil/crypto3/algebra/curves/params/multiexp/mnt6.hpp>
#include <nil/crypto3/algebra/curves/params/wnaf/mnt6.hpp>

#include <nil/crypto3/zk/components/disjunction.hpp>
#include <nil/crypto3/zk/components/conjunction.hpp>
#include <nil/crypto3/zk/components/comparison.hpp>
#include <nil/crypto3/zk/components/inner_product.hpp>
#include <nil/crypto3/zk/components/loose_multiplexing.hpp>

#include <nil/crypto3/zk/snark/schemes/ppzksnark/r1cs_gg_ppzksnark.hpp>

#include "verify_r1cs_scheme.hpp"

using namespace nil::crypto3;
using namespace nil::crypto3::zk;
using namespace nil::crypto3::algebra;

template<typename CurveType>
void test_disjunction_component(std::size_t w) {

    using field_type = typename CurveType::scalar_field_type;
    using curve_type = CurveType;

    std::size_t n = std::log2(w) +
        ((w > (1ul << std::size_t(std::log2(w))))? 1 : 0);

    components::blueprint<field_type> bp;
    components::blueprint_variable<field_type> output;
    output.allocate(bp);

    bp.set_input_sizes(1);

    components::blueprint_variable_vector<field_type> inputs;
    inputs.allocate(bp, n);

    components::disjunction<field_type> d(bp, inputs, output);
    d.generate_r1cs_constraints();

    for (std::size_t j = 0; j < n; ++j) {
        bp.val(inputs[j]) = typename field_type::value_type((w & (1ul << j)) ? 1 : 0);
    }

    d.generate_r1cs_witness();

    BOOST_CHECK(bp.val(output) == (w ? field_type::value_type::one() : field_type::value_type::zero()));
    BOOST_CHECK(bp.is_satisfied());

    BOOST_CHECK(verify_component<curve_type>(bp));
}

template<typename CurveType>
void test_conjunction_component(std::size_t w) {

    using field_type = typename CurveType::scalar_field_type;
    using curve_type = CurveType;

    std::size_t n = std::log2(w) +
        ((w > (1ul << std::size_t(std::log2(w))))? 1 : 0);

    components::blueprint<field_type> bp;

    components::blueprint_variable<field_type> output;
    output.allocate(bp);

    bp.set_input_sizes(1);

    components::blueprint_variable_vector<field_type> inputs;
    inputs.allocate(bp, n);

    components::conjunction<field_type> c(bp, inputs, output);
    c.generate_r1cs_constraints();

    for (std::size_t j = 0; j < n; ++j) {
        bp.val(inputs[j]) = (w & (1ul << j)) ? field_type::value_type::one() : field_type::value_type::zero();
    }

    c.generate_r1cs_witness();

    BOOST_CHECK(bp.val(output) ==
                (w == (1ul << n) - 1 ? field_type::value_type::one() : field_type::value_type::zero()));
    BOOST_CHECK(bp.is_satisfied());

    BOOST_CHECK(verify_component<curve_type>(bp));
}

template<typename CurveType>
void test_comparison_component(std::size_t a, std::size_t b) {

    using field_type = typename CurveType::scalar_field_type;
    using curve_type = CurveType;

    components::blueprint<field_type> bp;

    components::blueprint_variable<field_type> A, B, less, less_or_eq;
    A.allocate(bp);
    B.allocate(bp);
    less.allocate(bp);
    less_or_eq.allocate(bp);

    bp.set_input_sizes(1);
    std::size_t n = std::log2(std::max(a, b)) +
        ((std::max(a, b) > (1ul << std::size_t(std::log2(std::max(a, b)))))? 1 : 0);

    components::comparison<field_type> cmp(bp, n, A, B, less, less_or_eq);
    cmp.generate_r1cs_constraints();

    bp.val(A) = typename field_type::value_type(a);
    bp.val(B) = typename field_type::value_type(b);

    cmp.generate_r1cs_witness();

    BOOST_CHECK(bp.val(less) == (a < b ? field_type::value_type::one() : field_type::value_type::zero()));
    BOOST_CHECK(bp.val(less_or_eq) == (a <= b ? field_type::value_type::one() : field_type::value_type::zero()));
    BOOST_CHECK(bp.is_satisfied());

    BOOST_CHECK(verify_component<curve_type>(bp));
}

BOOST_AUTO_TEST_SUITE(basic_components_test_suite)

BOOST_AUTO_TEST_CASE(basic_components_disjunction_r1cs_gg_ppzksnark_test) {
    std::cout << "Disjunction component test started" << std::endl;
    std::cout << "Started for bls12<381>" << std::endl;
    test_disjunction_component<curves::bls12<381>>(10);
    std::cout << "Started for mnt4<298>" << std::endl;
    test_disjunction_component<curves::mnt4<298>>(10);
    std::cout << "Started for mnt6<298>" << std::endl;
    test_disjunction_component<curves::mnt6<298>>(10);
}

BOOST_AUTO_TEST_CASE(basic_components_conjunction_r1cs_gg_ppzksnark_test) {
    std::cout << "Conjunction component test started" << std::endl;
    std::cout << "Started for bls12<381>" << std::endl;
    test_conjunction_component<curves::bls12<381>>(10);
    std::cout << "Started for mnt4<298>" << std::endl;
    test_conjunction_component<curves::mnt4<298>>(10);
    std::cout << "Started for mnt6<298>" << std::endl;
    test_conjunction_component<curves::mnt6<298>>(10);
}

BOOST_AUTO_TEST_CASE(basic_components_comparison_r1cs_gg_ppzksnark_test) {
    std::cout << "Comparison component r1cs_gg_ppzksnark test started" << std::endl;
    std::cout << "Started for bls12<381>" << std::endl;
    test_comparison_component<curves::bls12<381>>(1, 4);
    std::cout << "Started for mnt4<298>" << std::endl;
    test_comparison_component<curves::mnt4<298>>(1, 4);
    std::cout << "Started for mnt6<298>" << std::endl;
    test_comparison_component<curves::mnt6<298>>(1, 4);
}

BOOST_AUTO_TEST_SUITE_END()


// All content Copyright (C) 2018 Genomics plc
#include "mapping/hashMapper.hpp"
#include "utils/exceptions.hpp"

#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#include <string>
#include <map>
#include <numeric>

using namespace wecall::mapping;

int hashValueForSequence( const wecall::utils::BasePairSequence & sequence )
{
    const auto length = sequence.size();
    HashFunction hashFunction( sequence, length );
    return hashFunction.next( sequence.at( length - 1 ) );
}

BOOST_AUTO_TEST_CASE( test_hash_function_raises_for_kmer_of_length_1 )
{
    wecall::utils::BasePairSequence sequence( "A" );
    BOOST_CHECK_THROW( HashFunction( sequence, 1 ), wecall::utils::wecall_exception );
}

BOOST_AUTO_TEST_CASE( test_hash_function_raises_if_sequence_length_less_than_kmer_size )
{
    wecall::utils::BasePairSequence sequence( "AAAAAAAAAAA" );
    BOOST_CHECK_THROW( HashFunction( sequence, sequence.size() + 1 ), wecall::utils::wecall_exception );
}

BOOST_AUTO_TEST_CASE( test_hash_value_is_bijective_length_8 )
{
    const auto length = 8;
    const std::vector< char > letters = {'A', 'T', 'C', 'G'};
    const auto alphabetSize = letters.size();

    std::map< int, std::vector< std::string > > hashToKmers;

    const int numberOfKmers = int( pow( alphabetSize, length ) );
    for ( auto kmerIndex = 0; kmerIndex < numberOfKmers; ++kmerIndex )
    {
        std::string kmer( length, ' ' );
        for ( auto letterIndex = 0; letterIndex < length; ++letterIndex )
        {
            auto index = ( kmerIndex / ( int( pow( alphabetSize, letterIndex ) ) ) ) % alphabetSize;
            kmer[letterIndex] = letters[index];
        }

        const auto hashValueThisKmer = hashValueForSequence( kmer );

        for ( auto letter : letters )
        {
            const auto preAppendedKmer = std::string( 1, letter ) + kmer;
            HashFunction hashFunction( preAppendedKmer, length );
            hashFunction.next( preAppendedKmer[preAppendedKmer.size() - 2] );
            BOOST_REQUIRE_EQUAL( hashValueThisKmer, hashFunction.next( preAppendedKmer[preAppendedKmer.size() - 1] ) );
        }

        hashToKmers[hashValueThisKmer].push_back( kmer );
    }

    BOOST_CHECK_EQUAL( hashToKmers.size(), numberOfKmers );

    for ( const auto & pair : hashToKmers )
    {
        BOOST_REQUIRE_EQUAL( pair.second.size(), 1 );
    }
}

BOOST_AUTO_TEST_CASE( test_next_hash_value_repetitive_sequence )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "AAAAAAAA" );

    HashFunction hashFunction( sequence, kmerSize );

    auto firstHash = hashFunction.next( 'A' );
    BOOST_CHECK_EQUAL( firstHash, 0 );

    for ( int i = 0; i < kmerSize * 2; ++i )
    {
        auto nextHash = hashFunction.next( 'A' );
        BOOST_CHECK_EQUAL( firstHash, nextHash );
    }
}

BOOST_AUTO_TEST_CASE( test_maps_substring_back_to_correct_start_position_in_non_repetitive_sequence )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "ACGAAATTTAAAAACCCCCTCCGCCCAAAAACATGACCCAATTGG" );

    HashMapper hasher( sequence, kmerSize, 0 );

    for ( std::size_t start = 0; start < sequence.size() - kmerSize; ++start )
    {
        const wecall::utils::BasePairSequence read = sequence.substr( start, kmerSize );
        auto mappingPositions = hasher.mapSequence( read, boost::none );

        BOOST_REQUIRE_EQUAL( mappingPositions.size(), 1 );
        BOOST_CHECK_EQUAL( mappingPositions[0], start );
    }
}

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

BOOST_AUTO_TEST_CASE( test_fails_to_index_sequence_shorter_than_kmer_length )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( kmerSize - 1, 'C' );

    BOOST_CHECK_THROW( HashMapper( sequence, kmerSize, 0 ), wecall::utils::wecall_exception );
}

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

BOOST_AUTO_TEST_CASE( test_fails_to_map_read_shorter_than_kmer_length )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "ACGAAATTTAAAAACCCCCTCCGCCCAAAAACATGACCCAATTGG" );

    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read = sequence.substr( 0, kmerSize - 1 );
    BOOST_CHECK_THROW( hasher.mapSequence( read, 0ul ), wecall::utils::wecall_exception );
}

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

BOOST_AUTO_TEST_CASE( test_fails_to_index_sequence_longer_than_hash_table_size )
{
    const auto kmerSize = 8;
    const auto hashTableSize = std::pow( 4, kmerSize );
    const wecall::utils::BasePairSequence sequence( hashTableSize, 'A' );

    BOOST_CHECK_THROW( HashMapper( sequence, kmerSize, 0 ), wecall::utils::wecall_exception );
}

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

BOOST_AUTO_TEST_CASE( test_fails_to_map_read_longer_than_hash_table_size )
{
    const auto kmerSize = 8;
    const auto hashTableSize = std::pow( 4, kmerSize );
    const wecall::utils::BasePairSequence sequence( 100, 'A' );

    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( hashTableSize + 1, 'A' );
    BOOST_CHECK_THROW( hasher.mapSequence( read, 0ul ), wecall::utils::wecall_exception );
}

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

BOOST_AUTO_TEST_CASE( test_counts_kmer_matches_correctly_for_short_read_in_repetitive_sequence )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "AAAAAAAAAACCCCCCCCCCAAAAAAAAAA" );

    KmerMatches hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( kmerSize, 'A' );
    const auto kmerMatches = hasher.countKmerMatches( read );

    BOOST_CHECK_EQUAL( std::accumulate( kmerMatches.cbegin(), kmerMatches.cend(), 0 ), 6 );
}

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

BOOST_AUTO_TEST_CASE( test_counts_kmer_matches_correctly_for_complete_read_of_repetitive_sequence )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "AAAAAAAAAA" );

    KmerMatches hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read = sequence;
    auto counts = hasher.countKmerMatches( read );

    BOOST_CHECK_EQUAL( std::accumulate( counts.cbegin(), counts.cend(), 0 ), 3 );
}

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

BOOST_AUTO_TEST_CASE( test_exact_matching_sequence_has_higher_matching_kmer_count_than_sequence_with_snp )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "AAAAAAAAAA" );
    const wecall::utils::BasePairSequence match( "AAAAAAAAAA" );
    const wecall::utils::BasePairSequence haplotypeWithSNP( "AAAATAAAAA" );

    KmerMatches hasher( sequence, kmerSize, 0 );

    const auto refMatches = hasher.countKmerMatches( match );
    const auto varMatches = hasher.countKmerMatches( haplotypeWithSNP );

    BOOST_CHECK_GT( *std::max_element( refMatches.cbegin(), refMatches.cend() ),
                    *std::max_element( varMatches.cbegin(), varMatches.cend() ) );

    BOOST_REQUIRE_EQUAL( refMatches.size(), 1 );
    BOOST_CHECK_EQUAL( refMatches[0], 3 );
}

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

BOOST_AUTO_TEST_CASE(
    test_exact_matching_sequence_has_higher_matching_kmer_count_than_sequence_with_one_base_insertion )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "TTTTTTTTTAAAAAAAAAATTTTTTTTTTN" );
    const wecall::utils::BasePairSequence match( "TTTTTTTTTAAAAAAAAAATTTTTTTTTTN" );
    const wecall::utils::BasePairSequence haplotypeWithInsertion( "TTTTTTTTTAAAAAAAAAAATTTTTTTTTT" );

    KmerMatches hasher( sequence, kmerSize, 0 );

    const auto refMatches = hasher.countKmerMatches( match );
    const auto varMatches = hasher.countKmerMatches( haplotypeWithInsertion );

    BOOST_CHECK_GT( *std::max_element( refMatches.cbegin(), refMatches.cend() ),
                    *std::max_element( varMatches.cbegin(), varMatches.cend() ) );
}

BOOST_AUTO_TEST_CASE( test_should_only_report_one_match_if_read_and_sequence_the_same )
{
    const wecall::utils::BasePairSequence sequence( "AA" );
    const auto kmerSize = 2;
    KmerMatches matcher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AA" );
    const auto matches = matcher.countKmerMatches( read );

    const auto expectedNumMatches = sequence.size() - read.size() + 1;
    BOOST_REQUIRE_EQUAL( matches.size(), expectedNumMatches );
    BOOST_CHECK_EQUAL( matches[0], 1 );
}

BOOST_AUTO_TEST_CASE( test_should_report_no_matches_if_read_and_sequence_are_different )
{
    const wecall::utils::BasePairSequence sequence( "AA" );
    const auto kmerSize = 2;
    KmerMatches matcher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AT" );
    const auto matches = matcher.countKmerMatches( read );

    const auto expectedNumMatches = sequence.size() - read.size() + 1;
    BOOST_REQUIRE_EQUAL( matches.size(), expectedNumMatches );
    BOOST_CHECK_EQUAL( matches[0], 0 );
}

BOOST_AUTO_TEST_CASE( test_should_report_match_only_at_matching_index )
{
    const wecall::utils::BasePairSequence sequence( "AAT" );
    const auto kmerSize = 2;
    KmerMatches matcher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AT" );
    const auto matches = matcher.countKmerMatches( read );

    const auto expectedNumMatches = sequence.size() - read.size() + 1;
    BOOST_REQUIRE_EQUAL( matches.size(), expectedNumMatches );
    BOOST_CHECK_EQUAL( matches[0], 0 );
    BOOST_CHECK_EQUAL( matches[1], 1 );
}

BOOST_AUTO_TEST_CASE( test_should_report_number_of_matching_subkmers )
{
    const wecall::utils::BasePairSequence sequence( "AAAAT" );
    const auto kmerSize = 2;
    KmerMatches matcher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AAAA" );
    const auto matches = matcher.countKmerMatches( read );

    const auto expectedNumMatches = sequence.size() - read.size() + 1;
    BOOST_REQUIRE_EQUAL( matches.size(), expectedNumMatches );
    BOOST_CHECK_EQUAL( matches[0], 3 );
    BOOST_CHECK_EQUAL( matches[1], 2 );
}

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

BOOST_AUTO_TEST_CASE( test_exact_matching_sequence_has_higher_matching_kmer_count_than_sequence_with_one_base_deletion )
{
    const auto kmerSize = 8;
    const wecall::utils::BasePairSequence sequence( "TTTTTTTTTAAAAAAAAAATTTTTTTTTT" );
    const wecall::utils::BasePairSequence match( "TTTTTTTTTAAAAAAAAAATTTTTTTTTT" );
    const wecall::utils::BasePairSequence haplotypeWithDeletion( "TTTTTTTTTAAAAAAAAATTTTTTTTTT" );

    KmerMatches hasher( sequence, kmerSize, 0 );

    const auto refMatches = hasher.countKmerMatches( match );
    const auto varMatches = hasher.countKmerMatches( haplotypeWithDeletion );

    BOOST_CHECK_GT( *std::max_element( refMatches.cbegin(), refMatches.cend() ),
                    *std::max_element( varMatches.cbegin(), varMatches.cend() ) );
}

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

BOOST_AUTO_TEST_CASE( test_finds_multiple_mapping_positions_in_repetitive_sequence )
{
    const auto kmerSize = 8;
    //                                                0         1         2
    //                                                012345678901234567890123456789
    const wecall::utils::BasePairSequence sequence( "AAAAAAAAAACCCCCCCCCCAAAAAAAAAA" );

    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read = sequence.substr( 0, kmerSize );
    auto mappingPositions = hasher.mapSequence( read, boost::none );

    BOOST_CHECK_EQUAL( mappingPositions.size(), 6 );
    BOOST_CHECK_EQUAL( mappingPositions[0], 22 );
    BOOST_CHECK_EQUAL( mappingPositions[1], 21 );
    BOOST_CHECK_EQUAL( mappingPositions[2], 20 );
    BOOST_CHECK_EQUAL( mappingPositions[3], 2 );
    BOOST_CHECK_EQUAL( mappingPositions[4], 1 );
    BOOST_CHECK_EQUAL( mappingPositions[5], 0 );
}

BOOST_AUTO_TEST_CASE( test_precondition_on_kmer_size )
{
    const wecall::utils::BasePairSequence sequence( "A" );
    const auto kmerSize = 1;
    BOOST_CHECK_THROW( HashMapper( sequence, kmerSize, 0 ), wecall::utils::wecall_exception );
}

BOOST_AUTO_TEST_CASE( test_sanity1 )
{
    const wecall::utils::BasePairSequence sequence( "AAA" );
    const auto kmerSize = 2;
    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AA" );
    auto mappingPositions = hasher.mapSequence( read, boost::none );

    BOOST_CHECK_EQUAL( mappingPositions.size(), 2 );
    BOOST_CHECK_EQUAL( mappingPositions[0], 1 );
    BOOST_CHECK_EQUAL( mappingPositions[1], 0 );
}

BOOST_AUTO_TEST_CASE( test_sanity2 )
{
    const wecall::utils::BasePairSequence sequence( "AAT" );
    const auto kmerSize = 2;
    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AA" );
    auto mappingPositions = hasher.mapSequence( read, boost::none );

    BOOST_CHECK_EQUAL( mappingPositions.size(), 1 );
    BOOST_CHECK_EQUAL( mappingPositions[0], 0 );
}

BOOST_AUTO_TEST_CASE( test_behaviour_when_no_decision_can_be_made )
{
    const wecall::utils::BasePairSequence sequence( "AAA" );
    const auto kmerSize = 2;
    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "TT" );
    auto mappingPositions = hasher.mapSequence( read, boost::none );

    BOOST_REQUIRE_EQUAL( mappingPositions.size(), 0 );
}

BOOST_AUTO_TEST_CASE( test_sanity3 )
{
    const wecall::utils::BasePairSequence sequence( "ATAT" );
    const auto kmerSize = 2;
    HashMapper hasher( sequence, kmerSize, 0 );

    const wecall::utils::BasePairSequence read( "AT" );
    auto mappingPositions = hasher.mapSequence( read, boost::none );

    BOOST_CHECK_EQUAL( mappingPositions.size(), 2 );
    BOOST_CHECK_EQUAL( mappingPositions[0], 2 );
    BOOST_CHECK_EQUAL( mappingPositions[1], 0 );
}


#ifndef TVMTL_TVMINIMIZER_IRLS3D_HPP
#define TVMTL_TVMINIMIZER_IRLS3D_HPP

//System includes
#include <iostream>
#include <map>
#include <vector>
#include <chrono>

#ifdef TVMTL_TVMIN_DEBUG
    #include <string>
#endif

//Eigen includes
#include <Eigen/Sparse>
#include <Eigen/Core>

//CGAL includes For linear Interpolation
#include <CGAL/Exact_predicates_inexact_constructions_kernel.h>
#include <CGAL/Delaunay_triangulation_3.h>
#include <CGAL/Interpolation_traits_2.h>
#include <CGAL/natural_neighbor_coordinates_3.h>
#include <CGAL/interpolation_functions.h>

//vpp includes
#include <vpp/vpp.hh>

namespace tvmtl {

template <class FUNCTIONAL, class MANIFOLD, class DATA, enum PARALLEL PAR>
    class TV_Minimizer< IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3> {

	public:
	    // Manifold typedefs
	    typedef typename MANIFOLD::scalar_type scalar_type;
	    typedef typename MANIFOLD::value_type value_type;
	    typedef typename MANIFOLD::tm_base_type tm_base_type;

	    // Functional typedefs
	    typedef typename FUNCTIONAL::gradient_type gradient_type;
	    typedef typename FUNCTIONAL::sparse_hessian_type hessian_type;
	    typedef typename FUNCTIONAL::tm_base_mat_type tm_base_mat_type;
	    typedef typename Eigen::SparseSelfAdjointView<hessian_type, Eigen::Upper> sa_hessian_type;

	    // Parameters from traits class
	    typedef algo_traits< MANIFOLD::MyType > AT;
	    static const int runtime = AT::max_runtime;

	    typedef double newton_error_type;

	    // Constructor
	    TV_Minimizer(FUNCTIONAL& func, DATA& dat):
		func_(func),
		data_(dat)
	    {
		max_runtime_=AT::max_runtime;
		max_irls_steps_=AT::max_irls_steps;
		max_newton_steps_=AT::max_newton_steps;
		tolerance_=AT::tolerance;
		sparse_pattern_analyzed_ = false;
	    }

	    void first_guess();
	    void smoothening(int smooth_steps);

	    newton_error_type newton_step();
	    void minimize();
	    void output() { std::cout << "OUTPUT TEST" << std::endl; }

	    void setMax_runtime(int t) { max_runtime_ = t; }
	    void setMax_irls_steps(int n) { max_irls_steps_ = n; }
	    void setMax_newton_steps(int n) { max_newton_steps_ = n; }
	    void setTolerance(double t) {tolerance_ =t; }

	    int max_runtime(int t) const { return max_runtime_; }
	    int  max_irls_steps(int n) const { return max_irls_steps_; }
	    int max_newton_steps(int n) const { return max_newton_steps_; }
	    int tolerance(double t) const { return tolerance_; }

	private:
	    FUNCTIONAL& func_;
	    DATA& data_;

	    typename AT::template solver< hessian_type > solver_;
	    bool sparse_pattern_analyzed_;

	    int irls_step_;
	    int newton_step_;

	    int max_runtime_;
	    int max_irls_steps_;
	    int max_newton_steps_;
	    double tolerance_;

	    std::vector< std::chrono::duration<double> > Ts_;
	    std::vector< typename FUNCTIONAL::result_type > Js_;

    };

/*----- IMPLEMENTATION IRLS------*/

//First Guess
template <class FUNCTIONAL, class MANIFOLD, class DATA, enum PARALLEL PAR>
void TV_Minimizer<IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3>::first_guess(){

    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();

    std::cout << "Starting interpolation of damaged Area" << std::endl;

    typedef CGAL::Exact_predicates_inexact_constructions_kernel K;
    typedef CGAL::Delaunay_triangulation_3<K> Delaunay_triangulation;
    typedef CGAL::Interpolation_traits_2<K> Traits;
    typedef K::FT Coord_type;
    typedef K::Point_3 Point;

    typedef typename DATA::inp_type inp_type;
    int value_dim = FUNCTIONAL::value_dim;
    int value_rows = value_type::RowsAtCompileTime;
    int value_cols = value_type::ColsAtCompileTime;

    /*
    if(MANIFOLD::non_isometric_embedding)
        pixel_wise3d([&] (value_type& i){ MANIFOLD::interpolation_preprocessing(i); }, data_.img_);

    for(int R=0; R<value_rows; R++)
	for(int C=0; C<value_cols; C++){
	    std::cout << "\t Channel " << value_cols*R+C+1 << " of " << value_dim << "..." << std::endl;
	    Delaunay_triangulation T;
	    std::map<Point, Coord_type, K::Less_xy_2> function_values;
	    typedef CGAL::Data_access< std::map<Point, Coord_type, K::Less_xy_2 > >  Value_access;
	    #ifdef TVMTL_TVMIN_DEBUG
		std::cout << "NZ-Entries in inpainting matrix:" << sum3d(data_.inp_) << std::endl;
	    #endif

	    int numnodes=0;

	    // Add Interpolation nodes
	    for(int s = 0; s < ns; ++s){
		for(int r = 0; r < nr; ++r){
		// Start of row pointers
		const inp_type* inp = &data_.inp_(s, r, 0);
		const value_type* i = &data_.img_(s, r, 0);
		for(int c = 0; c < nc; ++c)
		    if(!inp[c]){
			Point p(s, r, c);
			T.insert(p);
			function_values.insert(std::make_pair(p,i[c](R,C)));
			numnodes++;
		    }
		}
	    }

	    std::cout << "\t\tNumber of Nodes: " << numnodes << std::endl;

	    int numdampix=0;
	    // Interpolate missing nodes
	    for(int s = 0; s < ns; ++s){
		for(int r = 0; r < nr; ++r){
		    // Start of row pointers
		    const inp_type* inp = &data_.inp_(s, r, 0);
		    const value_type* i = &data_.img_(s, r, 0);
		    for(int c = 0; c < nc; ++c)
			if(inp[c]){
			    Point p(s, r, c);
		            std::vector< std::pair< Point, Coord_type > > coords;
			    Coord_type norm =  CGAL::natural_neighbor_coordinates_2(T, p,std::back_inserter(coords)).second;
			    Coord_type res = CGAL::linear_interpolation(coords.begin(), coords.end(), norm,Value_access(function_values));
			    i[c](R,C)=static_cast<scalar_type>(res);
			    numdampix++;
			}
		}
	    }

	    std::cout << "\t\tNumber of interpolated Pixels: " << numdampix << std::endl;
	}

	if(MANIFOLD::non_isometric_embedding)
	    pixel_wise3d([&] (value_type& i){ MANIFOLD::interpolation_postprocessing(i); }, data_.img_);

	pixel_wise3d([&] (value_type& i){ MANIFOLD::projector(i); }, data_.img_);
*/
}

//Smoothening
template <class FUNCTIONAL, class MANIFOLD, class DATA, enum PARALLEL PAR>
void TV_Minimizer<IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3>::smoothening(int smooth_steps){
    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();

    std::cout << "Start Smoothening with max_steps = " << smooth_steps << std::endl;

    //TODO: Add Manifold dependent tranformation Before and after smoothening
    typename FUNCTIONAL::result_type Jnew = func_.evaluateJ();
    typename FUNCTIONAL::result_type Jold = Jnew + 1;
    int step = 0;

    typename FUNCTIONAL::img_type temp_img(data_.img_.domain(), vpp::_border=1);
    //    typename FUNCTIONAL::nbh_type N(temp_img);

    std::cout << "Initial functional value J=" << Jnew << "\n\n" << std::endl;

    while(Jnew<Jold && step < smooth_steps){
	if(MANIFOLD::non_isometric_embedding)
	    pixel_wise3d([&] (value_type& i){ MANIFOLD::interpolation_preprocessing(i); }, data_.img_);

	clone3d(data_.img_, temp_img);
	//vpp::fill_border_closest(temp_img);

	std::cout << "\tSmoothen step #" << step+1 << std::endl;
	std::cout << "\t Value of Functional J: " << Jnew << std::endl;
	Jold = Jnew;

	const value_type* t = &temp_img(0,0,0);
        for(int s = 0; s < ns; ++s){
	    for(int r = 0; r < nr; ++r){
	     // Start of row pointers
	        value_type* i = &data_.img_(s, r, 0);
	        for(int c = 0; c < nc; ++c){
		    i[c] = (6 * t[s,r,c] + t[s,r,c-1] + t[s,r,c+1] + t[s,r-1,c] + t[s,r+1,c] + t[s-1,r,c] + t[s+1,r,c]) / 12.0;
		}
	    }
	}

	if(MANIFOLD::non_isometric_embedding)
	    pixel_wise3d([&] (value_type& i){ MANIFOLD::interpolation_postprocessing(i); }, data_.img_);

	Jnew = func_.evaluateJ();
	step++;
    }

    std::cout << "Smoothening completed with J=" << Jnew << "\n\n" << std::endl;

}

template <class FUNCTIONAL, class MANIFOLD, class DATA, enum PARALLEL PAR>
typename TV_Minimizer<IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3>::newton_error_type TV_Minimizer<IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3>::newton_step(){

    int ns = data_.img_.nslices();
    int nr = data_.img_.nrows();
    int nc = data_.img_.ncols();
    int value_dim = FUNCTIONAL::value_dim;
    int manifold_dim = FUNCTIONAL::manifold_dim;

    // Calculate the gradient and hessian
    #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\n\t\t...Calculate Gradient" << std::endl;
    #endif
    func_.evaluateDJ();

    #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\t\t...Calculate Hessian" << std::endl;
    #endif
    func_.evaluateHJ();

    // Set up the sparse Linear system
    gradient_type x;
    const gradient_type& b = func_.getDJ();
    const hessian_type& A = func_.getHJ();//.template selfadjointView<Eigen::Upper>();

    #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\t\t\t...Gradient Size: " << b.size() << std::endl;
	std::cout << "\t\t\t...Hessian Non-Zeros: " << A.nonZeros() << std::endl;
	std::cout << "\t\t\t...Hessian Rows: " << A.rows() << std::endl;
	std::cout << "\t\t\t...Hessian Cols: " << A.cols() << std::endl;
	std::cout << "\n\t\t...Analyze Sparse Pattern" << std::endl;
    #endif
    if (!sparse_pattern_analyzed_){
	solver_.analyzePattern(A);
	sparse_pattern_analyzed_ =  true;
    }

    #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\t\t...Solve System" << std::endl;
    #endif

    // Solve the System
    solver_.factorize(A);
    x = solver_.solve(b);

    // Apply Newton correction to picture
    // - Change VectorXd to something parametrized with scalar_type
    #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\t\t...Apply Newton Correction" << std::endl;
    #endif
    const tm_base_mat_type& T = func_.getT();

    for(int s = 0; s < ns; ++s){
	#pragma omp parallel for
	for(int r = 0; r < nr; ++r){
	    // Start of row pointers
	    value_type* i = &data_.img_(s, r, 0);
	    const tm_base_type* t = &T(s, r, 0);
	    for(int c = 0; c < nc; ++c){
		Eigen::VectorXd v = -t[c]*x.segment(manifold_dim * (c + nc * r + nr * nc * s), manifold_dim); // row_wise coordinates
		MANIFOLD::exp(i[c], Eigen::Map<value_type>(v.data()), i[c]);
	    }
	}
    }
       // Compute the Error
     #ifdef TVMTL_TVMIN_DEBUG_VERBOSE
	std::cout << "\t\t...Compute Newton error" << std::endl;
    #endif
    newton_error_type error = x.norm();

    return error;
}

template <class FUNCTIONAL, class MANIFOLD, class DATA, enum PARALLEL PAR>
void TV_Minimizer<IRLS, FUNCTIONAL, MANIFOLD, DATA, PAR, 3>::minimize(){
    std::cout << "Starting IRLS Algorithm with..." << std::endl;
    std::cout << "\t Lambda = \t" << func_.getlambda() << std::endl;
    std::cout << "\t eps^2 = \t" << func_.geteps2() << std::endl;
    std::cout << "\t Tolerance = \t" <<  tolerance_ << std::endl;
    std::cout << "\t Max Steps IRLS= \t" << max_irls_steps_ << std::endl;
    std::cout << "\t Max Steps Newton = \t" << max_newton_steps_ << std::endl;

    irls_step_ = 0;

    std::chrono::time_point<std::chrono::system_clock> start, end;
    std::chrono::duration<double> t = std::chrono::duration<double>::zero();
    start = std::chrono::system_clock::now();

    // IRLS Iteration Loop
    while(irls_step_ < max_irls_steps_ && t.count() < max_runtime_){

	std::cout << "IRLS Step #" << irls_step_+1 << std::endl;
	// NOTE: evaluation of J automatically calls updateWeights(): separate eventually
	typename FUNCTIONAL::result_type J = func_.evaluateJ();
	Js_.push_back(J);
	std::cout << "\t Value of Functional J: " << J << std::endl;

	newton_step_ = 0;
	newton_error_type error = tolerance_ + 1;

	// Newton Iteration Loop
	while(tolerance_ < error && t.count() < max_runtime_ && newton_step_ < max_newton_steps_){
	    std::cout << "\t Newton step #" << newton_step_+1;
	    error = newton_step();
	    #ifdef TVMTL_TVMIN_DEBUG
		    std::string fname("step_img.csv");
		    fname = std::to_string(irls_step_) + "." + std::to_string(newton_step_) + fname;
	    #endif

	    std::cout << "\t Error: " << error << std::endl;
	    newton_step_++;
	}

	end = std::chrono::system_clock::now();
	t = end - start;
	std::cout << "\t Elapsed time: " << t.count() << " seconds." << std::endl;
	irls_step_++;
	Ts_.push_back(t);
    }

    std::cout << "Minimization in " << t.count() << " seconds." << std::endl;
}

} // end namespace tvmtl

#endif


#define BOOST_TEST_MODULE example
#include <boost/test/included/unit_test.hpp>
#include "pagerank.hpp"
#include "vector.hpp"
#include <cstdio>

// #define PRECISION double

void clean_pagerank_files(const graphee::Properties& props){
    for(int i = 0; i < props.nslices; i++){
        for(int j = 0; j < props.nslices; j++){
            std::remove((props.name+"_adj_dmatblk_"+std::to_string(i)+"_"+std::to_string(j)+".gpe").c_str());
            std::remove((props.name+"_adj_tmpblk_"+std::to_string(i)+"_"+std::to_string(j)+".gpe").c_str());
        }
        std::remove((props.name+"_ob_dvecslc_"+std::to_string(i)+".gpe").c_str());
        std::remove((props.name+"_prp1_dvecslc_"+std::to_string(i)+".gpe").c_str());
        std::remove((props.name+"_pr_dvecslc_"+std::to_string(i)+".gpe").c_str());
    }
}

void pagerank_routine(graphee::Properties& props,
    std::vector<std::string>& filenames,
    int iters) {
  graphee::DiskSparseMatrix<graphee::SparseBMatrixCSR> adjacency_matrix(
      &props, // graph properties
      "adj");

  adjacency_matrix.load_edgelist(filenames);

    graphee::Pagerank<graphee::DiskSparseMatrix<graphee::SparseBMatrixCSR>>
        pagerank(&props,            // graph properties
                 &adjacency_matrix, // give the adress to the adjacency matrix
                 0.85); // damping factor of the Pagerank (original value)

    /**
     * Compute the Pagerank with 100 iterations
     */
    pagerank.compute_pagerank(iters);
}

BOOST_AUTO_TEST_CASE( test_dense_web )
/* Compare with void free_test_function() */
{
  graphee::Properties props(
      std::string("test_pagerank_web"),            // name of your graph
      325729,                              // number of nodes
      5,                         // number of slices
      4,                              // number of threads
      5 * graphee::Properties::GB,    // max RAM value
      32 * graphee::Properties::MB); // max size of sorting vector

  std::vector<std::string> filenames;
  filenames.push_back("test/ressources/web-NotreDame.txt.gz");
  pagerank_routine(props, filenames, 100);

  long n=0;
  double score_sum=0;
  graphee::Vector<double> vec(&props);
  for(uint64_t slice_i=0; slice_i<props.nslices; slice_i++){
      vec.load("test_pagerank_web_pr_dvecslc_"+std::to_string(slice_i)+".gpe");
      for(double score : vec){
          score_sum+=score;
          n++;
      }
  }
  BOOST_CHECK(abs(score_sum-1.)<0.001);
  std::cout<<"SCORE SUM : "<<score_sum<<std::endl;
  clean_pagerank_files(props);
  std::cout<<"NVERTEX : "<<n<<std::endl;
}

BOOST_AUTO_TEST_CASE( test_smallGraph )
/* Compare with void free_test_function() */
{
  graphee::Properties props(
      std::string("test_smallGraph"),            // name of your graph
      6,                              // number of nodes
      2,                         // number of slices
      1,                              // number of threads
      5 * graphee::Properties::GB,    // max RAM value
      128 * graphee::Properties::MB); // max size of sorting vector

  graphee::DiskSparseMatrix<graphee::SparseBMatrixCSR> adjacency_matrix(
      &props, // graph properties
      "adj");

  std::vector<std::string> filenames;
  filenames.push_back("test/ressources/test_smallGraph.txt.gz");
  pagerank_routine(props, filenames, 10);

  long n=0;
  double score_sum=0;
  graphee::Vector<double> vec(&props);
  double expected[] = {0.21495,0.15189,0.03953,0.26713,0.22387,0.10260};
  for(uint64_t slice_i=0; slice_i<props.nslices; slice_i++){
      vec.load("test_smallGraph_pr_dvecslc_"+std::to_string(slice_i)+".gpe");
      for(double score : vec){
          std::cout<<n<<"\t"<<score<<std::endl;
          score_sum+=score;
          BOOST_CHECK(abs(score-expected[n])<0.00001);
          n++;
      }
  }
  std::cout<<"SCORE SUM : "<<score_sum<<std::endl;
  clean_pagerank_files(props);
}

BOOST_AUTO_TEST_CASE( test_oneEdge )
/* Compare with void free_test_function() */
{
  graphee::Properties props(
      std::string("test_oneEdge"),            // name of your graph
      2,                              // number of nodes
      2,                         // number of slices
      1,                              // number of threads
      5 * graphee::Properties::GB,    // max RAM value
      128 * graphee::Properties::MB); // max size of sorting vector

  std::vector<std::string> filenames;
  filenames.push_back("test/ressources/oneEdge.txt.gz");
  pagerank_routine(props, filenames, 3);

  long n=0;
  double score_sum=0;
  graphee::Vector<double> vec(&props);
  for(uint64_t slice_i=0; slice_i<props.nslices; slice_i++){
      vec.load("test_oneEdge_pr_dvecslc_"+std::to_string(slice_i)+".gpe");
      for(double score : vec){
          std::cout<<n<<"\t"<<score<<std::endl;
          score_sum+=score;
          n++;
      }
  }
  std::cout<<"SCORE SUM : "<<score_sum<<std::endl;
  BOOST_CHECK(abs(score_sum-1)<0.00001);
  clean_pagerank_files(props);
}

BOOST_AUTO_TEST_CASE( test_smallGraph_sum_columns )
/* Compare with void free_test_function() */
{
  graphee::Properties props(
      std::string("test_sum_columns"),            // name of your graph
      1000,                              // number of nodes
      20,                         // number of slices
      4,                              // number of threads
      10 * graphee::Properties::GB,    // max RAM value
      1 * graphee::Properties::MB); // max size of sorting vector

  graphee::DiskSparseMatrix<graphee::SparseBMatrixCSR> adjacency_matrix(
      &props, // graph properties
      "adj");

  std::vector<std::string> filenames;
  filenames.push_back("test/ressources/test_sum_columns.txt.gz");
  adjacency_matrix.load_edgelist(filenames);

    graphee::DiskVector<graphee::Vector<double>> out_bounds = std::move(graphee::DiskVector<graphee::Vector<double>>(&props, "ob", 0.));
out_bounds.dmat_columns_sum(adjacency_matrix);
double n = 0;
for(int slice = 0 ; slice < props.nslices ; slice++){
      graphee::Vector<double>& vec = out_bounds.get_slice(slice);
      for(int i = 0;i<vec.get_lines();i++){
        BOOST_CHECK_EQUAL(vec.at(i),n);
        n+=1;
      }
    }
    clean_pagerank_files(props);
}


#include <boost/test/unit_test.hpp>

#include "math/floating_point.h" // header to test
#include <map>
#include "tools/log.h"

using namespace biosim;

BOOST_AUTO_TEST_SUITE(suite_floating_point)

BOOST_AUTO_TEST_CASE(floating_point_ctor) {
  math::floating_point<float> f;
  BOOST_CHECK(f.value == 0.0);
  f = math::floating_point<float>(1.2);
  BOOST_CHECK(f.value == float(1.2));

  math::floating_point<double> d;
  BOOST_CHECK(d.value == 0.0);
  d = math::floating_point<double>(1.2);
  BOOST_CHECK(d.value == 1.2);
}

BOOST_AUTO_TEST_CASE(floating_point_cmp) {
  std::map<int, bool> fp_distance_equality{
      {0, true}, {1, true}, {2, true}, {3, true}, {4, true}, {5, false}, {10, false}};
  std::list<float> float_numbers{0, 1, 1e4, 1e30};
  std::list<double> double_numbers{0, 1, 1e4, 1e30};

  BOOST_CHECK(math::floating_point<float>::get_max_ulp_tolerance() == 4);

  BOOST_CHECK(math::almost_equal((float)0.0, (float)0.0) == true);

  for(auto p : fp_distance_equality) {
    for(auto n : float_numbers) {
      BOOST_CHECK(math::almost_equal(n, boost::math::float_advance(n, p.first)) == p.second);
      BOOST_CHECK(math::almost_equal(n, boost::math::float_advance(n, -p.first)) == p.second);
    } // for
  } // for

  math::floating_point<float> f0, f1;
  f1.sign = 1;
  BOOST_CHECK(math::almost_equal(f0, f1) == true);

  f0.value = 0;
  f0.exponent = (1 << math::floating_point<float>::get_exponent_bit_count()) - 1;
  f0.mantissa = 1;
  f1 = f0;
  f1.sign = 1;
  math::floating_point<float> f2(f1);
  f2.mantissa = 10;
  BOOST_CHECK(math::almost_equal(f0, f0) == false);
  BOOST_CHECK(math::almost_equal(f0, f1) == false);
  BOOST_CHECK(math::almost_equal(f0, f2) == false);
  BOOST_CHECK(math::almost_equal(f1, f2) == false);
  BOOST_CHECK(math::almost_equal(f0, math::floating_point<float>(1.0)) == false);

  f0.value = 0;
  f0.exponent = (1 << math::floating_point<float>::get_exponent_bit_count()) - 1;
  f0.mantissa = 0;
  f1 = f0;
  f1.sign = 1;
  f2 = f0;
  --f2.representation;
  math::floating_point<float> f3(f2);
  f3.sign = 1;
  BOOST_CHECK(math::almost_equal(f0, f0) == true);
  BOOST_CHECK(math::almost_equal(f0, f2) == true);
  BOOST_CHECK(math::almost_equal(f1, f1) == true);
  BOOST_CHECK(math::almost_equal(f1, f3) == true);
  BOOST_CHECK(math::almost_equal(f0, f1) == false);
  BOOST_CHECK(math::almost_equal(f0, f3) == false);
  BOOST_CHECK(math::almost_equal(f1, f2) == false);

  BOOST_CHECK(math::floating_point<double>::get_max_ulp_tolerance() == 4);

  BOOST_CHECK(math::almost_equal((double)0.0, (double)0.0) == true);

  for(auto p : fp_distance_equality) {
    for(auto n : double_numbers) {
      BOOST_CHECK(math::almost_equal(n, boost::math::float_advance(n, p.first)) == p.second);
      BOOST_CHECK(math::almost_equal(n, boost::math::float_advance(n, -p.first)) == p.second);
    } // for
  } // for

  math::floating_point<double> d0, d1;
  d1.sign = 1;
  BOOST_CHECK(math::almost_equal(d0, d1) == true);

  d0.value = 0;
  d0.exponent = (1 << math::floating_point<double>::get_exponent_bit_count()) - 1;
  d0.mantissa = 1;
  d1 = d0;
  d1.sign = 1;
  math::floating_point<double> d2(d1);
  d2.mantissa = 10;
  BOOST_CHECK(math::almost_equal(d0, d0) == false);
  BOOST_CHECK(math::almost_equal(d0, d1) == false);
  BOOST_CHECK(math::almost_equal(d0, d2) == false);
  BOOST_CHECK(math::almost_equal(d1, d2) == false);
  BOOST_CHECK(math::almost_equal(d0, math::floating_point<double>(1.0)) == false);

  d0.value = 0;
  d0.exponent = (1 << math::floating_point<double>::get_exponent_bit_count()) - 1;
  d0.mantissa = 0;
  d1 = d0;
  d1.sign = 1;
  d2 = d0;
  --d2.representation;
  math::floating_point<double> d3(d2);
  d3.sign = 1;
  BOOST_CHECK(math::almost_equal(d0, d0) == true);
  BOOST_CHECK(math::almost_equal(d0, d2) == true);
  BOOST_CHECK(math::almost_equal(d1, d1) == true);
  BOOST_CHECK(math::almost_equal(d1, d3) == true);
  BOOST_CHECK(math::almost_equal(d0, d1) == false);
  BOOST_CHECK(math::almost_equal(d0, d3) == false);
  BOOST_CHECK(math::almost_equal(d1, d2) == false);
}

BOOST_AUTO_TEST_SUITE_END()


/****************************************************************************
 * Copyright (c) 2012-2019 by the DataTransferKit authors                   *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the DataTransferKit library. DataTransferKit is     *
 * distributed under a BSD 3-clause license. For the licensing terms see    *
 * the LICENSE file in the top-level directory.                             *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#include <DTK_DetailsTreeVisualization.hpp>
#include <DTK_LinearBVH.hpp>

#include <Kokkos_Core.hpp>
#include <Kokkos_DefaultNode.hpp>

#include <boost/program_options.hpp>

#include <algorithm>
#include <fstream>
#include <random>

#include <point_clouds.hpp>

template <typename View>
void printPointCloud( View points, std::ostream &os )
{
    auto const n = points.extent_int( 0 );
    for ( int i = 0; i < n; ++i )
        os << "\\node[leaf] at (" << points( i )[0] << "," << points( i )[1]
           << ") {\\textbullet};\n";
}

template <typename TreeType>
void viz( std::string const &prefix, std::string const &infile,
          int n_neighbors )
{
    using DeviceType = typename TreeType::device_type;
    using ExecutionSpace = typename DeviceType::execution_space;
    Kokkos::View<DataTransferKit::Point *, DeviceType> points( "points" );
    loadPointCloud( infile, points );

    TreeType bvh( points );

    using TreeVisualization =
        typename DataTransferKit::Details::TreeVisualization<DeviceType>;
    using TikZVisitor = typename TreeVisualization::TikZVisitor;
    using GraphvizVisitor = typename TreeVisualization::GraphvizVisitor;

    int const n_queries = bvh.size();
    if ( n_neighbors < 0 )
        n_neighbors = bvh.size();
    Kokkos::View<DataTransferKit::Nearest<DataTransferKit::Point> *, DeviceType>
        queries( "queries", n_queries );
    Kokkos::parallel_for( Kokkos::RangePolicy<ExecutionSpace>( 0, n_queries ),
                          KOKKOS_LAMBDA( int i ) {
                              queries( i ) = DataTransferKit::nearest(
                                  points( i ), n_neighbors );
                          } );
    Kokkos::fence();

    auto performQueries = [&bvh, &queries]( std::string const &p,
                                            std::string const &s ) {
        std::ofstream fout;
        for ( int i = 0; i < queries.extent_int( 0 ); ++i )
        {
            std::string const fname = p + std::to_string( i ) + s;
            fout.open( fname, std::fstream::out );
            TreeVisualization::visit( bvh, queries( i ),
                                      GraphvizVisitor{fout} );
            fout.close();
        }
    };

    std::fstream fout;

    // Print the point cloud
    fout.open( prefix + "points.tex", std::fstream::out );
    printPointCloud( points, fout );
    fout.close();

    // Print the bounding volume hierarchy
    fout.open( prefix + "bounding_volumes.tex", std::fstream::out );
    TreeVisualization::visitAllIterative( bvh, TikZVisitor{fout} );
    fout.close();

    // Print the entire tree
    fout.open( prefix + "tree_all_nodes_and_edges.dot.m4", std::fstream::out );
    TreeVisualization::visitAllIterative( bvh, GraphvizVisitor{fout} );
    fout.close();

    std::string const suffix = "_nearest_traversal.dot.m4";
    performQueries( prefix + "untouched_", suffix );

    // Shuffle the queries
    std::random_device rd;
    std::mt19937 g( rd() );
    std::shuffle( queries.data(), queries.data() + queries.size(), g );
    performQueries( prefix + "shuffled_", suffix );

    // Sort them
    auto permute = DataTransferKit::Details::BatchedQueries<
        DeviceType>::sortQueriesAlongZOrderCurve( bvh.bounds(), queries );
    queries =
        DataTransferKit::Details::BatchedQueries<DeviceType>::applyPermutation(
            permute, queries );
    performQueries( prefix + "sorted_", suffix );
}

// FIXME version of Kokkos in the CI/dev base image is too old and does not have
// scope guards so we define our own...
struct KokkosScopeGuard
{
    KokkosScopeGuard( Kokkos::InitArguments const &args )
    {
        Kokkos::initialize( args );
    }
    ~KokkosScopeGuard() { Kokkos::finalize(); }
};

int main( int argc, char *argv[] )
{
    Kokkos::InitArguments args;
    args.disable_warnings = true;
    KokkosScopeGuard guard( args );

    std::string prefix;
    std::string infile;
    int n_neighbors;
    boost::program_options::options_description desc( "Allowed options" );
    // clang-format off
    desc.add_options()
        ( "help", "produce help message" )
        ( "prefix", boost::program_options::value<std::string> (&prefix)->default_value("viz_"), "set prefix for output files" )
        ( "infile", boost::program_options::value<std::string> (&infile)->default_value("leaf_cloud.txt"), "set input point cloud file" )
        ( "neighbors", boost::program_options::value<int> (&n_neighbors)->default_value(5), "set the number of neighbors to search for (negative value means all)" )
    ;
    // clang-format on

    boost::program_options::variables_map vm;
    boost::program_options::store(
        boost::program_options::parse_command_line( argc, argv, desc ), vm );
    boost::program_options::notify( vm );

    if ( vm.count( "help" ) )
    {
        std::cout << desc << "\n";
        return 1;
    }

    using Serial = Kokkos::Compat::KokkosSerialWrapperNode::device_type;
    using Tree = DataTransferKit::BVH<Serial>;
    viz<Tree>( prefix, infile, n_neighbors );

    return 0;
}