//////////////////////////////////////////////////////////////////////////////
// random::poisson::devroye::detail::int_mean.hpp                        	//
//                                                                          //
//                                                                          //
//  (C) Copyright 2010 Erwann Rogard                                        //
//  Use, modification and distribution are subject to the                   //
//  Boost Software License, Version 1.0. (See accompanying file             //
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)        //
//////////////////////////////////////////////////////////////////////////////
#ifndef BOOST_RANDOM_POISSON_EXT_DEVROYE_DETAIL_INT_MEAN_HPP_ER_2010
#define BOOST_RANDOM_POISSON_EXT_DEVROYE_DETAIL_INT_MEAN_HPP_ER_2010
#include <boost/range.hpp>
#include <boost/math/special_functions/modf.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/random/exponential_distribution.hpp>

#include <boost/random/poisson_ext/devroye/detail/math.hpp>
#include <boost/random/poisson_ext/devroye/detail/parameters.hpp>

namespace boost{
namespace random{
namespace poisson{
namespace devroye{
namespace detail{

    // Samples from the poisson distribution with an integer mean
    //
    // Source : The computer generation of poisson random variables. L. Devroye,
    // Computing 26, Springer-Verlag, 1981
    //
    // Parameter     Description
    // Step4         A class that models the concept by the same name
    // Int           An integer type
    // T             A float type
    // P             An error handling policy
	//
    // Expression                 Requirement
    // int_mean<Step4,Int,T,P>    Public base of Step4
	//
    // Concept Step4:
    // int_mean<Step4,Int,T,P> is a public base, and
    // Expression    Result type   Side effect
    //  accept()     bool          if true, y is set to the value to be returned
    //
    // The numbers on the right are those of the lines of the Fortran listing
    // given in the reference above.
    template<typename Step4,typename Int,typename T,typename P>
    class int_mean : public devroye::detail::parameters<Int,T,P>
        {

		typedef devroye::detail::math<Int,T,P> ma_;
        typedef devroye::detail::parameters<Int,T,P> super1_;

		public:

        typedef T input_type;
        typedef Int result_type;

        input_type mean() const { return ma_::to_float(this->int_mean_); }
        void reset() { }

		int_mean():super1_(){}
        explicit int_mean(const result_type& m)
        	:super1_(m),int_mean_(m){
            	BOOST_ASSERT(this->int_mean_val()>=0);
            }

        int_mean (const int_mean& that)
        	:super1_(that),
            int_mean_(that.int_mean_),
			u_(that.u_),
        	x_(that.x_),
        	y_(that.y_),
        	v_(that.v_)
            {}

        int_mean& operator=(const int_mean& that){
        	if(that!=this){
				static_cast<super1_&>(*this) = that;
            	this->int_mean_ = (that.int_mean_);
				this->u_ = (that.u_);
        		this->x_ = (that.x_);
        		this->y_ = (that.y_);
        		this->v_ = (that.v_);
            }
			return (*this);
		}

  		template<class U>
  		result_type operator()(U& urng)
  		{
			result_type result =  this->step1(urng);
            BOOST_ASSERT(result >= 0);
            return result;
        }

            // --- I/O --- //

            template<class CharT, class Traits>
            friend std::basic_ostream<CharT,Traits>&
            operator<<(
                std::basic_ostream<CharT,Traits>& os,
                const int_mean& pd
            )
            {
                os 	<< "devroye("
            		<< pd.int_mean_val()
            		<< ','
            		<< pd.u()
            		<< ','
            		<< pd.x()
            		<< ','
            		<< pd.y()
            		<< ','
            		<< pd.v()
            		<< ')';
                return os;
            }

            template<class CharT, class Traits>
            friend std::basic_istream<CharT,Traits>&
            operator>>(
                std::basic_istream<CharT,Traits>& is,
                int_mean& pd
            )
            {
            	result_type new_mean;
                is	>> std::ws
            		>> new_mean
            		>> std::ws
                	>> std::ws 	// u
            		>> std::ws
                	>> std::ws	// x
            		>> std::ws
                	>> std::ws	// y
            		>> std::ws
                	>> std::ws	// v
            		>> std::ws;
				static_cast<Step4&>(pd) = Step4(new_mean);
            	return is;
            }

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

        protected:

		template<typename U>
		inline result_type step1(U& urng){
        	// TODO max recursion policy

            // line 10 : case mean = 0;

            this->u_ = urng();												//13
            if(this->u()>this->p1()){
            	// Reconciliation with Fortran listing :
                // ptail = p2 + p3 = 1-p1.
                // u' ~ 1-u. (u'<ptail) <=> (1-u<p2+p3) <=> (u > p1)
            	return this->step3(urng);
            }else{
                return this->step2(urng);
            }
		}

		template<typename U>
		inline result_type step2(U& urng){
        	input_type z = this->random_z(urng); 							//15
        	this->x_ = ( this->sd1() * z ) + super1_::loc1();
            if(
            	(this->x()>this->delta())
                	||  (this->x()< (-this->m1()))
            ){
            	return this->step1(urng);
            }else{
            	// y is rounded away from zero
                if(ma_::is_strictly_negative(this->x())){
                    this->y_ = ma_::floor(this->x());
                }else{
                    this->y_ = ma_::ceil(this->x());
                }
                input_type e  = this->random_exp1(urng);
                this->v_ = -( e + ma_::pow(z,2)/ma_::to_float(2) ) + this->c1();
                return this->step4(urng);
            }
        }

		template<typename U>
		inline result_type step3(U& urng){

            // Reconciliation with Fortran listing :
            // pbody = p3 = 1-(p1+p2)
            // u' ~ 1-u. (u'<pbody) <=> (1-u<p3) <=> (u > p1+p2)
        	if(this->u()>this->p1()+this->p2()){
            	this->y_ = this->int_mean_val();
                return this->y();
            }else{
            	input_type e1 = this->random_exp1(urng);
            	input_type e2 = this->random_exp1(urng);
            	this->x_ = this->delta() + e1 / this->shape2();				//51
            	this->y_ = ma_::to_int( ma_::ceil(this->x()) );
            	this->v_ = - (
                	e2 + this->shape2() * (ma_::to_float(1) + this->x())
                );
            	return this->step4(urng);
            }
        }

		template<typename U>
		inline result_type step4(U& urng){
            Step4& derived = static_cast<Step4&>(*this);
            if(derived.accept()){
            	return this->y();
            }else{
				return this->step1(urng);
            }
		}

		template<typename U>
		input_type random_z(U& urng)const{
        	typedef boost::normal_distribution<input_type> d_;
            static d_ d = d_(ma_::to_float(0),ma_::to_float(1));
            return d(urng);
        }

		template<typename U>
		input_type random_exp1(U& urng)const{
        	typedef boost::exponential_distribution<input_type> d_;
            static d_ d = d_(ma_::to_float(1));
            return d(urng);
        }

		// Parameter
		result_type int_mean_;

        // These are quantities that vary throughout sampling
		input_type	u_;
        input_type 	x_;
        result_type y_;
        input_type 	v_;

        const result_type& 	int_mean_val()const{ return this->int_mean_; }
        const input_type& 	u()const{ return this->u_; }
        const input_type& 	x()const{ return this->x_; }
        const result_type& 	y()const{ return this->y_; }
        const input_type& 	v()const{ return this->v_; }

	};

}// detail
}// devroye
}// poisson
}// random
}// boost

#endif


#include <vector>

#include <boost/shared_ptr.hpp>
#include <gflags/gflags.h>
#include <glog/logging.h>

#include <cmath>

#include "caffe/blob.hpp"
#include "caffe/common.hpp"
#include "caffe/flow_transformer_layer.hpp"
#include "caffe/util/math_functions.hpp"

namespace caffe {

template <typename Dtype>
void FlowTransformerLayer<Dtype>::LayerSetUp(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {

	string prefix = "\t\tFlow Transformer Layer:: LayerSetUp: \t";

	std::cout<<prefix<<"Getting output_H_ and output_W_"<<std::endl;
	output_H_ = bottom[0]->shape(2);
	output_W_ = bottom[0]->shape(3);
	std::cout<<prefix<<"output_H_ = "<<output_H_<<", output_W_ = "<<output_W_<<std::endl;

	std::cout<<prefix<<"Getting pre-defined parameters"<<std::endl;

	// check the validation for the parameter theta
	CHECK(bottom[1]->shape(2) == bottom[0]->shape(2)) << "The third dimension (height) of the flow field and " <<
			"U should be the same" << std::endl;
	CHECK(bottom[1]->shape(3) == bottom[0]->shape(3)) << "The forth dimension (width) of the flow field and " <<
			"U should be the same" << std::endl;

	// initialize the matrix for output grid
	std::cout<<prefix<<"Initializing the matrix for output grid"<<std::endl;

	//2 channel image, 1 channel each for U and V of flow field
	vector<int> shape_output(4);
	shape_output[0] = bottom[0]->shape(0);
	shape_output[1] = 2;
	shape_output[2] = output_H_;
	shape_output[3] = output_W_;
	output_grid.Reshape(shape_output);

	Dtype* data = output_grid.mutable_cpu_data();
	for(int n=0;n<bottom[0]->shape(0);n++){
	for(int r=0;r<output_H_;r++){
		for(int c=0;c<output_W_;c++){
			data[output_grid.offset(n,0,r,c)] = c;
			data[output_grid.offset(n,1,r,c)] = r;
		}
	}
	}

	// initialize the matrix for input grid
	std::cout<<prefix<<"Initializing the matrix for input grid"<<std::endl;

	vector<int> shape_input(4);
	shape_input[0] = bottom[0]->shape(0);
	shape_input[1] = 2;
	shape_input[2] = output_H_;
	shape_input[3] = output_W_;
	input_grid.Reshape(shape_input);

	std::cout<<prefix<<"Initialization finished."<<std::endl;
}

template <typename Dtype>
void FlowTransformerLayer<Dtype>::Reshape(const vector<Blob<Dtype>*>& bottom,
      const vector<Blob<Dtype>*>& top) {
	string prefix = "\t\tFlow Transformer Layer:: Reshape: \t";

	if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

	N = bottom[0]->shape(0);
	C = bottom[0]->shape(1);
	H = bottom[0]->shape(2);
	W = bottom[0]->shape(3);

	// reshape V
	vector<int> shape(4);

	shape[0] = N;
	shape[1] = C;
	shape[2] = output_H_;
	shape[3] = output_W_;

	top[0]->Reshape(shape);

	// reshape dTheta_tmp
	vector<int> dTheta_tmp_shape(4);

	dTheta_tmp_shape[0] = N;
	dTheta_tmp_shape[1] = 2;
	dTheta_tmp_shape[2] = 3;
	dTheta_tmp_shape[3] = output_H_ * output_W_ * C;

	dTheta_tmp.Reshape(dTheta_tmp_shape);

	// init all_ones_2
	vector<int> all_ones_2_shape(1);
	all_ones_2_shape[0] = output_H_ * output_W_ * C;
	all_ones_2.Reshape(all_ones_2_shape);

	// reshape full_theta
	vector<int> full_theta_shape(2);
	full_theta_shape[0] = N;
	full_theta_shape[1] = 6;
	full_theta.Reshape(full_theta_shape);

	if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
Dtype FlowTransformerLayer<Dtype>::transform_forward_cpu(const Dtype* pic, Dtype px, Dtype py) {
	/*
	so i think this does the bilinear interpolation for a single pixel
	*/

	bool debug = false;

	string prefix = "\t\tFlow Transformer Layer:: transform_forward_cpu: \t";

	if(debug) std::cout<<prefix<<"Starting!\t"<<std::endl;
	if(debug) std::cout<<prefix<<"(px, py) = ("<<px<<", "<<py<<")"<<std::endl;

	Dtype res = (Dtype)0.;

	Dtype x = px; Dtype y = py;

	if(debug) std::cout<<prefix<<"(x, y) = ("<<x<<", "<<y<<")"<<std::endl;

	int m, n; Dtype w;

	m = floor(x); n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"1: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x) + 1; n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"2: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x); n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"3: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	m = floor(x) + 1; n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"4: (m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));
		res += w * pic[m * W + n];
		if(debug) std::cout<<prefix<<"w = "<<w<<", pic[m, n] = "<<pic[m * W + n]<<std::endl;
	}

	if(debug) std::cout<<prefix<<"Finished. \tres = "<<res<<std::endl;

	return res;
}

template <typename Dtype>
void FlowTransformerLayer<Dtype>::Forward_cpu(const vector<Blob<Dtype>*>& bottom,
    const vector<Blob<Dtype>*>& top) {

	string prefix = "\t\tFlow Transformer Layer:: Forward_cpu: \t";

	/*
	CHECK(false) << "Don't use the CPU implementation! If you really want to, delete the" <<
			" CHECK in st_layer.cpp file. Line number: 240-241." << std::endl;
	*/

	if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

	const Dtype* U = bottom[0]->cpu_data();
	const Dtype* flows = bottom[1]->cpu_data();

	const Dtype* output_grid_data = output_grid.cpu_data();

	Dtype* input_grid_data = input_grid.mutable_cpu_data();
	Dtype* V = top[0]->mutable_cpu_data();

	caffe_set(input_grid.count(), (Dtype)0, input_grid_data);
	caffe_set(top[0]->count(), (Dtype)0, V);

	//combine flow with grid
	caffe_add(input_grid.count(),output_grid_data,flows,input_grid_data);

	// for each input
	for(int i = 0; i < N; ++i) {
		//coordinates is a matrix of coordinates on the original image
		//to sample pixel values from
		//Dtype* coordinates = input_grid_data + (output_H_ * output_W_ * 2) * i;
		//Dtype* coordinates = input_grid_data + input_grid.offset(i,0,0,0);

		//do affine transformation
		//caffe_cpu_gemm<Dtype>(CblasNoTrans, CblasTrans, output_H_ * output_W_, 2, 3, (Dtype)1.,
		//      output_grid_data, theta + 6 * i, (Dtype)0., coordinates);

		Dtype px, py;

		//for each pixel on the output, find the interpolated pixel value
		for(int s = 0; s < output_H_; ++s){
			for(int t = 0; t < output_W_; ++t) {
				//px,py should be the sample coordinates on the source img

				py = input_grid_data[input_grid.offset(i,0,s,t)];
				px = input_grid_data[input_grid.offset(i,1,s,t)];

				for(int j = 0; j < C; ++j){
					//do interpolation
					V[top[0]->offset(i, j, s, t)] = transform_forward_cpu(
							U + bottom[0]->offset(i, j, 0, 0), px, py);
				}
			}
		}
	}

	if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
void FlowTransformerLayer<Dtype>::transform_backward_cpu(Dtype dV, const Dtype* U, const Dtype px,
		const Dtype py, Dtype* dU, Dtype& dpx, Dtype& dpy) {
	/*
	U is the input image
	V is out output image, dV is local gradient from output
	*/

	bool debug = false;

	string prefix = "\t\tFlow Transformer Layer:: transform_backward_cpu: \t";

	if(debug) std::cout<<prefix<<"Starting!"<<std::endl;

	//Dtype x = (px + 1) / 2 * H; Dtype y = (py + 1) / 2 * W;
	Dtype x = px; Dtype y = py;
	if(debug) std::cout<<prefix<<"(x, y) = ("<<x<<", "<<y<<")"<<std::endl;

	int m, n; Dtype w;

	m = floor(x); n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		if(to_compute_dU_) dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x) + 1; n = floor(y); w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		if(to_compute_dU_) dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x); n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		if(to_compute_dU_) dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	m = floor(x) + 1; n = floor(y) + 1; w = 0;
	if(debug) std::cout<<prefix<<"(m, n) = ("<<m<<", "<<n<<")"<<std::endl;

	if(m >= 0 && m < H && n >= 0 && n < W) {
		w = max(0, 1 - abs(x - m)) * max(0, 1 - abs(y - n));

		if(to_compute_dU_) dU[m * W + n] += w * dV;

		if(abs(x - m) < 1) {
			if(m >= x) {
				dpx += max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx += "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			} else {
				dpx -= max(0, 1 - abs(y - n)) * U[m * W + n] * dV;// * H / 2;
				if(debug) std::cout<<prefix<<"dpx -= "<<max(0, 1 - abs(y - n))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<H / 2<<std::endl;
			}
		}

		if(abs(y - n) < 1) {
			if(n >= y) {
				dpy += max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy += "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			} else {
				dpy -= max(0, 1 - abs(x - m)) * U[m * W + n] * dV;// * W / 2;
				if(debug) std::cout<<prefix<<"dpy -= "<<max(0, 1 - abs(x - m))<<" * "<<U[m * W + n]<<" * "<<dV<<" * "<<W / 2<<std::endl;
			}
		}
	}

	if(debug) std::cout<<prefix<<"Finished."<<std::endl;
}

template <typename Dtype>
void FlowTransformerLayer<Dtype>::Backward_cpu(const vector<Blob<Dtype>*>& top,
    const vector<bool>& propagate_down,
    const vector<Blob<Dtype>*>& bottom) {

		string prefix = "\t\tFlow Transformer Layer:: Backward_cpu: \t";

		//CHECK(false) << "Don't use the CPU implementation! If you really want to, delete the" <<
		//		" CHECK in st_layer.cpp file. Line number: 420-421." << std::endl;

		if(global_debug) std::cout<<prefix<<"Starting!"<<std::endl;

		const Dtype* dV = top[0]->cpu_diff();
		Dtype* input_grid_data = input_grid.mutable_cpu_data();
		//const Dtype* output_grid_data = output_grid.cpu_data(); // remove if you don't need to recalc the input grid
		const Dtype* U = bottom[0]->cpu_data();

		Dtype* dU = bottom[0]->mutable_cpu_diff(); //we won't be setting this
		//Dtype* dTheta = bottom[1]->mutable_cpu_diff();
		//const Dtype* flows = bottom[1]->cpu_data(); // remove if you don't need to recalc the input grid
		Dtype* flow_diff = bottom[1]->mutable_cpu_diff();

		caffe_set(bottom[0]->count(), (Dtype)0, dU);
		//caffe_set(bottom[1]->count(), (Dtype)0, dTheta);
		caffe_set(input_grid.count(), (Dtype)0, flow_diff);

		//calc input grid data, recalc incase blob is reset between forward and backward
		//caffe_add(input_grid.count(),output_grid_data,flows,input_grid_data);

		for(int i = 0; i < N; ++i) {

			//const Dtype* coordinates = input_grid_data + (output_H_ * output_W_ * 2) * i;
			//Dtype* coordinates_diff = input_grid_diff + (output_H_ * output_W_ * 2) * i;

			const Dtype* curInputGrid = &input_grid_data[input_grid.offset(i,0,0,0)];
			//since its just element wise addition, we are
			//just going to set the diff directly to the flow
			Dtype* curFlowDiff = &flow_diff[bottom[1]->offset(i,0,0,0)];

			//int row_idx;
			Dtype px, py, delta_dpx, delta_dpy;

			for(int s = 0; s < output_H_; ++s)
				for(int t = 0; t < output_W_; ++t) {

					//row_idx = output_W_ * s + t;	//this means fix!

					px = curInputGrid[input_grid.offset(0,1,s,t)];
					py = curInputGrid[input_grid.offset(0,0,s,t)];

					for(int j = 0; j < C; ++j) {

						delta_dpx = delta_dpy = (Dtype)0.;

						transform_backward_cpu(dV[top[0]->offset(i, j, s, t)], U + bottom[0]->offset(i, j, 0, 0),
								px, py, dU + bottom[0]->offset(i, j, 0, 0), delta_dpx, delta_dpy);

						curFlowDiff[bottom[1]->offset(0,1,s,t)] += delta_dpx;
						curFlowDiff[bottom[1]->offset(0,0,s,t)] += delta_dpy;
					}

					/*
					dpx = curFlowDiff[row_idx * 2];
					dpy = curFlowDiff[row_idx * 2 + 1];

					dTheta[6 * i] += dpx * (s * 1.0 / output_H_ * 2 - 1);
					dTheta[6 * i + 1] += dpx * (t * 1.0 / output_W_ * 2 - 1);
					dTheta[6 * i + 2] += dpx;
					dTheta[6 * i + 3] += dpy * (s * 1.0 / output_H_ * 2 - 1);
					dTheta[6 * i + 4] += dpy * (t * 1.0 / output_W_ * 2 - 1);
					dTheta[6 * i + 5] += dpy;
					*/
				}
		}

		if(global_debug) std::cout<<prefix<<"Finished."<<std::endl;
}

#ifdef CPU_ONLY
STUB_GPU(FlowTransformerLayer);
#endif

INSTANTIATE_CLASS(FlowTransformerLayer);
REGISTER_LAYER_CLASS(FlowTransformer);

}  // namespace caffe


#pragma once

#include <numeric>

#include <boost/simd/constant/zero.hpp>
#include <boost/simd/function/cumsum.hpp>
#include <boost/simd/function/splat.hpp>
#include <boost/simd/function/sum.hpp>

#include "source.hpp"

typedef boost::simd::pack<double> pack_double;

template<typename tag>
double balanced_sum(source_1d<tag> &f)
{
    double res = 0.;
    while (f.can_procure(1, res) && !f.is_aligned()) {
        f.procure(res);
    }

    pack_double stack[62];
    size_t p = 0;
    for (size_t iteration = 0; f.can_procure(4, stack[0]); ++iteration) {
        pack_double v = boost::simd::Zero<pack_double>();
        f.procure(v);
        f.procure(v);
        pack_double w = boost::simd::Zero<pack_double>();
        f.procure(w);
        f.procure(w);
        v += w;
        size_t bitmask = 1;
        for (; iteration & bitmask; bitmask <<= 1, --p) {
            v += stack[p - 1];
        }
        stack[p++] = v;
    }
    pack_double vsum = boost::simd::Zero<pack_double>();
    for (size_t i = p; i > 0; --i) {
        vsum += stack[i - 1];
    }
    res = std::accumulate(vsum.begin(), vsum.end(), res);
    while (f.can_procure(1, res)) {
        f.procure(res);
    }
    return res;
}

template<typename tag>
double twin_balanced_sum(source_2d<tag>& f)
{
    double stat1 = 0.;
    double stat2 = 0.;
    while (f.can_procure(1, stat1) && !f.is_aligned()) {
        f.procure(stat1, stat2);
    }

    pack_double stack1[62];
    pack_double stack2[62];
    size_t p = 0;
    for (size_t iteration = 0; f.can_procure(4, stack1[0]); ++iteration) {
        pack_double v1 = boost::simd::Zero<pack_double>();
        pack_double v2 = boost::simd::Zero<pack_double>();
        f.procure(v1, v2);
        f.procure(v1, v2);
        pack_double w1 = boost::simd::Zero<pack_double>();
        pack_double w2 = boost::simd::Zero<pack_double>();
        f.procure(w1, w2);
        f.procure(w1, w2);
        v1 += w1;
        v2 += w2;
        size_t bitmask = 1;
        for (; iteration & bitmask; bitmask <<= 1, --p) {
            v1 += stack1[p - 1];
            v2 += stack2[p - 1];
        }
        stack1[p] = v1;
        stack2[p++] = v2;
    }
    pack_double vsum1 = boost::simd::Zero<pack_double>();
    pack_double vsum2 = boost::simd::Zero<pack_double>();
    for (size_t i = p; i > 0; --i) {
        vsum1 += stack1[i - 1];
        vsum2 += stack2[i - 1];
    }
    stat1 = std::accumulate(vsum1.begin(), vsum1.end(), stat1);
    stat2 = std::accumulate(vsum2.begin(), vsum2.end(), stat2);
    while (f.can_procure(1, stat1)) {
        f.procure(stat1, stat2);
    }
    return f.result(stat1, stat2);
}

inline void kahan_update(double &accumulator, double &compensator, double value)
{
    double const new_accumulator = accumulator + value;
    double const first_option = (accumulator - new_accumulator) + value;
    double const second_option = (value - new_accumulator) + accumulator;
    if (std::abs(accumulator) > std::abs(value)) {
        compensator += first_option;
    } else {
        compensator += second_option;
    }
    accumulator = new_accumulator;
}

template<typename tag>
void cum_sum(source_1d<tag> &f)
{
    double accumulator = 0.;
    double compensator = 0.;
    double value = 0.;
    while (f.can_procure(1, value) && !f.is_aligned()) {
        f.procure(value);
        kahan_update(accumulator, compensator, value);
        f.feed(accumulator + compensator);
    }
    pack_double pack_value1 = boost::simd::Zero<pack_double>();
    pack_double pack_value2 = boost::simd::Zero<pack_double>();
    while (f.can_procure(2, pack_value1)) {
        f.procure(pack_value1);
        f.procure(pack_value2);
        pack_double const pack_cum_sum1 = boost::simd::cumsum(pack_value1);
        pack_double const pack_cum_sum2 = boost::simd::cumsum(pack_value2);
        pack_double const pack_feed_slice1 = pack_cum_sum1 + (accumulator + compensator);
        f.feed(pack_feed_slice1);
        kahan_update(accumulator, compensator, pack_cum_sum1[pack_double::static_size - 1]);
        pack_double const pack_feed_slice2 = pack_cum_sum2 + (accumulator + compensator);
        f.feed(pack_feed_slice2);
        kahan_update(accumulator, compensator, pack_cum_sum2[pack_double::static_size - 1]);
    }
    value = 0.;
    while (f.can_procure(1, accumulator)) {
        f.procure(value);
        kahan_update(accumulator, compensator, value);
        f.feed(accumulator + compensator);
    }
}


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

#include "matrix_utils.hpp"

#include <cmath>
#include <vector>

#include <boost/dynamic_bitset.hpp>
#include <boost/format.hpp>

#include <xtensor/xio.hpp>
#include <xtensor-blas/xlinalg.hpp>

#include <core/utils/bitset_utils.hpp>
#include <core/utils/terminal.hpp>
#include <reversible/pauli_tags.hpp>
#include <reversible/target_tags.hpp>
#include <reversible/simulation/simple_simulation.hpp>

namespace cirkit
{

using namespace std::complex_literals;

/******************************************************************************
 * Types                                                                      *
 ******************************************************************************/

/******************************************************************************
 * Private functions                                                          *
 ******************************************************************************/

xt::xarray<complex_t> get_2by2_matrix( complex_t a, complex_t b, complex_t c, complex_t d )
{
  xt::xarray<complex_t> matrix(std::vector<size_t>{2, 2});
  matrix[{0,0}] = a;
  matrix[{0,1}] = b;
  matrix[{1,0}] = c;
  matrix[{1,1}] = d;
  return matrix;
}

xt::xarray<complex_t> kron( const std::vector<xt::xarray<complex_t>>& ms )
{
  assert( !ms.empty() );

  auto r = ms[0u];

  for ( auto i = 1u; i < ms.size(); ++i )
  {
    r = xt::linalg::kron( r, ms[i] );
  }

  return r;
}

xt::xarray<complex_t> identity( unsigned dimension )
{
  xt::xarray<complex_t> matrix(std::vector<size_t>{dimension, dimension});
  for ( auto i = 0u; i < dimension; ++i )
  {
    matrix[{i, i}] = 1.0;
  }
  return matrix;
}

xt::xarray<complex_t> identity_padding( const xt::xarray<complex_t>& matrix, unsigned from, unsigned to, unsigned lines )
{
  std::vector<xt::xarray<complex_t>> ms;

  if ( to + 1u < lines )
  {
    ms.push_back( identity( 1 << ( lines - to - 1u ) ) );
  }

  ms.push_back( matrix );

  if ( from > 0u )
  {
    ms.push_back( identity( 1 << from ) );
  }

  return kron( ms );
}

/******************************************************************************
 * Public functions                                                           *
 ******************************************************************************/

xt::xarray<complex_t> matrix_from_clifford_t_circuit( const circuit& circ, bool progress )
{
  xt::xarray<complex_t> matrix_X = get_2by2_matrix( 0.0, 1.0, 1.0, 0.0 );
  xt::xarray<complex_t> matrix_H = 1.0 / sqrt( 2.0 ) * get_2by2_matrix( 1.0, 1.0, 1.0, -1.0 );
  xt::xarray<complex_t> matrix_Z = get_2by2_matrix( 1.0, 0.0, 0.0, -1.0 );
  xt::xarray<complex_t> matrix_S = get_2by2_matrix( 1.0, 0.0, 0.0, 1i );
  xt::xarray<complex_t> matrix_Sdag = get_2by2_matrix( 1.0, 0.0, 0.0, -1i );
  xt::xarray<complex_t> matrix_T = get_2by2_matrix( 1.0, 0.0, 0.0, exp( 1i * M_PI / 4.0 ) );
  xt::xarray<complex_t> matrix_Tdag = get_2by2_matrix( 1.0, 0.0, 0.0, exp( -1i * M_PI / 4.0 ) );
  xt::xarray<complex_t> matrix_zc = get_2by2_matrix( 1.0, 0.0, 0.0, 0.0 );
  xt::xarray<complex_t> matrix_oc = get_2by2_matrix( 0.0, 0.0, 0.0, 1.0 );

  const auto n = circ.lines();

  std::vector<xt::xarray<complex_t>> gates;

  for ( const auto& g : circ )
  {
    const auto target = g.targets().front();

    if ( is_toffoli( g ) && g.controls().size() == 0u )
    {
      gates.push_back( identity_padding( matrix_X, target, target, n ) );
    }
    else if ( is_toffoli( g ) && g.controls().size() == 1u && g.controls().front().polarity() )
    {
      const auto control = g.controls().front().line();

      if ( control < target )
      {
        const auto act = target - control;
        const auto gate = xt::linalg::kron( identity( 1 << act ), matrix_zc ) + kron( {matrix_X, identity( 1 << ( act - 1 ) ), matrix_oc} );
        gates.push_back( identity_padding( gate, control, target, n ) );
      }
      else
      {
        const auto act = control - target;
        const auto gate = xt::linalg::kron( matrix_zc, identity( 1 << act ) ) + kron( {matrix_oc, identity( 1 << ( act - 1 ) ), matrix_X} );
        gates.push_back( identity_padding( gate, target, control, n ) );
      }
    }
    else if ( is_hadamard( g ) )
    {
      gates.push_back( identity_padding( matrix_H, target, target, n ) );
    }
    else if ( is_pauli( g ) )
    {
      const auto pauli = boost::any_cast<pauli_tag>( g.type() );
      switch ( pauli.axis )
      {
      case pauli_axis::X:
        if ( pauli.root == 1u )
        {
          gates.push_back( identity_padding( matrix_X, target, target, n ) );
        }
        else
        {
          std::cout << "[w] unsupported X gate" << std::endl;
        }
        break;
      case pauli_axis::Z:
        if ( pauli.root == 1u )
        {
          gates.push_back( identity_padding( matrix_Z, target, target, n ) );
        }
        else if ( pauli.root == 2u )
        {
          gates.push_back( identity_padding( pauli.adjoint ? matrix_Sdag : matrix_S, target, target, n ) );
        }
        else if ( pauli.root == 4u )
        {
          gates.push_back( identity_padding( pauli.adjoint ? matrix_Tdag : matrix_T, target, target, n ) );
        }
        else
        {
          std::cout << "[w] unsupported Z gate with root " << pauli.root << std::endl;
        }
        break;
      default:
        std::cout << "[w] unsupported Pauli gate" << std::endl;
        break;
      }
    }
    else
    {
      std::cout << "[w] unsupported gate" << std::endl;
    }
  }

  if ( gates.empty() )
  {
    return identity( 1 << n );
  }
  else
  {
    progress_line pline( boost::str( boost::format( "[i] (matrix_from_clifford_t_circuit) gate %%d / %d" ) % gates.size() ), progress );

    auto r = gates.front();
    for ( auto i = 1u; i < gates.size(); ++i )
    {
      pline( i + 1u );
      r = xt::linalg::dot( r, gates[i] );
    }

    return r;
  }
}

xt::xarray<complex_t> matrix_from_reversible_circuit( const circuit& circ )
{
  const size_t N = 1 << circ.lines();

  xt::xarray<complex_t> matrix(std::vector<size_t>{N, N});
  foreach_bitset( circ.lines(), [&circ, &matrix]( const boost::dynamic_bitset<>& input ) {
      boost::dynamic_bitset<> output;
      simple_simulation( output, circ, input );
      matrix[{input.to_ulong(), output.to_ulong()}] = 1.0;
    } );

  return matrix;
}

}

// Local Variables:
// c-basic-offset: 2
// eval: (c-set-offset 'substatement-open 0)
// eval: (c-set-offset 'innamespace 0)
// End:


/*===========================================================================*\

Author: Matthias W. Smith
Email:  mwsmith2@uw.edu
Date:   11/02/14

Detail: This is a new test program for my Fid libraries

\*===========================================================================*/

//--- std includes ----------------------------------------------------------//
#include <iostream>
#include <fstream>
#include <vector>
#include <cmath>
using std::vector;
using std::cout;
using std::endl;

//--- other includes --------------------------------------------------------//
#include <armadillo>

//--- project includes ------------------------------------------------------//
#include "fid.h"
using namespace fid;

int main(int argc, char** argv)
{
  // set precision
  cout.precision(10);
  cout.setf(std::ios::fixed, std::ios::floatfield);

  // declare variables
  int fid_length = 5000;
  double ti = -1.0;
  double dt = 0.001;
  double ftruth = 23.0;

  vector<double> wf;
  vector<double> tm;
  wf.reserve(fid_length);
  tm.reserve(fid_length);

  std::ofstream out;
  out.precision(10);

  for (int i = 0; i < fid_length; i++){
    tm.push_back(i * dt + ti);
  }

  FidFactory ff;
  ff.SetLarmorFreq(ftruth + sim::mixdown_freq);
  ff.IdealFid(wf, tm);

  for (int i = 0; i < wf.size(); ++i) {
    wf[i] = sin(40 * tm[i]);
  }

  out.open("wvd_test_wf.txt");
  for (auto it = wf.begin(); it != wf.end(); ++it) {
    out << *it << ", ";
  }
  out.close();

  auto wf_im = dsp::hilbert(wf);

  out.open("wvd_test_wf_im.txt");
  for (auto it = wf_im.begin(); it != wf_im.end(); ++it) {
    out << *it << ", ";
  }
  out.close();

  auto res = dsp::wvd(wf);
  res.quiet_save("wvd_test.txt", arma::csv_ascii);

  return 0;
}


#include "forward_dynamics.h"

#include <Eigen/Cholesky>
#include <iit/rbd/robcogen_commons.h>

using namespace iit::rbd;

// Initialization of static-const data
const iit::popi::dyn::ForwardDynamics::ExtForces
    iit::popi::dyn::ForwardDynamics::zeroExtForces(Force::Zero());

iit::popi::dyn::ForwardDynamics::ForwardDynamics(InertiaProperties& inertia, MotionTransforms& transforms) :
    inertiaProps( & inertia ),
    motionTransforms( & transforms )
{
    EpauleAVD_v.setZero();
    EpauleAVD_c.setZero();
    HJambeAVD_v.setZero();
    HJambeAVD_c.setZero();
    BJambeAVD_v.setZero();
    BJambeAVD_c.setZero();
    EpauleAVG_v.setZero();
    EpauleAVG_c.setZero();
    HJambeAVG_v.setZero();
    HJambeAVG_c.setZero();
    BJambeAVG_v.setZero();
    BJambeAVG_c.setZero();
    EpauleARD_v.setZero();
    EpauleARD_c.setZero();
    HJambeARD_v.setZero();
    HJambeARD_c.setZero();
    BJambeARD_v.setZero();
    BJambeARD_c.setZero();
    EpauleARG_v.setZero();
    EpauleARG_c.setZero();
    HJambeARG_v.setZero();
    HJambeARG_c.setZero();
    BJambeARG_v.setZero();
    BJambeARG_c.setZero();

    vcross.setZero();
    Ia_r.setZero();

}

void iit::popi::dyn::ForwardDynamics::fd(
    JointState& qdd,
    Acceleration& base_a,
    const Velocity& base_v,
    const Acceleration& g,
    const JointState& qd,
    const JointState& tau,
    const ExtForces& fext/* = zeroExtForces */)
{

    base_AI = inertiaProps->getTensor_base();
    base_p = - fext[BASE];
    EpauleAVD_AI = inertiaProps->getTensor_EpauleAVD();
    EpauleAVD_p = - fext[EPAULEAVD];
    HJambeAVD_AI = inertiaProps->getTensor_HJambeAVD();
    HJambeAVD_p = - fext[HJAMBEAVD];
    BJambeAVD_AI = inertiaProps->getTensor_BJambeAVD();
    BJambeAVD_p = - fext[BJAMBEAVD];
    EpauleAVG_AI = inertiaProps->getTensor_EpauleAVG();
    EpauleAVG_p = - fext[EPAULEAVG];
    HJambeAVG_AI = inertiaProps->getTensor_HJambeAVG();
    HJambeAVG_p = - fext[HJAMBEAVG];
    BJambeAVG_AI = inertiaProps->getTensor_BJambeAVG();
    BJambeAVG_p = - fext[BJAMBEAVG];
    EpauleARD_AI = inertiaProps->getTensor_EpauleARD();
    EpauleARD_p = - fext[EPAULEARD];
    HJambeARD_AI = inertiaProps->getTensor_HJambeARD();
    HJambeARD_p = - fext[HJAMBEARD];
    BJambeARD_AI = inertiaProps->getTensor_BJambeARD();
    BJambeARD_p = - fext[BJAMBEARD];
    EpauleARG_AI = inertiaProps->getTensor_EpauleARG();
    EpauleARG_p = - fext[EPAULEARG];
    HJambeARG_AI = inertiaProps->getTensor_HJambeARG();
    HJambeARG_p = - fext[HJAMBEARG];
    BJambeARG_AI = inertiaProps->getTensor_BJambeARG();
    BJambeARG_p = - fext[BJAMBEARG];
    // ---------------------- FIRST PASS ---------------------- //
    // Note that, during the first pass, the articulated inertias are really
    //  just the spatial inertia of the links (see assignments above).
    //  Afterwards things change, and articulated inertias shall not be used
    //  in functions which work specifically with spatial inertias.

    // + Link EpauleAVD
    //  - The spatial velocity:
    EpauleAVD_v = (motionTransforms-> fr_EpauleAVD_X_fr_base) * base_v;
    EpauleAVD_v(AZ) += qd(RF_HAA_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(EpauleAVD_v, vcross);
    EpauleAVD_c = vcross.col(AZ) * qd(RF_HAA_JOINT);

    //  - The bias force term:
    EpauleAVD_p += vxIv(EpauleAVD_v, EpauleAVD_AI);

    // + Link HJambeAVD
    //  - The spatial velocity:
    HJambeAVD_v = (motionTransforms-> fr_HJambeAVD_X_fr_EpauleAVD) * EpauleAVD_v;
    HJambeAVD_v(AZ) += qd(RF_HFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(HJambeAVD_v, vcross);
    HJambeAVD_c = vcross.col(AZ) * qd(RF_HFE_JOINT);

    //  - The bias force term:
    HJambeAVD_p += vxIv(HJambeAVD_v, HJambeAVD_AI);

    // + Link BJambeAVD
    //  - The spatial velocity:
    BJambeAVD_v = (motionTransforms-> fr_BJambeAVD_X_fr_HJambeAVD) * HJambeAVD_v;
    BJambeAVD_v(AZ) += qd(RF_KFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(BJambeAVD_v, vcross);
    BJambeAVD_c = vcross.col(AZ) * qd(RF_KFE_JOINT);

    //  - The bias force term:
    BJambeAVD_p += vxIv(BJambeAVD_v, BJambeAVD_AI);

    // + Link EpauleAVG
    //  - The spatial velocity:
    EpauleAVG_v = (motionTransforms-> fr_EpauleAVG_X_fr_base) * base_v;
    EpauleAVG_v(AZ) += qd(LF_HAA_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(EpauleAVG_v, vcross);
    EpauleAVG_c = vcross.col(AZ) * qd(LF_HAA_JOINT);

    //  - The bias force term:
    EpauleAVG_p += vxIv(EpauleAVG_v, EpauleAVG_AI);

    // + Link HJambeAVG
    //  - The spatial velocity:
    HJambeAVG_v = (motionTransforms-> fr_HJambeAVG_X_fr_EpauleAVG) * EpauleAVG_v;
    HJambeAVG_v(AZ) += qd(LF_HFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(HJambeAVG_v, vcross);
    HJambeAVG_c = vcross.col(AZ) * qd(LF_HFE_JOINT);

    //  - The bias force term:
    HJambeAVG_p += vxIv(HJambeAVG_v, HJambeAVG_AI);

    // + Link BJambeAVG
    //  - The spatial velocity:
    BJambeAVG_v = (motionTransforms-> fr_BJambeAVG_X_fr_HJambeAVG) * HJambeAVG_v;
    BJambeAVG_v(AZ) += qd(LF_KFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(BJambeAVG_v, vcross);
    BJambeAVG_c = vcross.col(AZ) * qd(LF_KFE_JOINT);

    //  - The bias force term:
    BJambeAVG_p += vxIv(BJambeAVG_v, BJambeAVG_AI);

    // + Link EpauleARD
    //  - The spatial velocity:
    EpauleARD_v = (motionTransforms-> fr_EpauleARD_X_fr_base) * base_v;
    EpauleARD_v(AZ) += qd(RH_HAA_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(EpauleARD_v, vcross);
    EpauleARD_c = vcross.col(AZ) * qd(RH_HAA_JOINT);

    //  - The bias force term:
    EpauleARD_p += vxIv(EpauleARD_v, EpauleARD_AI);

    // + Link HJambeARD
    //  - The spatial velocity:
    HJambeARD_v = (motionTransforms-> fr_HJambeARD_X_fr_EpauleARD) * EpauleARD_v;
    HJambeARD_v(AZ) += qd(RH_HFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(HJambeARD_v, vcross);
    HJambeARD_c = vcross.col(AZ) * qd(RH_HFE_JOINT);

    //  - The bias force term:
    HJambeARD_p += vxIv(HJambeARD_v, HJambeARD_AI);

    // + Link BJambeARD
    //  - The spatial velocity:
    BJambeARD_v = (motionTransforms-> fr_BJambeARD_X_fr_HJambeARD) * HJambeARD_v;
    BJambeARD_v(AZ) += qd(RH_KFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(BJambeARD_v, vcross);
    BJambeARD_c = vcross.col(AZ) * qd(RH_KFE_JOINT);

    //  - The bias force term:
    BJambeARD_p += vxIv(BJambeARD_v, BJambeARD_AI);

    // + Link EpauleARG
    //  - The spatial velocity:
    EpauleARG_v = (motionTransforms-> fr_EpauleARG_X_fr_base) * base_v;
    EpauleARG_v(AZ) += qd(LH_HAA_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(EpauleARG_v, vcross);
    EpauleARG_c = vcross.col(AZ) * qd(LH_HAA_JOINT);

    //  - The bias force term:
    EpauleARG_p += vxIv(EpauleARG_v, EpauleARG_AI);

    // + Link HJambeARG
    //  - The spatial velocity:
    HJambeARG_v = (motionTransforms-> fr_HJambeARG_X_fr_EpauleARG) * EpauleARG_v;
    HJambeARG_v(AZ) += qd(LH_HFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(HJambeARG_v, vcross);
    HJambeARG_c = vcross.col(AZ) * qd(LH_HFE_JOINT);

    //  - The bias force term:
    HJambeARG_p += vxIv(HJambeARG_v, HJambeARG_AI);

    // + Link BJambeARG
    //  - The spatial velocity:
    BJambeARG_v = (motionTransforms-> fr_BJambeARG_X_fr_HJambeARG) * HJambeARG_v;
    BJambeARG_v(AZ) += qd(LH_KFE_JOINT);

    //  - The velocity-product acceleration term:
    motionCrossProductMx<Scalar>(BJambeARG_v, vcross);
    BJambeARG_c = vcross.col(AZ) * qd(LH_KFE_JOINT);

    //  - The bias force term:
    BJambeARG_p += vxIv(BJambeARG_v, BJambeARG_AI);

    // + The floating base body
    base_p += vxIv(base_v, base_AI);

    // ---------------------- SECOND PASS ---------------------- //
    Matrix66 IaB;
    Force pa;

    // + Link BJambeARG
    BJambeARG_u = tau(LH_KFE_JOINT) - BJambeARG_p(AZ);
    BJambeARG_U = BJambeARG_AI.col(AZ);
    BJambeARG_D = BJambeARG_U(AZ);

    compute_Ia_revolute(BJambeARG_AI, BJambeARG_U, BJambeARG_D, Ia_r);  // same as: Ia_r = BJambeARG_AI - BJambeARG_U/BJambeARG_D * BJambeARG_U.transpose();
    pa = BJambeARG_p + Ia_r * BJambeARG_c + BJambeARG_U * BJambeARG_u/BJambeARG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_BJambeARG_X_fr_HJambeARG, IaB);
    HJambeARG_AI += IaB;
    HJambeARG_p += (motionTransforms-> fr_BJambeARG_X_fr_HJambeARG).transpose() * pa;

    // + Link HJambeARG
    HJambeARG_u = tau(LH_HFE_JOINT) - HJambeARG_p(AZ);
    HJambeARG_U = HJambeARG_AI.col(AZ);
    HJambeARG_D = HJambeARG_U(AZ);

    compute_Ia_revolute(HJambeARG_AI, HJambeARG_U, HJambeARG_D, Ia_r);  // same as: Ia_r = HJambeARG_AI - HJambeARG_U/HJambeARG_D * HJambeARG_U.transpose();
    pa = HJambeARG_p + Ia_r * HJambeARG_c + HJambeARG_U * HJambeARG_u/HJambeARG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_HJambeARG_X_fr_EpauleARG, IaB);
    EpauleARG_AI += IaB;
    EpauleARG_p += (motionTransforms-> fr_HJambeARG_X_fr_EpauleARG).transpose() * pa;

    // + Link EpauleARG
    EpauleARG_u = tau(LH_HAA_JOINT) - EpauleARG_p(AZ);
    EpauleARG_U = EpauleARG_AI.col(AZ);
    EpauleARG_D = EpauleARG_U(AZ);

    compute_Ia_revolute(EpauleARG_AI, EpauleARG_U, EpauleARG_D, Ia_r);  // same as: Ia_r = EpauleARG_AI - EpauleARG_U/EpauleARG_D * EpauleARG_U.transpose();
    pa = EpauleARG_p + Ia_r * EpauleARG_c + EpauleARG_U * EpauleARG_u/EpauleARG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_EpauleARG_X_fr_base, IaB);
    base_AI += IaB;
    base_p += (motionTransforms-> fr_EpauleARG_X_fr_base).transpose() * pa;

    // + Link BJambeARD
    BJambeARD_u = tau(RH_KFE_JOINT) - BJambeARD_p(AZ);
    BJambeARD_U = BJambeARD_AI.col(AZ);
    BJambeARD_D = BJambeARD_U(AZ);

    compute_Ia_revolute(BJambeARD_AI, BJambeARD_U, BJambeARD_D, Ia_r);  // same as: Ia_r = BJambeARD_AI - BJambeARD_U/BJambeARD_D * BJambeARD_U.transpose();
    pa = BJambeARD_p + Ia_r * BJambeARD_c + BJambeARD_U * BJambeARD_u/BJambeARD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_BJambeARD_X_fr_HJambeARD, IaB);
    HJambeARD_AI += IaB;
    HJambeARD_p += (motionTransforms-> fr_BJambeARD_X_fr_HJambeARD).transpose() * pa;

    // + Link HJambeARD
    HJambeARD_u = tau(RH_HFE_JOINT) - HJambeARD_p(AZ);
    HJambeARD_U = HJambeARD_AI.col(AZ);
    HJambeARD_D = HJambeARD_U(AZ);

    compute_Ia_revolute(HJambeARD_AI, HJambeARD_U, HJambeARD_D, Ia_r);  // same as: Ia_r = HJambeARD_AI - HJambeARD_U/HJambeARD_D * HJambeARD_U.transpose();
    pa = HJambeARD_p + Ia_r * HJambeARD_c + HJambeARD_U * HJambeARD_u/HJambeARD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_HJambeARD_X_fr_EpauleARD, IaB);
    EpauleARD_AI += IaB;
    EpauleARD_p += (motionTransforms-> fr_HJambeARD_X_fr_EpauleARD).transpose() * pa;

    // + Link EpauleARD
    EpauleARD_u = tau(RH_HAA_JOINT) - EpauleARD_p(AZ);
    EpauleARD_U = EpauleARD_AI.col(AZ);
    EpauleARD_D = EpauleARD_U(AZ);

    compute_Ia_revolute(EpauleARD_AI, EpauleARD_U, EpauleARD_D, Ia_r);  // same as: Ia_r = EpauleARD_AI - EpauleARD_U/EpauleARD_D * EpauleARD_U.transpose();
    pa = EpauleARD_p + Ia_r * EpauleARD_c + EpauleARD_U * EpauleARD_u/EpauleARD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_EpauleARD_X_fr_base, IaB);
    base_AI += IaB;
    base_p += (motionTransforms-> fr_EpauleARD_X_fr_base).transpose() * pa;

    // + Link BJambeAVG
    BJambeAVG_u = tau(LF_KFE_JOINT) - BJambeAVG_p(AZ);
    BJambeAVG_U = BJambeAVG_AI.col(AZ);
    BJambeAVG_D = BJambeAVG_U(AZ);

    compute_Ia_revolute(BJambeAVG_AI, BJambeAVG_U, BJambeAVG_D, Ia_r);  // same as: Ia_r = BJambeAVG_AI - BJambeAVG_U/BJambeAVG_D * BJambeAVG_U.transpose();
    pa = BJambeAVG_p + Ia_r * BJambeAVG_c + BJambeAVG_U * BJambeAVG_u/BJambeAVG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_BJambeAVG_X_fr_HJambeAVG, IaB);
    HJambeAVG_AI += IaB;
    HJambeAVG_p += (motionTransforms-> fr_BJambeAVG_X_fr_HJambeAVG).transpose() * pa;

    // + Link HJambeAVG
    HJambeAVG_u = tau(LF_HFE_JOINT) - HJambeAVG_p(AZ);
    HJambeAVG_U = HJambeAVG_AI.col(AZ);
    HJambeAVG_D = HJambeAVG_U(AZ);

    compute_Ia_revolute(HJambeAVG_AI, HJambeAVG_U, HJambeAVG_D, Ia_r);  // same as: Ia_r = HJambeAVG_AI - HJambeAVG_U/HJambeAVG_D * HJambeAVG_U.transpose();
    pa = HJambeAVG_p + Ia_r * HJambeAVG_c + HJambeAVG_U * HJambeAVG_u/HJambeAVG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_HJambeAVG_X_fr_EpauleAVG, IaB);
    EpauleAVG_AI += IaB;
    EpauleAVG_p += (motionTransforms-> fr_HJambeAVG_X_fr_EpauleAVG).transpose() * pa;

    // + Link EpauleAVG
    EpauleAVG_u = tau(LF_HAA_JOINT) - EpauleAVG_p(AZ);
    EpauleAVG_U = EpauleAVG_AI.col(AZ);
    EpauleAVG_D = EpauleAVG_U(AZ);

    compute_Ia_revolute(EpauleAVG_AI, EpauleAVG_U, EpauleAVG_D, Ia_r);  // same as: Ia_r = EpauleAVG_AI - EpauleAVG_U/EpauleAVG_D * EpauleAVG_U.transpose();
    pa = EpauleAVG_p + Ia_r * EpauleAVG_c + EpauleAVG_U * EpauleAVG_u/EpauleAVG_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_EpauleAVG_X_fr_base, IaB);
    base_AI += IaB;
    base_p += (motionTransforms-> fr_EpauleAVG_X_fr_base).transpose() * pa;

    // + Link BJambeAVD
    BJambeAVD_u = tau(RF_KFE_JOINT) - BJambeAVD_p(AZ);
    BJambeAVD_U = BJambeAVD_AI.col(AZ);
    BJambeAVD_D = BJambeAVD_U(AZ);

    compute_Ia_revolute(BJambeAVD_AI, BJambeAVD_U, BJambeAVD_D, Ia_r);  // same as: Ia_r = BJambeAVD_AI - BJambeAVD_U/BJambeAVD_D * BJambeAVD_U.transpose();
    pa = BJambeAVD_p + Ia_r * BJambeAVD_c + BJambeAVD_U * BJambeAVD_u/BJambeAVD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_BJambeAVD_X_fr_HJambeAVD, IaB);
    HJambeAVD_AI += IaB;
    HJambeAVD_p += (motionTransforms-> fr_BJambeAVD_X_fr_HJambeAVD).transpose() * pa;

    // + Link HJambeAVD
    HJambeAVD_u = tau(RF_HFE_JOINT) - HJambeAVD_p(AZ);
    HJambeAVD_U = HJambeAVD_AI.col(AZ);
    HJambeAVD_D = HJambeAVD_U(AZ);

    compute_Ia_revolute(HJambeAVD_AI, HJambeAVD_U, HJambeAVD_D, Ia_r);  // same as: Ia_r = HJambeAVD_AI - HJambeAVD_U/HJambeAVD_D * HJambeAVD_U.transpose();
    pa = HJambeAVD_p + Ia_r * HJambeAVD_c + HJambeAVD_U * HJambeAVD_u/HJambeAVD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_HJambeAVD_X_fr_EpauleAVD, IaB);
    EpauleAVD_AI += IaB;
    EpauleAVD_p += (motionTransforms-> fr_HJambeAVD_X_fr_EpauleAVD).transpose() * pa;

    // + Link EpauleAVD
    EpauleAVD_u = tau(RF_HAA_JOINT) - EpauleAVD_p(AZ);
    EpauleAVD_U = EpauleAVD_AI.col(AZ);
    EpauleAVD_D = EpauleAVD_U(AZ);

    compute_Ia_revolute(EpauleAVD_AI, EpauleAVD_U, EpauleAVD_D, Ia_r);  // same as: Ia_r = EpauleAVD_AI - EpauleAVD_U/EpauleAVD_D * EpauleAVD_U.transpose();
    pa = EpauleAVD_p + Ia_r * EpauleAVD_c + EpauleAVD_U * EpauleAVD_u/EpauleAVD_D;
    ctransform_Ia_revolute(Ia_r, motionTransforms-> fr_EpauleAVD_X_fr_base, IaB);
    base_AI += IaB;
    base_p += (motionTransforms-> fr_EpauleAVD_X_fr_base).transpose() * pa;

    // + The acceleration of the floating base base, without gravity
    Eigen::LLT<Matrix66d> llt(base_AI);
    base_a = - llt.solve(base_p);  // base_a = - IA^-1 * base_p

    // ---------------------- THIRD PASS ---------------------- //
    EpauleAVD_a = (motionTransforms-> fr_EpauleAVD_X_fr_base) * base_a + EpauleAVD_c;
    qdd(RF_HAA_JOINT) = (EpauleAVD_u - EpauleAVD_U.dot(EpauleAVD_a)) / EpauleAVD_D;
    EpauleAVD_a(AZ) += qdd(RF_HAA_JOINT);

    HJambeAVD_a = (motionTransforms-> fr_HJambeAVD_X_fr_EpauleAVD) * EpauleAVD_a + HJambeAVD_c;
    qdd(RF_HFE_JOINT) = (HJambeAVD_u - HJambeAVD_U.dot(HJambeAVD_a)) / HJambeAVD_D;
    HJambeAVD_a(AZ) += qdd(RF_HFE_JOINT);

    BJambeAVD_a = (motionTransforms-> fr_BJambeAVD_X_fr_HJambeAVD) * HJambeAVD_a + BJambeAVD_c;
    qdd(RF_KFE_JOINT) = (BJambeAVD_u - BJambeAVD_U.dot(BJambeAVD_a)) / BJambeAVD_D;
    BJambeAVD_a(AZ) += qdd(RF_KFE_JOINT);

    EpauleAVG_a = (motionTransforms-> fr_EpauleAVG_X_fr_base) * base_a + EpauleAVG_c;
    qdd(LF_HAA_JOINT) = (EpauleAVG_u - EpauleAVG_U.dot(EpauleAVG_a)) / EpauleAVG_D;
    EpauleAVG_a(AZ) += qdd(LF_HAA_JOINT);

    HJambeAVG_a = (motionTransforms-> fr_HJambeAVG_X_fr_EpauleAVG) * EpauleAVG_a + HJambeAVG_c;
    qdd(LF_HFE_JOINT) = (HJambeAVG_u - HJambeAVG_U.dot(HJambeAVG_a)) / HJambeAVG_D;
    HJambeAVG_a(AZ) += qdd(LF_HFE_JOINT);

    BJambeAVG_a = (motionTransforms-> fr_BJambeAVG_X_fr_HJambeAVG) * HJambeAVG_a + BJambeAVG_c;
    qdd(LF_KFE_JOINT) = (BJambeAVG_u - BJambeAVG_U.dot(BJambeAVG_a)) / BJambeAVG_D;
    BJambeAVG_a(AZ) += qdd(LF_KFE_JOINT);

    EpauleARD_a = (motionTransforms-> fr_EpauleARD_X_fr_base) * base_a + EpauleARD_c;
    qdd(RH_HAA_JOINT) = (EpauleARD_u - EpauleARD_U.dot(EpauleARD_a)) / EpauleARD_D;
    EpauleARD_a(AZ) += qdd(RH_HAA_JOINT);

    HJambeARD_a = (motionTransforms-> fr_HJambeARD_X_fr_EpauleARD) * EpauleARD_a + HJambeARD_c;
    qdd(RH_HFE_JOINT) = (HJambeARD_u - HJambeARD_U.dot(HJambeARD_a)) / HJambeARD_D;
    HJambeARD_a(AZ) += qdd(RH_HFE_JOINT);

    BJambeARD_a = (motionTransforms-> fr_BJambeARD_X_fr_HJambeARD) * HJambeARD_a + BJambeARD_c;
    qdd(RH_KFE_JOINT) = (BJambeARD_u - BJambeARD_U.dot(BJambeARD_a)) / BJambeARD_D;
    BJambeARD_a(AZ) += qdd(RH_KFE_JOINT);

    EpauleARG_a = (motionTransforms-> fr_EpauleARG_X_fr_base) * base_a + EpauleARG_c;
    qdd(LH_HAA_JOINT) = (EpauleARG_u - EpauleARG_U.dot(EpauleARG_a)) / EpauleARG_D;
    EpauleARG_a(AZ) += qdd(LH_HAA_JOINT);

    HJambeARG_a = (motionTransforms-> fr_HJambeARG_X_fr_EpauleARG) * EpauleARG_a + HJambeARG_c;
    qdd(LH_HFE_JOINT) = (HJambeARG_u - HJambeARG_U.dot(HJambeARG_a)) / HJambeARG_D;
    HJambeARG_a(AZ) += qdd(LH_HFE_JOINT);

    BJambeARG_a = (motionTransforms-> fr_BJambeARG_X_fr_HJambeARG) * HJambeARG_a + BJambeARG_c;
    qdd(LH_KFE_JOINT) = (BJambeARG_u - BJambeARG_U.dot(BJambeARG_a)) / BJambeARG_D;
    BJambeARG_a(AZ) += qdd(LH_KFE_JOINT);

    // + Add gravity to the acceleration of the floating base
    base_a += g;
}


#include <iostream>
#include <string>
#include <sstream>
#include <NTL/ZZ_pX.h>
#include <NTL/mat_ZZ_p.h>
#include <NTL/vec_ZZ_p.h>

using namespace NTL;
std::stringstream outstream;
std::string tStr;

static const char *ret_stream(void);
static char *prep_string(char *s);
extern "C" {
    const char *convert_matrix(char *in);
    const char *compute_laurent(char *t_yP, char *t_r, char *t_s);
    void clear_stringstream(void);
    void init(char *p);
}

void init(char *p) {
    ZZ_p::init(conv<ZZ>(p));
}

void clear_stringstream(void) {
    outstream.clear();
    outstream.str(std::string());
}

static const char *ret_stream(void) {
    tStr = std::move(outstream.str());
    for (auto it = tStr.begin(); it != tStr.end(); it++) {
        if ((*it == '\n') || (*it == ' ')) {
            *it = ',';
        }
    }
    return tStr.c_str();
}

static char *prep_string(char *s) {
    for (unsigned i = 0; s[i] != '\0'; i++) {
        if ((s[i] == ',') || (s[i] == 'L')) {
            s[i] = ' ';
        }
    }
    return s;
}

const char *convert_matrix(char *in) {
    clear_stringstream();

    Mat<ZZ_p> mat{conv<Mat<ZZ_p>>(prep_string(in))};
    outstream << mat;
    return ret_stream();
}

const char *compute_laurent(char *t_yP, char *t_r, char *t_s) {
    clear_stringstream();

    // convert
    Mat<ZZ_p> rT, rPT;
    unsigned n;
#ifdef EXTRATESTS
    Vec<ZZ_p> uv;
#endif
    {
        Vec<ZZ_p> yP{conv<Vec<ZZ_p>>(prep_string(t_yP))};
        Mat<ZZ_p> r{conv<Mat<ZZ_p>>(prep_string(t_r))};
        Mat<ZZ_p> s{conv<Mat<ZZ_p>>(prep_string(t_s)) * 2};

        // compute r' = r o yP + 2s
        Mat<ZZ_p> rP{r};
        for (unsigned i = 0; i < rP.NumRows(); i++) {
            for (unsigned j = 0; j < rP.NumCols(); j++) {
                rP[i][j] *= yP[j];
            }
            rP[i] += s[i];
        }
        rT = transpose(r);
        rPT = transpose(rP);
        n = yP.length();
#ifdef EXTRATESTS
        ZZ_p tmp;
        uv.SetLength(2*r.NumRows() - 1);
        for (unsigned ui = 0; ui < r.NumRows(); ui++) {
            for (unsigned vi = 0; vi < r.NumRows(); vi++) {
                InnerProduct(tmp, r[ui], rP[vi]);
                uv[ui + vi] += tmp;
            }
        }
#endif
    }

    // now compute r * r' using polymul
    ZZ_pX fg;
    {
        ZZ_pX f, g, tmp;
        for (unsigned j = 0; j < n; j++) {
            f.rep = std::move(rT[j]);
            g.rep = std::move(rPT[j]);
            tmp = f * g;
            fg += tmp;
        }
    }

#ifdef EXTRATESTS
    {
        bool correct = true;
        for (unsigned j = 0; j < uv.length(); j++) {
            correct &= uv[j] == fg[j];
        }
        if (!correct) {
            std::cerr << "ERROR: mismatched results\n" << uv << std::endl << fg << std::endl;
        }
    }
#endif

    outstream << fg;
    return ret_stream();
}


#include "drake/solvers/sos_basis_generator.h"

#include <vector>

#include <Eigen/Core>

#include "drake/solvers/integer_inequality_solver.h"
namespace drake {
namespace solvers {
namespace {
// Anonymous namespace containing collection of utility functions
using Variable = symbolic::Variable;
using Monomial = symbolic::Monomial;
using Variables = symbolic::Variables;
using MonomialVector = VectorX<symbolic::Monomial>;
using Exponent = Eigen::RowVectorXi;
using ExponentList = Eigen::Matrix<int, -1, -1, Eigen::RowMajor>;

// Given a list of exponents and variables, returns a vector of monomials.
// Ex: if exponents = [0, 1;1, 2], and vars = [x(0), x(1)], then the vector
// [x(1); x(0)* x(1)²] is returned.
MonomialVector ExponentsToMonomials(const ExponentList& exponents,
                                    const drake::VectorX<Variable>& vars) {
  MonomialVector monomials(exponents.rows());
  for (int i = 0; i < exponents.rows(); i++) {
    monomials(i) = Monomial(vars, exponents.row(i));
  }
  return monomials;
}

// Returns a list of all exponents that appear in a polynomial p.
// E.g., given p = 1 + 2x₀² + 3x₀*x₁², returns [0, 0; 2, 0; 1, 2];
ExponentList GetPolynomialExponents(const drake::symbolic::Polynomial& p) {
  const Variables& indeterminates{p.indeterminates()};
  ExponentList exponents(p.monomial_to_coefficient_map().size(),
                         indeterminates.size());
  int row = 0;
  for (const auto& m : p.monomial_to_coefficient_map()) {
    int col = 0;
    for (const auto& var : indeterminates) {
      exponents(row, col++) = m.first.degree(var);
    }
    row++;
  }
  return exponents;
}

ExponentList VerticalStack(const ExponentList& A, const ExponentList& B) {
  DRAKE_ASSERT(A.cols() == B.cols());
  if (A.rows() == 0) {
    return B;
  }
  if (B.rows() == 0) {
    return A;
  }
  ExponentList Y(A.rows() + B.rows(), B.cols());
  Y << A, B;
  return Y;
}

ExponentList PairwiseSums(const ExponentList& exponents) {
  int n = exponents.rows();
  ExponentList sums((n * n - n) / 2, exponents.cols());
  int cnt = 0;
  for (int i = 0; i < n; i++) {
    // Note: counter starts at i+1 to omit a+b when a = b.
    for (int j = i + 1; j < n; j++) {
      sums.row(cnt++) = exponents.row(i) + exponents.row(j);
    }
  }
  return sums;
}

// Returns true if the first num_rows of A contains B.
bool ContainsExponent(const ExponentList& A, int num_rows, const Exponent& B) {
  DRAKE_ASSERT(B.rows() == 1 && B.cols() == A.cols());
  DRAKE_ASSERT((num_rows >= 0) && (num_rows <= A.rows()));
  for (int i = 0; i < num_rows; i++) {
    if (A.row(i) == B) {
      return true;
    }
  }
  return false;
}

/* Intersection(A, B) removes duplicate rows from B and any row that doesn't
 * also appear in A.  For example, given A = [1, 0; 0, 1; 1, 1] and B = [1, 0;
 * 1, 1; 1, 1;], it overwrites B with [1, 0; 1, 1]. */
void Intersection(const ExponentList& A, ExponentList* B) {
  DRAKE_ASSERT(A.cols() == B->cols());
  int index = 0;
  for (int i = 0; i < B->rows(); i++) {
    if ((ContainsExponent(A, A.rows(), B->row(i))) &&
        !(ContainsExponent(*B, index, B->row(i)))) {
      B->row(index++) = B->row(i);
    }
  }
  B->conservativeResize(index, Eigen::NoChange);
}

/* Removes exponents of the monomials that aren't diagonally-consistent with
 * respect to the polynomial p and the given monomial basis.  A monomial is
 * diagonally-consistent if its square appears in p, or its square equals a
 * product of monomials in the basis; see, e.g., "Pre- and Post-Processing
 * Sum-of-Squares Programs in Practice Johan Löfberg, IEEE Transactions on
 * Automatic Control, 2009." After execution, all exponents of inconsistent
 * monomials are removed from exponents_of_basis.
*/
void RemoveDiagonallyInconsistentExponents(const ExponentList& exponents_of_p,
                                           ExponentList* exponents_of_basis) {
  while (1) {
    int num_exponents = exponents_of_basis->rows();

    ExponentList valid_squares =
        VerticalStack(PairwiseSums(*exponents_of_basis), exponents_of_p);

    (*exponents_of_basis) = (*exponents_of_basis) * 2;
    Intersection(valid_squares, exponents_of_basis);
    (*exponents_of_basis) = (*exponents_of_basis) / 2;

    if (exponents_of_basis->rows() == num_exponents) {
      break;
    }
  }
  return;
}

struct Hyperplanes {
  Eigen::MatrixXi normal_vectors;  // Each row contains a normal vector.
  Eigen::VectorXi max_dot_product;
  Eigen::VectorXi min_dot_product;
};

// Finding random supporting hyperplanes of 1/2 P, where P is the Newton
// polytope of the polynomial p (i.e., the convex hull of its exponents).
Hyperplanes RandomSupportingHyperplanes(const ExponentList& exponents_of_p,
                                        unsigned int seed) {
  Hyperplanes H;

  // get_random() samples uniformly between normal_vector_component_min/max.
  // Current values of min and max set heuristically.
  const int normal_vector_component_min = -10;
  const int normal_vector_component_max = 10;
  std::default_random_engine generator(seed);
  std::uniform_int_distribution<int> distribution(normal_vector_component_min,
                                                  normal_vector_component_max);
  auto get_random = [&]() { return distribution(generator); };

  // Number of hyperplanes currently picked heuristically.
  int num_hyperplanes = 10 * exponents_of_p.cols();

  H.normal_vectors = Eigen::MatrixXi(num_hyperplanes, exponents_of_p.cols());
  for (int i = 0; i < H.normal_vectors.cols(); i++) {
    H.normal_vectors.col(i)
        << Eigen::VectorXi::NullaryExpr(num_hyperplanes, get_random);
  }

  Eigen::MatrixXi dot_products = H.normal_vectors * exponents_of_p.transpose();
  H.max_dot_product = dot_products.rowwise().maxCoeff() / 2;
  H.min_dot_product = dot_products.rowwise().minCoeff() / 2;

  return H;
}

//  Generates the supporting hyperplanes of the Newton polytope that
//  are induced by the total degree ordering.
Hyperplanes DegreeInducedHyperplanes(const ExponentList& exponents_of_p) {
  Hyperplanes H;

  // The hyperplane for total degree.
  H.normal_vectors.resize(1, exponents_of_p.cols());
  H.normal_vectors.setConstant(1);

  Eigen::MatrixXi dot_products = H.normal_vectors * exponents_of_p.transpose();
  H.max_dot_product = dot_products.rowwise().maxCoeff() / 2;
  H.min_dot_product = dot_products.rowwise().minCoeff() / 2;

  return H;
}

ExponentList EnumerateInitialSet(const ExponentList& exponents_of_p) {
  Eigen::VectorXi lower_bounds = exponents_of_p.colwise().minCoeff() / 2;
  Eigen::VectorXi upper_bounds = exponents_of_p.colwise().maxCoeff() / 2;
  Hyperplanes hyperplanes = DegreeInducedHyperplanes(exponents_of_p);

  // We check the inequalities in two batches to allow for internal
  // infeasibility propagation inside of EnumerateIntegerSolutions,
  // which is done only if A has a column that is elementwise nonnegative
  // or nonpositive. (This condition never holds if we check the
  // inequalities in one batch, since A = [normal_vectors;-normal_vectors].)
  ExponentList basis_exponents_1 = drake::solvers::EnumerateIntegerSolutions(
      hyperplanes.normal_vectors, hyperplanes.max_dot_product, lower_bounds,
      upper_bounds);

  ExponentList basis_exponents = drake::solvers::EnumerateIntegerSolutions(
      -hyperplanes.normal_vectors, -hyperplanes.min_dot_product, lower_bounds,
      upper_bounds);

  Intersection(basis_exponents_1, &basis_exponents);
  return basis_exponents;
}

//  This function removes an element alpha from "basis" if a randomly generated
//  hyperplane separates 2*alpha from the Newton polytope of the polynomial p.
//  Note that this function is actually deterministic since the seed for
//  the random number generator is set to predetermined constants.
void RemoveWithRandomSeparatingHyperplanes(const ExponentList& exponents_of_p,
                                           ExponentList* basis) {
  // Declare this outside the main loop to avoid repeated dynamic memory
  // allocation.
  Eigen::MatrixXi dot_products;
  int random_seed = 0;

  while (1) {
    int next_basis_size = 0;
    int current_basis_size = basis->rows();

    auto H = RandomSupportingHyperplanes(exponents_of_p, random_seed++);

    // Remove monomials that the hyperplanes separate from the
    // Newton polytope.
    dot_products = (*basis) * H.normal_vectors.transpose();
    for (int i = 0; i < current_basis_size; i++) {
      bool keep_monomial = true;
      for (int j = 0; j < dot_products.cols(); j++) {
        if (dot_products(i, j) > H.max_dot_product(j) ||
            dot_products(i, j) < H.min_dot_product(j)) {
          keep_monomial = false;
          break;
        }
      }

      if (keep_monomial) {
        basis->row(next_basis_size++) = basis->row(i);
      }
    }

    basis->conservativeResize(next_basis_size, basis->cols());

    // Quit if the basis is now empty or if its size wasn't reduced
    // enough.
    constexpr double kMinimumPercentReduction = .1;
    if (next_basis_size >
            current_basis_size * (1.0 - kMinimumPercentReduction) ||
        next_basis_size == 0) {
      break;
    }
  }
  return;
}

ExponentList ConstructMonomialBasis(const ExponentList& exponents_of_p) {
  auto basis_exponents = EnumerateInitialSet(exponents_of_p);
  RemoveWithRandomSeparatingHyperplanes(exponents_of_p, &basis_exponents);
  RemoveDiagonallyInconsistentExponents(exponents_of_p, &basis_exponents);
  return basis_exponents;
}

}  // namespace

MonomialVector ConstructMonomialBasis(const drake::symbolic::Polynomial& p) {
  const Variables& indeterminates{p.indeterminates()};
  drake::VectorX<Variable> vars(indeterminates.size());
  int cnt = 0;
  for (auto& var : indeterminates) {
    vars(cnt++) = var;
  }

  auto polynomial_exponents = GetPolynomialExponents(p);
  auto basis_exponents = ConstructMonomialBasis(polynomial_exponents);
  auto monomial_basis = ExponentsToMonomials(basis_exponents, vars);
  return monomial_basis;
}
}  // namespace solvers
}  // namespace drake


#include "needleman.h"

#include "ereal.h"
#include <cassert>
#include <algorithm>
#include <blitz/tinyvec-et.h>
#include <btl/logspace.h>

// We assume the BLASTN default scoring scheme with Juke's/Cantor background. All logs are in base sqrt(2).
// Defaults          Adjusted to remove/include the odds portion.
// Match:       1   4
// Mismatch:   -3   -12
// Gap open:   -5   -20
// Gap extend: -2   -8
void needleman::init_pr() {
   // match/mismatch probabilities
   // pno_open should be rolled in here, but it's value is -0.006 which is
   // barely representable by the precisian we use here. Ignore for now.
   s =   4.0, 0.015625, 0.015625, 0.015625,
         0.015625, 4.0, 0.015625, 0.015625,
         0.015625, 0.015625, 4.0, 0.015625,
         0.015625, 0.015625, 0.015625, 4.0;

   // gap open
   popen = 0.001043033;

   // gap extend
   pext = 0.0625;
   pno_ext = 1.0-0.0625;
}

void needleman::set_parameters( nw_model_parameters const &mp ){
	set_mean_gap_length( mp.mean_gap_length );
	set_s( mp.pr_open, mp.p, mp.q );
}

// set mean gap length to n
void needleman::set_mean_gap_length( double n ) {
   // Since a gap is length 1 by the virtue of opening it, we adjust n down by 1.
   // The remaining gap extension probability corresponds to the mean of a
   // geometric distribution allowing zero extensions.
   pext =  1.0-1.0/n;
   pno_ext =  1.0/n;
	}

// computes score matrix from match probabilities and background
void needleman::set_s( double new_popen, blitz::Array<double,2> const &p, blitz::TinyVector<double,4> const &q ) {
   using namespace std;
   assert( p.rows() == 4 );
   assert( p.cols() == 4 );

   popen = new_popen;
	double pno_open = 1.0-2*new_popen;
   for( int i = 0; i < 4; i++ ) for( int j = 0; j < 4; j++ ) s(i,j) = pno_open*p(i,j)/q(i)/q(j);
}

void needleman::maindp_init_borders( dparray &dp, dna_sequence_region_data &a, dna_sequence_region_data &b ) {
   int n = (int) a.size(), m = (int)b.size();
   assert( (n>=1) && (m>=1) );

   // top corner is the only valid starting point
   init_cell(dp(0,0).M); init_cell(dp(0,0).Ga); init_cell(dp(0,0).Gb);
   dp(0,0).M.s = 1.0;

   // Initialize borders. Scores represent opening a big gap.
   init_cell(dp(1,0).M); init_cell(dp(1,0).Ga); init_cell(dp(1,0).Gb);
   dp(1,0).Gb.from = M_state; dp(1,0).Gb.s = popen*pno_ext;
   for( int i = 2; i <= n; i++ ) {
      init_cell(dp(i,0).M); init_cell(dp(i,0).Ga); init_cell(dp(i,0).Gb);
      dp(i,0).Gb.from = Gb_state; dp(i,0).Gb.s = dp(i-1,0).Gb.s*pext;
   }
   init_cell(dp(0,1).M); init_cell(dp(0,1).Ga); init_cell(dp(0,1).Gb);
   dp(0,1).Ga.from = M_state; dp(0,1).Ga.s = popen*pno_ext;
   for( int j = 2; j <= m; j++ ) {
      init_cell(dp(0,j).M); init_cell(dp(0,j).Ga); init_cell(dp(0,j).Gb);
      dp(0,j).Ga.from = Ga_state; dp(0,j).Ga.s = dp(0,j-1).Ga.s*pext;
   }
}

void needleman::maindp_filldp( dparray &dp, dna_sequence_region_data &a, dna_sequence_region_data &b ) {
   int n = (int) a.size(), m = (int)b.size();
	ereal score, match_cost;
   for( int i = 1; i <= n; i++ ) {
   for( int j = 1; j <= m; j++ ) {
      // first consider a match, default is come from match
      match_cost = s((int)a[(uint)i-1],(int)b[(uint)j-1]);
      dp(i,j).M.s = dp(i-1,j-1).M.s*match_cost;
      dp(i,j).M.from = M_state;

      // consider coming from a gap in A
      score = dp(i-1,j-1).Ga.s*match_cost;
      if( score > dp(i,j).M.s ) {
         dp(i,j).M.s = score;
         dp(i,j).M.from = Ga_state;
      }

      // consider coming from a gap in B
      score = dp(i-1,j-1).Gb.s*match_cost;
      if( score > dp(i,j).M.s ) {
         dp(i,j).M.s = score;
         dp(i,j).M.from = Gb_state;
      }

      // now consider extending or opening a gap in A, default is open a gap
      dp(i,j).Ga.s = dp(i,j-1).M.s*popen*pno_ext;
      dp(i,j).Ga.from = M_state;

      // extend a gap
      score = dp(i,j-1).Ga.s*pext;
      if( score > dp(i,j).Ga.s ) {
         dp(i,j).Ga.s = score;
         dp(i,j).Ga.from = Ga_state;
      }

      // now consider extending or opening a gap in B, default is open a gap
      dp(i,j).Gb.s = dp(i-1,j).M.s*popen*pno_ext;
      dp(i,j).Gb.from = M_state;

      // extend a gap
      score = dp(i-1,j).Gb.s*pext;
      if( score > dp(i,j).Gb.s ) {
         dp(i,j).Gb.s = score;
         dp(i,j).Gb.from = Gb_state;
      }
   } // for j
   } // for i
}

pairwise_dna_alignment needleman::maindp_traceback( dparray &dp, dna_sequence_region_data &a, dna_sequence_region_data &b ) {
   using namespace std;
   int n = (int) a.size(), m = (int)b.size();

   // alignment sequences
   dna_alignment_sequence_ptr newa = new_dna_alignment_sequence(), newb = new_dna_alignment_sequence();

   // determine starting block
   int i = n, j = m, next_i = n, next_j = m;
   uint8_t type = M_state, next_type = no_state;
   if( state_score(dp,i,j,type) < dp(i,j).Ga.s ) type = Ga_state;
   if( state_score(dp,i,j,type) < dp(i,j).Gb.s ) type = Gb_state;
   ereal alignment_score = state_score(dp,i,j,type);

   while( (i > 0) || (j > 0) ) {
      assert( type != no_state );
      assert( type < num_states );
      switch( type ) {
         case M_state:
            next_i = i - 1;
            next_j = j - 1;
            next_type = dp(i,j).M.from;
            newa->data.push_front( a[(uint)i-1] );
            newb->data.push_front( b[(uint)j-1] );
            break;
         case Ga_state:
            next_i = i;
            next_j = j - 1;
            next_type = dp(i,j).Ga.from;
            newa->data.push_front( dna_alignment_alpha::GAP );
            newb->data.push_front( b[(uint)j-1] );
            break;
         case Gb_state:
            next_i = i - 1;
            next_j = j;
            next_type = dp(i,j).Gb.from;
            newa->data.push_front( a[(uint)i-1] );
            newb->data.push_front( dna_alignment_alpha::GAP );
      }
      i = next_i;
      j = next_j;
      type = next_type;
   }
   return pairwise_dna_alignment( newa, newb, alignment_score.as_base() );
}

pairwise_dna_alignment needleman::maindp( dna_sequence_region_data &a, dna_sequence_region_data &b ) {
   int n = (int) a.size(), m = (int)b.size();
   dparray dp(n+1,m+1);
   maindp_init_borders( dp, a, b );
   maindp_filldp( dp, a, b );
   return maindp_traceback( dp, a, b );
}

// only need to initialize score scheme
needleman::needleman() {
   s.resize(dna_alpha::SIZE,dna_alpha::SIZE);
   init_pr();
}

pairwise_dna_alignment needleman::align( dna_sequence_region &seqa, dna_sequence_region &seqb ) {
   return maindp( seqa.data, seqb.data );
}

std::pair<nw_model_parameters_ptr,ereal > needleman::estimate( dna_sequence_region &seqa, dna_sequence_region &seqb ) {
	using namespace std;
	pairwise_dna_alignment aln = align(seqa,seqb);
	counts cnt;
	do_statistics(aln,cnt);
	nw_model_parameters_ptr nwp = update_probabilities(cnt);
	return make_pair(nwp,aln.score);
}

void needleman::do_statistics( pairwise_dna_alignment &align, needleman::counts &cnt ) {
	assert( align.a->data.size() == align.b->data.size() );

	int gap_a = -1, gap_b = -1;
	for( int c = 0; c < (int)align.a->data.size(); ++c ) {
		dna_alignment_alpha::symbol c_a = align.a->data[c], c_b = align.b->data[c];

		// found a pair!
		if( c_a != dna_alignment_alpha::GAP && c_b != dna_alignment_alpha::GAP ) {
			int base_a = dna_alpha::symbol(c_a).index();
			int base_b = dna_alpha::symbol(c_b).index();
			cnt.pairs( base_a, base_b) += 1;
			cnt.bases_a(base_a) += 1;
			cnt.bases_b(base_b) += 1;

			// closing a gap in A
			if( gap_a >= 0 ) {
				cnt.gap_lengths_a += c - gap_a + 1;
				gap_a = -1;
			}

			// closing a gap in B
			if( gap_b >= 0 ) {
				cnt.gap_lengths_b += c - gap_b + 1;
				gap_b = -1;
			}

			// found a gap in A
		} else if( c_a == dna_alignment_alpha::GAP && c_b != dna_alignment_alpha::GAP ) {
			cnt.bases_b((int)dna_alpha::symbol(c_b)) += 1;

			// gap is not open
			if( gap_a == -1 ) {
				gap_a = c;
				cnt.gaps_a += 1;
			} // otherwise nothing to do

			// closing a gap in B
			if( gap_b >= 0 ) {
				cnt.gap_lengths_b += c - gap_b + 1;
				gap_b = -1;
			}
			// found a gap in B
		} else if( c_a != dna_alignment_alpha::GAP && c_b == dna_alignment_alpha::GAP ) {
			cnt.bases_a((int)dna_alpha::symbol(c_a)) += 1;

			// closing a gap in A
			if( gap_a >= 0 ) {
				cnt.gap_lengths_a += c - gap_a + 1;
				gap_a = -1;
			}

			// gap is not open
			if( gap_b == -1 ) {
				gap_b = c;
				cnt.gaps_b += 1;
			} // otherwise nothing to do
		} else {
			using namespace std;
			cerr << c << "\t" << c_a << "\t" << c_b << endl;
			assert( 0 && "corrupt alignment" );
		}

	} // for c

	// finished region, count any last gaps
	// closing a gap in A
	if( gap_a >= 0 ) {
		cnt.gap_lengths_a += align.a->data.size() - gap_a;
		gap_a = -1;
	}

	// closing a gap in B
	if( gap_b >= 0 ) {
		cnt.gap_lengths_b += align.a->data.size() - gap_b;
		gap_b = -1;
	}

}

nw_model_parameters_ptr needleman::update_probabilities( needleman::counts &cnt ) {
	using namespace std;
	using namespace blitz;

	// get combined base count for both sequences
	base_count all_bases;
	all_bases = cnt.bases_a + cnt.bases_b;
	int total_bases = sum( all_bases );

	nw_model_parameters_ptr newp( new nw_model_parameters() );

	// compute p
	int total_pairs = sum( cnt.pairs(Range::all(),Range::all()) );
	for( int i = 0; i < dna_alpha::SIZE; ++i ) {
		for( int j = 0; j < dna_alpha::SIZE; ++j ) {
			newp->p(i,j) = (double) cnt.pairs(i,j)/(double)total_pairs;
		} // j
	} // i

	for( int i = 0; i < dna_alpha::SIZE; ++i ) {
		newp->q(i) = (double) all_bases(i)/(double)total_bases;
	}

	// gap open
	int total_gaps = cnt.gaps_a + cnt.gaps_b;
	newp->pr_open = (double)total_gaps/(double)total_bases;

	// mean gap length
	newp->mean_gap_length = (double)(cnt.gap_lengths_a+cnt.gap_lengths_b)/(double)total_gaps;
	return newp;
}


// Copyright 2018 The Simons Foundation, Inc. - All Rights Reserved.
//
// 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.

#ifndef NETKET_RMSPROP_HPP
#define NETKET_RMSPROP_HPP

#include <Eigen/Core>
#include <Eigen/Dense>
#include <cassert>
#include <cmath>
#include <complex>
#include <iostream>
#include "abstract_optimizer.hpp"

namespace netket {

class RMSProp : public AbstractOptimizer {
  int npar_;

  double eta_;
  double beta_;

  Eigen::VectorXd st_;

  double epscut_;

  const Complex I_;

 public:
  explicit RMSProp(double eta = 0.001, double beta = 0.9,
                   double epscut = 1.0e-7)
      : eta_(eta), beta_(beta), epscut_(epscut), I_(0, 1) {
    npar_ = -1;
    PrintParameters();
  }

  // TODO remove
  // Json constructor
  explicit RMSProp(const json &pars) : I_(0, 1) {
    npar_ = -1;

    from_json(pars);
    PrintParameters();
  }

  void PrintParameters() {
    InfoMessage() << "RMSProp optimizer initialized with these parameters :"
                  << std::endl;
    InfoMessage() << "Learning Rate = " << eta_ << std::endl;
    InfoMessage() << "Beta = " << beta_ << std::endl;
    InfoMessage() << "Epscut = " << epscut_ << std::endl;
  }

  void Init(const Eigen::VectorXd &pars) override {
    npar_ = pars.size();
    st_.setZero(npar_);
  }

  void Init(const Eigen::VectorXcd &pars) override {
    npar_ = 2 * pars.size();
    st_.setZero(npar_);
  }

  void Update(const Eigen::VectorXd &grad, Eigen::VectorXd &pars) override {
    assert(npar_ > 0);

    st_ = beta_ * st_ + (1. - beta_) * grad.cwiseAbs2();

    for (int i = 0; i < npar_; i++) {
      pars(i) -= eta_ * grad(i) / (std::sqrt(st_(i)) + epscut_);
    }
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXd &pars) override {
    Update(Eigen::VectorXd(grad.real()), pars);
  }

  void Update(const Eigen::VectorXcd &grad, Eigen::VectorXcd &pars) override {
    assert(npar_ == 2 * pars.size());

    for (int i = 0; i < pars.size(); i++) {
      st_(2 * i) =
          beta_ * st_(2 * i) + (1. - beta_) * std::pow(grad(i).real(), 2);
      st_(2 * i + 1) =
          beta_ * st_(2 * i + 1) + (1. - beta_) * std::pow(grad(i).imag(), 2);
      pars(i) -= eta_ * grad(i).real() / (std::sqrt(st_(2 * i)) + epscut_);
      pars(i) -=
          eta_ * I_ * grad(i).imag() / (std::sqrt(st_(2 * i + 1)) + epscut_);
    }
  }

  void Reset() override { st_ = Eigen::VectorXd::Zero(npar_); }

  // TODO remove
  void from_json(const json &pars) {
    // DEPRECATED (to remove for v2.0.0)
    std::string section = "Optimizer";
    if (!FieldExists(pars, section)) {
      section = "Learning";
    }
    eta_ = FieldOrDefaultVal(pars[section], "LearningRate", 0.001);
    beta_ = FieldOrDefaultVal(pars[section], "Beta", 0.9);
    epscut_ = FieldOrDefaultVal(pars[section], "Epscut", 1.0e-7);
  }
};

}  // namespace netket

#endif


#ifndef __SPARSEINDEXSET__CLASS__
#define __SPARSEINDEXSET__CLASS__

// Eigen Library
#include <Eigen/Dense>

#include "CONSTANTS.hpp"

class SparseIndexSet {
 public:
  /**     \brief initializes the sparse anisotropic index set
  *       \param[in] q is the maximum level for the sparse index set
  *       \param[in] dim is the dimension of the sparse index set
  *       \param[in] cpFun is the comparison Function. sparse index set is
  *                  constructed such that cpFun(alpha) <= q for all alpha in
  *                  the sparse index set
  *
  * 	    Additionally, the set of multiindices _alpha and the weights _cw
  *     	 are provided globally.
  *
  * 	    The function comp_indexSet computes a sparse multi index set
  * 	    _alpha\subset\N^n by the function call
  *       sparseIndexSet::Yw_alpha(0, &n, cpFun)
  * 	    and corresponding weights _cw \in \Z^n to each multi index.
  * 	    To determine the weights, we have to check for each multi index
  *       alpha contained in _alpha if alpha + beta is in _alpha for each beta
  *       in {0,1}^{n}. If this is the case, we update the weight by
  *       (-1)^|beta|. Of course, the condition is always fulfilled for the
  *       multiindex beta = 0. Therefore, we skip this multiindex and initialize
  *       the weight cw with 1 which is associated with the second variable in
  *       the function call to cw_alpha.
  * 	    For a multi index alpha, the corresponding weight has only to be
  *       computed if the multiindex alpha+1 is not contained in _alpha.
  *       Otherwise, the corresponding weight is 0. With the index set, the
  *       weights and the unidirectional quadrature formulas at hand, the
  *       integral I(f) is approximated by the combination technique formula
  * 	    I(f) \approx \sum_{i =0}^{_alpha.cols-1} _cw(i)*Q_{alpha.col(i)}.
  * 	    Herein, _alpha.cols determines the number of multiindices in _alpha
  *       and alpha.col(i) gives the i-th index in _alpha. The tensor product
  *       quadrature formula Q_{alpha.col(i)} is specified by the unidirectional
  *       quadrature formulas.
  */
  template <class Compare>
  void computeIndexSet(int q, int dim, const Compare &cpFun) {
    int k = 0;
    Eigen::VectorXi currInd(dim);
    // set max level for the sparse index set
    _q = q;
    // set dimension of the sparse grid
    _dim = dim;
    // allocate memory for _alpha and _cw to keep the overhead at a low level
    // blocks of size __MEMCHUNKSIZE__ are allocated at once
    _alpha.resize(_dim, __MEMCHUNKSIZE__);
    _cw.resize(__MEMCHUNKSIZE__);
    // set memory to 0
    _alpha.setZero();
    _cw.setZero();
    // get a vector with all ones to reduce the overhead for the constructor of
    // Eigen::VectorXi::Ones(_dim);
    _myOnes = Eigen::VectorXi::Ones(_dim);
    // start from multi index 0\in Xw_alpha
    currInd.setZero();
    k = 0;
    // test if 0 is in Xw_alpha, else index set empty due to downward closedness
    if (cpFun(currInd) <= _q) {
      // check if _cw(0)\neq 0
      if (cpFun(_myOnes) > _q) _cw(0) = combiWeights(0, 1, 1, currInd, cpFun);
      if (_cw(0)) ++k;
      // compute all other indices in Yw_alpha recursively
      combiIndexSet(0, &k, cpFun, currInd);
    }
    // crop memory for _alpha and _cw to actual size
    _alpha.conservativeResize(_dim, k);
    _cw.conservativeResize(k);
  }

  /**   \brief make an educated guess, what this function does...
  *
  */
  const Eigen::MatrixXi &get_alpha(void) const { return _alpha; };

  /**   \brief make an educated guess, what this function does...
  *
  */
  const Eigen::VectorXi &get_cw(void) const { return _cw; };

 protected:
  /**     \brief computes indices in weighted sparse grid space (recursive)
  *	    \param[in] maxBit position of the digit of the current multiindex
  *                  such that only digits with position greater or equal to
  *                  maxBit are considered in the recursion.
  *	 	 \param[in] k counter for the number of multiindices is passed
  *                  as pointer
  *	    \param[in] cpFun function/functor/whatever which decides whether a
  *                  multiindex is in the indexset or not.
  *
  *	    Additionally, the level of the sparse grid _q, the dimension of the
  *	    multiindices _dim and the memory _alpha for the indexset is
  *	    provided globally and referred to.
  *
  *	    The function Yw_alpha recursively computes the indexset of a
  *	    generalized sparse grid which is described by a level _q and a
  *	    function cpFun, i.e. it finds all multiindices alpha such that
  *	    cpFun(alpha)<= _q. To that end, we travel the set Xw_alpha and
  *       at a multi index only if the corresponding _cw is non zero.
  *	    The function needs to be called with input parameters
  *	    maxBit=0, k=0 defined as a reference, predefined
  *	    function cpFun which maps from \mathbb{N}^{_dim}\to \mathbb{R}.
  *	    Since cpFun(0)=0, the zero multiindex is always the first index
  *	    which is included in the sparse grid on default and serves as the
  *	    root node in the tree structure of the index set.
  *	    On level l of the tree, all indices alpha with
  *	    cpFun(alpha)<= _q and |alpha| = l-1 are computed.
  *	    Moreover, the tree is structured such that all indices alpha' which
  *	    are sons of a multiindex alpha fulfill alpha'>=alpha elementwise.
  *	    This structure immediately allows to exploit the downward closeness
  *	    of the indexset since we do not need a further recursion level
  *	    at an edge where the corresponding index does not belong to the
  *       indexset.
  *	    Moreover, the integer maxBit is necessary to avoid multiple
  *       repetition of multiindices. In every step maxBit is the position of
  *       the digit which is modified to obtain the current multiindex. In the
  *       subtree corresponding to this index only digits greater or equal to
  *       maxBit will be considered for modification. A simple example for the
  *       tree structure for _q=2, _dim=2 and cpFun(alpha) = alpha_1+alpha_2 is
  *       depicted below.
  *
  *								(0,0;mB=0;*k=0)
  *								/            \
  *						(1,0;mB=0;*k=1) (0,1;mB=1;*k=4)
  *						/		      \               \
  *			(2,0;mB=0;*k=2) (1,1;mB=1;*k=3)	(0,2;mB=1;*k=5)
  *
  */
  template <class Compare>
  void combiIndexSet(int maxBit, int *k, const Compare &cpFun,
                     Eigen::VectorXi &currInd) {
    int cw = 0;
    // successively increase all entries in the current index
    for (int i = maxBit; i < _dim; ++i) {
      ++currInd(i);
      if (cpFun(currInd) <= _q) {
        // if the index is in Xw_alpha, check if its weight is non zero
        if (cpFun(currInd + _myOnes) > _q)
          cw = combiWeights(0, 1, 1, currInd, cpFun);
        else
          cw = 0;
        // if its weight is non zero, add it to the index set
        if (cw) {
          if (_alpha.cols() <= *k) {
            _alpha.conservativeResize(_dim, _alpha.cols() + __MEMCHUNKSIZE__);
            _cw.conservativeResize(_cw.size() + __MEMCHUNKSIZE__);
          }
          _alpha.col(*k) = currInd;
          _cw(*k) = cw;
          ++(*k);
        }
        // check son indices only if father index is in Xw. this is possible
        // due to the downward closedness assumption
        combiIndexSet(i, k, cpFun, currInd);
      }
      --currInd(i);
    }
  }

  /**     \brief computes weights for the tensor product quadrature (recursive)
  *	    \param[in] maxBit position of the digit of the current multiindex
  *                  such that only digits with position greater to maxBit are
  *	 		         considered in the recursion.
  *       \param[in] cw is initialize by 1 due to alpha+0\in Xw if alpha in Xw
  *	    \param[in] lvl current level-1 in the recursion tree.
  *	    \param[in] ind current multi index alpha
  *	    \param[in] cpFun function/functor/whatever which decides whether a
  *                  multi index is in the indexset or not.
  *
  *	    Additionally, the level of the sparse grid _q, the dimension of the
  *	    multiindices _dim and the memory _alpha for the indexset is
  *	    provided globally.
  *
  *	    The function cw_alpha computes the coefficient of the indth
  *	    multi index in the combination technique formula corresponding to
  *	    the generalized sparse quadrature. This coefficient is given by
  *	    cw(alpha) = sum_{beta in {0,1}^_dim, alpha+beta in _alpha}
  *       (-1)^|beta|. Hence, in order to determine the value of the
  *       coefficient, we have to check for each beta in {0,1}^_dim if
  *       alpha+beta belongs to the index set _alpha. For beta=(0,...,0),
  *       alpha+beta belongs to _alpha and, hence, we
  *       initialize cw=1 and do not have to consider this multiindex.
  *	    The downward closeness of the index set can be exploited once more
  *	    which results in a similar tree structure as in the function
  *       Yw_alpha. The only difference of the tree is that once a digit of the
  *       multiindex alpha is modified, we do not have to consider further
  *       modification of this digit due to the condition beta in {0,1}^_dim.
  *       This is realized in the recursive algorithm consider only digits which
  *       are greater than maxBit in contrast to the "greater than and equal"
  *       condition in Yw_alpha.
  *	    In the recursive algorithm, the multiindices beta do not need to be
  *       computed explicitly, but the multiindex alpha is modified at each node
  *       of the tree. Of course, this modification has to be revoked when the
  *       recursion goes back to the parent node in the tree.
  *	    The value of lvl=|beta| is important since it determines whether the
  *	    coefficient is modified by +1 or -1.
  *	    A simple example of the algorithm for _q=2, _dim=2, ind = (1,0) and
  *	    cpFun(alpha) = alpha_1+alpha_2 is depicted below.
  *	    The further initial values are maxBit=mB=0, cw = 1, lvl=0.
  *
  *								(1,0;mB=0;lvl=0;cw=1)
  *								/                   \
  *					(2,0;mB=0;lvl=1;cw=0)   (1,1;mB=1;lvl=1;cw=-1)
  *							/
  *			[ (2,1;mB=1;lvl=2;cw=0) ] <-- is not contained in the
  *                                   indexset, no modification of cw
  *
  *	    The algorithm returns the final value of cw. In the example above
  *       this would be the correct value cw(1,0)=-1.
  */
  template <class Compare>
  int combiWeights(int maxBit, int cw, int lvl, Eigen::VectorXi &ind,
                   const Compare &cpFun) {
    // successively check all bits of the multiindex is contained
    for (int i = maxBit; i < _dim; ++i) {
      ++ind(i);
      if (cpFun(ind) <= _q) {
        if (lvl % 2)
          --cw;
        else
          ++cw;
        // again, exploit downward closedness and perform recursion only
        // if father is contained in Xw_alpha
        cw = combiWeights(i + 1, cw, lvl + 1, ind, cpFun);
      }
      --ind(i);
    }

    return cw;
  }

  // Member variables
  Eigen::MatrixXi _alpha;
  Eigen::VectorXi _cw;
  Eigen::VectorXi _myOnes;
  int _dim;
  int _q;
};
#endif