/**
 *    \file include/geombd/io/parser.hpp
 *    \author Keny Ordaz
 *    \version 1.0
 *    \date 2021
 *
 *    Class to parse a multibody system from URDF
 *    Copyright (c) 2021 Cinvestav
 *    This library is distributed under the MIT License.
 */

#ifndef HEADER_IO_PARSER_HPP
#define HEADER_IO_PARSER_HPP
#include <iostream>
#include <iomanip>
#include <vector>
#include <fstream>
#include <sstream>
#include <optional>
#include <memory>
#include <unordered_map>
//#include <chrono>
#include <Eigen/Core>
//#include "pugixml.hpp"

namespace pugi {
    class xml_node;
}
using namespace Eigen;

// support {{{1
RowVectorXd string_to_vector(const std::string& vec, int size=3);

class Pose
{
public:
    Pose(): xyz(0, 0, 0), rpy(0, 0, 0) {}

    void init_rpy(double roll, double pitch, double yaw)
    {
        auto phi = roll / 2.0;
        auto the = pitch / 2.0;
        auto psi = yaw / 2.0;
        auto sphi = sin(phi);  auto cphi = cos(phi);
        auto sthe = sin(the);  auto cthe = cos(the);
        auto spsi = sin(psi);  auto cpsi = cos(psi);

        this->x = sphi * cthe * cpsi - cphi * sthe * spsi;
        this->y = cphi * sthe * cpsi + sphi * cthe * spsi;
        this->z = cphi * cthe * spsi - sphi * sthe * cpsi;
        this->w = cphi * cthe * cpsi + sphi * sthe * spsi;

        this->normalize();
    };
    void normalize()
    {
        double s = sqrt(this->x * this->x +
                this->y * this->y +
                this->z * this->z +
                this->w * this->w);
        if (s == 0.0) {
            this->x = 0.0;
            this->y = 0.0;
            this->z = 0.0;
            this->w = 1.0;
        } else {
            this->x /= s;
            this->y /= s;
            this->z /= s;
            this->w /= s;
        }
    };
    double x, y, z, w;
    RowVector3d xyz;
    RowVector3d rpy;
};

class Axis
{
public:
    Axis(): xyz(1.0, 0., 0.) {}

    RowVector3d xyz;
};

class Inertia { // optional zero inertia
public:
    Inertia(double ixx=0.0, double ixy=0.0, double ixz=0.0, double iyy=0.0, double iyz=0.0, double izz=0.0)
        : ixx(ixx), ixy(ixy), ixz(ixz), iyy(iyy), iyz(iyz), izz(izz)
        { setMatrix(ixx, ixy, ixz, iyy, iyz, izz); }
    void setMatrix(double ixx=0.0, double ixy=0.0, double ixz=0.0, double iyy=0.0, double iyz=0.0, double izz=0.0)
        { value << ixx, ixy, ixz, ixy, iyy, iyz, ixz, iyz, izz; }

    double ixx;
    double ixy;
    double ixz;
    double iyy;
    double iyz;
    double izz;
    Matrix3d value;
};

class Inertial {
public:
    Inertial(): origin(), mass(0.0), inertia(), spatial() {}
    Inertial(const pugi::xml_node & node);
    std::string to_string() const {
        static IOFormat CommaInitFmt(StreamPrecision, DontAlignCols, ", ", ", ", "", "", " << ", ";");
        static IOFormat OctaveFmt(StreamPrecision, 0, ", ", ";\n", "", "", "[", "]");
        std::ostringstream out;
        out << "mass (" << mass << "),\n\t origin[xyz] (" << origin.xyz.format(CommaInitFmt);
        out << "), origin[rpy] (" << origin.rpy.format(CommaInitFmt) << "); inertia tensor:\n";
        out << inertia.value.format(OctaveFmt);
        return out.str();

    }
    Pose origin;
    double mass;
    Inertia inertia;
    Matrix<double, 6, 6> spatial;
};
// support }}}1

class Element // {{{2
{
public:
    Element(const std::string& name, int id): name(name), id(id), tree_id(-1) {}
    std::string name;
    int id;
    int tree_id;
    //~Element() { std::cout << __PRETTY_FUNCTION__ << "\n"; }
};  // }}}2

class Body : public Element  // {{{1
{
public:
    Body(const std::string& name="unassigned", int id=-1);
    std::string to_string() const;
    void addChild(int child_id);
    void addParent(int parent_id);
    void addJoint(int joint_id);
    void setInertialProperties(const Inertial& value);
    bool has_inertial;
    Inertial inertial_properties;
    std::vector<int> child_links;
    std::vector<int> parent_links;
    std::vector<int> parent_joints;
    int parent;
    std::vector<int> Pre;
    std::vector<int> Suc;
};  // }}}1

class Joint : public Element // {{{1
{
    // MAke enum for joint types
public:
    enum joint_type {
        revolute = 'r',  /*!< The type of the joint is revolute. */
        prismatic = 'p', /*!< The type of the joint is prismatic. */
        fixed = 'f',      /*!< The type of the joint is glue. */
        notype = 'n'      /*!< There is no information about the type. */
    };

    Joint(const std::string& name="unassigned", int id=-1): nq(0), Element(name, id) {}
    Joint(const std::string& name, unsigned int id,
            const std::string& parent, const std::string& child, const std::string& type)
        : /*name(name), id(id),*/ predecessor(parent), successor(child)
        , axis(1.0, 0, 0), Element(name, id) {
            if (type == "revolute") {
                this->type = Joint::joint_type::revolute;
            } else if (type == "prismatic") {
                this->type = Joint::joint_type::prismatic;
            } else if (type == "fixed") {
                this->type = Joint::joint_type::fixed;
            } else {
                this->type = Joint::joint_type::notype;

            }

            switch(this->type) {
            case joint_type::revolute:
            case joint_type::prismatic:
                // TODO: check limits
                setLimit();
                nq = 1;
                break;
            default:
                nq = 0;// Do nothing?;
            }
        }
    //~Joint() { std::cout << __PRETTY_FUNCTION__ << "\n"; }

    std::string to_string() {
        std::ostringstream out;
        out << "joint " << name << ", id " << id << ", type " << std::string(1, type);
        out << ", axis " << axis;
        if (type == joint_type::revolute || type == joint_type::prismatic) {
            out << ", limits(" << upper << ", " << lower << ", " << effort << ", " << velocity << ") ";
        }
        out << joint_axis_indicator();
        out << "\t/" << successor << "/" << predecessor << "/ ";
        return out.str();
    }

    std::string joint_axis_indicator() {
        std::ostringstream out;
        out << "[";
        switch(type) {
        case Joint::joint_type::revolute:
            out << "⚙";
            break;
        case Joint::joint_type::prismatic:
            out << "↕";
            break;
        case Joint::joint_type::fixed:
            out << "⚓";
            break;
        case Joint::joint_type::notype:
            out << "⦻";
        }
        if (type == joint_type::revolute || type == joint_type::prismatic) {
            if (axis == RowVector3d(1.0, 0.0, 0.0))
                out << "+X";
            else if (axis == RowVector3d(0.0, 1.0, 0.0))
                out << "+Y";
            else if (axis == RowVector3d(-1.0, 0.0, 0.0))
                out << "-X";
            else if (axis == RowVector3d(0.0, -1.0, 0.0))
                out << "-Y";
            else if (axis == RowVector3d(0.0, 0.0, 1.0))
                out << "+Z";
            else if (axis == RowVector3d(0.0, 0.0, -1.0))
                out << "-Z";
            else
                out << "⦻";
        }

        out << "]";
        return out.str();
    }

    void setLimit(double upper=0.0, double lower=0.0, double effort=0.0, double velocity=0.0) {
        this->upper = upper; this->lower = lower; this->effort = effort; this->velocity = velocity;
    }

    void setOrigin(const RowVector3d& xyz, const RowVector3d& rpy) {
        origin.xyz = xyz;
        origin.rpy = rpy;
    }

    void setDampingAndFriction(double damping=0.0, double friction=0.0) {} //FIXME: If required
    std::string predecessor;
    std::string successor;
    joint_type type;
    RowVector3d axis;
    Pose origin;
    double upper, lower, effort, velocity;
    double damping, friction;
    int nq;
};  // }}}1

class Robot { // {{{1
public:
    Robot(const std::string& name="unknown");
    void addJoint(const std::string& name, unsigned int id,
            const std::string& parent, const std::string& child, const std::string& type);
    void addBody(const std::string& name, unsigned int id);
    void dump(std::ostream & os);
    void print_featherstone(std::ostream & os);
    void print_tree(std::ostream & os, const std::string & prefix, int id, int parent_id=-1);
    void print_tree(std::ostream & os);
    int getBodyID(const std::string& frame_name) const;
    int getJointID(const std::string& name) const;
    std::shared_ptr<Joint> & operator()(int i) { // robot(j)
        return joints[i];
    }
    std::shared_ptr<Body> & operator[](int i) { // robot[i]
        return bodies[i];
    }
    int getRootBodyID() const;
    const std::string& getName() const { return name; }
    static std::optional<std::shared_ptr<Robot>> build_model(const std::string& urdf_file);
    void add_pre(std::shared_ptr<Body> body, std::shared_ptr<Body> body_s);
    void add_suc(std::shared_ptr<Body> body, std::shared_ptr<Body> body_p);
    void build_PSvector( );
    void build_spatial_inertias( );
protected:
    bool load_from_urdf(const pugi::xml_node & node);
    void visit_spanning_tree(int root_id);
    void build_graph();
protected:
    std::string name;
    std::vector<std::shared_ptr<Joint>> joints;
    std::vector<std::shared_ptr<Body>> bodies;
    std::unordered_map<std::string,int> joint_dico;
    std::unordered_map<std::string,int> joint_pc_dico;
    std::unordered_map<std::string,int> link_dico;
public:
    std::vector<int> body_l;
    std::vector<int> body_m;
    std::vector<int> joint_p;
    std::vector<int> joint_s;
    Eigen::RowVectorXi tree_ids;
    Eigen::RowVectorXi j_tree_ids;
public:
    int nq;
private:
    int _count;
};  // }}}1

#endif


/**************************************************************************\
|
|    Copyright (C) 2009 Marc Stevens
|
|    This program is free software: you can redistribute it and/or modify
|    it under the terms of the GNU General Public License as published by
|    the Free Software Foundation, either version 3 of the License, or
|    (at your option) any later version.
|
|    This program is distributed in the hope that it will be useful,
|    but WITHOUT ANY WARRANTY; without even the implied warranty of
|    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
|    GNU General Public License for more details.
|
|    You should have received a copy of the GNU General Public License
|    along with this program.  If not, see <http://www.gnu.org/licenses/>.
|
\**************************************************************************/

#ifndef MAIN_HPP
#define MAIN_HPP

#include <iostream>
#include <vector>
#include <string>
#include <stdexcept>

#include <boost/filesystem/operations.hpp>
#include <boost/thread.hpp>

#include <hashclash/saveload_gz.hpp>
#include <hashclash/sdr.hpp>
#include <hashclash/differentialpath.hpp>
#include <hashclash/booleanfunction.hpp>
#include <hashclash/timer.hpp>

using namespace hashclash;
using namespace std;

extern boost::mutex mut;
extern std::string workdir;
class path_container;
void dostep(path_container& container);
bool check_path_collfind(const differentialpath& diffpath, const uint32 mdiff[16]);

struct connect_bitdata {
	uint32 dQt;
	uint32 dQtp1;
	uint32 dFt;
	uint32 dFtp1;
	uint32 dFtp2;
	uint32 dFtp3;

	inline bool operator== (const connect_bitdata& r) const {
		return dQt==r.dQt && dQtp1==r.dQtp1 && dFt==r.dFt && dFtp1==r.dFtp1 && dFtp2==r.dFtp2 && dFtp3==r.dFtp3;
	}
	inline bool operator!= (const connect_bitdata& r) const {
		return !(*this == r);
	}
	inline bool operator< (const connect_bitdata& r) const {
		if (dQt < r.dQt) return true;
		if (dQt > r.dQt) return false;
		if (dQtp1 < r.dQtp1) return true;
		if (dQtp1 > r.dQtp1) return false;
		if (dFt < r.dFt) return true;
		if (dFt > r.dFt) return false;
		if (dFtp1 < r.dFtp1) return true;
		if (dFtp1 > r.dFtp1) return false;
		if (dFtp2 < r.dFtp2) return true;
		if (dFtp2 > r.dFtp2) return false;
		if (dFtp3 < r.dFtp3) return true;
		return false;
	}
};

struct md5_connect_thread {
	md5_connect_thread(): sw(true), countb(33,0), countbaborted(33,0), countbdepth(33,0)
	{}
	void md5_connect(const vector<differentialpath>& lowerpaths, const differentialpath& upperpath, path_container& container);
	timer sw/*(true)*/;
	vector<unsigned> countb/*(33, 0)*/;
	vector<unsigned> countbaborted/*(33, 0)*/;
	vector<unsigned> countbdepth/*(33, 0)*/;
	uint64 count, countall;
	vector<unsigned char> isgood;
	vector<int> lowerpathsmaxtunnel;

	unsigned md5_connect_bits(const vector<differentialpath>& lowers, unsigned index, const differentialpath& upper, path_container& container);
	connect_bitdata startbit0;
	vector<connect_bitdata> bitdataresults[33];
	vector<connect_bitdata> bitdatastart[32];
	vector<connect_bitdata> bitdataend[32];
	vector<byteconditions>  bitdatanewcond[32];
	differentialpath newpath2, tmppath;
	unsigned bindex[32];

	void connectbits2(const connect_bitdata& in, unsigned b, const differentialpath& lower, const differentialpath& upper);
	bf_outcome bfo0, bfo1, bfo2, bfo3;
	bf_conditions bfc0, bfc1, bfc2, bfc3;

	void connectbits(const connect_bitdata& in, vector<connect_bitdata>& out, unsigned b,
		const differentialpath& lower, const differentialpath& upper, vector<byteconditions>* newconds = 0);
	connect_bitdata result;
	//bf_outcome bfo0, bfo1, bfo2, bfo3;
	//bf_conditions bfc0, bfc1, bfc2, bfc3;
	bitcondition Qt, Qtp1;
	byteconditions newcond;
	bool lastdFp1, lastdFp2, lastdFp3;

	unsigned t;
	booleanfunction* Ft;
	booleanfunction* Ftp1;
	booleanfunction* Ftp2;
	booleanfunction* Ftp3;
	vector<uint32> dFt, dFtp1, dFtp2, dFtp3;
	uint32 dmt, dmtp1, dmtp2, dmtp3;
	uint32 dQtp1, dQtp2, dQtp3, dQtp4;
	differentialpath newpath;

	inline bool isinrange(const differentialpath& needle, const differentialpath& haystack, unsigned b)
	{
		for (unsigned j = t-2; j <= t; ++j)
		{
			const wordconditions& nt = needle[j];
			const wordconditions& ht = haystack[j];
			unsigned k = 0;
			while (k+7 <= b)
			{
				if (nt.bytes[k>>3] != ht.bytes[k>>3])
					return false;
				k += 8;
			}
			for (; k <= b; ++k)
				if (nt[k] != ht[k])
					return false;
		}
		return true;
	}
	inline unsigned binary_search_lower_paths(const vector<differentialpath>& lowerpaths,
									   unsigned i, unsigned b)
	{
		const differentialpath& needle = lowerpaths[i];
		unsigned l = i, u = lowerpaths.size()-1;
		while (l < u)
		{
			unsigned j = (l+u+1)>>1;
			if (isinrange(needle, lowerpaths[j], b)) {
				l = j;
			} else {
				u = j-1;
			}
		}
		return l;
	}
};

extern vector<uint32> lowdQt, lowdQtm1, lowdQtm2, lowdQtm3;

class path_container {
public:
	path_container()
		: modn(1), modi(0), inputfilelow(), inputfilehigh()
		, t(0), noverify(false), noenhancepath(false)
		, verified(0), verifiedbad(0), bestpathcond(1<<20)
		, bestmaxtunnel(0), showstats(false), bestmaxcomp(-1000), threads(1)
	{
		for (unsigned k = 0; k < 16; ++k)
			m_diff[k] = 0;
	}

	~path_container()
	{
		cout << "Best path: totcompl=" << bestmaxcomp << " tottunnel=" << bestmaxtunnel << ", totcond=" << bestpathcond << endl;
		if (!noverify)
			cerr << "Verified: " << verifiedbad << " bad out of " << verified << endl;
	}

	void push_back(const differentialpath& fullpath)
	{
		differentialpath pathback = fullpath;
		//pathback = fullpath;
		try {
		cleanup(pathback);
		} catch (std::exception& e) {
			cerr << "hashclash::cleanup(differentialpath&): unknown exception!:" << endl << e.what() << endl;
			show_path(pathback, m_diff);
		}catch (...) {
			cerr << "hashclash::cleanup(differentialpath&): unknown exception!:" << endl;
			show_path(pathback, m_diff);
		}

//		++verified;
		if (!noverify && !test_path_fast(pathback, m_diff))
		{
			mut.lock();
			++verified;
			++verifiedbad;
			mut.unlock();
			return;
		}

		unsigned tunnel = totaltunnelstrength(pathback);
		int tuncompl = tunnel;
		for (int k = pathback.tbegin(); k < pathback.tend() && k < 64; ++k)
			if (k >= Qcondstart)
				tuncompl -= int(pathback[k].hw());
		if (tuncompl < bestmaxcomp)
			return;
		if (tunnel < bestmaxtunnel)
			return;
		unsigned cond = 0;
		for (int k = pathback.tbegin(); k < pathback.tend(); ++k)
			cond += pathback[k].hw();
		if (!(tuncompl >= bestmaxcomp || tunnel >= bestmaxtunnel || cond <= bestpathcond))
			return;
		try {
			if (!noenhancepath) {
				enhancepath(pathback, m_diff);
				tunnel = totaltunnelstrength(pathback);
				tuncompl = tunnel;
				for (int k = pathback.tbegin(); k < pathback.tend() && k < 64; ++k)
					if (k >= Qcondstart)
						tuncompl -= int(pathback[k].hw());
				if (tuncompl < bestmaxcomp)
					return;
				if (tunnel < bestmaxtunnel)
					return;
				cond = 0;
				for (int k = pathback.tbegin(); k < pathback.tend(); ++k)
					cond += pathback[k].hw();
				if (!(tuncompl >= bestmaxcomp || tunnel >= bestmaxtunnel || cond <= bestpathcond))
					return;
			}
		} catch (std::exception&) {
			return;
		} catch (...) {
			return;
		}

		if (!check_path_collfind(pathback, m_diff))
			return;

		if (tuncompl == bestmaxcomp && tunnel == bestmaxtunnel && cond == bestpathcond) {
			mut.lock();
			++verified;
			bestpaths.push_back(pathback);
			if (hw(uint32(bestpaths.size()))==1) {
				save_gz(bestpaths, workdir + "/bestpaths", binary_archive);
				cout << "Best paths: " << bestpaths.size() << endl;
			}
			mut.unlock();
			return;
		}
		if (tuncompl < bestmaxcomp) return;
		if (tunnel < bestmaxtunnel) return;
		if (cond > bestpathcond) return;
		mut.lock();
		++verified;
		bestpaths.clear();
		bestpaths.push_back(pathback);

		bestmaxcomp = tuncompl;
		bestpathcond = cond;
		bestpath = pathback;
		bestmaxtunnel = tunnel;
		show_path(pathback, m_diff);
		double p = test_path(pathback, m_diff);
		cout << "Best path: totcompl=" << bestmaxcomp << " tottunnel=" << bestmaxtunnel << ", totcond=" << bestpathcond << ", p=" << p << endl;
		save_gz(pathback, workdir + "/bestpath_t" + boost::lexical_cast<string>(tunnel) + "_c" + boost::lexical_cast<string>(cond), binary_archive);
		save_gz(pathback, workdir + "/bestpath_new", binary_archive);
		save_gz(bestpaths, workdir + "/bestpaths_new", binary_archive);
		try { boost::filesystem::rename(workdir + "/bestpath.bin.gz", workdir + "/bestpath_old.bin.gz"); } catch (...) {}
		try { boost::filesystem::rename(workdir + "/bestpaths.bin.gz", workdir + "/bestpaths_old.bin.gz"); } catch (...) {}
		try { boost::filesystem::rename(workdir + "/bestpath_new.bin.gz", workdir + "/bestpath.bin.gz"); } catch (...) {}
		try { boost::filesystem::rename(workdir + "/bestpaths_new.bin.gz", workdir + "/bestpaths.bin.gz"); } catch (...) {}
		mut.unlock();
	}

	uint32 m_diff[16];

	unsigned t;
	int Qcondstart;

	bool showstats;
	bool noenhancepath;
	bool noverify;
	unsigned verified, verifiedbad;

	unsigned modn;
	unsigned modi;
	std::string inputfilelow, inputfilehigh;
	bool showinputpaths;

	differentialpath bestpath;
	vector<differentialpath> bestpaths;
	volatile unsigned bestpathcond;
	volatile unsigned bestmaxtunnel;
	volatile int bestmaxcomp;
	int threads;
};

struct diffpathlower_less
	: public std::binary_function<differentialpath, differentialpath, bool>
{
	bool operator()(const differentialpath& _Left, const differentialpath& _Right) const
	{
		if (_Left.tend() < _Right.tend() || _Left.tbegin() > _Right.tbegin())
			return true;
		if (_Left.tend() > _Right.tend() || _Left.tbegin() < _Right.tbegin())
			return false;
		if (_Left.path.size() < 3)
			throw std::runtime_error("Lower differential paths of insufficient size");
		unsigned t = _Left.tend() - 1;
		uint32 LdQt   = _Left[t].diff();
		uint32 RdQt   = _Right[t].diff();
		for (unsigned b = 0; b < 32; ++b)
		{
			if ((LdQt & (1<<b)) < (RdQt & (1<<b)))
				return true;
			if ((LdQt & (1<<b)) > (RdQt & (1<<b)))
				return false;
			if (_Left[t-2][b] < _Right[t-2][b])
				return true;
			if (_Left[t-2][b] > _Right[t-2][b])
				return false;
			if (_Left[t-1][b] < _Right[t-1][b])
				return true;
			if (_Left[t-1][b] > _Right[t-1][b])
				return false;
		}
		return false;
	}
};

struct diffpathupper_less
	: public std::binary_function<differentialpath, differentialpath, bool>
{
	bool operator()(const differentialpath& _Left, const differentialpath& _Right) const
	{
		if (_Left.tend() < _Right.tend() || _Left.tbegin() > _Right.tbegin())
			return true;
		if (_Left.tend() > _Right.tend() || _Left.tbegin() < _Right.tbegin())
			return false;
		if (_Left.path.size() < 3)
			throw std::runtime_error("Lower differential paths of insufficient size");
		unsigned t = _Left.tbegin() - 1;
		uint32 LdQtp1   = _Left[t+1].diff();
		uint32 RdQtp1   = _Right[t+1].diff();
		if (_Left[t+1].hw() > _Right[t+1].hw()) return true;
		if (_Left[t+1].hw() < _Right[t+1].hw()) return false;
		for (unsigned b = 0; b < 32; ++b)
		{
			if ((LdQtp1 & (1<<b)) < (RdQtp1 & (1<<b)))
				return true;
			if ((LdQtp1 & (1<<b)) > (RdQtp1 & (1<<b)))
				return false;
			if (_Left[t+2][b] < _Right[t+2][b])
				return true;
			if (_Left[t+2][b] > _Right[t+2][b])
				return false;
			if (_Left[t+3][b] < _Right[t+3][b])
				return true;
			if (_Left[t+3][b] > _Right[t+3][b])
				return false;
		}
		return false;
	}
};

#endif // MAIN_HPP


#pragma once
/******************************************************************************
*
*   Copyright (c) 2019 AT&T Intellectual Property.
*   Copyright (c) 2018-2019 Nokia.
*
*   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.
*
******************************************************************************/

// Standard Includes: ANSI C/C++, MSA, and Third-Party Libraries
#include <cmath>
#include <boost/integer/static_log2.hpp>

// Local Includes: Application specific classes, functions, and libraries
#include "asn/per/common.hpp"
#include "asn/per/binary_integer.hpp"

namespace asn {
namespace per {

/***************************************************************************************
* Encoding of a constrained whole number (X.691 10.5)
***************************************************************************************/

template<bound_t R, class E = void>
struct length_determinant;

template<bound_t R>
struct length_determinant<R, std::enable_if_t< R == 0 > > { static constexpr bound_t value = 0; };

template<bound_t R>
struct length_determinant<R, std::enable_if_t< R == 1 > > { static constexpr bound_t value = 1; };

template<bound_t R>
struct length_determinant<R, std::enable_if_t< (R  > 1) > > { static constexpr bound_t value = boost::static_log2<(R - 1)>::value + 1; };

/***************************************************************************************
***************************************************************************************/

template <class Range, class V, class Enable = void>
struct ConstrainedWholeNumber;

//Bit-field case
template <class Range, class V>
struct ConstrainedWholeNumber<Range, V, std::enable_if_t<(Range::upper_bound < (Range::lower_bound + 255))> >
{
	//non-negative-binary-integer X.691 10.5
	static void inline run(EncoderCtx& ctx, const V& val)
	{
		Tools::bit_accessor::put(static_cast<u8>(val - Range::lower_bound),
			length_determinant<(Range::upper_bound - Range::lower_bound + 1)>::value,
			ctx.refBuffer());
	}
	static V inline run(DecoderCtx& ctx)
	{
		return Range::lower_bound + Tools::bit_accessor::get(
			length_determinant<(Range::upper_bound - Range::lower_bound + 1)>::value,
			ctx.refBuffer());
	}
};

//One octet case
template <class Range, class V>
struct ConstrainedWholeNumber<Range, V, std::enable_if_t<(Range::upper_bound == (Range::lower_bound + 255))> >
{
	//non-negative-binary-integer X.691 10.5
	static void inline run(EncoderCtx& ctx, const V& val)
	{
		Tools::bit_accessor::padByte(ctx.refBuffer());
		ctx.refBuffer().putByte(static_cast<u8>(val - Range::lower_bound));
	}
	static V inline run(DecoderCtx& ctx)
	{
		V rval = 0;
		Tools::bit_accessor::padByte(ctx.refBuffer());
		u8 const* data = ctx.refBuffer().getBytes(1);
		if (data)
			rval = Range::lower_bound + data[0];
		return rval;
	}
};

//Two octets case
template <class Range, class V>
struct ConstrainedWholeNumber<Range, V, std::enable_if_t<(Range::upper_bound > (Range::lower_bound + 255)) && (Range::upper_bound <= (Range::lower_bound + 65535))> >
{
	//non-negative-binary-integer X.691 10.5
	static void inline run(EncoderCtx& ctx, const V& v)
	{
		u64 val = static_cast<u64>(v - Range::lower_bound);
		Tools::bit_accessor::padByte(ctx.refBuffer());
		ctx.refBuffer().putByte((u8)(val >> 8));
		ctx.refBuffer().putByte((u8)val);
	}
	static V inline run(DecoderCtx& ctx)
	{
		V rval = 0;
		Tools::bit_accessor::padByte(ctx.refBuffer());
		u8 const* data = ctx.refBuffer().getBytes(2);
		if (data) {
			rval = data[0];
			rval = rval << 8;
			rval |= data[1];
			rval += Range::lower_bound;
		}
		return rval;
	}
};

//Indefinite case
template <class Range, class V>
struct ConstrainedWholeNumber<Range, V, std::enable_if_t< (Range::upper_bound > Range::lower_bound + 65535) > >
{
	struct NormalizedValueRange
	{
		using boundary_type = typename Range::boundary_type;
		static constexpr bool extended = Range::extended;
		static constexpr boundary_type lower_bound = 0;
		static constexpr boundary_type upper_bound = Range::upper_bound - Range::lower_bound;
	};

	//non-negative-binary-integer X.691 10.5
	static void inline run(EncoderCtx& ctx, const V& val)
	{
		NonnegativeBinaryInteger<NormalizedValueRange>::run(val - Range::lower_bound, ctx);
	}
	static V inline run(DecoderCtx& ctx)
	{
		V rval = 0;
		NonnegativeBinaryInteger<NormalizedValueRange>::run(rval, ctx, false);
		return rval + Range::lower_bound;
	}
};

} //namespace per
} //namespace asn


// Boost.Geometry

// Copyright (c) 2020, Oracle and/or its affiliates.

// Contributed and/or modified by Adam Wulkiewicz, on behalf of Oracle

// Licensed under the Boost Software License version 1.0.
// http://www.boost.org/users/license.html

#ifndef BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_HPP
#define BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_HPP

#include <boost/geometry/strategies/area/geographic.hpp>
#include <boost/geometry/strategies/convex_hull/geographic.hpp>
#include <boost/geometry/strategies/envelope/geographic.hpp>
#include <boost/geometry/strategies/expand/geographic.hpp>

namespace boost { namespace geometry
{

namespace strategies
{

template
<
    typename FormulaPolicy = strategy::andoyer,
    std::size_t SeriesOrder = strategy::default_order<FormulaPolicy>::value,
    typename Spheroid = srs::spheroid<double>,
    typename CalculationType = void
>
class geographic : strategies::detail::geographic_base<Spheroid>
{
    using base_t = strategies::detail::geographic_base<Spheroid>;

public:
    geographic()
        : base_t()
    {}

    explicit geographic(Spheroid const& spheroid)
        : base_t(spheroid)
    {}

    // area

    template <typename Geometry>
    auto area(Geometry const&) const
    {
        return strategy::area::geographic
            <
                FormulaPolicy, SeriesOrder, Spheroid, CalculationType
            >(base_t::m_spheroid);
    }

    // envelope

    template <typename Geometry, typename Box>
    static auto envelope(Geometry const&, Box const&,
                         typename util::enable_if_point_t<Geometry> * = nullptr)
    {
        return strategy::envelope::spherical_point();
    }

    template <typename Geometry, typename Box>
    static auto envelope(Geometry const&, Box const&,
                         typename util::enable_if_multi_point_t<Geometry> * = nullptr)
    {
        return strategy::envelope::spherical_multipoint();
    }

    template <typename Geometry, typename Box>
    static auto envelope(Geometry const&, Box const&,
                         typename util::enable_if_box_t<Geometry> * = nullptr)
    {
        return strategy::envelope::spherical_box();
    }

    template <typename Geometry, typename Box>
    auto envelope(Geometry const&, Box const&,
                  typename util::enable_if_segment_t<Geometry> * = nullptr) const
    {
        return strategy::envelope::geographic_segment
            <
                FormulaPolicy, Spheroid, CalculationType
            >(base_t::m_spheroid);
    }

    template <typename Geometry, typename Box>
    auto envelope(Geometry const&, Box const&,
                  typename util::enable_if_polysegmental_t<Geometry> * = nullptr) const
    {
        return strategy::envelope::geographic
            <
                FormulaPolicy, Spheroid, CalculationType
            >(base_t::m_spheroid);
    }

    // expand

    template <typename Box, typename Geometry>
    static auto expand(Box const&, Geometry const&,
                       typename util::enable_if_point_t<Geometry> * = nullptr)
    {
        return strategy::expand::spherical_point();
    }

    template <typename Box, typename Geometry>
    static auto expand(Box const&, Geometry const&,
                       typename util::enable_if_box_t<Geometry> * = nullptr)
    {
        return strategy::expand::spherical_box();
    }

    template <typename Box, typename Geometry>
    auto expand(Box const&, Geometry const&,
                typename util::enable_if_segment_t<Geometry> * = nullptr) const
    {
        return strategy::expand::geographic_segment
            <
                FormulaPolicy, Spheroid, CalculationType
            >(base_t::m_spheroid);
    }
};

} // namespace strategies

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_STRATEGIES_GEOGRAPHIC_HPP


#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <string>
#include <vector>
#include <unistd.h>
#include <time.h>
#include <stdio.h>
#include <set>
#include <boost/graph/adjacency_list.hpp>
#include <boost/pending/disjoint_sets.hpp>
#include <boost/graph/incremental_components.hpp>
#include <boost/graph/graph_traits.hpp>
#include <boost/graph/connected_components.hpp>
#include <boost/graph/filtered_graph.hpp>

using namespace std;

/* CROCO - Complete Redundant Orthology Clusters */
namespace boost {
	enum vertex_component_t { vertex_component = 111 };
	BOOST_INSTALL_PROPERTY(vertex, component);
}
using namespace boost;

typedef map<string, int> GeneID;
typedef map<string, int> stringcounts;
typedef map<int, string> invGeneID;
typedef set<string> setstring;
typedef set<int> setint;
template <typename ComponentMap>
struct vertexComponent {

	vertexComponent() {}

	vertexComponent(ComponentMap component, int f_component) : m_component(component), m_f_component(f_component) {}

	template <typename Vertex>
	bool operator()(const Vertex& v) const {
		return (get(m_component, v) == m_f_component);
	}

	ComponentMap m_component;
	int m_f_component;
};

/* MAIN */
int main(int argc, char* argv[]) {
	/* USAGE */
	if (argc != 2) {
		std::cerr << "Usage: " << argv[0] << " genelinks (results to STDOUT)" << std::endl;
		exit(1);
	}

	/* DECLARATIONS */
	char      line[200];
	char      scanL[20];
	char      scanR[20];

	int lico = 0;
	GeneID ids;
	invGeneID name;
	int runningid = 0;

	typedef adjacency_list<vecS, vecS, undirectedS, property<vertex_component_t, int> > Graph;

	typedef property_map<Graph, vertex_component_t>::type ComponentMap;
	typedef filtered_graph<Graph, keep_all,	vertexComponent<ComponentMap> > FilteredGraph;

	graph_traits < Graph >::vertex_iterator vi, vi_end;
	graph_traits < Graph >::out_edge_iterator oei, oei_end;
	graph_traits < FilteredGraph >::vertex_iterator fvi, fvi_end;
	graph_traits < FilteredGraph >::out_edge_iterator foei, foei_end;

	Graph * g;
	FilteredGraph * fg;

	g = new Graph;

	// READ LINK FILE and add edges
	std::ifstream linkFile(argv[1]);
	if (linkFile.fail()) {
		std::cerr << "Error: could not read from file " << argv[1] << std::endl;
		exit(1);
	}
	std::cerr << "Reading gene links " << argv[1] << std::endl;
	while (!linkFile.eof()) {
		linkFile.getline(line, 200);
		if (linkFile.eof()) break;

		if (++lico % 1000 == 0) std::cerr << "\r" << lico;
		sscanf(line, "%s%s", scanL, scanR);
		string scanLstring = string(scanL);
		string scanRstring = string(scanR);

		if (ids.find(scanLstring) == ids.end()) {
			ids.insert(pair<string,int>(scanLstring,++runningid));
			name.insert(pair<int,string>(runningid,scanLstring));
		}
		if (ids.find(scanRstring) == ids.end()) {
			ids.insert(pair<string,int>(scanRstring,++runningid));
			name.insert(pair<int,string>(runningid,scanRstring));
		}
		// add edge
		add_edge(ids.find(scanLstring)->second, ids.find(scanRstring)->second, *g);
	}
	linkFile.close();
	cerr << "\r" << lico << " gene pairs read" << std::endl;

	//// SHOW GRAPH
	/*	std::cerr<<std::endl<<"Graph out-edges:"<<std::endl;
	 for (tie(vi, vi_end) = vertices(*g); vi != vi_end; ++vi) {
	 std::cerr<<name[*vi] << " " << *vi <<" outedges - ";
	 for(tie(oei, oei_end)=out_edges(*vi, *g); oei != oei_end; ++oei) {
	 std::cerr<<name[source(*oei, *g)]<<"<->"<<name[target(*oei, *g)]<<"  ";
	 }
	 std::cerr<<std::endl;
	 }
	 std::cerr<<std::endl;
	 */	//// SHOW GRAPH

	int numComponents = connected_components(*g, get(vertex_component, *g));

	std::cerr<<"Graph has "<<numComponents<<" components."<<std::endl;

	string familyprefix = "GF";

	int connex = 0;
	int compo = numComponents;
	int clusterNumber = 0;
	for(int i=0; i<numComponents; i++) {
		if (--compo % 100 == 0) std::cerr << "\r" << compo << " ";
		keep_all efilter;
		vertexComponent<ComponentMap> vfilter(get(vertex_component, *g), i);
		fg = new FilteredGraph(*g, efilter, vfilter);

		//std::cerr<<"Filtered graph (component "<<i<<") ";

		setstring Genes;
		int elements = 0;
		for (tie(fvi, fvi_end) = vertices(*fg); fvi != fvi_end; ++fvi) {
			Genes.insert(name[*fvi]);
			++elements;
		}
		// printing
		++connex;
		if (elements > 1 && elements == Genes.size()) {
			++clusterNumber;
			for (setstring::const_iterator it = Genes.begin(); it != Genes.end(); ++it) {
				std::cout << familyprefix << clusterNumber << "\t" << *it << std::endl;
			}
		}
		else std::cerr << "D'oh!" << endl;

		delete fg;
		fg = NULL;
	}

	std::cerr << "\rFound " << clusterNumber << " (" << connex << ") connected components (gene families)."<<std::endl;

	delete g;
	g = NULL;

	return EXIT_SUCCESS;

} // END OF MAIN


//
// SPDX-License-Identifier: BSD-3-Clause
// Copyright Contributors to the OpenEXR Project.
//

// clang-format off

#include <Python.h>
#include <boost/python.hpp>
#include <boost/python/make_constructor.hpp>
#include <boost/format.hpp>
#include <ImathVec.h>
#include "PyImath.h"
#include "PyImathMathExc.h"
#include "PyImathEuler.h"
#include "PyImathDecorators.h"
#include "PyImathExport.h"
#include "PyImathOperators.h"

// XXX incomplete array wrapping, docstrings missing

namespace PyImath {
template<> const char *PyImath::EulerfArray::name() { return "EulerfArray"; }
template<> const char *PyImath::EulerdArray::name() { return "EulerdArray"; }
}

namespace PyImath {
using namespace boost::python;
using namespace IMATH_NAMESPACE;

template <class T> struct EulerName { static const char *value; };
template<> const char *EulerName<float>::value  = "Eulerf";
template<> const char *EulerName<double>::value = "Eulerd";

template <class T>
static std::string nameOfOrder(typename IMATH_NAMESPACE::Euler<T>::Order order)
{
    switch(order)
    {
        case IMATH_NAMESPACE::Euler<T>::XYZ:
            return "EULER_XYZ";
        case IMATH_NAMESPACE::Euler<T>::XZY:
            return "EULER_XZY";
        case IMATH_NAMESPACE::Euler<T>::YZX:
            return "EULER_YZX";
        case IMATH_NAMESPACE::Euler<T>::YXZ:
            return "EULER_YXZ";
        case IMATH_NAMESPACE::Euler<T>::ZXY:
            return "EULER_ZXY";
        case IMATH_NAMESPACE::Euler<T>::ZYX:
            return "EULER_ZYX";
        case IMATH_NAMESPACE::Euler<T>::XZX:
            return "EULER_XZX";
        case IMATH_NAMESPACE::Euler<T>::XYX:
            return "EULER_XYX";
        case IMATH_NAMESPACE::Euler<T>::YXY:
            return "EULER_YXY";
        case IMATH_NAMESPACE::Euler<T>::YZY:
            return "EULER_YZY";
        case IMATH_NAMESPACE::Euler<T>::ZYZ:
            return "EULER_ZYZ";
        case IMATH_NAMESPACE::Euler<T>::ZXZ:
            return "EULER_ZXZ";
        case IMATH_NAMESPACE::Euler<T>::XYZr:
            return "EULER_XYZr";
        case IMATH_NAMESPACE::Euler<T>::XZYr:
            return "EULER_XZYr";
        case IMATH_NAMESPACE::Euler<T>::YZXr:
            return "EULER_YZXr";
        case IMATH_NAMESPACE::Euler<T>::YXZr:
            return "EULER_YXZr";
        case IMATH_NAMESPACE::Euler<T>::ZXYr:
            return "EULER_ZXYr";
        case IMATH_NAMESPACE::Euler<T>::ZYXr:
            return "EULER_ZYXr";
        case IMATH_NAMESPACE::Euler<T>::XZXr:
            return "EULER_XZXr";
        case IMATH_NAMESPACE::Euler<T>::XYXr:
            return "EULER_XYXr";
        case IMATH_NAMESPACE::Euler<T>::YXYr:
            return "EULER_YXYr";
        case IMATH_NAMESPACE::Euler<T>::YZYr:
            return "EULER_YZYr";
        case IMATH_NAMESPACE::Euler<T>::ZYZr:
            return "EULER_ZYZr";
        case IMATH_NAMESPACE::Euler<T>::ZXZr:
            return "EULER_ZXZr";
        default:
            break;
    }

    return "";
}

template <class T>
static std::string Euler_str(const Euler<T> &e)
{
    std::stringstream stream;
    stream << EulerName<T>::value << "(" << e.x << ", " << e.y << ", " << e.z << ", "
           << nameOfOrder<T> (e.order()) << ")";
    return stream.str();
}

// Non-specialized repr is same as str
template <class T>
static std::string Euler_repr(const Euler<T> &e)
{
    return Euler_str(e);
}

// Specialization for float to full precision
template <>
std::string Euler_repr(const Euler<float> &e)
{
    return (boost::format("%s(%.9g, %.9g, %.9g, %s)")
                        % EulerName<float>::value
                        % e.x % e.y % e.z
                        % nameOfOrder<float>(e.order()).c_str()).str();
}

// Specialization for double to full precision
template <>
std::string Euler_repr(const Euler<double> &e)
{
    return (boost::format("%s(%.17g, %.17g, %.17g, %s)")
                        % EulerName<double>::value
                        % e.x % e.y % e.z
                        % nameOfOrder<double>(e.order()).c_str()).str();
}

template <class T>
static bool
equal(const Euler<T> &e0, const Euler<T> &e1)
{
    if(e0.x == e1.x && e0.y == e1.y && e0.z == e1.z && (e0.order())==(e1.order()))
        return true;
    else
        return false;
}

template <class T>
static bool
notequal(const Euler<T> &e0, const Euler<T> &e1)
{
    if(e0.x != e1.x || e0.y != e1.y || e0.z != e1.z || (e0.order()) != (e1.order()))
    {
        return true;
    }
    else
        return false;
}

template <class T>
static IMATH_NAMESPACE::Vec3 <int> getAngleOrder(Euler <T> &euler)
{
    int i, j, k;
    euler.angleOrder(i, j, k);
    return IMATH_NAMESPACE::Vec3 <int> (i, j, k);
}

template <class T>
static void
setXYZTuple(Euler<T> &euler, const tuple &t)
{
    MATH_EXC_ON;
    Vec3<T> v;
    if(t.attr("__len__")() == 3)
    {
        v.x = extract<T>(t[0]);
        v.y = extract<T>(t[1]);
        v.z = extract<T>(t[2]);

        euler.setXYZVector(v);
    }
    else
        throw std::invalid_argument ("Color3 expects tuple of length 3");
}

// needed to convert Eulerf::Order to Euler<T>::Order
template <class T>
static typename Euler<T>::Order interpretOrder(typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = Euler<T>::XYZ;
    switch(order)
    {
        case IMATH_NAMESPACE::Eulerf::XYZ:
        {
            o = Euler<T>::XYZ;
        }break;
        case IMATH_NAMESPACE::Eulerf::XZY:
        {
            o = Euler<T>::XZY;
        }break;
        case IMATH_NAMESPACE::Eulerf::YZX:
        {
            o = Euler<T>::YZX;
        }break;
        case IMATH_NAMESPACE::Eulerf::YXZ:
        {
            o = Euler<T>::YXZ;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZXY:
        {
            o = Euler<T>::ZXY;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZYX:
        {
            o = Euler<T>::ZYX;
        }break;
        case IMATH_NAMESPACE::Eulerf::XZX:
        {
            o = Euler<T>::XZX;
        }break;
        case IMATH_NAMESPACE::Eulerf::XYX:
        {
            o = Euler<T>::XYX;
        }break;
        case IMATH_NAMESPACE::Eulerf::YXY:
        {
            o = Euler<T>::YXY;
        }break;
        case IMATH_NAMESPACE::Eulerf::YZY:
        {
            o = Euler<T>::YZY;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZYZ:
        {
            o = Euler<T>::ZYZ;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZXZ:
        {
            o = Euler<T>::ZXZ;
        }break;
        case IMATH_NAMESPACE::Eulerf::XYZr:
        {
            o = Euler<T>::XYZr;
        }break;
        case IMATH_NAMESPACE::Eulerf::XZYr:
        {
            o = Euler<T>::XZYr;
        }break;
        case IMATH_NAMESPACE::Eulerf::YZXr:
        {
            o = Euler<T>::YZXr;
        }break;
        case IMATH_NAMESPACE::Eulerf::YXZr:
        {
            o = Euler<T>::YXZr;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZXYr:
        {
            o = Euler<T>::ZXYr;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZYXr:
        {
            o = Euler<T>::ZYXr;
        }break;
        case IMATH_NAMESPACE::Eulerf::XZXr:
        {
            o = Euler<T>::XZXr;
        }break;
        case IMATH_NAMESPACE::Eulerf::XYXr:
        {
            o = Euler<T>::XYXr;
        }break;
        case IMATH_NAMESPACE::Eulerf::YXYr:
        {
            o = Euler<T>::YXYr;
        }break;
        case IMATH_NAMESPACE::Eulerf::YZYr:
        {
            o = Euler<T>::YZYr;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZYZr:
        {
            o = Euler<T>::ZYZr;
        }break;
        case IMATH_NAMESPACE::Eulerf::ZXZr:
        {
            o = Euler<T>::ZXZr;
        }break;
        default:
            break;
    }

    return o;
}

// needed to convert Eulerf::Axis to Euler<T>::Axis
template <class T>
static typename Euler<T>::Axis interpretAxis(typename IMATH_NAMESPACE::Eulerf::Axis axis)
{
    if (axis == IMATH_NAMESPACE::Eulerf::X)
        return Euler<T>::X;
    else if (axis == IMATH_NAMESPACE::Eulerf::Y)
        return Euler<T>::Y;
    else
        return Euler<T>::Z;
}

template <class T>
static Euler<T> *
eulerConstructor1(const Vec3<T> &v, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    return new Euler<T>(v, o);
}

template <class T>
static Euler<T> *
eulerConstructor1a(const Vec3<T> &v)
{
    return eulerConstructor1 (v, IMATH_NAMESPACE::Eulerf::Default);
}

template <class T>
static Euler<T> *
eulerConstructor1b(const Vec3<T> &v, int iorder)
{
    typename Euler<T>::Order o = typename Euler<T>::Order (iorder);
    return new Euler<T>(v, o);
}

template <class T>
static Euler<T> *
eulerConstructor2(T i, T j, T k, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    return new Euler<T>(i, j, k, o);
}

template <class T>
static Euler<T> *
eulerConstructor2a(T i, T j, T k)
{
    return eulerConstructor2 (i, j, k, IMATH_NAMESPACE::Eulerf::Default);
}

template <class T>
static Euler<T> *
eulerConstructor2b(T i, T j, T k, int iorder)
{
    typename Euler<T>::Order o = typename Euler<T>::Order (iorder);
    return new Euler<T>(i, j, k, o);
}

template <class T>
static Euler<T> *
eulerConstructor3(const Matrix33<T> &mat, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    return new Euler<T>(mat, o);
}

template <class T>
static Euler<T> *
eulerConstructor3a(const Matrix33<T> &mat)
{
    return eulerConstructor3 (mat, IMATH_NAMESPACE::Eulerf::Default);
}

template <class T>
static Euler<T> *
eulerConstructor3b(const Matrix33<T> &mat, int iorder)
{
    typename Euler<T>::Order o = typename Euler<T>::Order (iorder);
    return new Euler<T>(mat, o);
}

template <class T>
static Euler<T> *
eulerConstructor4(const Matrix44<T> &mat, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    return new Euler<T>(mat, o);
}

template <class T>
static Euler<T> *
eulerConstructor4a(const Matrix44<T> &mat)
{
    return eulerConstructor4 (mat, IMATH_NAMESPACE::Eulerf::Default);
}

template <class T>
static Euler<T> *
eulerConstructor4b(const Matrix44<T> &mat, int iorder)
{
    typename Euler<T>::Order o = typename Euler<T>::Order (iorder);
    return new Euler<T>(mat, o);
}

template <class T>
static Euler<T> *
eulerConstructor5(typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    return new Euler<T>(o);
}

template <class T>
static Euler<T> *
eulerConstructor5a()
{
    typename Euler<T>::Order o = interpretOrder<T>(IMATH_NAMESPACE::Eulerf::Default);
    return new Euler<T>(o);
}

template <class T>
static Euler<T> *
eulerConstructor5b(int iorder)
{
    typename Euler<T>::Order o = typename Euler<T>::Order (iorder);
    return new Euler<T>(o);
}

template <class T>
static Euler<T> *
eulerConstructor6(T x, T y, T z)
{
    return new Euler<T>(Vec3<T>(x,y,z));
}

template <class T>
static Euler<T> *
eulerConstructor7(const Quat<T> &quat, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    Euler<T> *e = eulerConstructor5<T>(order);
    e->extract(quat);
    return e;
}

template <class T>
static Euler<T> *
eulerConstructor7a(const Quat<T> &quat)
{
    return eulerConstructor7(quat, IMATH_NAMESPACE::Eulerf::Default);
}

template <class T>
static Euler<T> *
eulerConstructor7b(const Quat<T> &quat, int iorder)
{
    Euler<T> *e = eulerConstructor5b<T>(iorder);
    e->extract(quat);
    return e;
}

template <class T, class S>
static Euler<T> *
eulerConversionConstructor(const Euler<S> &euler)
{
    MATH_EXC_ON;
    Euler<T> *e = new Euler<T>;
    *e = euler;
    return e;
}

template <class T>
static void
eulerMakeNear(Euler<T> &euler, Euler<T> &target)
{
    MATH_EXC_ON;
    euler.makeNear (target);
}

template <class T>
static void
eulerSetOrder(Euler<T> &euler, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    typename Euler<T>::Order o = interpretOrder<T>(order);
    euler.setOrder (o);
}

template <class T>
static void
eulerSet(Euler<T> &euler, IMATH_NAMESPACE::Eulerf::Axis axis, int relative, int parityEven, int firstRepeats)
{
    MATH_EXC_ON;
    typename Euler<T>::Axis a = interpretAxis<T>(axis);
    euler.set (a, relative, parityEven, firstRepeats);
}

template <class T>
static void
extract1(Euler<T> &euler, const Matrix33<T> &m)
{
    MATH_EXC_ON;
    euler.extract(m);
}

template <class T>
static void
extract2(Euler<T> &euler, const Matrix44<T> &m)
{
    MATH_EXC_ON;
    euler.extract(m);
}

template <class T>
static void
extract3(Euler<T> &euler, const Quat<T> &q)
{
    MATH_EXC_ON;
    euler.extract(q);
}

template <class T>
static Matrix33<T>
toMatrix33(Euler<T> &euler)
{
    MATH_EXC_ON;
    return euler.toMatrix33();
}

template <class T>
static Matrix44<T>
toMatrix44(Euler<T> &euler)
{
    MATH_EXC_ON;
    return euler.toMatrix44();
}

template <class T>
static Quat<T>
toQuat(Euler<T> &euler)
{
    MATH_EXC_ON;
    return euler.toQuat();
}

template <class T>
static Vec3<T>
toXYZVector(Euler<T> &euler)
{
    MATH_EXC_ON;
    return euler.toXYZVector();
}

template <class T>
class_<Euler<T>,bases<IMATH_NAMESPACE::Vec3<T> > >
register_Euler()
{
    class_<Euler<T>,bases<Vec3<T> > > euler_class(EulerName<T>::value,EulerName<T>::value,init<Euler<T> >("copy construction"));
    euler_class
        .def(init<>("imath Euler default construction"))
        .def("__init__", make_constructor(eulerConstructor1<T>))
        .def("__init__", make_constructor(eulerConstructor1a<T>))
        .def("__init__", make_constructor(eulerConstructor1b<T>))
        .def("__init__", make_constructor(eulerConstructor2<T>))
        .def("__init__", make_constructor(eulerConstructor2a<T>))
        .def("__init__", make_constructor(eulerConstructor2b<T>))
        .def("__init__", make_constructor(eulerConstructor3<T>),
             "Euler-from-matrix construction assumes, but does\n"
             "not verify, that the matrix includes no shear or\n"
             "non-uniform scaling.  If necessary, you can fix\n"
             "the matrix by calling the removeScalingAndShear()\n"
             "function.\n")
        .def("__init__", make_constructor(eulerConstructor3a<T>))
        .def("__init__", make_constructor(eulerConstructor3b<T>))
        .def("__init__", make_constructor(eulerConstructor4<T>))
        .def("__init__", make_constructor(eulerConstructor4a<T>))
        .def("__init__", make_constructor(eulerConstructor4b<T>))
        .def("__init__", make_constructor(eulerConstructor5<T>))
        .def("__init__", make_constructor(eulerConstructor5a<T>))
        .def("__init__", make_constructor(eulerConstructor5b<T>))
        .def("__init__", make_constructor(eulerConstructor6<T>))
        .def("__init__", make_constructor(eulerConstructor7<T>))
        .def("__init__", make_constructor(eulerConstructor7a<T>))
        .def("__init__", make_constructor(eulerConstructor7b<T>))
        .def("__init__", make_constructor(eulerConversionConstructor<T, float>))
        .def("__init__", make_constructor(eulerConversionConstructor<T, double>))

        .def("angleOrder", &getAngleOrder<T>, "angleOrder() set the angle order")

        .def("frameStatic", &Euler<T>::frameStatic,
             "e.frameStatic() -- returns true if the angles of e\n"
             "are measured relative to a set of fixed axes,\n"
             "or false if the angles of e are measured relative to\n"
             "each other\n")

        .def("initialAxis", &Euler<T>::initialAxis,
             "e.initialAxis() -- returns the initial rotation\n"
             "axis of e (EULER_X_AXIS, EULER_Y_AXIS, EULER_Z_AXIS)")

        .def("initialRepeated", &Euler<T>::initialRepeated,
             "e.initialRepeated() -- returns 1 if the initial\n"
             "rotation axis of e is repeated (for example,\n"
             "e.order() == EULER_XYX); returns 0 if the initial\n"
             "rotation axis is not repeated.\n")

        .def("makeNear", &eulerMakeNear<T>,
             "e.makeNear(t) -- adjusts Euler e so that it\n"
             "represents the same rotation as before, but the\n"
             "individual angles of e differ from the angles of\n"
             "t by as little as possible.\n"
             "This method might not make sense if e.order()\n"
             "and t.order() are different\n")

        .def("order", &Euler<T>::order,
             "e.order() -- returns the rotation order in e\n"
             "(EULER_XYZ, EULER_XZY, ...)")

        .def("parityEven", &Euler<T>::parityEven,
             "e.parityEven() -- returns the parity of the\n"
             "axis permutation of e\n")

        .def("set", &eulerSet<T>,
             "e.set(i,r,p,f) -- sets the rotation order in e\n"
             "according to the following flags:\n"
             "\n"
             "   i   initial axis (EULER_X_AXIS,\n"
             "       EULER_Y_AXIS or EULER_Z_AXIS)\n"
             "\n"
             "   r   rotation angles are measured relative\n"
             "       to each other (r == 1), or relative to a\n"
             "       set of fixed axes (r == 0)\n"
             "\n"
             "   p   parity of axis permutation is even (r == 1)\n"
             "       or odd (r == 0)\n"
             "\n"
             "   f   first rotation axis is repeated (f == 1)\n"
             "	or not repeated (f == 0)\n")

        .def("setOrder", &eulerSetOrder<T>,
             "e.setOrder(o) -- sets the rotation order in e\n"
             "to o (EULER_XYZ, EULER_XZY, ...)")

        .def("setXYZVector", &Euler<T>::setXYZVector,
             "e.setXYZVector(v) -- sets the three rotation\n"
             "angles in e to v[0], v[1], v[2]")
        .def("setXYZVector", &setXYZTuple<T>)

        .def("extract", &extract1<T>,
             "e.extract(m) -- extracts the rotation component\n"
             "from 3x3 matrix m and stores the result in e.\n"
             "Assumes that m does not contain shear or non-\n"
             "uniform scaling.  If necessary, you can fix m\n"
             "by calling m.removeScalingAndShear().")

        .def("extract", &extract2<T>,
             "e.extract(m) -- extracts the rotation component\n"
             "from 4x4 matrix m and stores the result in e.\n"
             "Assumes that m does not contain shear or non-\n"
             "uniform scaling.  If necessary, you can fix m\n"
             "by calling m.removeScalingAndShear().")

        .def("extract", &extract3<T>,
             "e.extract(q) -- extracts the rotation component\n"
             "from quaternion q and stores the result in e")

        .def("toMatrix33", &toMatrix33<T>, "e.toMatrix33() -- converts e into a 3x3 matrix\n")

        .def("toMatrix44", &toMatrix44<T>, "e.toMatrix44() -- converts e into a 4x4 matrix\n")

        .def("toQuat", &toQuat<T>, "e.toQuat() -- converts e into a quaternion\n")

        .def("toXYZVector", &toXYZVector<T>,
             "e.toXYZVector() -- converts e into an XYZ\n"
             "rotation vector")
        .def("__str__", &Euler_str<T>)
        .def("__repr__", &Euler_repr<T>)

        .def("__eq__", &equal<T>)
        .def("__ne__", &notequal<T>)
        ;

    // fill in the Euler scope
    {
        scope euler_scope(euler_class);
        enum_<typename Euler<T>::Order> euler_order("Order");
        euler_order
            .value("XYZ",Euler<T>::XYZ)
            .value("XZY",Euler<T>::XZY)
            .value("YZX",Euler<T>::YZX)
            .value("YXZ",Euler<T>::YXZ)
            .value("ZXY",Euler<T>::ZXY)
            .value("ZYX",Euler<T>::ZYX)
            .value("XZX",Euler<T>::XZX)
            .value("XYX",Euler<T>::XYX)
            .value("YXY",Euler<T>::YXY)
            .value("YZY",Euler<T>::YZY)
            .value("ZYZ",Euler<T>::ZYZ)
            .value("ZXZ",Euler<T>::ZXZ)
            .value("XYZr",Euler<T>::XYZr)
            .value("XZYr",Euler<T>::XZYr)
            .value("YZXr",Euler<T>::YZXr)
            .value("YXZr",Euler<T>::YXZr)
            .value("ZXYr",Euler<T>::ZXYr)
            .value("ZYXr",Euler<T>::ZYXr)
            .value("XZXr",Euler<T>::XZXr)
            .value("XYXr",Euler<T>::XYXr)
            .value("YXYr",Euler<T>::YXYr)
            .value("YZYr",Euler<T>::YZYr)
            .value("ZYZr",Euler<T>::ZYZr)
            .value("ZXZr",Euler<T>::ZXZr)

            // don't export these, they're not really part of the public interface
            //.value("Legal",Euler<T>::Legal)
            //.value("Min",Euler<T>::Min)
            //.value("Max",Euler<T>::Max)

            // handle Default seperately since boost sets up a 1-1 mapping for enum values
            //.value("Default",Euler<T>::Default)
            .export_values()
            ;
        // just set it to the XYZ value manually
        euler_scope.attr("Default") = euler_scope.attr("XYZ");

        enum_<typename Euler<T>::Axis>("Axis")
            .value("X",Euler<T>::X)
            .value("Y",Euler<T>::Y)
            .value("Z",Euler<T>::Z)
            .export_values()
            ;

        enum_<typename Euler<T>::InputLayout>("InputLayout")
            .value("XYZLayout",Euler<T>::XYZLayout)
            .value("IJKLayout",Euler<T>::IJKLayout)
            .export_values()
            ;
    }

    decoratecopy(euler_class);

    return euler_class;
}

// XXX fixme - template this
// really this should get generated automatically...

/*
template <class T,int index>
static FixedArray<T>
EulerArray_get(FixedArray<IMATH_NAMESPACE::Euler<T> > &qa)
{
    return FixedArray<T>( &(qa[0].r)+index, qa.len(), 4*qa.stride());
}
*/

template <class T>
static FixedArray<IMATH_NAMESPACE::Euler<T> > *
EulerArray_eulerConstructor7a(const FixedArray<IMATH_NAMESPACE::Quat<T> > &q)
{
    MATH_EXC_ON;
    size_t len = q.len();
    FixedArray<IMATH_NAMESPACE::Euler<T> >* result = new FixedArray<IMATH_NAMESPACE::Euler<T> >(len);
    for (size_t i = 0; i < len; ++i) {
        (*result)[i].extract(q[i]);
    }
    return result;
}

template <class T>
static FixedArray<IMATH_NAMESPACE::Euler<T> > *
EulerArray_eulerConstructor8a(const FixedArray<IMATH_NAMESPACE::Vec3<T> >& v)
{
    MATH_EXC_ON;
    size_t len = v.len();
    FixedArray<IMATH_NAMESPACE::Euler<T> >* result = new FixedArray<IMATH_NAMESPACE::Euler<T> >(len);

    for (size_t i = 0; i < len; ++i)
        (*result)[i] = Euler<T>(v[i]);

    return result;
}

template <class T>
static FixedArray<IMATH_NAMESPACE::Euler<T> > *
EulerArray_eulerConstructor9a(const FixedArray<IMATH_NAMESPACE::Vec3<T> >& v, typename IMATH_NAMESPACE::Eulerf::Order order)
{
    MATH_EXC_ON;
    size_t len = v.len();
    FixedArray<IMATH_NAMESPACE::Euler<T> >* result = new FixedArray<IMATH_NAMESPACE::Euler<T> >(len);

    typename Euler<T>::Order o = interpretOrder<T>(order);
    for (size_t i = 0; i < len; ++i)
        (*result)[i] = Euler<T>(v[i], o);

    return result;
}

template <class T>
static FixedArray<IMATH_NAMESPACE::Vec3<T> >
EulerArray_toXYZVector(const FixedArray<IMATH_NAMESPACE::Euler<T> >& e)
{
    MATH_EXC_ON;
    size_t len = e.len();
    FixedArray<IMATH_NAMESPACE::Vec3<T> > result(len, UNINITIALIZED);
    for (size_t i = 0; i < len; ++i)
        result[i] = e[i].toXYZVector();
    return result;
}

template <class T>
static FixedArray<IMATH_NAMESPACE::Quat<T> >
EulerArray_toQuat(const FixedArray<IMATH_NAMESPACE::Euler<T> >& e)
{
    MATH_EXC_ON;
    size_t len = e.len();
    FixedArray<IMATH_NAMESPACE::Quat<T> > result(len, UNINITIALIZED);
    for (size_t i = 0; i < len; ++i)
        result[i] = e[i].toQuat();
    return result;
}

template <class T>
class_<FixedArray<IMATH_NAMESPACE::Euler<T> > >
register_EulerArray()
{
    class_<FixedArray<IMATH_NAMESPACE::Euler<T> > > eulerArray_class = FixedArray<IMATH_NAMESPACE::Euler<T> >::register_("Fixed length array of IMATH_NAMESPACE::Euler");
    eulerArray_class
        //.add_property("x",&EulerArray_get<T,1>)
        //.add_property("y",&EulerArray_get<T,2>)
        //.add_property("z",&EulerArray_get<T,3>)
        .def("__init__", make_constructor(EulerArray_eulerConstructor7a<T>))
        .def("__init__", make_constructor(EulerArray_eulerConstructor8a<T>))
        .def("__init__", make_constructor(EulerArray_eulerConstructor9a<T>))
        .def("toXYZVector", EulerArray_toXYZVector<T>)
        .def("toQuat", EulerArray_toQuat<T>)
        ;

    add_comparison_functions(eulerArray_class);
    PyImath::add_explicit_construction_from_type<IMATH_NAMESPACE::Matrix33<T> >(eulerArray_class);
    PyImath::add_explicit_construction_from_type<IMATH_NAMESPACE::Matrix44<T> >(eulerArray_class);
    return eulerArray_class;
}

template PYIMATH_EXPORT class_<IMATH_NAMESPACE::Euler<float>,bases<IMATH_NAMESPACE::Vec3<float> > > register_Euler<float>();
template PYIMATH_EXPORT class_<IMATH_NAMESPACE::Euler<double>,bases<IMATH_NAMESPACE::Vec3<double> > > register_Euler<double>();

template PYIMATH_EXPORT class_<FixedArray<IMATH_NAMESPACE::Euler<float> > > register_EulerArray<float>();
template PYIMATH_EXPORT class_<FixedArray<IMATH_NAMESPACE::Euler<double> > > register_EulerArray<double>();

template<> PYIMATH_EXPORT IMATH_NAMESPACE::Euler<float> FixedArrayDefaultValue<IMATH_NAMESPACE::Euler<float> >::value() { return IMATH_NAMESPACE::Euler<float>(); }
template<> PYIMATH_EXPORT IMATH_NAMESPACE::Euler<double> FixedArrayDefaultValue<IMATH_NAMESPACE::Euler<double> >::value() { return IMATH_NAMESPACE::Euler<double>(); }

} // namespace PyImath


#include "boxsampler.hpp"
#include "gibbssampler.hpp"
#include "heliumsampler.hpp"
#include "importancesampler.hpp"
#include "metropolissampler.hpp"
#include "sampler.hpp"

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

namespace py = pybind11;

void init_sampler(py::module& main)
{
    auto m  = main.def_submodule("samplers");
    m.doc() = R"doc(
Samplers
--------

The sampling classes are the foundation of the Monte Carlo machinery that VMC
depends on. They provide a way of sampling system configurations (i.e.
positions of all particles) from a possibly unnormalized probability amplitude
(i.e. wavefunctions), which in turn allows us to efficiently evaluate
multidimensional integrals.
)doc";

    py::class_<Sampler>(m, "Sampler", R"doc(
The :class:`Sampler` class provides a unified abstraction for the generation of
successive system instances drawn from some probability distribution (i.e. a
wavefunction).

All other code expecting a sampler instance will take a :class:`Sampler` reference.

The subclasses of :class:`Sampler` will differ only in how they generate new
samples, according to their respective algorithms. As such they have different
parameters that can be set at initialization.
)doc")
        .def("next_configuration", &Sampler::next_configuration, R"doc(
Return a newly sampled system configuration.

There is no guarantee that this will always differ between successive calls,
but with a sufficient number of calls the distribution of outputs,
:math:`P(\mathbf{X})` should approximate

.. math::

    P(\mathbf{X}) \simeq  |\Psi(\mathbf{X})|^2

)doc")

        .def("thermalize", &Sampler::thermalize, R"doc(
Generate a given number of samples, discarding all.

This is useful to ensure that the sampler has reached a stable point where
samples are representative of the underlying distribution.
)doc")

        .def_property_readonly("acceptance_rate", &Sampler::get_acceptance_rate, R"doc(
Rate at which newly proposed samples where accepted by the algorithm.
)doc")

        .def_property_readonly("current_system", &Sampler::get_current_system);

    py::class_<MetropolisSampler, Sampler>(m, "MetropolisSampler", R"doc(
Implementation of the standard Metropolis-Hastings algorithm, producing a
Markov Chain of system configurations. [1]_

Each configuration is a _random_ perturbation of its predecessor if accepted, or
a copy if rejected. The perturbation is defined as a :math:`D` dimensional
vector of uniform random numbers in :math:`\frac{1}{2}[-s, s]` for some scale
parameter :math:`s`.

Notes
-----
This implementation only changes the coordinates of one particle at a time. This
means that two successive calls to `next_configuration()` will at most differ
in one row of the output. This can be trusted to always be the case, and
potential optimizations can be made based on this knowledge. For instance, if
caching of relative distances is employed, only the distances corresponding to
the changed particle would have to be recalculated.

References
----------
.. [1] W. K. Hastings; Monte Carlo sampling methods using Markov chains and
       their applications, Biometrika, Volume 57, Issue 1, 1 April 1970, Pages 97–109,
       https://doi.org/10.1093/biomet/57.1.97

Examples
--------

    >>> import numpy as np
    >>> from qflow.samplers import MetropolisSampler
    >>> from qflow.wavefunctions import SimpleGaussian
    >>> psi = SimpleGaussian()
    >>> sampler = MetropolisSampler(np.zeros((2, 3)), psi, step_size=1)

We can now get samples of the desired size on demand:

    >>> sampler.next_configuration()
    array([[-0.61370758,  0.08936709,  0.15872668],
           [ 0.05973557,  0.07445129, -0.29230947]])

Thermalize the sampler (equivalent to running `next_configuration` the given number of times, only faster):

    >>> sampler.thermalize(100)

Inspect the acceptance rate:

    >>> sampler.acceptance_rate
    0.7326732673267327
)doc")
        .def(py::init<const System&, Wavefunction&, Real>(),
             py::arg("system"),
             py::arg("wavefunction"),
             py::arg("step_size") = 1,
             R"doc(
Construct a sampler that uses the standard Metropolis algorithm. The
`step_size` determines how different successive configurations will be, and
should be tuned such that the acceptance rate remains high at all times.

Arguments
---------
 )doc");

    py::class_<ImportanceSampler, Sampler>(m, "ImportanceSampler", R"doc(
Modified version of the Metropolis-Hastings algorithm, which employs a smarter
way of generating new samples. [2]_ The variance of integrals computed using
this algorithm tends to be significantly lower compared to
:class:`MetropolisSampler`, at the expense of higher run-time cost.

.. warning::
    Importance sampling may only be used with wavefunctions that implements the
    :class:`Wavefunction.drift_force` method. If this is not fulfilled, the
    program will halt immediately without any way of catching an exception from
    Python.

Notes
-----
This algorithm differs from Metropolis-Hastings in the perturbations made, with
a corresponding change in the acceptance probability. In our case,

.. math::

    X_k^{(i+1)} = X_k^{(i)} + \sqrt{t}\mathcal{n} + t \frac{1}{\Psi}\nabla_k \Psi,

where :math:`X_k^{(i)}` is the coordinates of particle :math:`k` at time step
:math:`i`, :math:`\mathcal{n}` is a random number drawn from the standard normal
distribution and :math:`t` is the step size parameter used to tune how
different successive samples are.

Similarly to :class:`MetropolisSampler`, only one particle is perturbed at a time.

References
----------
.. [2] Reiher, W. (1966), Hammersley, J. M., D. C. Handscomb: Monte Carlo
       Methods. Methuen & Co., London, and John Wiley & Sons, New York, 1964. VII +
       178 S., Preis: 25 s. Biom. J., 8: 209-209. doi:10.1002/bimj.19660080314

Examples
--------
This is used exactly like :class:`MetropolisSampler`, with the only exception being the
meaning of the `step_size`.

    >>> import numpy as np
    >>> from qflow.samplers import ImportanceSampler
    >>> from qflow.wavefunctions import SimpleGaussian
    >>> psi = SimpleGaussian()
    >>> sampler = ImportanceSampler(np.zeros((2, 3)), psi, step_size=0.1)
    >>> sampler.next_configuration()
    array([[ 0.56938477,  0.25037102, -0.50411809],
           [-0.8038079 ,  0.15799471, -0.06645576]])
    >>> sampler.thermalize(100)
    >>> sampler.acceptance_rate
    0.9603960396039604

)doc")
        .def(py::init<const System&, Wavefunction&, Real>(),
             py::arg("system"),
             py::arg("wavefunction"),
             py::arg("step_size") = 0.1);

    py::class_<GibbsSampler, Sampler>(m, "GibbsSampler")
        .def(py::init<const System&, RBMWavefunction&>(),
             py::arg("system"),
             py::arg("rbm_wavefunction"));

    py::class_<BoxMetropolisSampler, MetropolisSampler>(m, "BoxMetropolisSampler")
        .def(py::init<const System&, Wavefunction&, Real, Real>(),
             py::arg("system"),
             py::arg("wavefunction"),
             py::arg("box_size"),
             py::arg("step_size") = 1)
        .def("initialize_from_system", &BoxMetropolisSampler::initialize_from_system);

    py::class_<BoxImportanceSampler, ImportanceSampler>(m, "BoxImportanceSampler")
        .def(py::init<const System&, Wavefunction&, Real, Real>(),
             py::arg("system"),
             py::arg("wavefunction"),
             py::arg("box_size"),
             py::arg("step_size") = 0.1);

    py::class_<HeliumSampler, Sampler>(m, "HeliumSampler")
        .def(py::init<const System&, Wavefunction&, Real, Real>(),
             py::arg("system"),
             py::arg("wavefunction"),
             py::arg("step_size"),
             py::arg("box_size"));
}


// Copyright (c) 2015-2018 Daniel Cooke
// Use of this source code is governed by the MIT license that can be found in the LICENSE file.

#include "variational_bayes_mixture_mixture_model.hpp"

#include <utility>
#include <cmath>
#include <iterator>
#include <algorithm>
#include <numeric>
#include <iostream>

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

#include "utils/maths.hpp"

namespace octopus { namespace model {

VariationalBayesMixtureMixtureModel::VariationalBayesMixtureMixtureModel(Options options)
: options_ {std::move(options)}
{}

VariationalBayesMixtureMixtureModel::Inferences
VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genotype_log_priors,
                                              const HaplotypeLikelihoodMatrix& log_likelihoods,
                                              const GroupConcentrationVector& group_concentrations,
                                              const MixtureConcentrationArray& mixture_concentrations,
                                              std::vector<LogProbabilityVector> seeds) const
{
    Inferences result {};
    result.approx_log_evidence = std::numeric_limits<double>::lowest();
    for (auto& seed : seeds) {
        auto inferences = evaluate(genotype_log_priors, log_likelihoods, group_concentrations, mixture_concentrations, seed);
        if (inferences.approx_log_evidence > result.approx_log_evidence) {
            result = std::move(inferences);
        }
    }
    return result;
}

VariationalBayesMixtureMixtureModel::Inferences
VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genotype_log_priors,
                                              const HaplotypeLikelihoodMatrix& log_likelihoods,
                                              const double group_concentration, const double mixture_concentration,
                                              std::vector<LogProbabilityVector> seeds) const
{
    const auto num_samples = log_likelihoods.size();
    const auto num_groups = log_likelihoods.front().front().size();
    std::vector<unsigned> group_mixture_sizes {};
    group_mixture_sizes.reserve(num_groups);
    for (const auto& genotype : log_likelihoods.front().front()) {
        group_mixture_sizes.push_back(genotype.size());
    }
    const GroupConcentrationVector group_concentrations(num_groups, group_concentration);
    MixtureConcentrationArray mixture_concentrations(num_samples, MixtureConcentrationVector(num_groups));
    for (auto& sample_concentrations : mixture_concentrations) {
        unsigned t {0};
        for (auto& sample_group_sample_concentrations : sample_concentrations) {
            sample_group_sample_concentrations.assign(group_mixture_sizes[t++], mixture_concentration);
        }
    }
    return evaluate(genotype_log_priors, log_likelihoods, group_concentrations, mixture_concentrations, std::move(seeds));
}

// Private methods

namespace {

VariationalBayesMixtureMixtureModel::ProbabilityVector&
exp(const LogProbabilityVector& log_probabilities, ProbabilityVector& result) noexcept
{
    std::transform(std::cbegin(log_probabilities), std::cend(log_probabilities), std::begin(result),
                   [] (const auto lp) noexcept { return std::exp(lp); });
    return result;
}

auto exp(const VariationalBayesMixtureMixtureModel::LogProbabilityVector& log_probabilities)
{
    VariationalBayesMixtureMixtureModel::ProbabilityVector result(log_probabilities.size());
    return exp(log_probabilities, result);
}

template <typename T>
auto sum(const std::vector<T>& values) noexcept
{
    return std::accumulate(std::cbegin(values), std::cend(values), T {0});
}

auto sum(const VBReadLikelihoodArray& likelihoods) noexcept
{
    using T = VBReadLikelihoodArray::BaseType::value_type;
    return std::accumulate(std::cbegin(likelihoods), std::cend(likelihoods), T {0});
}

template <typename T1, typename T2>
auto inner_product(const T1& lhs, const T2& rhs) noexcept
{
    assert(std::distance(std::cbegin(lhs), std::cend(lhs)) == std::distance(std::cbegin(rhs), std::cend(rhs)));
    using T = typename T1::value_type;
    return std::inner_product(std::cbegin(lhs), std::cend(lhs), std::cbegin(rhs), T {0});
}

template <typename T>
inline auto digamma_diff(const T a, const T b)
{
    using boost::math::digamma;
    return digamma(a) - digamma(b);
}

template <typename T>
auto dirichlet_expectation_log(const std::vector<T>& concentrations)
{
    std::vector<T> result(concentrations.size());
    const auto alpha_0= sum(concentrations);
    std::transform(std::cbegin(concentrations), std::cend(concentrations), std::begin(result),
                   [=] (auto alpha) { return digamma_diff(alpha, alpha_0); });
    return result;
}

template <typename T>
T log_sum_exp(const std::vector<T>& logs)
{
    return maths::log_sum_exp(logs);
}

} // namespace

VariationalBayesMixtureMixtureModel::Inferences
VariationalBayesMixtureMixtureModel::evaluate(const LogProbabilityVector& genotype_log_priors,
                                              const HaplotypeLikelihoodMatrix& log_likelihoods,
                                              const GroupConcentrationVector& prior_group_concentrations,
                                              const MixtureConcentrationArray& prior_mixture_concentrations,
                                              LogProbabilityVector& genotype_log_posteriors) const
{
    auto genotype_posteriors = exp(genotype_log_posteriors);
    auto posterior_group_concentrations = prior_group_concentrations;
    auto posterior_mixture_concentrations = prior_mixture_concentrations;
    auto group_responsabilities = init_responsabilities(prior_group_concentrations, prior_mixture_concentrations,
                                                        genotype_posteriors, log_likelihoods);
    auto component_responsabilities = init_responsabilities(prior_group_concentrations, prior_mixture_concentrations,
                                                            genotype_posteriors, group_responsabilities, log_likelihoods);
    auto prev_evidence = std::numeric_limits<double>::lowest();
    for (unsigned i {0}; i < options_.max_iterations; ++i) {
        update_genotype_log_posteriors(genotype_log_posteriors, genotype_log_priors,
                                       group_responsabilities, component_responsabilities,
                                       log_likelihoods);
        exp(genotype_log_posteriors, genotype_posteriors);
        update_group_concentrations(posterior_group_concentrations, prior_group_concentrations, group_responsabilities);
        update_mixture_concentrations(posterior_mixture_concentrations, prior_mixture_concentrations,
                                      group_responsabilities, component_responsabilities);
        auto curr_evidence = calculate_evidence(prior_group_concentrations, posterior_group_concentrations,
                                                prior_mixture_concentrations, posterior_mixture_concentrations,
                                                genotype_log_priors, genotype_log_posteriors, genotype_posteriors,
                                                group_responsabilities, component_responsabilities,
                                                log_likelihoods);
//        assert(curr_evidence + options_.epsilon >= prev_evidence);
        if (curr_evidence <= prev_evidence || (curr_evidence - prev_evidence) < options_.epsilon) {
            prev_evidence = curr_evidence;
            break;
        }
        prev_evidence = curr_evidence;
        update_responsabilities(group_responsabilities, posterior_group_concentrations, posterior_mixture_concentrations,
                                genotype_posteriors, component_responsabilities, log_likelihoods);
        update_responsabilities(component_responsabilities, posterior_group_concentrations, posterior_mixture_concentrations,
                                genotype_posteriors, group_responsabilities, log_likelihoods);
    }
    Inferences result {};
    result.genotype_log_posteriors = std::move(genotype_log_posteriors);
    result.genotype_posteriors = std::move(genotype_posteriors);
    result.group_concentrations = std::move(posterior_group_concentrations);
    result.mixture_concentrations = std::move(posterior_mixture_concentrations);
    result.group_responsabilities = std::move(group_responsabilities);
    result.approx_log_evidence = prev_evidence;
    return result;
}

VariationalBayesMixtureMixtureModel::GroupResponsabilityVector
VariationalBayesMixtureMixtureModel::init_responsabilities(const GroupConcentrationVector& group_concentrations,
                                                           const MixtureConcentrationArray& mixture_concentrations,
                                                           const ProbabilityVector& genotype_priors,
                                                           const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    const auto S = log_likelihoods.size();
    const auto T = group_concentrations.size();
    GroupResponsabilityVector result(S, Sigma(T));
    const auto G = genotype_priors.size();
    const auto ln_ex_psi = dirichlet_expectation_log(group_concentrations);
    std::size_t max_K {0};
    for (std::size_t s {0}; s < S; ++s) {
        for (std::size_t t {0}; t < T; ++t) {
            max_K = std::max(max_K, mixture_concentrations[s][t].size());
        }
    }
    for (std::size_t s {0}; s < S; ++s) {
        const auto N = log_likelihoods[s][0][0][0].size();
        const auto tau = 1.0 / max_K;
        const auto tau_sum = N * tau;
        for (std::size_t t {0}; t < T; ++t) {
            result[s][t] = ln_ex_psi[t];
            const auto ln_ex_pi = dirichlet_expectation_log(mixture_concentrations[s][t]);
            const auto K = mixture_concentrations[s][t].size();
            for (std::size_t k {0}; k < K; ++k) {
                result[s][t] += ln_ex_pi[k] * tau_sum;
                for (std::size_t g {0}; g < G; ++g) {
                    result[s][t] += genotype_priors[g] * tau * sum(log_likelihoods[s][g][t][k]);
                }
            }
        }
        maths::normalise_exp(result[s]);
    }
    return result;
}

void
VariationalBayesMixtureMixtureModel::update_responsabilities(GroupResponsabilityVector& result,
                                                             const GroupConcentrationVector& group_concentrations,
                                                             const MixtureConcentrationArray& mixture_concentrations,
                                                             const ProbabilityVector& genotype_posteriors,
                                                             const ComponentResponsabilityMatrix& component_responsabilities,
                                                             const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    const auto T = group_concentrations.size();
    const auto S = log_likelihoods.size();
    const auto G = genotype_posteriors.size();
    const auto ln_ex_psi = dirichlet_expectation_log(group_concentrations);
    for (std::size_t s {0}; s < S; ++s) {
        std::vector<double> component_responsability_sums(component_responsabilities[s].size());
        std::transform(std::cbegin(component_responsabilities[s]), std::cend(component_responsabilities[s]),
                       std::begin(component_responsability_sums), [] (const auto& taus) { return sum(taus); });
        for (std::size_t t {0}; t < T; ++t) {
            result[s][t] = ln_ex_psi[t];
            const auto ln_ex_pi = dirichlet_expectation_log(mixture_concentrations[s][t]);
            const auto K = ln_ex_pi.size();
            for (std::size_t k {0}; k < K; ++k) {
                result[s][t] += ln_ex_pi[k] * component_responsability_sums[k];
                for (std::size_t g {0}; g < G; ++g) {
                    result[s][t] += genotype_posteriors[g] * inner_product(component_responsabilities[s][k], log_likelihoods[s][g][t][k]);
                }
            }
        }
        maths::normalise_exp(result[s]);
    }
}

VariationalBayesMixtureMixtureModel::ComponentResponsabilityMatrix
VariationalBayesMixtureMixtureModel::init_responsabilities(const GroupConcentrationVector& group_concentrations,
                                                           const MixtureConcentrationArray& mixture_concentrations,
                                                           const ProbabilityVector& genotype_priors,
                                                           const GroupResponsabilityVector& group_responsabilities,
                                                           const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    const auto S = log_likelihoods.size();
    const auto T = group_concentrations.size();
    std::size_t K {0};
    for (std::size_t s {0}; s < S; ++s) {
        for (std::size_t t {0}; t < T; ++t) {
            K = std::max(K, mixture_concentrations[s][t].size());
        }
    }
    ComponentResponsabilityMatrix result(S, ComponentResponsabilityVector(K));
    for (std::size_t s {0}; s < S; ++s) {
        const auto N = log_likelihoods[s][0][0][0].size();
        for (std::size_t k {0}; k < K; ++k) {
            result[s][k].resize(N);
        }
    }
    update_responsabilities(result, group_concentrations, mixture_concentrations,
                            genotype_priors, group_responsabilities, log_likelihoods);
    return result;
}

namespace {

auto inner_product(const VariationalBayesMixtureMixtureModel::ProbabilityVector& genotype_posteriors,
                   const VariationalBayesMixtureMixtureModel::GenotypeCombinationLikelihoodVector& log_likelihoods,
                   const std::size_t t, const std::size_t k, const std::size_t n) noexcept
{
    ProbabilityVector::value_type result {0};
    const auto G = genotype_posteriors.size();
    for (std::size_t g {0}; g < G; ++g) {
        result += genotype_posteriors[g] * log_likelihoods[g][t][k][n];
    }
    return result;
}

} // namespace

void
VariationalBayesMixtureMixtureModel::update_responsabilities(ComponentResponsabilityMatrix& result,
                                                             const GroupConcentrationVector& group_concentrations,
                                                             const MixtureConcentrationArray& mixture_concentrations,
                                                             const ProbabilityVector& genotype_posteriors,
                                                             const GroupResponsabilityVector& group_responsabilities,
                                                             const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    const auto T = group_concentrations.size();
    const auto S = log_likelihoods.size();
    const auto max_K = result[0].size();
    for (std::size_t s {0}; s < S; ++s) {
        const auto N = log_likelihoods[s][0][0][0].size();
        for (std::size_t t {0}; t < T; ++t) {
            const auto ln_exp_pi = dirichlet_expectation_log(mixture_concentrations[s][t]);
            const auto K = ln_exp_pi.size();
            for (std::size_t k {0}; k < K; ++k) {
                for (std::size_t n {0}; n < N; ++n) {
                    if (t == 0) result[s][k][n] = 0;
                    result[s][k][n] += group_responsabilities[s][t] * (ln_exp_pi[k] + inner_product(genotype_posteriors, log_likelihoods[s], t, k, n));
                }
            }
        }
        std::vector<double> ln_rho(max_K);
        for (std::size_t n {0}; n < N; ++n) {
            for (std::size_t k {0}; k < max_K; ++k) {
                ln_rho[k] = result[s][k][n];
            }
            const auto ln_rho_norm = log_sum_exp(ln_rho);
            for (std::size_t k {0}; k < max_K; ++k) {
                result[s][k][n] = std::exp(ln_rho[k] - ln_rho_norm);
            }
        }
    }
}

void
VariationalBayesMixtureMixtureModel::update_genotype_log_posteriors(LogProbabilityVector& result,
                                                                    const LogProbabilityVector& genotype_log_priors,
                                                                    const GroupResponsabilityVector& group_responsabilities,
                                                                    const ComponentResponsabilityMatrix& component_responsabilities,
                                                                    const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    for (std::size_t g {0}; g < genotype_log_priors.size(); ++g) {
        result[g] = genotype_log_priors[g] + marginalise(group_responsabilities, component_responsabilities, log_likelihoods, g);
    }
    maths::normalise_logs(result);
}

VariationalBayesMixtureMixtureModel::LogProbability
VariationalBayesMixtureMixtureModel::marginalise(const GroupResponsabilityVector& group_responsabilities,
                                                 const ComponentResponsabilityMatrix& component_responsabilities,
                                                 const HaplotypeLikelihoodMatrix& log_likelihoods,
                                                 const std::size_t g) const noexcept
{
    const auto T = group_responsabilities.front().size();
    const auto S = component_responsabilities.size();
    LogProbability result {0};
    for (std::size_t s {0}; s < S; ++s) {
        for (std::size_t t {0}; t < T; ++t) {
            result += group_responsabilities[s][t] * marginalise(component_responsabilities[s], log_likelihoods[s][g][t]);
        }
    }
    return result;
}

VariationalBayesMixtureMixtureModel::LogProbability
VariationalBayesMixtureMixtureModel::marginalise(const ComponentResponsabilityVector& responsabilities,
                                                 const HaplotypeLikelihoodVector& log_likelihoods) const noexcept
{
    const auto K = log_likelihoods.size();
    LogProbability result {0};
    for (std::size_t k {0}; k < K; ++k) {
        result += inner_product(responsabilities[k], log_likelihoods[k]);
    }
    return result;
}

void
VariationalBayesMixtureMixtureModel::update_group_concentrations(GroupConcentrationVector& result,
                                                                 const GroupConcentrationVector& prior_group_concentrations,
                                                                 const GroupResponsabilityVector& group_responsabilities) const
{
    assert(result.size() == prior_group_concentrations.size());
    const auto T = prior_group_concentrations.size();
    const auto S = group_responsabilities.size();
    for (std::size_t t {0}; t < T; ++t) {
        result[t] = prior_group_concentrations[t];
        for (std::size_t s {0}; s < S; ++s) {
            result[t] += group_responsabilities[s][t];
        }
    }
}

void
VariationalBayesMixtureMixtureModel::update_mixture_concentrations(MixtureConcentrationArray& result,
                                                                   const MixtureConcentrationArray& prior_mixture_concentrations,
                                                                   const GroupResponsabilityVector& group_responsabilities,
                                                                   const ComponentResponsabilityMatrix& component_responsabilities) const
{
    const auto S = prior_mixture_concentrations.size();
    const auto T = group_responsabilities.front().size();
    for (std::size_t s {0}; s < S; ++s) {
        for (std::size_t t {0}; t < T; ++t) {
            const auto K = prior_mixture_concentrations[s][t].size();
            for (std::size_t k {0}; k < K; ++k) {
                result[s][t][k] = prior_mixture_concentrations[s][t][k] + group_responsabilities[s][t] * sum(component_responsabilities[s][k]);
            }
        }
    }
}

namespace {

template <typename Range1, typename Range2>
auto inner_inner_product(const Range1& lhs, const Range2& rhs) noexcept
{
    using T = decltype(inner_product(lhs[0], rhs[0]));
    const auto k = std::min(lhs.size(), rhs.size());
    return std::inner_product(std::cbegin(lhs), std::next(std::cbegin(lhs), k), std::cbegin(rhs), T {}, std::plus<> {},
                              [] (const auto& a, const auto& b) { return inner_product(a, b); });
}

template <typename T>
auto shannon_entropy(const std::vector<T>& probabilities) noexcept
{
    return -std::accumulate(std::cbegin(probabilities), std::cend(probabilities), T {0},
                            [] (const auto curr, const auto p) noexcept { return curr + (p > 0 ? p * std::log(p) : 0.0); });
}

template <typename T>
auto shannon_entropy(const std::vector<std::vector<T>>& probabilities) noexcept
{
    return std::accumulate(std::cbegin(probabilities), std::cend(probabilities), T {0},
                           [] (const auto curr, const auto& ps) noexcept { return curr + shannon_entropy(ps); });
}

} // namespace

double
VariationalBayesMixtureMixtureModel::calculate_evidence(const GroupConcentrationVector& prior_group_concentrations,
                                                        const GroupConcentrationVector& posterior_group_concentrations,
                                                        const MixtureConcentrationArray& prior_mixture_concentrations,
                                                        const MixtureConcentrationArray& posterior_mixture_concentrations,
                                                        const LogProbabilityVector& genotype_log_priors,
                                                        const LogProbabilityVector& genotype_log_posteriors,
                                                        const ProbabilityVector& genotype_posteriors,
                                                        const GroupResponsabilityVector& group_responsabilities,
                                                        const ComponentResponsabilityMatrix& component_responsabilities,
                                                        const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    const auto G = genotype_log_priors.size();
    const auto S = prior_mixture_concentrations.size();
    const auto T = group_responsabilities.front().size();
    double result {0};
    for (std::size_t g {0}; g < G; ++g) {
        auto w = genotype_log_priors[g] - genotype_log_posteriors[g];
        for (std::size_t s {0}; s < S; ++s) {
            for (std::size_t t {0}; t < T; ++t) {
                w += group_responsabilities[s][t] * inner_inner_product(component_responsabilities[s], log_likelihoods[s][g][t]);
            }
        }
        result += genotype_posteriors[g] * w;
    }
    for (std::size_t s {0}; s < S; ++s) {
        result += shannon_entropy(group_responsabilities[s]);
        result += shannon_entropy(component_responsabilities[s]);
        for (std::size_t t {0}; t < T; ++t) {
            result += maths::log_beta(posterior_mixture_concentrations[s][t]) - maths::log_beta(prior_mixture_concentrations[s][t]);
        }
    }
    result += maths::log_beta(posterior_group_concentrations) - maths::log_beta(prior_group_concentrations);
    return result;
}

void VariationalBayesMixtureMixtureModel::print_concentrations(const GroupConcentrationVector& concentrations) const
{
    for (std::size_t t {0}; t < concentrations.size(); ++t) {
        std::cout << "t: " << t << " = " << concentrations[t] << std::endl;
    }
}

void VariationalBayesMixtureMixtureModel::print_concentrations(const MixtureConcentrationArray& concentrations) const
{
    for (std::size_t s {0}; s < concentrations.size(); ++s) {
        for (std::size_t t {0}; t < concentrations[s].size(); ++t) {
            for (std::size_t k {0}; k < concentrations[s][t].size(); ++k) {
                std::cout << "s: " << s << " t: " << t << " k: " << k << " = " << concentrations[s][t][k] << std::endl;
            }
        }
    }
}

void VariationalBayesMixtureMixtureModel::print(const ProbabilityVector& probabilities) const
{
    for (std::size_t g {0}; g < probabilities.size(); ++g) {
        std::cout << "g: " << g << " = " << probabilities[g] << std::endl;
    }
}

void VariationalBayesMixtureMixtureModel::print(const GroupResponsabilityVector& responsabilities) const
{
    for (std::size_t s {0}; s < responsabilities.size(); ++s) {
        for (std::size_t t {0}; t < responsabilities[s].size(); ++t) {
            std::cout << "s: " << s << " t: " << t << " = " << responsabilities[s][t] << std::endl;
        }
    }
}

void VariationalBayesMixtureMixtureModel::print(const ComponentResponsabilityMatrix& responsabilities) const
{
    for (std::size_t s {0}; s < responsabilities.size(); ++s) {
        for (std::size_t k {0}; k < responsabilities[s].size(); ++k) {
            for (std::size_t n {0}; n < responsabilities[s][k].size(); ++n) {
                std::cout << "s: " << s <<  " k: " << k << " n: " << n << " = " << responsabilities[s][k][n] << std::endl;
            }
        }
    }
}

void VariationalBayesMixtureMixtureModel::print(const HaplotypeLikelihoodMatrix& log_likelihoods) const
{
    for (std::size_t s {0}; s < log_likelihoods.size(); ++s) {
        for (std::size_t g {0}; g < log_likelihoods[s].size(); ++g) {
            for (std::size_t t {0}; t < log_likelihoods[s][g].size(); ++t) {
                for (std::size_t k {0}; k < log_likelihoods[s][g][t].size(); ++k) {
                    for (std::size_t n {0}; n < log_likelihoods[s][g][t][k].size(); ++n) {
                        std::cout << "s: " << s <<  " g: " << g << " t: " << t << " k: " << k
                                  << " n: " << n << " = " << log_likelihoods[s][g][t][k][n] << std::endl;
                    }
                }
            }
        }
    }
}

} // namespace model
} // namespace octopus