#include <carma>
#include <armadillo>

#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <pybind11/numpy.h>
#include <cmath> // fabs
#include <chrono>
#include <limits>
#include "svd_3x3.h" // fast 3x3 svd
#include "carma_svd.h" // fast 3x3 svd using carma

using namespace std::chrono;
using std::vector;

using namespace pybind11::literals;
namespace py = pybind11;

// Compile command: g++ -O3 -Wall -shared -std=c++17 -fPIC -Wl,-undefined,dynamic_lookup $(python3 -m pybind11 --includes) example.cpp -o example$(python3-config --extension-suffix)

constexpr auto rank_comb2(size_t i, size_t j, size_t n) noexcept -> size_t {
  if (j < i){ std::swap(i,j); }
  return(size_t(n*i - i*(i+1)/2 + j - i - 1));
}

inline std::array< size_t, 2 > unrank_comb2(const size_t x, const size_t n) noexcept {
	auto i = static_cast< size_t >( (n - 2 - floor(sqrt(-8*x + 4*n*(n-1)-7)/2.0 - 0.5)) );
	auto j = static_cast< size_t >( x + i + 1 - n*(n-1)/2 + (n-i)*((n-i)-1)/2 );
	return (std::array< size_t, 2 >{ i, j });
}

// Require column-major layout (AoS - Fortran style)
using np_array_t = py::array_t< float, py::array::f_style | py::array::forcecast >;

// Evaluations a cost function quickly
struct StiefelLoss {
	using index_mat_map = std::unordered_map< size_t, arma::mat >;
	const size_t n; // number of points
	const size_t d; // intrinsic dimension
	const size_t D; // target dimension of coordinatization
	index_mat_map rotations; // stores the rotation matrices (Omega)

	arma::mat frames; // all column-stacked frames (dJ x dn) for some choice of iota, weighted by PoU
	arma::sp_mat frames_sparse; // same as frames, but as sp_mat. Only one should be used

	arma::sp_mat pou; // (J x n) partition of unity
	np_array_t output; // preallocated output for (A^T x Phi) => (D x dn)

	StiefelLoss(int n_points, int dim, int target_dim) : n(n_points), d(dim), D(target_dim) {
		output = np_array_t({ D, d*n });
	}

	// Assuming output contains the result of (A^* x Phi)
	auto gradient(const py::array_t< double >& At, bool normalize=false) -> std::list< np_array_t > {

		// First project the alignment to a D-dimensional space
		arma::Mat< double > A_star { carma::arr_to_mat< double >(At) };
		output = carma::mat_to_arr< double >(A_star * frames);

		// Calculate all the SVDs
		// const size_t J = frames.n_rows / d;
		auto nuclear_norm = (float) 0.0;
		auto G = arma::mat(D, d*n);
		auto i = size_t(0);
		arma::mat S = arma::zeros(D, d);
		fast_svd_stream(output, d, [this, &i, &G, &S, &nuclear_norm](np_array_t& u, np_array_t& s, np_array_t& vt){
			auto u_copy = py::array_t< double >(u);
			auto v_copy = py::array_t< double >(vt);
			auto s_copy = py::array_t< double >(s);
			arma::mat U { carma::arr_to_mat< double >(u_copy, true) };
			S.diag() = carma::arr_to_col< double >(s_copy, true);
			arma::mat V { carma::arr_to_mat< double >(v_copy, true) };
			// py::print(U.n_rows, U.n_cols, S.n_rows, S.n_cols, V.n_rows, V.n_cols);
			G(arma::span::all, arma::span(i, i+d-1)) = U * S * V; // TODO: transpose?
			nuclear_norm += arma::trace(S);
			i += d;
		});
		// Calculate gradient
		arma::mat GF = frames * G.t(); // (dJ x dn)*(dn x D) => (dJ x D)
		if (normalize){
			GF /= n;
			nuclear_norm /= n;
		}
		auto out = std::list< np_array_t >();
		out.push_back(carma::mat_to_arr(arma::Mat< float >(&nuclear_norm, 1, 1)));
		out.push_back(carma::mat_to_arr(GF));
		return(out);
	}

	// Populate the O(J^2)-sized hashmap mapping index pairs i,j \in J -> rotation matrices
	// omega_ is expected to contain the vertically-stacked d-d rotation matrices corresponding to each (i,j) pair
	void init_rotations(py::list I_ind, py::list J_ind, py::array_t< double > omega_, const size_t J){
		std::vector< size_t > I1 = py::cast< std::vector< size_t > >(I_ind);
		std::vector< size_t > I2 = py::cast< std::vector< size_t > >(J_ind);
		// np_array_t O = static_cast< np_array_t >(omega_);
		arma::mat omega = carma::arr_to_mat< double >(omega_, true);
		for (size_t j = 0; j < I1.size(); ++j){
			size_t ii = I1[j], jj = I2[j];
			size_t key = rank_comb2(ii,jj,J);
			arma::mat R = omega(arma::span(j*d, (j+1)*d-1), arma::span::all);
			rotations.emplace(key, R);
		}
	}

	auto get_rotation(const size_t i, const size_t j, const size_t J) -> py::array_t< double > {
		size_t key = rank_comb2(i,j,J);
		if (rotations.find(key) != rotations.end()){
			return(carma::mat_to_arr(rotations[key], true));
		}
		throw std::invalid_argument("Invalid key given");
	}

	// Generates a weighted (dJ x d) frame relative to some origin subset
	// This is equivalent to Phi_{origin}(x) where 'weights' are specific to 'x'
	auto generate_frame(const size_t origin, py::array_t< double > weights) -> py::array_t< double > {
		const size_t J = weights.size();
		arma::mat d_frame(d*J, d, arma::fill::zeros); // output
		vector< double > weights_vec = weights.cast< vector< double > >();
		generate_frame_(origin, weights_vec, d_frame);
		return(carma::mat_to_arr< double >(d_frame, true));
	}

	// This generates a given dense (dJ x d) frame with the weighted rotation matrices given in the 'rotations' table
	// Note: this applies sqrt to the phi weights!
	// TODO: should the weights be sqrt'ed on identity matrices? Assume yes.
	// TODO: consider filling a (d x dJ) matrix instead to be more cache-friendly
	void generate_frame_(const size_t origin, const vector< double >& weights, arma::mat& d_frame) {
		const size_t J = weights.size();
		arma::mat I = arma::eye(d, d);

		// This fills up the (dJ x d) frame one (d x d) matrix at a time
		for (size_t j = 0; j < J; ++j){
			auto r_rng = arma::span(j*d,(j+1)*d-1);
			if (j == origin){
				d_frame(r_rng, arma::span::all) = std::sqrt(weights[j]) * I;
				continue;
			} else if (weights[j] == 0.0){
				d_frame(r_rng, arma::span::all) = std::sqrt(weights[j]) * I;
				continue;
			} else {
				// If pair exists, load it up, otherwise use identity
				const size_t key = rank_comb2(origin, j, J);
				bool key_exists = rotations.find(key) != rotations.end();
				d_frame(r_rng, arma::span::all) = std::sqrt(weights[j])*(key_exists ? (origin < j ? rotations[key] : rotations[key].t()) : I);
			}
		}
	}

	// This generates a given dense (d x dJ) frame with the weighted rotation matrices given in the 'rotations' table
	// Note: this applies sqrt to the phi weights!
	void generate_frame_T(const size_t origin, const vector< double >& weights, arma::mat& d_frame) {
		const size_t J = weights.size();
		arma::mat I = arma::eye(d, d);

		// This fills up the (d x dJ) frame one (d x d) matrix at a time (column-wise)
		for (size_t j = 0; j < J; ++j){
			auto c_rng = arma::span(j*d,(j+1)*d-1);
			auto weight = std::sqrt(weights[j]);
			if (j == origin || weights[j] == 0.0){
				d_frame(arma::span::all, c_rng) = std::sqrt(weights[j]) * I;
				continue;
			} else {
				// If pair exists, load it up, otherwise use identity
				const size_t key = rank_comb2(origin, j, J);
				bool key_exists = rotations.find(key) != rotations.end();
				d_frame(arma::span::all, c_rng) = weight *  (key_exists ? (origin < j ? rotations[key] : rotations[key].t()) : I);
			}
		}
	}

	// TODO: come back to this to fix the populate_frames functions efficiency
	template< typename OutputIt1, typename OutputIt2 >
	inline void generate_frame_ijx(const size_t origin, const size_t j, const size_t J, const double weight, const size_t col_offset, OutputIt1 RC, OutputIt2 X){

		// For each cover set, detect whether there exists rotation matrices to transform between them
		if (j == origin){
			const size_t row_offset = (j*d);
			auto w = std::sqrt(weight);
			for (size_t ri = 0; ri < d; ++ri){
				// py::print("row: ", row_offset + ri, ", col: ", col_offset + ri, ", val: ", weight);
				*RC++ = row_offset + ri; // row index of non-zero element
				*RC++ = col_offset + ri; // column index of non-zero element
				*X++ = w;
			}
		} else { // j != origin and weights[j] > 0
			const size_t row_offset = (j*d), key = rank_comb2(origin, j, J);
			auto w = std::sqrt(weight);

			// If key exists, there's a non-empty intersection between the cover sets (i,j)
			// otherwise, just use the (weighted) identity matrix
			bool key_exists = rotations.find(key) != rotations.end();
			if (!key_exists){
				for (size_t ri = 0; ri < d; ++ri){
					// py::print("row: ", row_offset + ri, ", col: ", col_offset + ri, ", val: ", weight);
					*RC++ = row_offset + ri; // row index of non-zero element
					*RC++ = col_offset + ri; // column index of non-zero element
					*X++ = w;
				}
			} else {
				const arma::mat& omega = origin < j ? rotations[key] : rotations[key].t();
				arma::umat nz_ind = arma::ind2sub(arma::size(omega), arma::find(omega != 0.0));
				nz_ind.each_col([&](arma::uvec& a){
					// py::print("| row: ", row_offset + a[0], ", col: ", col_offset + a[1], ", val: ", omega(a[0], a[1]));
					*RC++ = row_offset + a[0];
					*RC++ = col_offset + a[1];
					*X++ = w*omega(a[0], a[1]);
				});
			}
		}
	}

	// Using the rotations from the omega map, initialize the phi matrix representing the concatenation
	// of the weighted frames for some choice of iota
	// py::array_t< double > iota, bool sparse = false
	void populate_frame(const size_t i, py::array_t< double > weights, bool sparse = false){
		// if (iota.size() != n || weights.size() != ){
		// 	throw std::invalid_argument("Invalid input. Must have one weight for each cover element.")
		// }
		const size_t J = weights.size();
		auto w = weights.unchecked< 1 >();

		// Find the reference frame
		size_t k = 0;
		double max_weight = 0.0;
		for (size_t j = 0; j < J; ++j){
			if (w(j) > max_weight){
				max_weight = w(j);
				k = j;
			}
		}

		if (sparse && frames_sparse.is_empty()){
			frames_sparse.resize(d*J, d*n);
		} else if (!sparse && frames.is_empty()){
			frames.resize(d*J, d*n);
		}

		py::array_t< double > _d_frame = generate_frame(k, weights);
		arma::mat d_frame = carma::arr_to_mat(_d_frame, true);

		// arma::mat d_frame(d*J, d); // output
		// arma::mat I = arma::eye(d, d);
		// for (size_t j = 0; j < J; ++j){
		// 	auto r_rng = arma::span(j*d,(j+1)*d-1);
		// 	if (j == k){ d_frame(r_rng, arma::span::all) = I; }

		// 	// If pair exists, load it up, otherwise use identity
		// 	size_t key = rank_comb2(k, j, J);
		// 	bool key_exists = rotations.find(key) != rotations.end();
		// 	d_frame(r_rng, arma::span::all) = double(w(j))*(key_exists ? (k < j ? rotations[key] : rotations[key].t()) : I);
		// }

		// Assign the frame to right position in the frames matrix
		if (sparse){
			frames_sparse(arma::span::all, arma::span(i*d, (i+1)*d-1)) = d_frame;
		} else {
			frames(arma::span::all, arma::span(i*d, (i+1)*d-1)) = d_frame;
		}
	}

	void setup_pou(py::object& P_csc){
		to_sparse(P_csc, pou);
	}

	auto extract_iota(){
		if (pou.is_empty()){ throw std::invalid_argument("Partition of unity matrix not populated."); }
		// weights.assign(J, 0.0);
		// auto ci = pou_.begin_col(i);
		// for (; ci != pou_.end_col(i); ++ci){ weights[ci.row()] = *ci; }
		// size_t iota_i = std::distance(weights.begin(), std::max_element(weights.begin(), weights.end()));
		arma::urowvec u = arma::index_max(pou,0);
		return(carma::row_to_arr< unsigned long long >(std::move(u)));
	}

	// Populate the 'frames_sparse' arma::sp_mat member variable using the iota bijection
	// Postcondition: self.frames_sparse is a sparse (dJ x dn) matrix representing the horizontal concatenation of all the Phi's for each x \in X
	void populate_frames_sparse(py::array_t< arma::uword >& iota){
		if (pou.n_cols != n){ throw std::invalid_argument("Invalid input. Must have one weight for each cover element."); }
		const size_t J = pou.n_rows;
		arma::uvec I = carma::arr_to_col(iota, true);

		// Prepare the vectors needed to construct the sparse matrix using COO-input
		vector< arma::uword > RC;
		vector< double > X;
		RC.reserve(2*d*d*n);
		X.reserve(d*d*n);

		// Output iterators
		auto rc_out = std::back_inserter(RC);
		auto x_out = std::back_inserter(X);

		// Generate all the frames using iota
		for (size_t i = 0; i < n; ++i){

			// Iterate through sparse columns
			auto ci = pou.begin_col(i);
			for (; ci != pou.end_col(i); ++ci){
				generate_frame_ijx(I[i], ci.row(), J, *ci, i*d, rc_out, x_out);
			}
		}

		// Assign to frames_sparse
		auto locations = arma::umat(RC.data(), 2, RC.size()/2, false, true);
		frames_sparse = arma::sp_mat(std::move(locations), arma::vec(std::move(X)), d*J, d*n);
	} // populate_frames_sparse

	// Using the rotations from the omega map, initialize the phi matrix representing the concatenation
	// of the weighted frames for some *fixed* choice of iota
	// iota := n-length vector of indices each in [0, J) indicating the most similar cover set
	// pou := (J x n) sparse csc_matrix representing the partition of unity
	// Note the transpose!
	// TODO: change this to construct the sparse matrix *not* using block assignments
	void populate_frames(const py::array_t< size_t >& iota, py::object& pou, bool sparse = false){
		if (iota.size() != size_t(n)){
			throw std::invalid_argument("Invalid input. Must have one weight for each cover element.");
		}

		// Convert partition of unity to arma
		arma::sp_mat pou_;
		to_sparse(pou, pou_);
		const size_t J = pou_.n_rows;

		if (sparse && frames_sparse.is_empty()){
			frames_sparse.resize(d*J, d*n);
		} else if (!sparse && frames.is_empty()){
			frames.resize(d*J, d*n);
		}

		// Generate all the frames using iota
		vector< double > weights(J, 0.0);
		arma::mat d_frame(d*J, d, arma::fill::zeros);
		for (size_t i = 0; i < n; ++i){

			// Fill the weight vector
			weights.assign(J, 0.0);
			auto ci = pou_.begin_col(i);
			for (; ci != pou_.end_col(i); ++ci){ weights[ci.row()] = *ci; }

			// Generate the current frame using iota to specify the origin
			generate_frame_(iota.at(i), weights, d_frame);

			// Assign the frame to right position in the frames matrix
			if (sparse){
				frames_sparse(arma::span::all, arma::span(i*d, (i+1)*d-1)) = d_frame;
			} else {
				frames(arma::span::all, arma::span(i*d, (i+1)*d-1)) = d_frame;
			}
		}

		// If sparse, make sure to clean up
		if (sparse){ frames_sparse.clean(std::numeric_limits< double >::epsilon()); }
	} // populate_frames

	py::tuple initial_guess(const size_t D, bool sparse=true){
		if (sparse){
			if (frames_sparse.empty()) { throw std::invalid_argument("Frames sparse matrix has not been populated."); }
			arma::vec eigval;
			arma::mat eigvec;
			arma::eigs_sym(eigval, eigvec, frames_sparse * frames_sparse.t(), D, "lm"); // largest first
			eigval = arma::reverse(eigval);
			eigvec = arma::fliplr(eigvec);
			return py::make_tuple(carma::col_to_arr(eigval), carma::mat_to_arr(eigvec));
		} else {
			if (frames.empty()) { throw std::invalid_argument("Frames matrix has not been populated."); }
			arma::vec eigval;
			arma::mat eigvec;
			arma::eig_sym(eigval, eigvec, frames * frames.t()); // largest first
			eigval = arma::reverse(eigval);
			eigvec = arma::fliplr(eigvec);
			return py::make_tuple(carma::col_to_arr(eigval), carma::mat_to_arr(eigvec));
		}
	}

	// Returns the i'th frame of the matrix
	auto get_frame(const size_t i) -> py::array_t< double > {
		if (i >= n){ throw std::invalid_argument("Invalid index supplied."); }
		arma::mat d_frame = frames(arma::span::all, arma::span(i*d, (i+1)*d-1));
		py::array_t< double > res = carma::mat_to_arr< double >(d_frame, true);
		return(res);
	}

	// Returns all the frames
	auto all_frames() -> py::array_t< double > {
		py::array_t< double > res = carma::mat_to_arr< double >(frames, true);
		return(res);
	}

	auto all_frames_sparse() -> py::tuple {
		// py::tuple shape = S.attr("shape").cast< py::tuple >();
		// const size_t nr = shape[0].cast< size_t >(), nc = shape[1].cast< size_t >();
		// S.attr("indices") = py::array_t< arma::uword >();
		// S.attr("indptr") = py::array_t< arma::uword >()
		// arma::uvec ind_ptr = carma::arr_to_col(.cast< py::array_t< arma::uword > >());
		// arma::vec data = carma::arr_to_col(S.attr("data").cast< py::array_t< double > >());
		// out = arma::sp_mat(std::move(ind), std::move(ind_ptr), std::move(data), nr, nc);

		// py::array_t< double > res = carma::mat_to_arr< double >(frames, true);
		// return(res);
		frames_sparse.sync();
		//const ptr_aux_mem, number_of_elements
		auto ri = carma::col_to_arr(arma::uvec(frames_sparse.row_indices, frames_sparse.n_nonzero));
		auto cp = carma::col_to_arr(arma::uvec(frames_sparse.col_ptrs, frames_sparse.n_cols+1));
		auto x = carma::col_to_arr(arma::vec(frames_sparse.values, frames_sparse.n_nonzero));
		// auto ri = carma::col_to_arr< const arma::uword >(*frames_sparse.row_indices, true);
		// auto out = py::dict(
		// 	"indices"= py::array(ri),
		// 	"indptr"=cp, "values"=x);
		return(py::make_tuple(ri, cp, x));
	}

	void embed(const py::array_t< double >& At){
		arma::Mat< double > A_star { carma::arr_to_mat< double >(At) };
		output = carma::mat_to_arr< double >(A_star * frames);
	}

	void benchmark_embedding(py::array_t< double >& At, const size_t m){
		if (frames_sparse.is_empty()){ frames_sparse = arma::sp_mat(frames); }
		arma::Mat< double > A_star { carma::arr_to_mat< double >(At, true) };

		size_t ms_dense = 0, ms_sparse = 0;

		auto start = high_resolution_clock::now();
		for (size_t i = 0; i < m; ++i){
			output = carma::mat_to_arr< double >(A_star * frames);
		}
		auto stop = high_resolution_clock::now();
		auto duration = duration_cast< milliseconds >(stop - start);
		ms_dense = duration.count();

		start = high_resolution_clock::now();
		for (size_t i = 0; i < m; ++i){
			output = carma::mat_to_arr< double >(A_star * frames_sparse);
		}
		stop = high_resolution_clock::now();
		duration = duration_cast< milliseconds >(stop - start);
		ms_sparse = duration.count();

		py::print("ms dense: ", ms_dense, "ms sparse: ", ms_sparse);
	}

	using index_list = vector< vector< size_t > >;
	using vec_mats = vector< arma::mat >;

	// Given sorted range [b,e), finds 'element' in log time, or throws an exception if not found
	template< typename Iter, typename T = typename Iter::value_type >
	auto find_index(Iter b, const Iter e, T element) -> size_t {
		auto lb = std::lower_bound(b, e, element);
		if (lb != e && (*lb == element)){
			return(std::distance(b, lb));
		}
		throw std::logic_error("Unable to find element in array.");
	}
	void to_sparse(const py::object& S, arma::sp_mat& out){
		py::tuple shape = S.attr("shape").cast< py::tuple >();
		const size_t nr = shape[0].cast< size_t >(), nc = shape[1].cast< size_t >();
		arma::uvec ind = carma::arr_to_col(S.attr("indices").cast< py::array_t< arma::uword > >());
		arma::uvec ind_ptr = carma::arr_to_col(S.attr("indptr").cast< py::array_t< arma::uword > >());
		arma::vec data = carma::arr_to_col(S.attr("data").cast< py::array_t< double > >());
		out = arma::sp_mat(std::move(ind), std::move(ind_ptr), std::move(data), nr, nc);
	}

	// A := (dJ x D) dense orthonormal matrix
	// pou := (J x n) sparse matrix representing the transpose of the PoU
	// cover_subsets := (J)-length vector of sorted cover sets
	// local_models := (J)-length vector of column-oriented euclidean coordinate models (point for each column)
	// T := (D x n) matrix of translation vectors
	// Note: with armadillo's csc impmenentations, (dense x sparse) is faster than (sparse x dense)
	// Output => (D x n) matrix of the assembled coordinates
	void fast_assembly2(const arma::mat& A, const arma::sp_mat& pou, const index_list& cover_subsets, const vec_mats& local_models, const arma::mat& T, arma::mat& assembly){
		// arma::mat assembly = arma::zeros(D, n);
		arma::vec coords = arma::zeros(D);
		const size_t J = pou.n_rows;

		// Variables to re-use/cache in the loop
		vector< double > phi_i(J); // the partition of unity weights for x_i
		arma::mat d_frame(d*J, d); // the current frame to populate
		arma::mat U, V;
		arma::vec s;

		// Build the assembly
		for (size_t i = 0; i < n; ++i){
			phi_i.assign(J, 0.0);
			coords.zeros();

			// Fill phi weight vector
			arma::sp_mat::const_col_iterator ci = pou.begin_col(i);
			for (; ci != pou.end_col(i); ++ci){ phi_i[ci.row()] = *ci; }

			// Compute the weighted average of the Fj's using the partition of unity
			ci = pou.begin_col(i);
			for (; ci != pou.end_col(i); ++ci){
				size_t j = ci.row();
				generate_frame_(j, phi_i, d_frame); 			// populates the (dJ x d) frame in d_frame
				svd_econ(U, s, V, A * (A.t() * d_frame)); // compute SVD of A A^T phi_j (U := dj x r, V := r x d, r <= d)

				// In cover set U_j, find the index of the local coordinate for point x_i, add the translation vector
				size_t jj = find_index(cover_subsets[j].begin(), cover_subsets[j].end(), i); // todo: remove this
				arma::vec local_coord = local_models[j].col(jj) + T.col(j); // should be column vector

				// Add to the current coordinate
				coords += ((*ci) * A.t() * U * V.t()) * local_coord; // lhs := (D x d), local coords := (d x 1)
			}
			assembly.col(i) = coords;
		}
	}

	// High dimensional (dJ) assembly
	// assembly := output (dJ x n) matrix
	void fast_assembly_high(const arma::sp_mat& pou, const index_list& cover_subsets, const vec_mats& local_models, const arma::mat& T, arma::mat& assembly){

		const size_t J = pou.n_rows;
		arma::vec coords = arma::zeros(d*J);

		// Variables to re-use/cache in the loop
		vector< double > phi_i(J); // the partition of unity weights for x_i
		arma::mat d_frame(d*J, d); // the current frame to populate

		// Build the assembly
		for (size_t i = 0; i < n; ++i){
			phi_i.assign(J, 0.0);
			coords.zeros();

			// Fill phi weight vector
			arma::sp_mat::const_col_iterator ci = pou.begin_col(i);
			for (; ci != pou.end_col(i); ++ci){ phi_i[ci.row()] = *ci; }

			// Compute the weighted average of the Fj's using the partition of unity
			ci = pou.begin_col(i);
			for (; ci != pou.end_col(i); ++ci){
				size_t j = ci.row();
				generate_frame_(j, phi_i, d_frame); 			// populates the (dJ x d) frame in d_frame
				size_t jj = find_index(cover_subsets[j].begin(), cover_subsets[j].end(), i); // todo: remove this
				arma::vec local_coord = local_models[j].col(jj) + T.col(j); // should be (d)-length column vector
				coords += (*ci) * d_frame * local_coord; // lhs := (dJ x 1), local coords := (d x 1)
			}
			assembly.col(i) = coords;
		}
	}

	// Uses the fast_assembly2() function.
	// All inputs as passed as-is to fast_assembly2; do not transpose anything here
	auto assemble_frames2(const py::array_t< double >& A, const py::object& pou, const py::list& cover_subsets, const py::list& local_models, const py::array_t< double >& T, bool high) -> py::array_t< double > {
		arma::mat A_ = carma::arr_to_mat(A);

		// Partition of unity
		arma::sp_mat pou_;
		to_sparse(pou, pou_);
		const size_t J = pou_.n_rows;

		// Convert cover subsets to C++ versions
		auto subsets = vector< vector< size_t > >();
		for (auto ind: cover_subsets){ subsets.push_back(ind.cast< vector< size_t > >()); }

		// Copy the local euclidean models (transposed)
		auto models = vector< arma::mat >();
		for (auto pts: local_models){
			py::array_t< double > pts_ = pts.cast< py::array_t< double > >();
			models.push_back(carma::arr_to_mat(pts_, true));
		}

		// Copy/move the translations
		arma::mat translations = carma::arr_to_mat(T);

		// Output assembly
		if (high){
			arma::mat assembly = arma::zeros(d*J, n);
			fast_assembly_high(pou_, subsets, models, translations, assembly);
			return(carma::mat_to_arr(assembly));
		}
		else {
			arma::mat assembly = arma::zeros(D, n);
			fast_assembly2(A_, pou_, subsets, models, translations, assembly);
			return(carma::mat_to_arr(assembly));
		}
	}

	// TODO: make pou and local_models transposed to access by column
	arma::mat fast_assembly(const arma::mat& A, const arma::sp_mat& pou, const index_list& cover_subsets, const vec_mats& local_models, const arma::mat& T){
		//py::array_t< double > phi_i = generate_frame(const size_t origin, py::array_t< double > weights);
		arma::mat assembly = arma::zeros(n, D);
		arma::vec coords = arma::zeros(D);
		const size_t J = pou.n_cols;

		// Variables to re-use/cache in the loop
		vector< double > phi_i(J); // the partition of unity weights for x_i
		arma::mat d_frame(d*J, d); // the current frame to populate
		arma::mat U, V;
		arma::vec s;

		// Build the assembly
		for (size_t i = 0; i < n; ++i){
			phi_i.assign(J, 0.0);
			coords.zeros();

			// Fill phi weight vector
			arma::sp_mat::const_row_iterator ri = pou.begin_row(i);
			for (; ri != pou.end_row(i); ++ri){ phi_i[ri.col()] = *ri; }

			// Compute the weighted average of the Fj's using the partition of unity
			ri = pou.begin_row(i);
			for (; ri != pou.end_row(i); ++ri){
				size_t j = ri.col();
				size_t jj = find_index(cover_subsets[j].begin(), cover_subsets[j].end(), i);
				generate_frame_(j, phi_i, d_frame); 			// populates d_frame
				svd_econ(U, s, V, A * (A.t() * d_frame)); // compute SVD of A A^T phi_j
				arma::rowvec local_coord = local_models[j].row(jj) + T.row(j); // should be column vector
				coords += ((*ri) * A.t() * U * V.t()) * local_coord.t(); // left-side should be (D x d)
			}
			assembly.row(i) = coords.t();
		}
		return(assembly);
	}

	// Wrapper for the fast_assembly above
	auto assemble_frames(py::array_t< double >& A, py::object& pou, py::list& cover_subsets, const py::list& local_models, py::array_t< double >& T ) -> py::array_t< double > {
		arma::mat A_ = carma::arr_to_mat(A, true);
		arma::sp_mat pou_;
		to_sparse(pou, pou_);
		auto subsets = vector< vector< size_t > >();
		for (auto ind: cover_subsets){
			subsets.push_back(ind.cast< vector< size_t > >());
		}
		auto models = vector< arma::mat >();
		for (auto pts: local_models){
			py::array_t< double > pts_ = pts.cast< py::array_t< double > >();
			models.push_back(carma::arr_to_mat(pts_, true));
		}
		arma::mat translations = carma::arr_to_mat(T, true);
		arma::mat assembly = fast_assembly(A_, pou_, subsets, models, translations);
		return(carma::mat_to_arr(assembly));
	}

};

// if len(translations) != len(cover): raise ValueError("There should be a translation vector for each subset of the cover.")
// assembly = np.zeros((stf.n, stf.D), dtype=np.float64)
// coords = np.zeros((1,stf.D), dtype=np.float64)
// index_set = list(local_models.keys())
// for i in range(stf.n):
// 	w_i = np.ravel(pou[i,:].todense())
// 	nz_ind = np.where(w_i > 0)[0]
// 	coords.fill(0)
// 	## Construct assembly functions F_j(x) for x_i
// 	for j in nz_ind:
// 		subset_j = cover[index_set[j]]
// 		relative_index = find_where(i, subset_j, True) ## This should always be true!
// 		u, s, vt = np.linalg.svd((A @ (A.T @ stf.generate_frame(j, w_i))), full_matrices=False, compute_uv=True)
// 		d_coords = local_models[index_set[j]][relative_index,:]
// 		coords += (w_i[j]*A.T @ (u @ vt) @ (d_coords + translations[j]).T).T
// 	assembly[i,:] = coords

// extern void slaed1(int* N, float* D, float* Q, int* LDQ, int* INDXQ, float* RHO, int* CUTPNT, float* WORK, int* IWORK, int* INFO);

// Computes U A U^T + sigma (u * u^T)
// SLAED1 computes the updated eigensystem of a diagonal matrix after modification by a rank-one symmetric matrix.
// void dpr1(py::array_t< float > D, py::array_t< float > V, float sigma, py::array_t< float > u){
// 	arma::fvec d = carma::arr_to_col(D, true);
// 	arma::fvec v = carma::arr_to_col(u, true);
// 	arma::fmat Q = carma::arr_to_mat(V, true);
// 	int N = d.size();
// 	int LDQ = Q.n_rows;
// 	vector< int > indxq(N);
// 	std::iota(indxq.begin(), indxq.end(), 0);
// 	int CUTPNT = N/2;
// 	vector< float > workspace(4*N + N*N);
// 	vector< int > iworkspace(4*N);
// 	int info = 0;
// 	slaed1(&N, d.memptr(), Q.memptr(), &LDQ, indxq.data(), &sigma, &CUTPNT, workspace.data(), iworkspace.data(), &info);
// 	py::print("Info: ", info);
// }

// diag( D ) +  RHO *  Z * Z_transpose.
// void dpr1(py::array_t< float > D, float rho, py::array_t< float > Z, int I){
// 	arma::fvec d = carma::arr_to_col(D, true);
// 	arma::fvec z = carma::arr_to_col(Z, true);
// 	int N = D.size(), info = 0;
// 	arma::fvec delta(N); // used for reconstructing eiegnvectors
// 	float lambda = 0; // output eigenvalue
// 	slaed4(&N, &I, d.memptr(), z.memptr(), delta.memptr(), &rho, &lambda, &info);
// 	py::print("Info: ", info, "updated ev: ", lambda);
// }

// slaed9(int* K, int* KSTART, int* KSTOP, int* N, float* D, float* Q, int* LDQ, float* rho, float* dlambda, float* W, float* S, int& lds, int& info);
// auto dpr1_ev(py::array_t< float > Q, py::array_t< float > D, float rho, py::array_t< float > Z) -> py::dict {
// 	arma::fmat q = carma::arr_to_mat(Q, true); // eigenvectors
// 	arma::fvec d = carma::arr_to_col(D, true); // diagonal entries / poles
// 	arma::fvec z = carma::arr_to_col(Z, true); // perturbation vector
// 	int K = d.size(), N = q.n_rows, info = 0;
// 	int KSTART = 1, KEND = K;
// 	arma::fvec lambda(K);
// 	arma::fmat S(q.n_rows, q.n_cols);
// 	slaed9(&K, &KSTART, &KEND, &N, lambda.memptr(), q.memptr(), &N, &rho, d.memptr(), z.memptr(), S.memptr(), &N, &info);
// 	// py::print("Info: ", info);
// 	py::dict output;
// 	output["info"] = info;
// 	output["eval"] = carma::col_to_arr(lambda, true);
// 	output["evec"] = carma::mat_to_arr(S, true);
// 	return(output);
// }

// T = Q(in) ( D(in) + RHO * Z*Z**T ) Q**T(in) = Q(out) * D(out) * Q**T(out)

PYBIND11_MODULE(fast_svd, m) {
	m.def("fast_svd", &fast_svd, "Yields the svd of a matrix of low dimension");
	//m.def("lapack_svd", &lapack_svd, "Yields the svd of a matrix of low dimension");
	//m.def("test_sparse", &test_sparse, "Test conversion to sparse matrix");
	//m.def("dpr1", &dpr1, "Diagonal + rank-1 matrix eigenvalue update");
	//m.def("dpr1_ev", &dpr1_ev, "Diagonal + rank-1 matrix eigenvalue update");
	py::class_<StiefelLoss>(m, "StiefelLoss")
		.def(py::init< int, int, int >())
		.def_readonly("d", &StiefelLoss::d)
		.def_readonly("n", &StiefelLoss::n)
		.def_readonly("D", &StiefelLoss::D)
		.def_readwrite("output", &StiefelLoss::output)
		// .def_readwrite("rotations", &StiefelLoss::rotations)
		//.def("benchmark_gradient", &StiefelLoss::benchmark_gradient)
		.def("gradient", &StiefelLoss::gradient)
		.def("init_rotations", &StiefelLoss::init_rotations)
		.def("get_rotation", &StiefelLoss::get_rotation)
		.def("populate_frame", &StiefelLoss::populate_frame)
		.def("populate_frames", &StiefelLoss::populate_frames)
		.def("populate_frames_sparse", &StiefelLoss::populate_frames_sparse)
		.def("generate_frame", &StiefelLoss::generate_frame)
		.def("get_frame", &StiefelLoss::get_frame)
		.def("all_frames", &StiefelLoss::all_frames)
		.def("all_frames_sparse", &StiefelLoss::all_frames_sparse)
		.def("embed", &StiefelLoss::embed)
		.def("extract_iota", &StiefelLoss::extract_iota)
		.def("setup_pou", &StiefelLoss::setup_pou)
		.def("benchmark_embedding", &StiefelLoss::benchmark_embedding)
		.def("assemble_frames", &StiefelLoss::assemble_frames)
		.def("assemble_frames2", &StiefelLoss::assemble_frames2)
		.def("initial_guess", &StiefelLoss::initial_guess)
		.def("__repr__",[](const StiefelLoss &stf) {
			return("Stiefel Loss w/ parameters n="+std::to_string(stf.n)+",d="+std::to_string(stf.d)+",D="+std::to_string(stf.D));
  	});
}

// double numpy_svd(){
// 	auto svd = py::module::import("numpy.linalg").attr("svd");
// 	py::buffer_info output_buffer = output.request();
// 	const size_t inc = D*d;
// 	double nuclear_norm = 0.0;
// 	for (int j = 0; j < n; ++j){
// 		np_array_t inp = np_array_t({ D, d }, output.data()+(j*inc));
// 		np_array_t sv = svd(inp, false, false, false);
// 		auto r = sv.unchecked< 1 >();
// 		for (size_t i = 0; i < sv.shape(0); i++){
// 			nuclear_norm += r(i);
// 		}
// 	}
// 	return(nuclear_norm);
// }


#include <Eigen/Core>
#include <Eigen/Geometry>
#include <cmath>
#include "UKF/Types.h"
#include "UKF/Integrator.h"
#include "UKF/StateVector.h"
#include "UKF/MeasurementVector.h"
#include "UKF/Core.h"
#include "cukf.h"

/*
This file implements the same UKF as was used for the 2014 UAV Challenge.
There are a number of efficiency and performance improvements that could be
made with the new library structure, but the main purpose of this is for
direct comparison of performance of the new library against the old
implementation.
*/

enum SFWA_States {
    LatLon,
    Altitude,
    Velocity,
    Acceleration,
    Attitude,
    AngularVelocity,
    AngularAcceleration,
    WindVelocity,
    GyroBias
};

using SFWA_StateVector = UKF::StateVector<
    UKF::Field<LatLon, UKF::Vector<2>>,                 /* Latitude and longitude (rad) */
    UKF::Field<Altitude, real_t>,                       /* Altitude above the WGS84 ellipsoid (m) */
    UKF::Field<Velocity, UKF::Vector<3>>,               /* Velocity (NED frame, m/s) */
    UKF::Field<Acceleration, UKF::Vector<3>>,           /* Acceleration (body frame, m/s^2) */
    UKF::Field<Attitude, UKF::Quaternion>,              /* Attitude as a quaternion (NED frame to body frame) */
    UKF::Field<AngularVelocity, UKF::Vector<3>>,        /* Angular velocity (body frame, rad/s) */
    UKF::Field<AngularAcceleration, UKF::Vector<3>>,    /* Angular acceleration (body frame, rad/s^2) */
    UKF::Field<WindVelocity, UKF::Vector<3>>,           /* Wind velocity (NED frame, m/s) */
    UKF::Field<GyroBias, UKF::Vector<3>>                /* Gyro bias (body frame, rad/s) */
>;

/* WGS84 reference ellipsoid constants. */
#define WGS84_A (6378137.0)
#define WGS84_B (6356752.314245)
#define WGS84_A2 (WGS84_A*WGS84_A)
#define WGS84_B2 (WGS84_B*WGS84_B)
#define WGS84_AB2 (WGS84_A2*WGS84_B2)

#define G_ACCEL (9.80665)
#define RHO (1.225)

/* SFWA vehicle dynamics model. Not used for this comparison. */
UKF::Vector<6> x8_dynamics_model(const SFWA_StateVector &state, const UKF::Vector<3> &control) {
    UKF::Vector<6> output;

    /* Cache state data for convenience */
    UKF::Quaternion attitude = state.get_field<Attitude>();
    real_t yaw_rate = state.get_field<AngularVelocity>()[2],
           pitch_rate = state.get_field<AngularVelocity>()[1],
           roll_rate = state.get_field<AngularVelocity>()[0];

    /* External axes */
    UKF::Vector<3> airflow = attitude * (state.get_field<WindVelocity>() - state.get_field<Velocity>());

    /*
    Calculate axial airflow
    */
    real_t airflow_x2, airflow_y2, airflow_z2, airflow_v2;
    airflow_x2 = airflow[0]*airflow[0];
    airflow_y2 = airflow[1]*airflow[1];
    airflow_z2 = airflow[2]*airflow[2];
    airflow_v2 = airflow_x2 + airflow_y2 + airflow_z2;

    /*
    Determine motor thrust and torque.
    */
    real_t rpm = control[0] * 12000.0, thrust,
           ve2 = (0.0025 * 0.0025) * rpm * rpm;
    /* 1 / 3.8kg times area * density of air */
    thrust = (ve2 - airflow_v2) *
             (0.26315789473684 * 0.5 * RHO * 0.02);

    /*
    Calculate airflow in the horizontal and vertical planes, as well as
    pressure
    */
    real_t v_inv, vertical_v, vertical_v_inv, qbar;

    qbar = (RHO * 0.5) * airflow_v2;
    v_inv = 1.0 / std::sqrt(std::max(1.0, airflow_v2));

    vertical_v = std::sqrt(airflow_x2 + airflow_z2);
    vertical_v_inv = 1.0 / std::max(1.0, vertical_v);

    /* Work out sin/cos of alpha and beta */
    real_t sin_alpha, cos_alpha, sin_beta, cos_beta, sin_cos_alpha;

    sin_beta = airflow[1] * v_inv;
    cos_beta = vertical_v * v_inv;

    sin_alpha = -airflow[2] * vertical_v_inv;
    cos_alpha = -airflow[0] * vertical_v_inv;

    sin_cos_alpha = sin_alpha * cos_alpha;

    /* Work out aerodynamic forces in wind frame */
    real_t lift, drag, side_force;

    /* 0.26315789473684 is the reciprocal of mass (3.8kg) */
    lift = (qbar * 0.26315789473684) * (0.8 * sin_cos_alpha + 0.18);
    drag = (qbar * 0.26315789473684) *
           (0.05 + 0.7 * sin_alpha * sin_alpha);
    side_force = (qbar * 0.26315789473684) * 0.2 * sin_beta * cos_beta;

    /* Convert aerodynamic forces from wind frame to body frame */
    real_t x_aero_f = lift * sin_alpha - drag * cos_alpha -
                             side_force * sin_beta,
           z_aero_f = lift * cos_alpha + drag * sin_alpha,
           y_aero_f = side_force * cos_beta;

    output.segment<3>(0) << UKF::Vector<3>(x_aero_f + thrust, y_aero_f, -z_aero_f) +
        (attitude * UKF::Vector<3>(0, 0, G_ACCEL));

    /* Determine moments */
    real_t pitch_moment, yaw_moment, roll_moment,
           left_aileron = control[1] - 0.5, right_aileron = control[2] - 0.5;
    pitch_moment = 0.0 - 0.0 * sin_alpha - 0.0 * pitch_rate -
                   0.1 * (left_aileron + right_aileron) * vertical_v * 0.1;
    roll_moment = 0.05 * sin_beta - 0.1 * roll_rate +
                  0.15 * (left_aileron - right_aileron) * vertical_v * 0.1;
    yaw_moment = -0.02 * sin_beta - 0.05 * yaw_rate -
                 0.02 * (std::abs(left_aileron) + std::abs(right_aileron)) *
                 vertical_v * 0.1;
    pitch_moment *= qbar;
    roll_moment *= qbar;
    yaw_moment *= qbar;

    /*
    Calculate angular acceleration (tau / inertia tensor).
    Inertia tensor is:
        0.3 0 -0.0334
        0 0.17 0
        -0.0334 0 0.405
    So inverse is:
        3.36422 0 0.277444
        0 5.88235 0
        0.277444 0 2.49202
    */
    output.segment<3>(3) << UKF::Vector<3>(
        (3.364222 * roll_moment + 0.27744448 * yaw_moment),
        10.8823528 * pitch_moment,
        (0.27744448 * roll_moment + 2.4920163 * yaw_moment));

    return output;
}

namespace UKF {

/* SFWA state vector process model. */
template <> template <>
SFWA_StateVector SFWA_StateVector::derivative<>() const {
    SFWA_StateVector output;

    /* Calculate the normal and meridional radii of curvature. */
    real_t lat = get_field<LatLon>()[0];
    real_t tempA = WGS84_A*std::cos(lat), tempB = WGS84_B*std::sin(lat),
           temp = tempA * tempA + tempB * tempB,
           temp_sqrt = std::sqrt(temp);
    real_t M = WGS84_AB2 / (temp_sqrt * temp);
    real_t N = WGS84_A2 / temp_sqrt;

    /*
    Calculate change in position. Using the small angle approximation, this
    becomes very simple – no trig required for latitude derivative, and one
    cosine function for longitude derivative.
    */
    UKF::Vector<3> vel = get_field<Velocity>();
    output.set_field<LatLon>(Vector<2>(
        vel[0] / (M + get_field<Altitude>()),
        (vel[1] / (N + get_field<Altitude>())) * std::cos(lat)));
    output.set_field<Altitude>(-vel[2]);

    /* Calculate change in velocity. */
    output.set_field<Velocity>(get_field<Attitude>().conjugate() * get_field<Acceleration>());

    /* Change in linear acceleration is zero. */
    output.set_field<Acceleration>(UKF::Vector<3>(0, 0, 0));

    /* Calculate change in attitude. */
    UKF::Quaternion omega_q;
    omega_q.vec() = get_field<AngularVelocity>() * 0.5;
    omega_q.w() = 0;
    output.set_field<Attitude>(omega_q.conjugate() * get_field<Attitude>());

    /* Calculate change in angular velocity (just angular acceleration). */
    output.set_field<AngularVelocity>(get_field<AngularAcceleration>());

    /* Change in angular acceleration is zero. */
    output.set_field<AngularAcceleration>(UKF::Vector<3>(0, 0, 0));

    /* Change in wind velocity is zero. */
    output.set_field<WindVelocity>(UKF::Vector<3>(0, 0, 0));

    /* Change in gyro bias is zero. */
    output.set_field<GyroBias>(UKF::Vector<3>(0, 0, 0));

    return output;
}

}

enum SFWA_Measurements {
    Accelerometer,
    Gyroscope,
    Magnetometer,
    GPS_Position,
    GPS_Velocity,
    Airspeed,
    PressureAltitude
};

using SFWA_MeasurementVector = UKF::DynamicMeasurementVector<
    UKF::Field<Accelerometer, UKF::Vector<3>>,
    UKF::Field<Gyroscope, UKF::Vector<3>>,
    UKF::Field<Magnetometer, UKF::Vector<3>>,
    UKF::Field<GPS_Position, UKF::Vector<3>>,
    UKF::Field<GPS_Velocity, UKF::Vector<3>>,
    UKF::Field<Airspeed, real_t>,
    UKF::Field<PressureAltitude, real_t>
>;

/*
Hard code the magnetic field vector to the WMM value for
-37.954690, 145.237575 for the purposes of this comparison. Value is in
microtesla.
*/
static UKF::Vector<3> local_mag_field = UKF::Vector<3>(21.2584, 4.4306, -55.9677);

using SFWA_UKF = UKF::Core<
    SFWA_StateVector,
    SFWA_MeasurementVector,
    UKF::IntegratorRK4
>;

namespace UKF {
/* SFWA measurement model. */
template <> template <>
UKF::Vector<3> SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, Accelerometer>(const SFWA_StateVector& state) {
    return state.get_field<Acceleration>() + state.get_field<Attitude>() * UKF::Vector<3>(0, 0, -G_ACCEL);
}

template <> template <>
UKF::Vector<3> SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, Gyroscope>(const SFWA_StateVector& state) {
    return state.get_field<AngularVelocity>() + state.get_field<GyroBias>();
}

template <> template <>
UKF::Vector<3> SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, Magnetometer>(const SFWA_StateVector& state) {
    return state.get_field<Attitude>() * local_mag_field;
}

template <> template <>
UKF::Vector<3> SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, GPS_Position>(const SFWA_StateVector& state) {
    return UKF::Vector<3>(state.get_field<LatLon>()[0], state.get_field<LatLon>()[1], state.get_field<Altitude>());
}

template <> template <>
UKF::Vector<3> SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, GPS_Velocity>(const SFWA_StateVector& state) {
    return state.get_field<Velocity>();
}

template <> template <>
real_t SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, Airspeed>(const SFWA_StateVector& state) {
    return (state.get_field<Attitude>() * (state.get_field<Velocity>() - state.get_field<WindVelocity>()))[0];
}

template <> template <>
real_t SFWA_MeasurementVector::expected_measurement
<SFWA_StateVector, PressureAltitude>(const SFWA_StateVector& state) {
    return state.get_field<Altitude>();
}

}
static SFWA_UKF ukf;
static SFWA_MeasurementVector meas;

/*
The following functions provide a ctypes-compatible interface for ease of
testing.
*/

void ukf_init() {
    ukf.state.set_field<LatLon>(UKF::Vector<2>(-0.662434, 2.534874));
    ukf.state.set_field<Altitude>(0);
    ukf.state.set_field<Velocity>(UKF::Vector<3>(0, 0, 0));
    ukf.state.set_field<Acceleration>(UKF::Vector<3>(0, 0, 0));
    ukf.state.set_field<Attitude>(UKF::Quaternion(1, 0, 0, 0));
    ukf.state.set_field<AngularVelocity>(UKF::Vector<3>(0, 0, 0));
    ukf.state.set_field<AngularAcceleration>(UKF::Vector<3>(0, 0, 0));
    ukf.state.set_field<WindVelocity>(UKF::Vector<3>(0, 0, 0));
    ukf.state.set_field<GyroBias>(UKF::Vector<3>(0, 0, 0));
    ukf.covariance = SFWA_StateVector::CovarianceMatrix::Zero();
    ukf.covariance.diagonal() <<
        1, 1, 10000, 1e-6, 1e-6, 1e-6, 1e-6, 1e-6, 1e-6, 10, 10, 10,
        1e-6, 1e-6, 1e-6, 1e-6, 1e-6, 1e-6, 100, 100, 100, 0.01, 0.01, 0.01;
    ukf.measurement_covariance <<1, 1, 1, 0.01, 0.01, 0.01, 10, 10, 10, 1e-20, 1e-20, 1, 1, 1, 1, 1, 1;

    local_mag_field = UKF::Vector<3>(21.2584, 4.4306, -55.9677);

    ukf.process_noise_covariance = SFWA_StateVector::CovarianceMatrix::Zero();
    ukf.process_noise_covariance.diagonal() <<
        1e-17, 1e-17, 1e-4,     /* lat, lon, alt */
        2e-3, 2e-3, 2e-3,       /* velocity N, E, D */
        2e-2, 2e-2, 2e-2,       /* acceleration x, y, z */
        7e-8, 7e-8, 7e-8,       /* attitude roll, pitch, yaw */
        1e-3, 1e-3, 1e-3,       /* angular velocity roll, pitch, yaw */
        1e-3, 1e-3, 1e-3,       /* angular acceleration roll, pitch, yaw */
        1e-7, 1e-7, 1e-7,       /* wind velocity N, E, D -- NOTE: from FCS armed mode */
        1e-12, 1e-12, 1e-12;    /* gyro bias x, y, z -- NOTE: from FCS armed mode */
}

void ukf_set_position(real_t lat, real_t lon, real_t alt) {
    ukf.state.set_field<LatLon>(UKF::Vector<2>(lat, lon));
    ukf.state.set_field<Altitude>(alt);
}

void ukf_set_velocity(real_t x, real_t y, real_t z) {
    ukf.state.set_field<Velocity>(UKF::Vector<3>(x, y, z));
}

void ukf_set_acceleration(real_t x, real_t y, real_t z) {
    ukf.state.set_field<Acceleration>(UKF::Vector<3>(x, y, z));
}

void ukf_set_attitude(real_t w, real_t x, real_t y, real_t z) {
    ukf.state.set_field<Attitude>(UKF::Quaternion(w, x, y, z));
}

void ukf_set_angular_velocity(real_t x, real_t y, real_t z) {
    ukf.state.set_field<AngularVelocity>(UKF::Vector<3>(x, y, z));
}

void ukf_set_angular_acceleration(real_t x, real_t y, real_t z) {
    ukf.state.set_field<AngularAcceleration>(UKF::Vector<3>(x, y, z));
}

void ukf_set_wind_velocity(real_t x, real_t y, real_t z) {
    ukf.state.set_field<WindVelocity>(UKF::Vector<3>(x, y, z));
}

void ukf_set_gyro_bias(real_t x, real_t y, real_t z) {
    ukf.state.set_field<GyroBias>(UKF::Vector<3>(x, y, z));
}

void ukf_get_state(struct ukf_state_t *in) {
    in->position[0] = ukf.state.get_field<LatLon>()[0];
    in->position[1] = ukf.state.get_field<LatLon>()[1];
    in->position[2] = ukf.state.get_field<Altitude>();
    in->velocity[0] = ukf.state.get_field<Velocity>()[0];
    in->velocity[1] = ukf.state.get_field<Velocity>()[1];
    in->velocity[2] = ukf.state.get_field<Velocity>()[2];
    in->acceleration[0] = ukf.state.get_field<Acceleration>()[0];
    in->acceleration[1] = ukf.state.get_field<Acceleration>()[1];
    in->acceleration[2] = ukf.state.get_field<Acceleration>()[2];
    in->attitude[0] = ukf.state.get_field<Attitude>().x();
    in->attitude[1] = ukf.state.get_field<Attitude>().y();
    in->attitude[2] = ukf.state.get_field<Attitude>().z();
    in->attitude[3] = ukf.state.get_field<Attitude>().w();
    in->angular_velocity[0] = ukf.state.get_field<AngularVelocity>()[0];
    in->angular_velocity[1] = ukf.state.get_field<AngularVelocity>()[1];
    in->angular_velocity[2] = ukf.state.get_field<AngularVelocity>()[2];
    in->angular_acceleration[0] = ukf.state.get_field<AngularAcceleration>()[0];
    in->angular_acceleration[1] = ukf.state.get_field<AngularAcceleration>()[1];
    in->angular_acceleration[2] = ukf.state.get_field<AngularAcceleration>()[2];
    in->wind_velocity[0] = ukf.state.get_field<WindVelocity>()[0];
    in->wind_velocity[1] = ukf.state.get_field<WindVelocity>()[1];
    in->wind_velocity[2] = ukf.state.get_field<WindVelocity>()[2];
    in->gyro_bias[0] = ukf.state.get_field<GyroBias>()[0];
    in->gyro_bias[1] = ukf.state.get_field<GyroBias>()[1];
    in->gyro_bias[2] = ukf.state.get_field<GyroBias>()[2];
}

void ukf_set_state(struct ukf_state_t *in) {
    ukf.state.set_field<LatLon>(UKF::Vector<2>(in->position[0], in->position[1]));
    ukf.state.set_field<Altitude>(in->position[2]);
    ukf.state.set_field<Velocity>(
        UKF::Vector<3>(in->velocity[0], in->velocity[1], in->velocity[2]));
    ukf.state.set_field<Acceleration>(
        UKF::Vector<3>(in->acceleration[0], in->acceleration[1], in->acceleration[2]));
    ukf.state.set_field<Attitude>(
        UKF::Quaternion(in->attitude[3], in->attitude[0], in->attitude[1], in->attitude[2]));
    ukf.state.set_field<AngularVelocity>(
        UKF::Vector<3>(in->angular_velocity[0], in->angular_velocity[1], in->angular_velocity[2]));
    ukf.state.set_field<AngularAcceleration>(
        UKF::Vector<3>(in->angular_acceleration[0], in->angular_acceleration[1], in->angular_acceleration[2]));
    ukf.state.set_field<WindVelocity>(
        UKF::Vector<3>(in->wind_velocity[0], in->wind_velocity[1], in->wind_velocity[2]));
    ukf.state.set_field<GyroBias>(
        UKF::Vector<3>(in->gyro_bias[0], in->gyro_bias[1], in->gyro_bias[2]));
}

void ukf_get_state_covariance(
        real_t state_covariance[SFWA_StateVector::covariance_size()*SFWA_StateVector::covariance_size()]) {
    Eigen::Map<typename SFWA_StateVector::CovarianceMatrix> covariance_map(state_covariance);
    covariance_map = ukf.covariance;
}

void ukf_get_state_covariance_diagonal(
        real_t state_covariance_diagonal[SFWA_StateVector::covariance_size()]) {
    Eigen::Map<UKF::Vector<SFWA_StateVector::covariance_size()>> covariance_map(state_covariance_diagonal);
    covariance_map = ukf.covariance.diagonal();
}

void ukf_get_state_error(real_t state_error[SFWA_StateVector::covariance_size()]) {
    Eigen::Map<typename SFWA_StateVector::StateVectorDelta> error_map(state_error);
    error_map = ukf.covariance.cwiseAbs().rowwise().sum().cwiseSqrt();
}

void ukf_sensor_clear() {
    meas = SFWA_MeasurementVector();
}

void ukf_sensor_set_accelerometer(real_t x, real_t y, real_t z) {
    meas.set_field<Accelerometer>(UKF::Vector<3>(x, y, z));
}

void ukf_sensor_set_gyroscope(real_t x, real_t y, real_t z) {
    meas.set_field<Gyroscope>(UKF::Vector<3>(x, y, z));
}

void ukf_sensor_set_magnetometer(real_t x, real_t y, real_t z) {
    meas.set_field<Magnetometer>(UKF::Vector<3>(x, y, z));
}

void ukf_sensor_set_gps_position(real_t lat, real_t lon, real_t alt) {
    meas.set_field<GPS_Position>(UKF::Vector<3>(lat, lon, alt));
}

void ukf_sensor_set_gps_velocity(real_t x, real_t y, real_t z) {
    meas.set_field<GPS_Velocity>(UKF::Vector<3>(x, y, z));
}

void ukf_sensor_set_pitot_tas(real_t tas) {
    meas.set_field<Airspeed>(tas);
}

void ukf_sensor_set_barometer_amsl(real_t amsl) {
    meas.set_field<PressureAltitude>(amsl);
}

void ukf_set_params(struct ukf_ioboard_params_t *in) {
    local_mag_field << in->mag_field[0], in->mag_field[1], in->mag_field[2];
    ukf.measurement_covariance <<
        in->accel_covariance[0], in->accel_covariance[1], in->accel_covariance[2],
        in->gyro_covariance[0], in->gyro_covariance[1], in->gyro_covariance[2],
        in->mag_covariance[0], in->mag_covariance[1], in->mag_covariance[2],
        in->gps_position_covariance[0], in->gps_position_covariance[1], in->gps_position_covariance[2],
        in->gps_velocity_covariance[0], in->gps_velocity_covariance[1], in->gps_velocity_covariance[2],
        in->pitot_covariance,
        in->barometer_amsl_covariance;
}

void ukf_choose_dynamics(enum ukf_model_t t) {

}

void ukf_set_custom_dynamics_model(ukf_model_function_t func) {

}

void ukf_iterate(float dt, real_t control_vector[UKF_CONTROL_DIM]) {
    ukf.step(dt, meas);
}

void ukf_set_process_noise(real_t process_noise_covariance[SFWA_StateVector::covariance_size()]) {
    Eigen::Map<typename SFWA_StateVector::StateVectorDelta> covariance_map(process_noise_covariance);
    ukf.process_noise_covariance = SFWA_StateVector::CovarianceMatrix::Zero();
    ukf.process_noise_covariance.diagonal() << covariance_map;
}

uint32_t ukf_config_get_state_dim() {
    return SFWA_StateVector::covariance_size();
}

uint32_t ukf_config_get_measurement_dim() {
    return SFWA_MeasurementVector::max_size();
}

uint32_t ukf_config_get_control_dim() {
    return UKF_CONTROL_DIM;
}

enum ukf_precision_t ukf_config_get_precision() {
    return UKF_PRECISION_DOUBLE;
}


// UNDER CONSTRUCTION
#include <utility>

#include "skeletons.h"
#include <boost/graph/connected_components.hpp>
#include <boost/graph/adjacency_list.hpp>
#include <boost/config.hpp>

#include <cctbx/error.h>

namespace cctbx { namespace maptbx
{

typedef boost::adjacency_list<boost::vecS,boost::vecS,boost::undirectedS>
  graph_t;

std::pair<size_t,std::vector<int> > skeleton_components(const joins_t &joins,
  std::size_t nv)
{
  namespace b = boost;
  graph_t g(nv);
  for( const auto & j : joins )
    b::add_edge(j.ilt, j.igt, g);
  size_t nvj = b::num_vertices(g);
  CCTBX_ASSERT( nv == nvj );
  std::vector<int> cm(nv);
  size_t n = b::connected_components(g,&cm[0]);
  return std::make_pair(n,std::move(cm));
}

std::vector<std::size_t> mask_components(marks_t &mask,
  const std::vector<int> &components)
{
  std::vector<std::size_t> sizes(components.size(),0);
  for( auto & m : mask )
  {
    if( m!=0 )
    {
      m = components.at(m);
      ++(sizes.at(m));
    }
  }
  return sizes;
}

void mask_density_map(asymmetric_map::data_ref_t map, const marks_t &mask,
  unsigned val)
{
  // CCTBX_ASSERT( map.accessor() == mask.accessor() );

  CCTBX_ASSERT( map.size() == mask.size() );
  for(size_t i=0; i<map.size(); ++i)
  {
    if( mask[i]!=val )
      map[i] = 0.;
  }
}

// consider useing boost::graph
shortest_paths dijkstra(const skeleton &skelet,
    const std::vector<std::size_t> &molids, std::size_t iatom)
{
  shortest_paths result;
  if( iatom> skelet.maximums.size() )
    throw std::runtime_error("wrong vertex");
  CCTBX_ASSERT( skelet.maximums.size() == molids.size() );
  std::size_t molid = molids[iatom];
  std::size_t nv = std::count(molids.begin(), molids.end(), molid);
  result.distances.resize(nv,0);
  result.predecessors.resize(nv,0);
  std::vector<std::size_t> Q(nv);
  for(std::size_t i=0; i<Q.size(); ++i)
    Q[i] = i;
  std::size_t j=0;
  std::size_t inv = std::numeric_limits<std::size_t>::max();
  do
  {
    std::size_t mloc = std::min_element(result.distances.begin(),
       result.distances.end()) - result.distances.begin();
    mloc = Q[mloc];
    if( mloc == inv )
      break;
    if( j>nv )
      throw std::logic_error("index is out of range");
    ++j;
    for(joins_t::const_iterator i=skelet.joins.begin();
        i!=skelet.joins.end(); ++i)
    {
      const join_t &join = *i; // kelet.joins[i];
      std::size_t adj = 0;
      if( join.ilt == mloc )
        adj = join.igt;
      if( join.igt == mloc )
        adj = join.ilt;
      if( adj!=0 )
      {
        if( result.distances[adj] > result.distances[mloc]+1 )
        {
          result.distances[adj] = result.distances[mloc] + 1;
          result.predecessors[adj] = mloc;
        }
      }
    }
    Q[mloc] = inv;
  } while(true);
  return result;
}

}}


#include "wrappers/qpWrapperBase.hpp"

#include <Eigen/Cholesky>

namespace cvx::internal
{

    QPWrapperBase::QPWrapperBase(OptimizationProblem &problem)
    {
        std::vector<Eigen::Triplet<Parameter>> A_coeffs, P_coeffs;
        std::vector<Parameter> l_coeffs, u_coeffs;

        // Build equality constraint parameters
        for (internal::EqualityConstraint &constraint : problem.equality_constraints)
        {
            constraint.affine.cleanUp();
            if (constraint.affine.isConstant())
            {
                continue;
            }

            for (Term &term : constraint.affine.terms)
            {
                addVariable(term.variable);
                A_coeffs.emplace_back(u_coeffs.size(),
                                      term.variable.getProblemIndex(),
                                      term.parameter);
            }
            l_coeffs.push_back(Parameter(-1.) * constraint.affine.constant);
            u_coeffs.push_back(Parameter(-1.) * constraint.affine.constant);
        }

        // Build positive constraint parameters
        for (internal::PositiveConstraint &constraint : problem.positive_constraints)
        {
            constraint.affine.cleanUp();
            if (constraint.affine.isConstant())
            {
                continue;
            }

            for (Term &term : constraint.affine.terms)
            {
                addVariable(term.variable);
                A_coeffs.emplace_back(u_coeffs.size(),
                                      term.variable.getProblemIndex(),
                                      term.parameter);
            }
            l_coeffs.push_back(Parameter(-1.) * constraint.affine.constant);
            u_coeffs.push_back(Parameter(std::numeric_limits<double>::max()));
        }

        // Build box constraint parameters
        for (internal::BoxConstraint &constraint : problem.box_constraints)
        {
            if (constraint.lower.isConstant() and constraint.upper.isConstant())
            {
                // lower <= middle <= upper
                constraint.middle.cleanUp();
                if (constraint.middle.isConstant())
                {
                    continue;
                }

                for (Term &term : constraint.middle.terms)
                {
                    addVariable(term.variable);
                    A_coeffs.emplace_back(u_coeffs.size(),
                                          term.variable.getProblemIndex(),
                                          term.parameter);
                }
                l_coeffs.push_back(constraint.lower.constant - constraint.middle.constant);
                u_coeffs.push_back(constraint.upper.constant - constraint.middle.constant);
            }
            else
            {
                // c_lower - c_middle <= middle - lower <= inf
                Affine middle_m_lower = constraint.middle - constraint.lower;
                middle_m_lower.cleanUp();

                if (middle_m_lower.isFirstOrder())
                {
                    for (Term &term : middle_m_lower.terms)
                    {
                        addVariable(term.variable);
                        A_coeffs.emplace_back(u_coeffs.size(),
                                              term.variable.getProblemIndex(),
                                              term.parameter);
                    }
                    l_coeffs.push_back(constraint.lower.constant - constraint.middle.constant);
                    u_coeffs.push_back(Parameter(std::numeric_limits<double>::max()));
                }

                // c_middle - c_upper <= upper - middle <= inf
                Affine upper_m_middle = constraint.upper - constraint.middle;
                upper_m_middle.cleanUp();

                if (upper_m_middle.isFirstOrder())
                {
                    for (Term &term : upper_m_middle.terms)
                    {
                        addVariable(term.variable);
                        A_coeffs.emplace_back(u_coeffs.size(),
                                              term.variable.getProblemIndex(),
                                              term.parameter);
                    }
                    l_coeffs.push_back(constraint.middle.constant - constraint.upper.constant);
                    u_coeffs.push_back(Parameter(std::numeric_limits<double>::max()));
                }
            }
        }

        // Build cost function
        if (problem.costFunction.getOrder() == 0 or problem.costFunction.isNorm())
        {
            throw std::runtime_error("QP cost functions must be linear or quadratic.");
        }

        // Linear part
        for (Term &term : problem.costFunction.affine.terms)
        {
            addVariable(term.variable);
            q_params(term.variable.getProblemIndex()) += term.parameter;
        }

        // Quadratic part
        for (Product &product : problem.costFunction.products)
        {
            for (Term &term1 : product.firstTerm().terms)
            {
                for (Term &term2 : product.secondTerm().terms)
                {
                    addVariable(term1.variable);
                    addVariable(term2.variable);

                    const std::pair<size_t, size_t> sorted = std::minmax(term1.variable.getProblemIndex(),
                                                                         term2.variable.getProblemIndex());

                    Parameter param = term1.parameter * term2.parameter;

                    // Explicitly double diagonal elements
                    if (sorted.first == sorted.second)
                    {
                        param *= Parameter(2.);
                    }

                    P_coeffs.emplace_back(sorted.first,
                                          sorted.second,
                                          param);
                }
            }

            // The linear parts from the multiplication
            if (not product.firstTerm().constant.isZero())
            {
                for (Term &term : product.secondTerm().terms)
                {
                    addVariable(term.variable);
                    q_params(term.variable.getProblemIndex()) += product.firstTerm().constant * term.parameter;
                }
            }
            if (not product.secondTerm().constant.isZero())
            {
                for (Term &term : product.firstTerm().terms)
                {
                    addVariable(term.variable);
                    q_params(term.variable.getProblemIndex()) += product.secondTerm().constant * term.parameter;
                }
            }
        }

        assert(l_coeffs.size() == u_coeffs.size());

        // Fill matrices and vectors
        A_params.resize(l_coeffs.size(), getNumVariables());
        P_params.resize(getNumVariables(), getNumVariables());

        A_params.setFromTriplets(A_coeffs.begin(), A_coeffs.end());
        P_params.setFromTriplets(P_coeffs.begin(), P_coeffs.end());

        l_params = Eigen::Map<VectorXp>(l_coeffs.data(), l_coeffs.size());
        u_params = Eigen::Map<VectorXp>(u_coeffs.data(), u_coeffs.size());

        solution->resize(getNumVariables());
    }

    size_t QPWrapperBase::getNumInequalityConstraints() const
    {
        return A_params.rows();
    }

    void QPWrapperBase::addVariable(Variable &variable)
    {
        const bool was_linked = variable.linkToSolver(solution, getNumVariables());
        if (was_linked)
        {
            variables.push_back(variable);
            q_params.conservativeResize(getNumVariables());
        }
    }

    std::ostream &operator<<(std::ostream &os, const QPWrapperBase &wrapper)
    {
        Eigen::MatrixXd A = eval(wrapper.A_params);
        Eigen::MatrixXd P = eval(wrapper.P_params);
        P += P.triangularView<Eigen::StrictlyUpper>().transpose();
        Eigen::VectorXd q = eval(wrapper.q_params);
        Eigen::VectorXd l = eval(wrapper.l_params);
        Eigen::VectorXd u = eval(wrapper.u_params);

        os << "Quadratic problem\n";
        os << "Minimize 0.5x'Px + q'x\n";
        os << "Subject to l <= Ax <= u\n";
        os << "With:\n\n";

        os << "P:\n"
           << P << "\n\n";
        os << "q:\n"
           << q << "\n\n";
        os << "A:\n"
           << A << "\n\n";
        os << "l:\n"
           << l << "\n\n";
        os << "u:\n"
           << u;

        return os;
    }

    bool QPWrapperBase::isConvex() const
    {
        if (P_params.nonZeros() == 0)
        {
            return true;
        }
        else
        {
            Eigen::LLT<Eigen::MatrixXd> llt(eval(P_params));
            return llt.info() != Eigen::NumericalIssue;
        }
    }

} // namespace cvx::internal


#include <iostream>
#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/cycle_canceling.hpp>
#include <boost/graph/push_relabel_max_flow.hpp>
#include <boost/graph/successive_shortest_path_nonnegative_weights.hpp>
#include <boost/graph/find_flow_cost.hpp>

// Graph Type with nested interior edge properties for Cost Flow Algorithms
typedef boost::adjacency_list_traits<boost::vecS, boost::vecS, boost::directedS> traits;
typedef boost::adjacency_list<boost::vecS, boost::vecS, boost::directedS, boost::no_property,
                              boost::property<boost::edge_capacity_t, long,
                                              boost::property<boost::edge_residual_capacity_t, long,
                                                              boost::property<boost::edge_reverse_t, traits::edge_descriptor,
                                                                              boost::property<boost::edge_weight_t, long>>>>>
    graph; // new! weightmap corresponds to costs
typedef boost::graph_traits<graph>::edge_descriptor edge_desc;
typedef boost::graph_traits<graph>::out_edge_iterator out_edge_it; // Iterator

// Custom edge adder class
class edge_adder
{
    graph &G;

public:
    explicit edge_adder(graph &G) : G(G) {}
    void add_edge(int from, int to, long capacity, long cost)
    {
        auto c_map = boost::get(boost::edge_capacity, G);
        auto r_map = boost::get(boost::edge_reverse, G);
        auto w_map = boost::get(boost::edge_weight, G); // new!
        const edge_desc e = boost::add_edge(from, to, G).first;
        const edge_desc rev_e = boost::add_edge(to, from, G).first;
        c_map[e] = capacity;
        c_map[rev_e] = 0; // reverse edge has no capacity!
        r_map[e] = rev_e;
        r_map[rev_e] = e;
        w_map[e] = cost;      // new assign cost
        w_map[rev_e] = -cost; // new negative cost
    }
};

using namespace std;

void solve()
{
    int buyers, sites, states;
    cin >> buyers >> sites >> states;

    graph G(states + sites + buyers);
    edge_adder adder(G);
    auto c_map = boost::get(boost::edge_capacity, G);
    auto r_map = boost::get(boost::edge_reverse, G);
    auto rc_map = boost::get(boost::edge_residual_capacity, G);

    auto source = boost::add_vertex(G);
    auto target = boost::add_vertex(G);

    int stateLimit;
    for (int state = 0; state < states; ++state) {
        cin >> stateLimit;
        adder.add_edge(source, state, stateLimit, 0);
    }

    int state;
    int offsetSites = states;
    for (int site = offsetSites; site < offsetSites + sites; ++site) {
        cin >> state; --state; // convert to 0 based
        adder.add_edge(state, site, 1, 0);
    }

    int bid;
    int maxBid = 100;
    int offsetBuyers = states + sites;
    for (int buyer = offsetBuyers; buyer < offsetBuyers + buyers; ++buyer) {
        for (int site = offsetSites; site < offsetSites + sites; ++site) {
            cin >> bid;
            adder.add_edge(site, buyer, 1, maxBid - bid); // Avoid negative costs
        }
        adder.add_edge(buyer, target, 1, 0);
    }

    boost::successive_shortest_path_nonnegative_weights(G, source, target);
    auto cost = boost::find_flow_cost(G);

    int flow = 0;
    out_edge_it e, eend;
    for (boost::tie(e, eend) = boost::out_edges(boost::vertex(source, G), G); e != eend; ++e)
    {
        flow += c_map[*e] - rc_map[*e];
    }
    cost = maxBid * flow - cost; // Convert back to actual cost

    cout << flow << " " << cost << endl;
}

int main()
{
    int t; cin >> t;
    for (int i = 0; i < t; ++i) {
        solve();
    }
    return 0;
}


/*
 * StokesM2L.cpp
 *
 *  Created on: Oct 12, 2016
 *      Author: wyan
 */

#include "SVD_pvfmm.hpp"

#include <Eigen/Dense>

#include <iomanip>
#include <iostream>

#define DIRECTLAYER 2

namespace StokesPVel1D3D {

using EVec3 = Eigen::Vector3d;
using EVec4 = Eigen::Vector4d;
using EMat3 = Eigen::Matrix3d;
using EMat4 = Eigen::Matrix4d;

/**
 * \brief Returns the coordinates of points on the surface of a cube.
 * \param[in] p Number of points on an edge of the cube is (n+1)
 * \param[in] c Coordinates to the centre of the cube (3D array).
 * \param[in] alpha Scaling factor for the size of the cube.
 * \param[in] depth Depth of the cube in the octree.
 * \return Vector with coordinates of points on the surface of the cube in the
 * format [x0 y0 z0 x1 y1 z1 .... ].
 */
template <class Real_t>
std::vector<Real_t> surface(int p, Real_t *c, Real_t alpha, int depth) {
    size_t n_ = (6 * (p - 1) * (p - 1) + 2); // Total number of points.

    std::vector<Real_t> coord(n_ * 3);
    coord[0] = coord[1] = coord[2] = -1.0;
    size_t cnt = 1;
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = -1.0;
            coord[cnt * 3 + 1] = (2.0 * (i + 1) - p + 1) / (p - 1);
            coord[cnt * 3 + 2] = (2.0 * j - p + 1) / (p - 1);
            cnt++;
        }
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = (2.0 * i - p + 1) / (p - 1);
            coord[cnt * 3 + 1] = -1.0;
            coord[cnt * 3 + 2] = (2.0 * (j + 1) - p + 1) / (p - 1);
            cnt++;
        }
    for (int i = 0; i < p - 1; i++)
        for (int j = 0; j < p - 1; j++) {
            coord[cnt * 3] = (2.0 * (i + 1) - p + 1) / (p - 1);
            coord[cnt * 3 + 1] = (2.0 * j - p + 1) / (p - 1);
            coord[cnt * 3 + 2] = -1.0;
            cnt++;
        }
    for (size_t i = 0; i < (n_ / 2) * 3; i++)
        coord[cnt * 3 + i] = -coord[i];

    Real_t r = 0.5 * pow(0.5, depth);
    Real_t b = alpha * r;
    for (size_t i = 0; i < n_; i++) {
        coord[i * 3 + 0] = (coord[i * 3 + 0] + 1.0) * b + c[0];
        coord[i * 3 + 1] = (coord[i * 3 + 1] + 1.0) * b + c[1];
        coord[i * 3 + 2] = (coord[i * 3 + 2] + 1.0) * b + c[2];
    }
    return coord;
}

// fx,fy,fz,trD -> p, vx,vy,vz
inline void Wkernel(const EVec3 &target, const EVec3 &source, EMat4 &answer) {
    auto rst = target - source;
    double rnorm = rst.norm();
    if (rnorm < 1e-13) {
        answer.setZero();
        return;
    }
    double rnorm3 = rnorm * rnorm * rnorm;

    answer.block<3, 3>(1, 0) = EMat3::Identity() / rnorm;
    answer.block<3, 3>(1, 0) += rst * rst.transpose() / rnorm3;

    answer(0, 0) = rst[0] / rnorm3;
    answer(0, 1) = rst[1] / rnorm3;
    answer(0, 2) = rst[2] / rnorm3;
    answer(0, 3) = 0;
    answer(1, 3) = -rst[0] / rnorm3;
    answer(2, 3) = -rst[1] / rnorm3;
    answer(3, 3) = -rst[2] / rnorm3;
    answer.row(0) *= (1 / (4 * M_PI));
    answer.block<3, 4>(1, 0) *= (1 / (8 * M_PI));
}

inline void WkernelFF(const EVec3 &target, const EVec3 &source, EMat4 &answer) {
    answer.setZero();
    const int imageN = 1000000; // images to sum
    for (int per = DIRECTLAYER + 1; per < imageN; per++) {
        EVec3 perVec(1.0 * per, 0, 0);
        EMat4 W1 = EMat4::Zero();
        EMat4 W2 = EMat4::Zero();
        Wkernel(target, source + perVec, W1);
        Wkernel(target, source - perVec, W2);
        answer += (W1 + W2);
    }
}

// calculate the M2L matrix of images from 2 to 1000
int main(int argc, char **argv) {
    Eigen::initParallel();
    Eigen::setNbThreads(1);

    const int pEquiv = atoi(argv[1]); // (8-1)^2*6 + 2 points
    const int pCheck = atoi(argv[1]);
    const double scaleEquiv = 1.05;
    const double scaleCheck = 2.95;
    const double pCenterEquiv[3] = {-(scaleEquiv - 1) / 2, -(scaleEquiv - 1) / 2, -(scaleEquiv - 1) / 2};
    const double pCenterCheck[3] = {-(scaleCheck - 1) / 2, -(scaleCheck - 1) / 2, -(scaleCheck - 1) / 2};
    auto pointMEquiv = surface(pEquiv, (double *)&(pCenterEquiv[0]), scaleEquiv,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0
    auto pointMCheck = surface(pCheck, (double *)&(pCenterCheck[0]), scaleCheck,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0

    auto pointLEquiv = surface(pEquiv, (double *)&(pCenterCheck[0]), scaleCheck,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0
    auto pointLCheck = surface(pCheck, (double *)&(pCenterEquiv[0]), scaleEquiv,
                               0); // center at 0.5,0.5,0.5, periodic box 1,1,1, scale 1.05, depth = 0

    // calculate the operator M2L with least square
    const int equivN = pointMEquiv.size() / 3;
    const int checkN = pointLCheck.size() / 3;
    Eigen::MatrixXd A(4 * checkN, 4 * equivN);
#pragma omp parallel for
    for (int k = 0; k < checkN; k++) {
        EMat4 W;
        W.setZero();
        Eigen::Vector3d Cpoint(pointLCheck[3 * k], pointLCheck[3 * k + 1], pointLCheck[3 * k + 2]);
        for (int l = 0; l < equivN; l++) {
            const Eigen::Vector3d Lpoint(pointLEquiv[3 * l], pointLEquiv[3 * l + 1], pointLEquiv[3 * l + 2]);
            Wkernel(Cpoint, Lpoint, W);
            A.block<4, 4>(4 * k, 4 * l) = W;
        }
    }
    Eigen::MatrixXd ApinvU(A.cols(), A.rows());
    Eigen::MatrixXd ApinvVT(A.cols(), A.rows());
    pinv(A, ApinvU, ApinvVT);

    Eigen::MatrixXd M2L(4 * equivN, 4 * equivN);
#pragma omp parallel for
    for (int i = 0; i < equivN; i++) {
        const Eigen::Vector3d Mpoint(pointMEquiv[3 * i], pointMEquiv[3 * i + 1], pointMEquiv[3 * i + 2]);
        Eigen::MatrixXd f(4 * checkN, 4);
        for (int k = 0; k < checkN; k++) {
            EMat4 temp = EMat4::Zero();
            Eigen::Vector3d Cpoint(pointLCheck[3 * k], pointLCheck[3 * k + 1], pointLCheck[3 * k + 2]);
            WkernelFF(Cpoint, Mpoint, temp);
            f.block<4, 4>(4 * k, 0) = temp;
        }
        M2L.block(0, 4 * i, 4 * equivN, 4) = (ApinvU.transpose() * (ApinvVT.transpose() * f));
    }

    // dump M2L
    for (int i = 0; i < 4 * equivN; i++) {
        for (int j = 0; j < 4 * equivN; j++) {
            std::cout << i << " " << j << " " << std::scientific << std::setprecision(18) << M2L(i, j) << std::endl;
        }
    }

    // Test
    // Sum of force and trD must be zero
    std::vector<EVec3, Eigen::aligned_allocator<EVec3>> forcePoint(3);
    std::vector<EVec4, Eigen::aligned_allocator<EVec4>> forceValue(3);
    forcePoint[0] = EVec3(0.5, 0.55, 0.2);
    forcePoint[1] = EVec3(0.5, 0.5, 0.5);
    forcePoint[2] = EVec3(0.7, 0.7, 0.7);
    forceValue[0] = EVec4(0.1, 0.2, 0.3, 0.4);
    forceValue[1] = EVec4(-0.1, -0.1, -0.3, -0.4);
    forceValue[2] = EVec4(0, -0.1, 0, 0);

    // solve M
    A.resize(4 * checkN, 4 * equivN);
    ApinvU.resize(A.cols(), A.rows());
    ApinvVT.resize(A.cols(), A.rows());
    Eigen::VectorXd f(4 * checkN);
#pragma omp parallel for
    for (int k = 0; k < checkN; k++) {
        Eigen::Vector3d Cpoint(pointMCheck[3 * k], pointMCheck[3 * k + 1], pointMCheck[3 * k + 2]);
        EVec4 temp = EVec4::Zero();
        for (int p = 0; p < forceValue.size(); p++) {
            EMat4 W = EMat4::Zero();
            Wkernel(Cpoint, forcePoint[p], W);
            temp += W * (forceValue[p]);
        }
        f.block<4, 1>(4 * k, 0) = temp;
        for (int l = 0; l < equivN; l++) {
            Eigen::Vector3d Mpoint(pointMEquiv[3 * l], pointMEquiv[3 * l + 1], pointMEquiv[3 * l + 2]);
            EMat4 W = EMat4::Zero();
            Wkernel(Cpoint, Mpoint, W);
            A.block<4, 4>(4 * k, 4 * l) = W;
        }
    }
    pinv(A, ApinvU, ApinvVT);
    Eigen::VectorXd Msource = (ApinvU.transpose() * (ApinvVT.transpose() * f));

    std::cout << "Msource: " << Msource << std::endl;
    std::cout << "Msource Sum: " << Msource.sum() << std::endl;
    Eigen::VectorXd M2Lsource = M2L * (Msource);

    {
        EVec3 samplePoint(0.5, 0.2, 0.8);
        // Compute: WFF from L, WFF from WkernelFF
        EVec4 WFFL = EVec4::Zero();
        EVec4 WFFK = EVec4::Zero();

        for (int k = 0; k < equivN; k++) {
            EVec3 Lpoint(pointLEquiv[3 * k], pointLEquiv[3 * k + 1], pointLEquiv[3 * k + 2]);
            EMat4 W = EMat4::Zero();
            Wkernel(samplePoint, Lpoint, W);
            WFFL += W * M2Lsource.block<4, 1>(4 * k, 0);
        }

        for (int k = 0; k < forceValue.size(); k++) {
            EMat4 W;
            WkernelFF(samplePoint, forcePoint[k], W);
            WFFK += W * forceValue[k];
        }
        std::cout << "WFF from Lequiv: " << WFFL << std::endl;
        std::cout << "WFF from Kernel: " << WFFK << std::endl;
        std::cout << "FF Error: " << WFFL - WFFK << std::endl;
    }

    return 0;
}

} // namespace StokesPVel1D3D

#undef DIRECTLAYER


//
// Created by Alex Beccaro on 18/01/18.
//

#define BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#include "../../src/problems/1-50/6/problem6.hpp"

BOOST_AUTO_TEST_SUITE( Problem6 )

    BOOST_AUTO_TEST_CASE( Example ) {
        auto res = problems::problem6::solve(10);
        BOOST_CHECK_EQUAL(res, 2640);
    }

    BOOST_AUTO_TEST_CASE( Solution ) {
        auto res = problems::problem6::solve();
        BOOST_CHECK_EQUAL(res, 25164150);
    }

BOOST_AUTO_TEST_SUITE_END()


#include "../include/Length.h"

#include <vector>
#include <fstream>
#include <iostream>

#include "../include/utility.h"
#include "../include/Timer.h"
#include "../include/exceptions.h"

#include <boost/lexical_cast.hpp>
#include <boost/iostreams/filtering_streambuf.hpp>
#include <boost/iostreams/copy.hpp>
#include <boost/iostreams/filter/zlib.hpp>
//#include "Model.cpp"

using namespace boost;

string Length::ending = "length";

Length::Length(Corpus* co) : Model<Length>(co)
{
    size1 = LMAX;
    size2 = LMAX;
    systemMessage("Creating Length Model");
    model = new float*[LMAX];
    float norm[LMAX];
    for(int i = 0; i < LMAX; i++)
    {
        model[i] = new float[LMAX];
        norm[i] = SMOOTH * LMAX;
        for(int j = 0; j < LMAX; j++)
        {
            model[i][j] = SMOOTH;
        }
    }

    vector<int*>::iterator eIt;
    vector<int*>::iterator fIt;

    int m = 0;
    int n = 0;

    for(eIt = co->getElines()->begin(), fIt = co->getFlines()->begin(); eIt != co->getElines()->end(), fIt != co->getFlines()->end(); eIt++, fIt++)
    {
        m = (*eIt)[0];
        n = (*fIt)[0];
        if(m < LMAX && n < LMAX)
        {
            model[n][m] += 1.0;
            norm[m] += 1.0;
        }
    }

    for(int i = 0; i < LMAX; i++)
    {
        for(int j = 0; j < LMAX; j++)
        {
            model[i][j] = model[i][j]/norm[j];
        }
    }
    saveModel();
}

float** Length::getLengthmodel()
{
    return model;
}


//==============================================================================
//         Copyright 2003 - 2011   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2013   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//         Copyright 2012 - 2013   MetaScale SAS
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#ifndef BOOST_SIMD_META_IS_POWER_OF_2_HPP_INCLUDED
#define BOOST_SIMD_META_IS_POWER_OF_2_HPP_INCLUDED

#include <cstddef>
#include <boost/mpl/bool.hpp>

namespace boost { namespace simd {  namespace meta
{
  /*!
    @brief Checks if an integral value is a power of two.

    is_power_of_2_c checks if any given integral @c N is a non-zero power of two.

    @par Semantic:
    For any given integral value @c N:

    @code
    typedef is_power_of_2_c<N>::type r;
    @endcode

    is equivalent to:

    @code
    typedef mpl::bool<(!(N & (N - 1)) && N)> r;
    @endcode

    @usage{meta/is_power_of_2_c.cpp}

    @tparam N Integral value to check
  **/
  template<std::size_t N> struct  is_power_of_2_c
#if !defined(DOXYGEN_ONLY)
        : boost::mpl::bool_<(!(N & (N - 1)) && N)>
#endif
  {};

  /*!
    @brief Checks if an @mplint is a power of two.

    is_power_of_2_c is a Boolean @metafunction that checks if any given
    @mplint @c N is a non-zero power of two.

    @par Semantic:
    For any given @mplint @c N:

    @code
    typedef is_power_of_2<N>::type r;
    @endcode

    is equivalent to:

    @code
    typedef is_power_of_2_c<N::value>::type r;
    @endcode

    @usage{meta/is_power_of_2.cpp}

    @tparam N @mplint to check
  **/
  template<class N> struct  is_power_of_2
#if !defined(DOXYGEN_ONLY)
        : boost::mpl::bool_<(!(N::value & (N::value - 1)) && N::value)>
#endif
  {};
} } }

#endif


//==============================================================================
//         Copyright 2003 - 2013   LASMEA UMR 6602 CNRS/Univ. Clermont II
//         Copyright 2009 - 2013   LRI    UMR 8623 CNRS/Univ Paris Sud XI
//
//          Distributed under the Boost Software License, Version 1.0.
//                 See accompanying file LICENSE.txt or copy at
//                     http://www.boost.org/LICENSE_1_0.txt
//==============================================================================
#include <nt2/trigonometric/include/functions/indeg.hpp>

#include <nt2/sdk/functor/meta/call.hpp>
#include <boost/simd/sdk/simd/native.hpp>
#include <nt2/sdk/unit/tests/type_expr.hpp>
#include <nt2/sdk/unit/tests/ulp.hpp>
#include <nt2/include/functions/splat.hpp>
#include <nt2/include/functions/unary_minus.hpp>
#include <nt2/sdk/unit/module.hpp>
#include <boost/simd/sdk/config.hpp>
#include <boost/simd/sdk/simd/io.hpp>

#include <nt2/include/constants/pi.hpp>
#include <nt2/include/constants/pio_2.hpp>
#include <nt2/include/constants/pio_4.hpp>
#include <nt2/include/constants/zero.hpp>
#include <nt2/include/constants/inf.hpp>
#include <nt2/include/constants/minf.hpp>
#include <nt2/include/constants/nan.hpp>
#include <nt2/include/constants/_45.hpp>
#include <nt2/include/constants/_90.hpp>
#include <nt2/include/constants/_180.hpp>

NT2_TEST_CASE_TPL ( indeg_real_1,  NT2_SIMD_REAL_TYPES)
{
  using nt2::indeg;
  using nt2::tag::indeg_;
  using boost::simd::native;
  typedef BOOST_SIMD_DEFAULT_EXTENSION  ext_t;
  typedef native<T,ext_t>                  vT;

  typedef typename nt2::meta::call<indeg_(vT)>::type r_t;
  typedef vT wished_r_t;

  // return type conformity test
  NT2_TEST_TYPE_IS(r_t, wished_r_t);

  // specific values tests
#ifndef BOOST_SIMD_NO_INVALIDS
  NT2_TEST_ULP_EQUAL(indeg(nt2::Inf<vT>()), nt2::Inf<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Minf<vT>()), nt2::Minf<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Nan<vT>()), nt2::Nan<r_t>(), 0.5);
#endif
  NT2_TEST_ULP_EQUAL(indeg(-nt2::Pi<vT>()), -nt2::_180<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(-nt2::Pio_2<vT>()), -nt2::_90<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(-nt2::Pio_4<vT>()), -nt2::_45<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Pi<vT>()), nt2::_180<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Pio_2<vT>()), nt2::_90<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Pio_4<vT>()), nt2::_45<r_t>(), 0.5);
  NT2_TEST_ULP_EQUAL(indeg(nt2::Zero<vT>()), nt2::Zero<r_t>(), 0.5);
}


//
// Created by riku on 2016/07/19.
//
#include "core/ctxt_util.hpp"
#include "core/literal.hpp"
#include "HElib/Ctxt.h"
#include "HElib/EncryptedArray.h"
#include "HElib/PAlgebra.h"
#include <NTL/ZZ.h>
#include <random>
#include <algorithm>
#include <sstream>

namespace core {
size_t number_bits(long a) {
    size_t bits = 0;
    while (a > 0) {
        bits += 1;
        a = a >> 1;
    }
    return bits;
}

static bool test_bit(long a, size_t i) {
    return (a & (1 << i)) != 0;
}

Ctxt repeat0(const Ctxt &c,
             const long n_slots, const long rep,
             const EncryptedArray &ea) {
    Ctxt repeated(c);
    Ctxt res(c.getPubKey()); /// zero-like
    long offset = n_slots;
    for (size_t i = 0; i < number_bits(rep); i++) {
        if (test_bit(rep, i)) {
            FHE_NTIMER_START(ea_rotate);
            ea.rotate(res, offset);
            FHE_NTIMER_STOP(ea_rotate);
            res += repeated;
        }
        Ctxt tmp(repeated);
        FHE_NTIMER_START(ea_rotate);
        ea.rotate(tmp, offset);
        FHE_NTIMER_STOP(ea_rotate);
        repeated += tmp;
        offset = offset << 1;
    }
    return res;
}

Ctxt repeat(const Ctxt &c,
            const long n_slots, const long rep,
            const EncryptedArray &ea) {
    assert(n_slots * rep <= ea.size() && "n_slots * rep > ea.size()");
    assert(n_slots >= 0 && rep >= 0 && "n_slots < 0 || rep < 0");
    Ctxt repeated(c);
    mask_first(repeated, n_slots, ea);
    return repeat0(repeated, n_slots, rep, ea);
}

// NOTE. This function might introduce too many noise.
void replicate(Ctxt *out, const int pos ,const int length, const EncryptedArray *ea) {
    ZZX mask;
    ea->encodeUnitSelector(mask, pos);
    out->multByConstant(mask);
    ea->rotate(*out, -pos);
    long offset = 0;
    long k = number_bits(length);
    Ctxt ctxt_orig(*out);
    long e = 1;
    // now process bits k-2 down to 0
    for (long j = k-2; j >= 0; j--) {
        // e -> 2*e
        Ctxt tmp = *out;
        ea->rotate(tmp, e);
        out->addCtxt(tmp);
        e <<= 1;

        long b = test_bit(length, j);
        // e -> e+b
        if (b) {
            ea->rotate(*out, 1);
            out->addCtxt(ctxt_orig);
            e++;
        }
    }
}

void mask_first(Ctxt &ctxt, size_t n, const EncryptedArray &ea) {
    std::vector<long> mask(ea.size(), 0);
    for (long i = 0; i < n; i++) mask[i] = 1;
    NTL::ZZX poly;
    ea.encode(poly, mask);
    ctxt.multByConstant(poly);
}

std::vector<std::vector<long>> random_permutation(long D, const EncryptedArray &ea) {
    size_t nr = (D + ea.size() - 1) / ea.size();
    std::vector<std::vector<long>> parts(nr, std::vector<long>(ea.size(), 0));
    std::vector<long> I(nr * ea.size(), 0);
    for (size_t i = 0; i < I.size(); i++) I.at(i) = i;

    std::random_shuffle(I.begin(), I.end());
    auto itr = I.begin();
    for (auto &part : parts) {
        for (auto &value : part) {
            value = *itr;
            itr++;
        }
    }
    return parts;
}

std::vector<std::vector<long>> randomness(long D, const EncryptedArray &ea) {
    size_t nr = (D + ea.size() - 1) / ea.size();
    std::vector<std::vector<long>> parts(nr, std::vector<long>(ea.size()));
    std::vector<long> I(nr * ea.size());
    long pr = ea.getAlMod().getPPowR();
    for (size_t i = 0; i < I.size(); i++) {
        do{
            I.at(i) = NTL::RandomBnd(pr);
        } while (I.at(i) == 0);
    }
    auto itr = I.begin();
    for (auto &part : parts) {
        for (auto &value : part) {
            value = *itr;
            itr++;
        }
    }
    return parts;
}

template<>
std::string conv(const Ctxt &obj) {
    std::stringstream sstream;
    sstream << obj;
    return sstream.str();
}

template<>
void conv(Ctxt &obj, const std::string &str) {
    std::stringstream sstream(str);
    sstream >> obj;
}

bool dumpCtxts(const std::vector<Ctxt> &ctxts,
               const std::string &outputDirPath) {
    auto makePath = [](std::string const& path, const long count) -> std::string {
        return path + literal::separator + "FILE_" + std::to_string(count);
    };

    long file_nr = 1;
    long ctx_dumped = 0;
    std::ofstream fout(makePath(outputDirPath, file_nr), std::ios::binary | std::ios::trunc);
    if (!fout.is_open()) {
        return false;
    }

    for (const Ctxt &ctx :ctxts) {
        if (ctx_dumped < core::core_setting.CTX_PER_FILE) {
            fout << ctx;
            ctx_dumped += 1;
        } else {
            fout.close();
            file_nr += 1;
            fout.open(makePath(outputDirPath, file_nr), std::ios::binary);
            if (!fout.is_open())
                return false;
            fout << ctx;
            ctx_dumped = 1;
        }
    }
    fout.close();
    return true;
}

}


//
// Copyright (c) 2019 INRIA
//

#include <pinocchio/autodiff/casadi.hpp>
#include <Eigen/Dense>

#include <boost/variant.hpp> // to avoid C99 warnings

#include <boost/test/unit_test.hpp>
#include <boost/utility/binary.hpp>

BOOST_AUTO_TEST_SUITE(BOOST_TEST_MODULE)

BOOST_AUTO_TEST_CASE(test_eigen)
{
  Eigen::Matrix<casadi::SX, 3, 3> A, B;
  Eigen::Matrix<casadi::SX, 3, 1> a, b;
  Eigen::Matrix<casadi::SX, 3, 1> c = A * a - B.transpose() * b;
}

// A function working with Eigen::Matrix'es parameterized by the Scalar type
template <typename Scalar, typename T1, typename T2, typename T3, typename T4>
Eigen::Matrix<Scalar, Eigen::Dynamic, 1>
eigenFun(Eigen::MatrixBase<T1> const& A,
         Eigen::MatrixBase<T2> const& a,
         Eigen::MatrixBase<T3> const& B,
         Eigen::MatrixBase<T4> const& b)
{
  Eigen::Matrix<Scalar, Eigen::Dynamic, 1> c(4);
  c.segment(1, 3) = A * a.segment(1, 3) - B.transpose() * b;
  c[0] = 0.;

  return c;
}

BOOST_AUTO_TEST_CASE(test_example)
{
  // Declare casadi symbolic matrix arguments
  casadi::SX cs_a = casadi::SX::sym("a", 4);
  casadi::SX cs_b = casadi::SX::sym("b", 3);

  // Declare Eigen matrices
  Eigen::Matrix<casadi::SX, 3, 3> A, B;
  Eigen::Matrix<casadi::SX, -1, 1> a (4), c (4);
  Eigen::Matrix<casadi::SX, 3, 1> b;

  // Let A, B be some numeric matrices
  for (Eigen::Index i = 0; i < A.rows(); ++i)
  {
    for (Eigen::Index j = 0; j < A.cols(); ++j)
    {
      A(i, j) = 10. * static_cast<double>(i) + static_cast<double>(j);
      B(i, j) = -10. * static_cast<double>(i) - static_cast<double>(j);
    }
  }

  // Let a, b be symbolic arguments of a function
  pinocchio::casadi::copy(cs_b, b);
  pinocchio::casadi::copy(cs_a, a);

  // Call the function taking Eigen matrices
  c = eigenFun<casadi::SX>(A, a, B, b);

  // Copy the result from Eigen matrices to casadi matrix
  casadi::SX cs_c = casadi::SX(casadi::Sparsity::dense(c.rows(), 1));
  pinocchio::casadi::copy(c, cs_c);

  // Display the resulting casadi matrix
  std::cout << "c = " << cs_c << std::endl;

  // Do some AD
  casadi::SX dc_da = jacobian(cs_c, cs_a);

  // Display the resulting jacobian
  std::cout << "dc/da = " << dc_da << std::endl;

  // Create a function which takes a, b and returns c and dc_da
  casadi::Function fun("fun", casadi::SXVector {cs_a, cs_b}, casadi::SXVector {cs_c, dc_da});
  std::cout << "fun = " << fun << std::endl;

  // Evaluate the function
  casadi::DMVector res = fun(casadi::DMVector {std::vector<double> {1., 2., 3., 4.}, std::vector<double> {-1., -2., -3.}});
  std::cout << "fun(a, b)=" << res << std::endl;
}

BOOST_AUTO_TEST_CASE(test_jacobian)
{
  casadi::SX cs_x = casadi::SX::sym("x", 3);

  casadi::SX cs_y = casadi::SX::sym("y", 1);
  cs_y(0) = cs_x(0) + cs_x(1) + cs_x(2);

  // Display the resulting expression
  std::cout << "y = " << cs_y << std::endl;

  // Do some AD
  casadi::SX dy_dx = jacobian(cs_x, cs_x);

  // Display the resulting jacobian
  std::cout << "dy/dx = " << dy_dx << std::endl;
}

BOOST_AUTO_TEST_CASE(test_copy_casadi_to_eigen)
{
  casadi::SX cs_mat = casadi::SX::sym("A", 3, 4);
  Eigen::Matrix<casadi::SX, 3, 4> eig_mat;

  pinocchio::casadi::copy(cs_mat, eig_mat);
  std::cout << eig_mat << std::endl;
}

BOOST_AUTO_TEST_CASE(test_copy_eigen_to_casadi)
{
  Eigen::Matrix<casadi::SX, 3, 4> eig_mat;
  pinocchio::casadi::sym(eig_mat, "A");

  casadi::SX cs_mat;

  pinocchio::casadi::copy(eig_mat, cs_mat);
  std::cout << cs_mat << std::endl;
}

BOOST_AUTO_TEST_CASE(test_casadi_codegen)
{
  casadi::SX x = casadi::SX::sym("x");
  casadi::SX y = casadi::SX::sym("y");
  casadi::Function fun("fun", casadi::SXVector {x, y}, casadi::SXVector {x + y});

  casadi::CodeGenerator gen("module");
  gen.add(fun);

  std::cout << gen.dump();
}

BOOST_AUTO_TEST_CASE(test_max)
{
  casadi::SX x = casadi::SX::sym("x");
  casadi::SX y = casadi::SX::sym("y");

  casadi::SX max_x_y = pinocchio::math::max(x,y);
  casadi::SX max_x_0 = pinocchio::math::max(x,0.);
  casadi::SX max_0_y = pinocchio::math::max(0.,y);
}

BOOST_AUTO_TEST_SUITE_END()


#include <math.h>
#include <stdio.h>

#include <boost/math/quaternion.hpp>

extern "C" void externalForces(double t, double *f, unsigned int size_z, double *z)
{
  f[0]=0;
  f[1]=0;
  f[2]=-10;
}
extern "C" void externalMoment(double t,double *m, unsigned int size_z, double *z)
{
  m[0]=0;
  m[1]=0;
  m[2]=0;
}
extern "C" void externalForcesB1(double t, double *f, unsigned int size_z,double *z)
{
  f[0]=0;
  f[1]=0;
  f[2]=-9.81*0.038;
  // printf("externalForcesB1 :\n");
  // printf("f[0] = %e\t f[1] = %e\t, f[2]=%e\n",f[0],f[1],f[2]);
}
extern "C" void internalForcesB1(double t, double *q, double *v, double *f, unsigned int size_z,double *z)
{
  // Simple spring in z direction.
  f[0]=0;
  f[1]=0;
  f[2]=1e4*q[2];
  // printf("internalForcesB1 :\n");
  // printf("f[0] = %e\t f[1] = %e\t, f[2]=%e\n",f[0],f[1],f[2]);
}

extern "C" void internalForcesB1_Jacq(double t, double *q, double *v, double *jac, unsigned int size_z,double *z)
{
  for(int i =0; i < 3; i++)
  {
    for(int j=0; j<7; j++)
      jac[i+j*3]=0.0;
  }
  jac[2+2*3]=1e4;
  // printf("internalForcesB1_Jacq :\n");
  // printf("jac[2+2*3] = %e\n", jac[2+2*3]);
}

extern "C" void internalMomentsB1(double t, double *q, double *v, double *m, unsigned int size_z,double *z)
{
  //  printf("internalMomentsB1 :\n");
  // Simple torsional spring around y axis
  // printf("q[3] = %e\n", q[3]);
  // printf("q[4] = %e\n", q[4]);
  // printf("q[5] = %e\n", q[5]);
  // printf("q[6] = %e\n", q[6]);

  double angle = 2*asin(q[5]);
  // printf("angle = %e\n", angle);
  m[0]=0.0;
  m[1]=1e3*(angle);
  m[2]=0.0;
  // printf("m[0] = %e\t m[1] = %e\t, m[2]=%e\n",m[0],m[1],m[2]);
}

extern "C" void internalMomentsB1_Jacq(double t, double *q, double *v, double *jac, unsigned int size_z,double *z)
{
  //printf("internalMomentsB1_Jacq :\n");
  for(int i =0; i < 3; i++)
  {
    for(int j=0; j<7; j++)
      jac[i+j*3]=0.0;
  }
  // printf("q[3] = %e\n", q[3]);
  // printf("q[4] = %e\n", q[4]);
  // printf("q[5] = %e\n", q[5]);
  // printf("q[6] = %e\n", q[6]);

  double angle = 2*asin(q[5]);

  // printf("angle = %e\n", angle);
  jac[1+5*3]=1e3 * 2.0 / sqrt(1 - q[5]*q[5]) ;
  // printf("jac[3+3*3] = %e\n", jac[3+3*3]);
}

extern "C" void externalForcesB2(double t,double *f, unsigned  int size_z,double *z)
{
  f[0]=0;
  f[1]=0;
  f[2]=-9.81*0.038;

}

extern "C" void externalForcesS(double t,double *f, unsigned int size_z, double *z)
{
  f[0]=0;
  f[1]=0;
  f[2]=-9.81*0.076;

}
extern "C" void externalForceG(double t,double *f, unsigned  int size_z, double *z)
{
  // std::cout << "externalForceG"<<std::endl;
  f[0]=0;
  f[1]=0;
  f[2]=-9.81;
}

extern "C" void prescribedvelocityB1(double time, unsigned int sizeofprescribedvelocity, double *pv)
{
  /* the plugin implements v(t) = C + A cos(omega *t) */

  double C = -150.0 ;
  double omega = M_PI / 2.0;
  double A = 10.0;

  //pv[0] =  A * cos(omega * time*100.0);
  pv[0] =  C;
  //printf("prescribed velocity = %e\n", pv[0]);
}


#include "Solver.h"
#include <algorithm> /* min, max */
#include <armadillo>
#include <assert.h>
#include <math.h>   /* pow */
#include <functional>
#include <stdlib.h> /* abs, drand48 */

using namespace arma;

double SmoSolver::SvmOutputOnPoint(int i) {
  vec point = this->x.row(i).t();
  double result = (kernel->*(kernel->KernelFunction))(this->theta, point);
  return result - this->b;
}

double SmoSolver::Predict(vec &x) { return dot(this->theta, x) - this->b; }

double SmoSolver::KernelCal(int i1, int i2) {
  vec point1 = this->x.row(i1).t();
  vec point2 = this->x.row(i2).t();
  return (kernel->*(kernel->KernelFunction))(point1, point2);
}

int SmoSolver::TakeStep(int i1, int i2) {
  double alpha1 = 0.0, alpha2 = 0.0;
  double a1 = 0.0, a2 = 0.0;
  int y1 = 0, y2 = 0;
  double e1 = 0.0, e2 = 0.0;
  int s = 0;
  double low = 0.0, high = 0.0;
  double k11 = 0.0, k12 = 0.0, k22 = 0.0, eta = 0.0;

  if (i1 == i2) {
    return 0;
  }

  alpha1 = this->lagrangeMultiplier[i1];
  alpha2 = this->lagrangeMultiplier[i2];
  y1 = this->y[i1];
  y2 = this->y[i2];

  if (alpha1 > 0 && alpha1 < this->C) {
    e1 = this->errorCache[i1];
  } else {
    e1 = this->SvmOutputOnPoint(i1) - y1;
  }
  if (alpha2 > 0 && alpha2 < this->C) {
    e2 = this->errorCache[i2];
  } else {
    e2 = this->SvmOutputOnPoint(i1) - y2;
  }

  s = y1 * y2;
  if (y1 != y2) {
    double temp = alpha2 - alpha1;
    low = std::max(0.0, temp);
    high = std::min(this->C, this->C + temp);
  } else {
    double temp = alpha2 + alpha1;
    low = std::max(0.0, temp - this->C);
    high = std::min(this->C, temp);
  }

  // check if low is equal to high
  if (abs(low - high) < 1.0e-7) {
    return 0;
  }
  k11 = this->KernelCal(i1, i1);
  k12 = this->KernelCal(i1, i2);
  k22 = this->KernelCal(i2, i2);
  eta = k11 + k22 - 2 * k12;

  if (eta > 0) {
    a2 = alpha2 + y2 * (e1 - e2) / eta;
    if (a2 < low) {
      a2 = low;
    } else if (a2 > high) {
      a2 = high;
    }
  } else {
    // In papaer 2.1 (19)
    double f1 = y1 * (e1 + this->b) - alpha1 * k11 - s * alpha2 * k12;
    double f2 = y2 * (e2 + this->b) - s * alpha1 * k12 - alpha2 * k12;
    double low1 = alpha1 + s * (alpha2 - low);
    double high1 = alpha1 + s * (alpha2 - high);
    double objLow = low1 * f1 + low * f2 + 0.5 * pow(low1, 2.0) * k11 +
                    0.5 * pow(low, 2.0) * k22 + s * low * low1 * k12;
    double objHigh = high1 * f1 + high * f2 + 0.5 * pow(high1, 2.0) * k11 +
                     0.5 * pow(high, 2.0) * k22 + s * high * high1 * k12;
    if (objLow < objHigh - this->eps) {
      a2 = low;
    } else if (objLow > objHigh + this->eps) {
      a2 = high;
    } else {
      a2 = alpha2;
    }
  }

  if (std::abs(a2 - alpha2) < this->eps * (a2 + alpha2 + this->eps)) {
    return 0;
  }
  a1 = alpha1 + s * (alpha2 - a2);

  // Update threshold to reflect change in Lagrange multipliers
  double b1 = 0.0;
  double b2 = 0.0;
  double bReal = 0.0;
  double temp1 = y1 * (a1 - alpha1);
  double temp2 = y2 * (a2 - alpha2);
  if (a1 > 0 && a1 < this->C) {
    bReal = e1 + temp1 * k11 + temp2 * k12 + b;
  } else if (a2 > 0 && a2 < this->C) {
    bReal = e2 + temp1 * k12 + temp2 * k22 + b;
  } else {
    b1 = e1 + temp1 * k11 + temp2 * k12 + b;
    b2 = e2 + temp1 * k12 + temp2 * k22 + b;
    bReal = (b1 + b2) / 2.0;
  }
  double bDiff = bReal - this->b;
  this->b = bReal;

  // Update weight vector (theta) to reflect change in al & a2, if SVM is
  // linear
  if (this->kernel->kernelType == LINEAR) {
    this->theta = this->theta + temp1 * this->x.row(i1).t() +
                  temp2 * this->x.row(i2).t();
  }
  // Update error cache using new Lagrange multipliers
  int exampleNum = this->ExampleNum();
  for (int i = 0; i < exampleNum; i++) {
    if (lagrangeMultiplier[i] > 0 && lagrangeMultiplier[i] < this->C) {
      this->errorCache[i] += temp1 * this->KernelCal(i1, i) +
                             temp2 * this->KernelCal(i2, i) - bDiff;
    }
  }

  this->errorCache[i1] = 0.0;
  this->errorCache[i2] = 0.0;

  // Store a1, a2 in the alpha array
  this->lagrangeMultiplier[i1] = a1;
  this->lagrangeMultiplier[i2] = a2;
  return 1;
}

int SmoSolver::ExamineExample(int i2) {
  double y2 = 0.0;
  double alpha2 = 0.0;
  double e2 = 0.0;
  double r2 = 0.0;

  alpha2 = this->lagrangeMultiplier[i2];
  y2 = this->y[i2];
  if (alpha2 > 0 && alpha2 < this->C) {
    e2 = this->errorCache[i2];
  } else {
    e2 = this->SvmOutputOnPoint(i2) - y2;
  }
  r2 = e2 * y2;

  int exampleNum = this->ExampleNum();
  //
  double tmax = 0.0;
  int i1 = 0;
  int k = 0;
  if ((r2 < -tol && alpha2 < this->C) || (r2 > tol && alpha2 > 0)) {
    for (i1 = -1, tmax = 0, k = 0; k < exampleNum; k++) {
      if (lagrangeMultiplier[k] > 0 && lagrangeMultiplier[k] < this->C) {
        double e1 = 0.0;
        double temp = 0.0;
        e1 = this->errorCache[k];
        temp = std::abs(e2 - e1);
        if (temp > tmax) {
          tmax = temp;
          i1 = k;
        }
      }
      if (i1 >= 0) {
        if (TakeStep(i1, i2)) {
          return 1;
        }
      }
    }

    for (int i = (int)(drand48() * exampleNum), k = i; k < exampleNum + i;
         k++) {
      i1 = k % exampleNum;
      if (lagrangeMultiplier[i1] > 0 && lagrangeMultiplier[i1] < C) {
        if (TakeStep(i1, i2)) {
          return 1;
        }
      }
    }
    for (int i = (int)(drand48() * exampleNum), k = i; k < exampleNum + i;
         k++) {
      i1 = k % exampleNum;
      if (TakeStep(i1, i2)) {
        return 1;
      }
    }
  }
  return 0;
}

int SmoSolver::ExampleNum() { return (int)this->x.n_rows; }

int SmoSolver::Train() {
  int exampleNum = this->ExampleNum();
  if (!trained) {
    this->b = 0.0;
    this->theta = zeros<vec>(exampleNum);
    this->errorCache = zeros<vec>(exampleNum);
    this->lagrangeMultiplier = zeros<vec>(exampleNum);
  }

  unsigned int numChanged = 0;
  unsigned int examineAll = 1;
  while (numChanged > 0 || examineAll) {
    numChanged = 0;
    if (examineAll) {
      for (unsigned int i = 0; i < exampleNum; i++) {
        numChanged += ExamineExample(i);
      }
    } else {
      for (unsigned int i = 0; i < exampleNum; i++) {
        if (lagrangeMultiplier[i] != 0 && lagrangeMultiplier[i] != C) {
          numChanged += ExamineExample(i);
        }
      }
    }
    if (examineAll == 1) {
      examineAll = 0;
    } else if (numChanged == 0) {
      examineAll = 1;
    }
  }
  this->trained = true;
  return 0;
}