#include "../../include/laser_door_detection/door_detection.h"
#include <algorithm>
#include <Eigen/Dense>
#include <iostream>
#include <fstream>
#include <math.h>

#include "ros/console.h"

namespace door_detection
{

/**
 * DoorDetection constroctour.
 */
DoorDetection::DoorDetection(Status &status)
{
    c_data_ = status.getCachedData();
    r_data_ = status.getRangeData();
    lines_ = status.getLines();
    params_line_ = status.getParamsLine();
}

/**
 * DoorDetection destructour.
 */
DoorDetection::~DoorDetection()
{
    c_data_.reset();
    r_data_.reset();
    lines_.reset();
    params_line_.reset();
}

/**
 * Main function called for door detection.
 */
void DoorDetection::detectDoors(std::vector<Door> &doors)
{
    doors_.clear();

    extractDoors();
    filterDoorsInliers();
    filterDoorsAngles();

    doors = doors_;
}

/**
 * Filter out lines not allows lines.
 */
std::vector<line_extraction::Line> DoorDetection::filterLines()
{
    std::vector<line_extraction::Line> tmp_lines;

    for (unsigned int i = 0; i < lines_->size(); i++)
    {
        // Filter small lines
        if (door_detection::euclideanDist(lines_->at(i).getStart(), lines_->at(i).getEnd()) > params_.min_line_door_length)
            tmp_lines.push_back(lines_->at(i));
    }
    return tmp_lines;
}

/**
 * Extract doors from lines.
 */
void DoorDetection::extractDoors()
{
    std::vector<line_extraction::Line> filtered_lines = filterLines();

    for (unsigned int i = 0; i < filtered_lines.size(); i++)
    {
        for (unsigned int j = 0; j < lines_->size(); j++)
        {
            // No doors w/ the same wall
            if (i == j)
                continue;

            // Create all possible door combinations given two walls
            Door door1(filtered_lines[i].getStart(), filtered_lines[j].getStart(), filtered_lines[i], filtered_lines[j]);
            Door door2(filtered_lines[i].getStart(), filtered_lines[j].getEnd(), filtered_lines[i], filtered_lines[j]);
            Door door3(filtered_lines[i].getEnd(), filtered_lines[j].getStart(), filtered_lines[i], filtered_lines[j]);
            Door door4(filtered_lines[i].getEnd(), filtered_lines[j].getEnd(), filtered_lines[i], filtered_lines[j]);
            std::vector<Door> doors = {door1, door2, door3, door4};

            int count = 0;
            boost::array<boost::array<double, 2>, 2> close_points;
            // Sort doors by their width
            std::sort(doors.begin(), doors.end());

            // Filter Doors by width
            // ---------------------
            for (unsigned int k = 0; k < doors.size(); k++)
            {
                // No more than two doors
                // between two walls.
                if (count >= 2)
                    break;

                // Check Not Possible Doors
                if (close_points[0] == doors[k].getStart() || close_points[1] == doors[k].getEnd())
                    continue;

                // Detect Not Possible Doors
                // If the distance between two extremities of two lines is minor of a threshold,
                // then there is no chance to have any door between one of these two points.
                if (doors[k].getWidth() < 0.8)
                {
                    // saving the walls' points
                    close_points[0] = doors[k].getStart();
                    close_points[1] = doors[k].getEnd();
                }

                // Possible Door Detected
                // TODO: Set param for thresholds
                if (params_.min_door_width < doors[k].getWidth() && doors[k].getWidth() < params_.max_door_width)
                {
                    doors_.push_back(doors[k]);
                    count++;
                }
            }
        }
    }
}

/**
 * Filter out doors that have inliers points and keep doors
 * with a gap between the walls. This will help to remove
 * fake doors detected along walls.
 */
void DoorDetection::filterDoorsInliers()
{
    std::vector<Door> tmp_doors;
    for (unsigned int i = 0; i < doors_.size(); i++)
    {
        unsigned int count = 0;
        for (std::vector<unsigned int>::const_iterator cit = c_data_->indices.begin();
             cit != c_data_->indices.end(); ++cit)
        {
            boost::array<double, 2> point = {r_data_->xs[*cit], r_data_->ys[*cit]};
            // Check if the bearing point is inlier the door
            if (doors_[i].isInlier(point))
                count++;
        }

        // Filter out the door if more than
        // two inline point were detected.
        if (count < params_.max_allowed_inliers)
            tmp_doors.push_back(doors_[i]);
    }
    doors_ = tmp_doors;
}

/**
 * Filter out doors that are generated with a
 * not allowed angle between the two walls.
 */
void DoorDetection::filterDoorsAngles()
{
    std::vector<Door> tmp_doors;
    for (unsigned int i = 0; i < doors_.size(); i++)
    {
        // Get the angle of the first and second wall that generated the door
        double ang1 = angleFromEndpoints(doors_[i].getLine1().getStart(), doors_[i].getLine1().getEnd());
        double ang2 = angleFromEndpoints(doors_[i].getLine2().getStart(), doors_[i].getLine2().getEnd());
        // Compute the angles diff and convert to degree
        double ang_diff = (ang1 - ang2) * (180 / M_PI);

        // Filter out doors with wrong walls angle
        if (params_.min_allowed_angle <= ang_diff && ang_diff <= params_.max_allowed_angle)
            tmp_doors.push_back(doors_[i]);
    }

    doors_ = tmp_doors;
}

/**
 * Set Parameters.
 * ---------------
 */
void DoorDetection::setMinLineDoorLength(double value)
{
    params_.min_line_door_length = value;
}
void DoorDetection::setMinDoorWidth(double value)
{
    params_.min_door_width = value;
}
void DoorDetection::setMaxDoorWidth(double value)
{
    params_.max_door_width = value;
}
void DoorDetection::setMaxAllowedInliers(int value)
{
    params_.max_allowed_inliers = value;
}
void DoorDetection::setMinAllowedAngle(double value)
{
    params_.min_allowed_angle = value;
}
void DoorDetection::setMaxAllowedAngle(double value)
{
    params_.max_allowed_angle = value;
}

} // namespace door_detection


// Petter Strandmark 2012.

#include <cstring>
#include <iostream>
#include <limits>
#include <memory>
#include <random>
#include <stdexcept>

// GNU 4.8.1 define _X on Cygwin.
// This breaks Eigen.
// http://eigen.tuxfamily.org/bz/show_bug.cgi?id=658
#ifdef _X
#undef _X
#endif
#include <Eigen/Dense>
#include <Eigen/Eigenvalues>
#include <Eigen/Sparse>

extern "C" {
	#include "matrix.h"
	#include "matrix2.h"
	#undef min
	#undef max
	#undef catch
	#undef SPARSE
}

#include <spii/spii.h>
#include <spii/solver.h>

namespace spii {

template<typename EigenMat>
void Eigen_to_Meschach(const EigenMat& eigen_matrix, MAT* A)
{
	auto m = eigen_matrix.rows();
	auto n = eigen_matrix.cols();

	for (int i = 0; i < m; ++i) {
		for (int j = 0; j < n; ++j) {
			A->me[i][j] = eigen_matrix(i, j);
		}
	}
}

struct FactorizationCacheInternal
{
	PERM* pivot;
	PERM* block;
	MAT* Hmat;
	VEC* x;
	VEC* b;
};

FactorizationCache::FactorizationCache(int n)
{
	auto cache = new FactorizationCacheInternal;
	cache->pivot = px_get(n);
	cache->block = px_get(n);
	cache->Hmat  = m_get(n, n);
	cache->x     = v_get(n);
	cache->b     = v_get(n);
	this->data = cache;
}

FactorizationCache::~FactorizationCache()
{
	v_free(this->data->b);
	v_free(this->data->x);
	m_free(this->data->Hmat);
	px_free(this->data->pivot);
	px_free(this->data->block);
	delete this->data;
}

void Solver::BKP_dense(const Eigen::MatrixXd& H,
                       const Eigen::VectorXd& g,
                       const FactorizationCache& cache_input,
                       Eigen::VectorXd* p,
                       SolverResults* results) const
{
	using namespace Eigen;
	double start_time = wall_time();

	auto cache = cache_input.data;
	auto n = H.rows();

	Eigen_to_Meschach(H, cache->Hmat);

	//m_foutput(stderr, cache->Hmat);

	BKPfactor(cache->Hmat, cache->pivot, cache->block);

	//m_foutput(stderr, cache->Hmat);

	MatrixXd B(n, n);
	MatrixXd Q(n, n);
	VectorXd tau(n);
	VectorXd lambda(n);
	B.setZero();
	Q.setZero();

	SelfAdjointEigenSolver<MatrixXd> eigensolver;

	double delta = 1e-12;

	int onebyone;
	for (int i = 0; i < n; i = onebyone ? i+1 : i+2 ) {
		onebyone = ( cache->block->pe[i] == i );
		if ( onebyone ) {
		    B(i, i) = m_entry(cache->Hmat, i, i);
			lambda(i) = B(i, i);
			if (lambda(i) >= delta) {
				tau(i) = 0;
			}
			else {
				tau(i) = delta - (1.0 + delta) * lambda(i);
			}
			Q(i, i) = 1;
		}
		else {
		    auto a11 = m_entry(cache->Hmat, i, i);
		    auto a22 = m_entry(cache->Hmat, i+1, i+1);
		    auto a12 = m_entry(cache->Hmat, i+1, i);
			B(i,   i)   = a11;
			B(i+1, i)   = a12;
			B(i,   i+1) = a12;
			B(i+1, i+1) = a22;
			eigensolver.compute(B.block(i, i, 2, 2));

			lambda(i)   = eigensolver.eigenvalues()(0);
			lambda(i+1) = eigensolver.eigenvalues()(1);
			for (int k = i; k <= i + 1; ++k) {
				if (lambda(k) >= delta) {
					tau(k) = 0;
				}
				else {
					tau(k) = delta - (1.0 + delta) * lambda(k);
				}
			}

			Q.block(i, i, 2, 2) = eigensolver.eigenvectors();
		}
	}

	//std::cerr << "B = \n" << B << "\n\n";
	//std::cerr << "F = \n" <<  Q * tau.asDiagonal() * Q.transpose() << "\n\n";

	B = B + Q * tau.asDiagonal() * Q.transpose();

	//std::cerr << "B = \n" << B << "\n\n";

	for (int i = 0; i < n; i = onebyone ? i+1 : i+2 ) {
		onebyone = ( cache->block->pe[i] == i );
		if ( onebyone ) {
		   m_entry(cache->Hmat, i, i) = B(i, i);
		}
		else {
		    m_entry(cache->Hmat, i,   i)   = B(i,   i);
		    m_entry(cache->Hmat, i+1, i+1) = B(i+1, i+1);
		    m_entry(cache->Hmat, i+1, i)   = B(i+1, i);
			m_entry(cache->Hmat, i,   i+1) = B(i,   i+1);
		}
	}

	results->matrix_factorization_time += wall_time() - start_time;
	start_time = wall_time();

	//m_foutput(stderr, cache->Hmat);

	for (int i = 0; i < n; ++i) {
		cache->b->ve[i] = -g(i);
	}
	BKPsolve(cache->Hmat, cache->pivot, cache->block, cache->b, cache->x);
	for (int i = 0; i < n; ++i) {
		(*p)(i) = cache->x->ve[i];
	}

	results->linear_solver_time += wall_time() - start_time;
	start_time = wall_time();

	results->matrix_factorization_time += wall_time() - start_time;
}

}  // namespace spii


#ifndef POLIMIDL_LAYERS_INTERNAL_POINTWISE_CONVOLUTION_INPLACE_HPP
#define POLIMIDL_LAYERS_INTERNAL_POINTWISE_CONVOLUTION_INPLACE_HPP

#include <algorithm>
#include <chrono>
#include <iterator>

#include <Eigen/Dense>

#include "../alignment.hpp"
#include "../layer.hpp"

namespace polimidl {
namespace layers {
namespace internal {
template <typename type_t, typename components>
class pointwise_convolution_inplace : public layer<type_t, components, true> {
 public:
  pointwise_convolution_inplace(const type_t* coeff) :
      coeff_(coeff),
      batch_size_(1) {}

  static std::size_t temporary_size(std::size_t rows, std::size_t columns,
                                    std::size_t number_of_workers) {
    const std::size_t cells = rows * columns;
    std::size_t batch_size = 1;
    // Ensure the alignment of the input Eigen Matrix.
    while ((batch_size * components::input) % alignment_t::type_t_alignment ||
           batch_size < std::max(rows, columns) / 2) {
      batch_size *= 2;
    }
    if (batch_size > cells) {
      batch_size = cells;
    }
    return batch_size * components::output * number_of_workers;
  }

  template <typename input_t, typename temporary_t, typename output_t,
            typename scheduler_t>
  void operator()(input_t input, temporary_t temporary, output_t output,
                  std::size_t rows, std::size_t columns,
                  const scheduler_t& scheduler) const {
    const std::size_t cells = rows * columns;
    const unsigned int number_of_workers = scheduler.number_of_workers();
    for (std::size_t start_cell = 0; start_cell < cells;
         start_cell += batch_size_) {
      const std::size_t end_cell = std::min(cells, start_cell + batch_size_);
      scheduler.schedule([=] (std::size_t worker) {
        const std::size_t cells_in_this_batch = end_cell - start_cell;
        auto temporary_worker = temporary.slice(worker, number_of_workers);
        auto i_input = &input[components::input * start_cell];
        matrix_input_t val(i_input, components::input, cells_in_this_batch);
        auto i_temp = &temporary_worker[0];
        matrix_output_t out(i_temp, components::output, cells_in_this_batch);
        out.noalias() = coeff_ * val;
        std::copy_n(i_temp, components::output * cells_in_this_batch, i_input);
      });
    }
    scheduler.wait();
    if constexpr (components::are_reducing) {
      if (input.begin() == output.begin() ) { // beginning of the buffer
        auto i_input = input.begin();
        auto i_output = output.begin();
        for (std::size_t start_cell = 0; start_cell < cells;
             start_cell += batch_size_) {
          const std::size_t end_cell =
              std::min(cells, start_cell + batch_size_);
          const std::size_t cells_in_this_batch = end_cell - start_cell;
          const std::size_t input_batch_size =
              components::input * cells_in_this_batch;
          const std::size_t output_batch_size =
              components::output * cells_in_this_batch;
          i_output = std::copy_n(i_input, output_batch_size, i_output);
          i_input += input_batch_size;
        }
      } else { // end of the buffer
        auto i_input = std::reverse_iterator(input.end());
        auto i_output = std::reverse_iterator(output.end());
        const std::size_t n_batches = cells / batch_size_ + 1;
        std::size_t cells_in_this_batch = cells % batch_size_;
        for (std::size_t batch = 0; batch < n_batches; ++batch) {
          const std::size_t input_batch_size =
              components::input * cells_in_this_batch;
          const std::size_t output_batch_size =
              components::output * cells_in_this_batch;
          i_input += input_batch_size - output_batch_size;
          i_output = std::copy_n(i_input, output_batch_size, i_output);
          i_input += output_batch_size;
          cells_in_this_batch = batch_size_;
        }
      }
    }
  }

  template <typename input_t, typename temporary_t, typename output_t,
            typename scheduler_t>
  void optimize_for(input_t input, temporary_t temporary, output_t output,
                    std::size_t rows, std::size_t columns,
                    const scheduler_t& scheduler) {
    const std::size_t cells = rows * columns;
    const std::size_t worker_temp_size =
        temporary.slice_size(scheduler.number_of_workers());
    const std::size_t max_batch_size =
        std::min(cells, worker_temp_size / components::output);

    auto best_duration = std::chrono::high_resolution_clock::duration::max();

    std::size_t best_batch_size = 1;
    auto test_and_maybe_set = [&]() {
      for (std::size_t iteration = 0; iteration < 3; ++iteration) {
        const auto start = std::chrono::high_resolution_clock::now();
        operator()(input, temporary, output, rows, columns,
                   scheduler);
        const auto current_duration =
            std::chrono::high_resolution_clock::now() - start;
        if (current_duration < best_duration) {
          best_duration = current_duration;
          best_batch_size = batch_size_;
        }
      }
    };
    batch_size_ = 1;
    // Ensure the alignment of the output Eigen Matrix.
    while (batch_size_ <= max_batch_size) {
      if ((batch_size_ * components::input) % alignment_t::type_t_alignment) {
        batch_size_ *= 2;
        continue;
      }
      test_and_maybe_set();
      batch_size_ *= 2;
    }
    if (cells == max_batch_size) {
      batch_size_ = cells;
      test_and_maybe_set();
    }
    batch_size_ = best_batch_size;
  }

 private:
  using alignment_t = alignment<type_t>;
  using matrix_input_t = Eigen::Map<
      Eigen::Matrix<type_t, components::input, Eigen::Dynamic>,
      alignment_t::eigen_alignment>;
  using matrix_output_t = Eigen::Map<
      Eigen::Matrix<type_t, components::output, Eigen::Dynamic>,
      alignment_t::eigen_alignment>;
  using matrix_coefficients_t = Eigen::Map<
      const Eigen::Matrix<type_t, components::output, components::input,
                          Eigen::RowMajor>,
      alignment_t::eigen_alignment>;

  matrix_coefficients_t coeff_;
  std::size_t batch_size_;
};
}
}
}

#endif  // POLIMIDL_LAYERS_INTERNAL_POINTWISE_CONVOLUTION_INPLACE_HPP


#include <algorithm>
#include <numeric>
#include <limits>
#include <iostream>
#include <cmath>
#include <cassert>

#include <boost/lambda/lambda.hpp>

#include "TreeCoverDecomposition.h"
#include "LinearProgrammingMAPInference.h"

using namespace boost::lambda;

namespace Grante {

LinearProgrammingMAPInference::LinearProgrammingMAPInference(
	const FactorGraph* fg, bool verbose)
	: InferenceMethod(fg), verbose(verbose), max_iter(100), conv_tol(1.0e-6),
		primal_best_energy(0.0), T(0) {
	// Check each variable has at least one unary factors attached
	size_t var_count = fg->Cardinalities().size();
	std::vector<int> has_unary(var_count, -1);
	has_unary.resize(var_count, -1);
	const std::vector<Factor*>& factors = fg->Factors();
	for (size_t fi = 0; fi < factors.size(); ++fi) {
		const Factor* fac = factors[fi];
		if (fac->Variables().size() > 1)
			continue;

		assert(fac->Variables().size() == 1);
		unsigned int var_index = fac->Variables()[0];
		if (has_unary[var_index] >= 0)
			continue;
		has_unary[var_index] = static_cast<int>(fi);
	}
	if (std::find(has_unary.begin(), has_unary.end(), -1) !=
		has_unary.end()) {
		std::cout << "Factor graph has variables without unary factors.  "
			<< "This is required for MAP-MRF LP inference." << std::endl;
		assert(false);
	}

	// Build a fast factor-to-variable lookup map for the unary factors
	std::unordered_map<unsigned int, unsigned int> unary_fi_to_var;
	for (unsigned int ui = 0; ui < has_unary.size(); ++ui)
		unary_fi_to_var[static_cast<unsigned int>(has_unary[ui])] = ui;

	primal_best.resize(var_count, 0);

	// Perform tree decomposition
	TreeCoverDecomposition tcov_decomp(fg);
	tcov_decomp.ComputeDecompositionGreedy(tree_factor_indices,
		factor_cover_count);
	T = tree_factor_indices.size();

	// Instantiate spanning trees
	trees.resize(T, 0);
	tree_var_to_factor_map.resize(T);
	tree_inf.resize(T, 0);
	for (size_t t = 0; t < T; ++t) {
		// Scale all factors by the inverse number of times they are covered
		std::vector<double> f_scale(tree_factor_indices[t].size());
		for (size_t tfi = 0; tfi < f_scale.size(); ++tfi) {
			// Factor tfi in tree t
			f_scale[tfi] = 1.0 /static_cast<double>(
				factor_cover_count[tree_factor_indices[t][tfi]]);
		}

		// Create a subgraph with the given set of factors
		trees[t] = new SubFactorGraph(fg, tree_factor_indices[t], f_scale);
		assert(trees[t]->FG()->Cardinalities().size() == var_count);

		// Build a map for "global variable index -> per-tree factor index"
		// lookups
		for (size_t tfi = 0; tfi < tree_factor_indices[t].size(); ++tfi) {
			// Is it a unary factor?
			unsigned int fi = tree_factor_indices[t][tfi];
			if (unary_fi_to_var.count(fi) == 0)
				continue;

			// It is, find the variable index
			unsigned int var_index = unary_fi_to_var[fi];
			tree_var_to_factor_map[t][var_index] =
				static_cast<unsigned int>(tfi);
		}

		// Tree inference object for this subgraph
		tree_inf[t] = new TreeInference(trees[t]->FG());
	}
}

LinearProgrammingMAPInference::~LinearProgrammingMAPInference() {
	for (size_t t = 0; t < T; ++t) {
		delete (tree_inf[t]);
		delete (trees[t]);
	}
}

InferenceMethod* LinearProgrammingMAPInference::Produce(
	const FactorGraph* new_fg) const {
	return (new LinearProgrammingMAPInference(new_fg));
}

void LinearProgrammingMAPInference::SetParameters(
	unsigned int max_iter, double conv_tol) {
	assert(conv_tol >= 0.0);
	this->max_iter = max_iter;
	this->conv_tol = conv_tol;
}

void LinearProgrammingMAPInference::PerformInference() {
	// Distribute energies uniformly over decomposed trees
	for (size_t t = 0; t < T; ++t)
		trees[t]->ForwardMap();

	// Initialize primal labeling
	const std::vector<unsigned int>& var_card = fg->Cardinalities();
	size_t var_count = fg->Cardinalities().size();
	primal_best.resize(var_count);
	std::fill(primal_best.begin(), primal_best.end(), 0);
	primal_best_energy = std::numeric_limits<double>::infinity();

	// Initialize relaxed solution
#if 0
	const std::vector<Factor*>& factors = fg->Factors();
	relaxed_sol.resize(factors.size());
	for (unsigned int fi = 0; fi < factors.size(); ++fi) {
		relaxed_sol[fi].resize(factors[fi]->Type()->ProdCardinalities());
		std::fill(relaxed_sol[fi].begin(), relaxed_sol[fi].end(), 0.0);
	}
	relaxed_sol_energy = -std::numeric_limits<double>::infinity();
#endif

	// \lambda_{t,vi,state} = 0
	std::vector<std::vector<double> > sol_avg(var_count);
	for (size_t vi = 0; vi < var_count; ++vi) {
		unsigned int vi_card = var_card[vi];
		sol_avg[vi].resize(vi_card, 0.0);
	}
	// Individual tree solutions
	std::vector<std::vector<unsigned int> > cur_sol_t(T);
	for (size_t t = 0; t < T; ++t)
		cur_sol_t[t].resize(var_count);

	// Stepsize control
	// The control mechanism is from Section 8.2, "Path-Based Incremental
	// Target Level Algorithm" in Bertsekas, Nedic, Ozdaglar, "Convex Analysis
	// and Optimization".
	double delta = -1.0;
	double B = 2.0;	// Travel-length bound
	double sigma = 0.0;	// Actual travel-length since last control
	double dual_obj_target = -std::numeric_limits<double>::infinity();
	double dual_obj_best = -std::numeric_limits<double>::infinity();

	// Iterate
	double sol_step = 1.0 / static_cast<double>(T);
	for (int iter = 1; max_iter == 0 || iter <= static_cast<int>(max_iter);
		++iter) {
		// Clear averaged solution
		for (unsigned int vi = 0; vi < var_count; ++vi)
			std::fill(sol_avg[vi].begin(), sol_avg[vi].end(), 0.0);

		// Perform inference for all submodels
		double dual_obj = 0.0;
		bool new_primal_best = false;
		for (unsigned int t = 0; t < T; ++t) {
			double t_obj = tree_inf[t]->MinimizeEnergy(cur_sol_t[t]);

			// Dual objective is simply the sum of all tree objectives
			dual_obj += t_obj;

			// Produce averaged solution (primal infeasible)
			for (unsigned int vi = 0; vi < var_count; ++vi)
				sol_avg[vi][cur_sol_t[t][vi]] += sol_step;

			// Identify best feasible integral labeling
			double t_primal_energy = fg->EvaluateEnergy(cur_sol_t[t]);
			if (t_primal_energy < primal_best_energy) {
				std::copy(cur_sol_t[t].begin(), cur_sol_t[t].end(),
					primal_best.begin());
				primal_best_energy = t_primal_energy;
				new_primal_best = true;
			}
		}

		if (dual_obj > dual_obj_best)
			dual_obj_best = dual_obj;

		// Initial delta: half the primal-dual gap
		if (delta < 0.0)
			delta = 0.5 * (primal_best_energy - dual_obj);

		bool sufficient_descent = false;
		bool oscillation = false;
		if (dual_obj >= (dual_obj_target + 0.5 * delta)) {
			// Sufficient descent because target level is reached
			sufficient_descent = true;
			sigma = 0.0;
			dual_obj_target = dual_obj_best;
		} else if (sigma > B) {
			// Oscillation detected
			oscillation = true;
			sigma = 0.0;
			delta *= 0.5;
			dual_obj_target = dual_obj_best;
		}

#if 0
		// Linear program primal solution recovery by subgradient averaged
		// solution (Anstreicher and Wolsey, MathProg 2009)
#if 1
		// Uniform average
		double new_sol_scale = 1.0 / static_cast<double>(T * iter);
		double old_sol_scale = static_cast<double>(iter - 1) /
			static_cast<double>(iter);
#endif
#if 0
		// Geometric average
		double vol_alpha = 0.005;
		double new_sol_scale = vol_alpha / static_cast<double>(T);
		double old_sol_scale = (1.0 - vol_alpha);
		if (iter == 1) {
			old_sol_scale = 0.0;
			new_sol_scale = 1.0;
		}
#endif
		relaxed_sol_energy = 0.0;
		for (unsigned int fi = 0; fi < factors.size(); ++fi) {
			std::transform(relaxed_sol[fi].begin(), relaxed_sol[fi].end(),
				relaxed_sol[fi].begin(), _1 * old_sol_scale);

			for (unsigned int t = 0; t < T; ++t) {
				unsigned int ei =
					factors[fi]->ComputeAbsoluteIndex(cur_sol_t[t]);
				assert(ei < relaxed_sol[fi].size());
				relaxed_sol[fi][ei] += new_sol_scale * 1.0;
			}
			relaxed_sol_energy += std::inner_product(
				relaxed_sol[fi].begin(), relaxed_sol[fi].end(),
				factors[fi]->Energies().begin(), 0.0);
		}
#endif

		// Output statistics
		if (verbose) {
			std::cout << "iter " << iter << ", primal " << primal_best_energy
				<< ", dual " << dual_obj << ", best dual " << dual_obj_best
				<< ", gap " << (primal_best_energy - dual_obj) << std::endl;
		}

		// Compute step size
		double subgradient_norm = 0.0;
		for (size_t t = 0; t < T; ++t) {
			for (unsigned int vi = 0; vi < var_count; ++vi) {
				// Obtain a unary factor of the variable
				unsigned int vi_card = var_card[vi];
				for (unsigned int vs = 0; vs < vi_card; ++vs) {
					subgradient_norm += std::pow(
						(cur_sol_t[t][vi] == vs ? 1.0 : 0.0)
						- sol_avg[vi][vs], 2.0);
				}
			}
		}
		//subgradient_norm = std::sqrt(subgradient_norm);
		double gamma = 1.95;
		double alpha = gamma * ((dual_obj_target + delta) - dual_obj) /
			subgradient_norm;
		double dgap = primal_best_energy - dual_obj;
		double convergence_measure = dgap / (fabs(dual_obj) + 1.0e-5);
		sigma += std::min(1.0, alpha * subgradient_norm);
		if (subgradient_norm <= 1.0e-5 || convergence_measure <= conv_tol) {
			if (verbose) {
				std::cout << "Converged, subg norm " << subgradient_norm
					<< ", conv " << convergence_measure << std::endl;
			}
			break;
		}

		// 4. Compute step size
#if 0
		alpha = 1.0 / (10.0 + static_cast<double>(iter));
#endif

		// Update Lagrange multipliers implicitly by directly updating the
		// energies of the trees in the decomposition
		for (unsigned int t = 0; t < T; ++t) {
			const std::vector<Factor*>& factors = trees[t]->FG()->Factors();
			for (unsigned int vi = 0; vi < var_count; ++vi) {
				// Obtain a unary factor of the variable
				unsigned int t_fi = tree_var_to_factor_map[t][vi];
				std::vector<double>& t_fi_energies = factors[t_fi]->Energies();

				// Modify the energies
				unsigned int vi_card = var_card[vi];
				for (unsigned int vs = 0; vs < vi_card; ++vs) {
					// Subgradient update
					t_fi_energies[vs] += alpha * (
						(cur_sol_t[t][vi] == vs ? 1.0 : 0.0) -
						sol_avg[vi][vs]);
				}
			}
		}
	}
}

void LinearProgrammingMAPInference::ClearInferenceResult() {
#if 0
	relaxed_sol.clear();
#endif
}

// Returns one part of the (approximate) relaxed solution
const std::vector<double>& LinearProgrammingMAPInference::Marginal(
	unsigned int factor_id) const {
	assert(factor_id < relaxed_sol.size());
	return (relaxed_sol[factor_id]);
}

// Returns the relaxed solution
const std::vector<std::vector<double> >&
LinearProgrammingMAPInference::Marginals() const {
	assert(0);
	return (relaxed_sol);
}

// XXX: not implemented
double LinearProgrammingMAPInference::LogPartitionFunction() const {
	assert(false);
	return (std::numeric_limits<double>::signaling_NaN());
}

// XXX: not implemented
void LinearProgrammingMAPInference::Sample(
	std::vector<std::vector<unsigned int> >& states,
	unsigned int sample_count) {
	assert(false);
}

// Obtain an approximate minimum energy state for the current factor graph
// energies.
double LinearProgrammingMAPInference::MinimizeEnergy(
	std::vector<unsigned int>& state) {
	PerformInference();
	state = primal_best;

	return (primal_best_energy);
}

}


#include <Eigen/Core>
#include <cstdint>
#include <string>

#include "pyinterp/axis.hpp"
#include "pyinterp/detail/math/frame.hpp"
#include "pyinterp/grid.hpp"

namespace pyinterp {

// Error thrown if it' s not possible to frame the value on the specified axis.
template <typename T>
auto index_error(const std::string& axis, T value, size_t n) -> void {
  throw std::invalid_argument(
      "Unable to frame the value " + std::to_string(value) + " with " +
      std::to_string(n) + " items of the " + axis + " axis");
}

/// Loads the interpolation frame into memory
template <typename DataType>
auto load_frame(const Grid2D<DataType>& grid, const double x, const double y,
                const axis::Boundary boundary, const bool bounds_error,
                detail::math::Frame2D& frame) -> bool {
  const auto& x_axis = *grid.x();
  const auto& y_axis = *grid.y();
  const auto y_indexes =
      y_axis.find_indexes(y, static_cast<uint32_t>(frame.ny()), boundary);
  const auto x_indexes =
      x_axis.find_indexes(x, static_cast<uint32_t>(frame.nx()), boundary);

  if (x_indexes.empty() || y_indexes.empty()) {
    if (bounds_error) {
      if (x_indexes.empty()) {
        index_error("x", x, frame.nx());
      }
      index_error("y", y, frame.ny());
    }
    return false;
  }

  auto x0 = x_axis(x_indexes[0]);

  for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
    frame.y(jx) = y_axis(y_indexes[jx]);
  }

  for (Eigen::Index ix = 0; ix < frame.x()->size(); ++ix) {
    const auto index = x_indexes[ix];

    frame.x(ix) = x_axis.is_angle()
                      ? detail::math::normalize_angle(x_axis(index), x0, 360.0)
                      : x_axis(index);

    for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
      frame.q(ix, jx) = static_cast<double>(grid.value(index, y_indexes[jx]));
    }
  }
  return frame.is_valid();
}

/// Loads the interpolation frame into memory
template <typename DataType, typename AxisType>
auto load_frame(const Grid3D<DataType, AxisType>& grid, const double x,
                const double y, const AxisType z, const axis::Boundary boundary,
                const bool bounds_error, detail::math::Frame3D<AxisType>& frame)
    -> bool {
  const auto& x_axis = *grid.x();
  const auto& y_axis = *grid.y();
  const auto& z_axis = *grid.z();
  const auto z_indexes =
      z_axis.find_indexes(z, static_cast<uint32_t>(frame.nz()), boundary);
  const auto y_indexes =
      y_axis.find_indexes(y, static_cast<uint32_t>(frame.ny()), boundary);
  const auto x_indexes =
      x_axis.find_indexes(x, static_cast<uint32_t>(frame.nx()), boundary);

  if (x_indexes.empty() || y_indexes.empty() || z_indexes.empty()) {
    if (bounds_error) {
      if (x_indexes.empty()) {
        index_error("x", x, frame.nx());
      } else if (y_indexes.empty()) {
        index_error("y", y, frame.ny());
      }
      index_error("z", z, frame.nz());
    }
    return false;
  }

  auto x0 = x_axis(x_indexes[0]);

  for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
    frame.y(jx) = y_axis(y_indexes[jx]);
  }

  for (Eigen::Index kx = 0; kx < frame.z().size(); ++kx) {
    frame.z(kx) = z_axis(z_indexes[kx]);
  }

  for (Eigen::Index ix = 0; ix < frame.x()->size(); ++ix) {
    const auto x_index = x_indexes[ix];

    frame.x(ix) = x_axis.is_angle() ? detail::math::normalize_angle(
                                          x_axis(x_index), x0, 360.0)
                                    : x_axis(x_index);

    for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
      const auto y_index = y_indexes[jx];

      for (Eigen::Index kx = 0; kx < frame.z().size(); ++kx) {
        frame.q(ix, jx, kx) =
            static_cast<double>(grid.value(x_index, y_index, z_indexes[kx]));
      }
    }
  }
  return frame.is_valid();
}

/// Loads the interpolation frame into memory
template <typename DataType, typename AxisType>
auto load_frame(const Grid4D<DataType, AxisType>& grid, const double x,
                const double y, const AxisType z, const double u,
                const axis::Boundary boundary, const bool bounds_error,
                detail::math::Frame4D<AxisType>& frame) -> bool {
  const auto& x_axis = *grid.x();
  const auto& y_axis = *grid.y();
  const auto& z_axis = *grid.z();
  const auto& u_axis = *grid.u();
  const auto u_indexes =
      u_axis.find_indexes(u, static_cast<uint32_t>(frame.nu()), boundary);
  const auto z_indexes =
      z_axis.find_indexes(z, static_cast<uint32_t>(frame.nz()), boundary);
  const auto y_indexes =
      y_axis.find_indexes(y, static_cast<uint32_t>(frame.ny()), boundary);
  const auto x_indexes =
      x_axis.find_indexes(x, static_cast<uint32_t>(frame.nx()), boundary);

  if (x_indexes.empty() || y_indexes.empty() || z_indexes.empty() ||
      u_indexes.empty()) {
    if (bounds_error) {
      if (x_indexes.empty()) {
        index_error("x", x, frame.nx());
      } else if (y_indexes.empty()) {
        index_error("y", y, frame.ny());
      } else if (z_indexes.empty()) {
        index_error("z", z, frame.nz());
      }
      index_error("u", u, frame.nu());
    }
    return false;
  }

  auto x0 = x_axis(x_indexes[0]);

  for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
    frame.y(jx) = y_axis(y_indexes[jx]);
  }

  for (Eigen::Index kx = 0; kx < frame.z().size(); ++kx) {
    frame.z(kx) = z_axis(z_indexes[kx]);
  }

  for (Eigen::Index lx = 0; lx < frame.u().size(); ++lx) {
    frame.u(lx) = u_axis(u_indexes[lx]);
  }

  for (Eigen::Index ix = 0; ix < frame.x()->size(); ++ix) {
    const auto x_index = x_indexes[ix];

    frame.x(ix) = x_axis.is_angle() ? detail::math::normalize_angle(
                                          x_axis(x_index), x0, 360.0)
                                    : x_axis(x_index);

    for (Eigen::Index jx = 0; jx < frame.y()->size(); ++jx) {
      const auto y_index = y_indexes[jx];

      for (Eigen::Index kx = 0; kx < frame.z().size(); ++kx) {
        const auto z_index = z_indexes[kx];

        for (Eigen::Index lx = 0; lx < frame.u().size(); ++lx) {
          frame.q(ix, jx, kx, lx) = static_cast<double>(
              grid.value(x_index, y_index, z_index, u_indexes[lx]));
        }
      }
    }
  }
  return frame.is_valid();
}

}  // namespace pyinterp


/*
 * @Description: Kalman Filter interface.
 * @Author: Ge Yao
 * @Date: 2020-11-12 15:14:07
 */

#ifndef LIDAR_LOCALIZATION_MODELS_KALMAN_FILTER_KALMAN_FILTER_HPP_
#define LIDAR_LOCALIZATION_MODELS_KALMAN_FILTER_KALMAN_FILTER_HPP_

#include <yaml-cpp/yaml.h>

#include <deque>

#include <Eigen/Core>
#include <Eigen/Dense>

#include "lidar_localization/sensor_data/imu_data.hpp"

namespace lidar_localization {

class KalmanFilter {
public:
  /**
   * @class MeasurementType
   * @brief enum for observation type
   */
  enum MeasurementType {
    POSE = 0,
    POSE_VEL,
    POSI,
    POSI_VEL,
    POSI_MAG,
    POSI_VEL_MAG,
    NUM_TYPES
  };

  /**
   * @class Measurement
   * @brief Kalman filter measurement data
   */
  struct Measurement {
    // timestamp:
    double time;
    // a. pose observation, lidar/visual frontend:
    Eigen::Matrix4d T_nb;
    // b. body frame velocity observation, odometer:
    Eigen::Vector3d v_b;
    // c. body frame angular velocity, needed by motion constraint:
    Eigen::Vector3d w_b;
    // d. magnetometer:
    Eigen::Vector3d B_b;
  };

  /**
   * @class Cov
   * @brief Kalman filter process covariance data
   */
  struct Cov {
    struct {
      double x;
      double y;
      double z;
    } pos;
    struct {
      double x;
      double y;
      double z;
    } vel;
    // here quaternion is used for orientation representation:
    struct {
      double w;
      double x;
      double y;
      double z;
    } ori;
    struct {
      double x;
      double y;
      double z;
    } gyro_bias;
    struct {
      double x;
      double y;
      double z;
    } accel_bias;
  };

  /**
   * @brief  init filter
   * @param  imu_data, input IMU measurements
   * @return true if success false otherwise
   */
  virtual void Init(const Eigen::Vector3d &vel, const IMUData &imu_data) = 0;

  /**
   * @brief  update state & covariance estimation, Kalman prediction
   * @param  imu_data, input IMU measurements
   * @return true if success false otherwise
   */
  virtual bool Update(const IMUData &imu_data) = 0;

  /**
   * @brief  correct state & covariance estimation, Kalman correction
   * @param  measurement_type, input measurement type
   * @param  measurement, input measurement
   * @return void
   */
  virtual bool Correct(const IMUData &imu_data,
                       const MeasurementType &measurement_type,
                       const Measurement &measurement) = 0;

  /**
   * @brief  get filter time
   * @return filter time as double
   */
  double GetTime(void) const { return time_; }

  /**
   * @brief  get odometry estimation
   * @param  pose, output pose
   * @param  vel, output vel
   * @return void
   */
  virtual void GetOdometry(Eigen::Matrix4f &pose, Eigen::Vector3f &vel) = 0;

  /**
   * @brief  get covariance estimation
   * @param  cov, output covariance
   * @return void
   */
  virtual void GetCovariance(Cov &cov) = 0;

  /**
   * @brief  update observability analysis
   * @param  time, measurement time
   * @param  measurement_type, measurement type
   * @return void
   */
  virtual void
  UpdateObservabilityAnalysis(const double &time,
                              const MeasurementType &measurement_type) = 0;

  /**
   * @brief  save observability analysis to persistent storage
   * @param  measurement_type, measurement type
   * @return void
   */
  virtual bool
  SaveObservabilityAnalysis(const MeasurementType &measurement_type) = 0;

protected:
  KalmanFilter() {}

  static void AnalyzeQ(const int DIM_STATE, const double &time,
                       const Eigen::MatrixXd &Q, const Eigen::VectorXd &Y,
                       std::vector<std::vector<double>> &data);

  static void WriteAsCSV(const int DIM_STATE,
                         const std::vector<std::vector<double>> &data,
                         const std::string filename);

  // time:
  double time_;

  // data buff:
  std::deque<IMUData> imu_data_buff_;

  // earth constants:
  Eigen::Vector3d g_;
  Eigen::Vector3d w_;
  Eigen::Vector3d b_;

  // observability analysis:
  struct {
    std::vector<double> time_;
    std::vector<Eigen::MatrixXd> Q_;
    std::vector<Eigen::VectorXd> Y_;
  } observability;

  // hyper-params:
  // a. earth constants:
  struct {
    double GRAVITY_MAGNITUDE;
    double ROTATION_SPEED;
    double LATITUDE;
    double LONGITUDE;
    struct {
      double B_E;
      double B_N;
      double B_U;
    } MAG;
  } EARTH;
  // b. prior state covariance, process & measurement noise:
  struct {
    struct {
      double POSI;
      double VEL;
      double ORI;
      double EPSILON;
      double DELTA;
    } PRIOR;
    struct {
      double GYRO;
      double ACCEL;
      double BIAS_ACCEL;
      double BIAS_GYRO;
      bool BIAS_FLAG;
    } PROCESS;
    struct {
      struct {
        double POSI;
        double ORI;
      } POSE;
      double POSI;
      double VEL;
      double ORI;
      double MAG;
    } MEASUREMENT;
  } COV;
  // c. motion constraint:
  struct {
    bool ACTIVATED;
    double W_B_THRESH;
  } MOTION_CONSTRAINT;
};

} // namespace lidar_localization

#endif // LIDAR_LOCALIZATION_MODELS_KALMAN_FILTER_KALMAN_FILTER_HPP_


/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
/*                                                                           */
/*  This file is part of the library KASKADE 7                               */
/*    see http://www.zib.de/projects/kaskade7-finite-element-toolbox         */
/*                                                                           */
/*  Copyright (C) 2002-2011 Zuse Institute Berlin                            */
/*                                                                           */
/*  KASKADE 7 is distributed under the terms of the ZIB Academic License.    */
/*    see $KASKADE/academic.txt                                              */
/*                                                                           */
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */

#ifndef UZAWA_H
#define UZAWA_H

#include <boost/fusion/sequence.hpp>

#include "dune/common/static_assert.hh"
#include "dune/istl/preconditioners.hh"
#include "dune/istl/solvers.hh"

#include "fem/linearspace.hh"

namespace Kaskade
{
  /**
   * Convenience function for simple application of
   * preconditioners. Calls pre(), apply(), and post() sequentially.
   */
  template <class X, class Y>
  void applyPreconditioner(Dune::Preconditioner<X,Y>& p, X& x, Y& y) {
    p.pre(x,y);
    p.apply(x,y);
    p.post(x);
  }

  /**
   * An inexact Uzawa solver for solving symmetric saddle point systems
   *
   * \f[ \begin{array}{cc} A & B^T \\ B & \end{array}\begin{matrix}{c} u \\ \lambda \end{matrix} =
   * \begin{matrix}{c} f \\ g \end{matrix} \f]
   *
   * It implements an inexact preconditioned conjugate gradient method for the
   * Schur complement \f$ BA^{-1}B^T\f$.
   *
   * Template parameters:
   * - X: the linear space of the primal variable \f$ u \f$
   * - Y: the linear space of the dual variable \f$ \lambda \f$
   */
  template<class X, class Y>
  class UzawaSolver : public Dune::InverseOperator<LinearProductSpace<typename X::field_type,boost::fusion::vector<X,Y> >,
  LinearProductSpace<typename X::field_type,boost::fusion::vector<X,Y> > > {
  public:
    typedef LinearProductSpace<typename X::field_type,boost::fusion::vector<X,Y> > domain_type;
    typedef domain_type                                                            range_type;
    typedef typename X::field_type                                                 field_type;
    typedef field_type Scalar;
    typedef domain_type Domain;
    typedef range_type Range;

    /**
     * Constructor.
     * \param opA_ the linear operator \f$ A \f$
     * \param solA_ the approximate inverse operator \f$ A^{-1} \f$, needs to be a subclass of Dune::InverseOperator<X,X>
     * \param opB_ the linear operaotr \f$ B \f$
     * \param opBt_ the linear operator \f$ B^T \f$
     * \param pb_ the Schur complement preconditioner
     * \param reduction_ the default required reduction in preconditioned Schur complement's residual
     * \param maxit_ the maximum number of iterations
     * \param verbose_ printing level to stdout (0=silent, 1=result, 2=every iteration)
     */
    template<class A, class SolA, class B, class Bt, class PB>
    UzawaSolver (A& opA_, SolA& solA_, B& opB_, Bt& opBt_, PB& pb_, double reduction_, int maxit_, int verbose_) :
    opA(opA_), solA(solA_), opB(opB_), opBt(opBt_), precB(pb_), reduction(reduction_), maxit(maxit_), verbose(verbose_) {
      using namespace Dune;
      dune_static_assert( static_cast<int>(PB::category) == static_cast<int>(SolverCategory::sequential) , "Preconditioner PB should be sequential.");
      dune_static_assert( static_cast<int>(A::category) == static_cast<int>(SolverCategory::sequential)  , "Operator A should be sequential");
      dune_static_assert( static_cast<int>(B::category) == static_cast<int>(SolverCategory::sequential)  , "Operator B should be sequential");
      dune_static_assert( static_cast<int>(Bt::category) == static_cast<int>(SolverCategory::sequential) , "Operator Bt should be sequential");
    }

    virtual void apply (Domain& x, Range& b, Dune::InverseOperatorResult& res) {
      using namespace boost::fusion;
      using namespace Dune;

      Timer watch;                // start a timer

      InverseOperatorResult resA;
      if (verbose>0) printf("=== UzawaSolver\n");

      // Declare variables
      X& u(at_c<0>(x.data));
      X& f(at_c<0>(b.data));
      X r(u), h(u), rSave(u);
      Y& lambda(at_c<1>(x.data));
      Y g(at_c<1>(b.data)), d(g), p(g);
      Y tmp(g); // WARNING: needed only for applying P twice!

      // initialize
      r = f;
        opBt.applyscaleadd(-1,lambda,r);  // r = f-B'*u
      solA.apply(u,r,resA);                    // Au = r
      opB.applyscaleadd(-1,u,g);               // g = g-Bu
      applyPreconditioner(precB,tmp,g);
      applyPreconditioner(precB,d,tmp); // d = Pg // WARNING: assumes precB^2 is preco
      double sigma = spY.dot(g,d);            // sigma = g'd
      double sigma0 = sigma;

      //std::cout << "sigma = " << sigma << "\n g=" << at_c<0>(g.data) << "d=" << at_c<0>(d.data);

      // Perform stationary iteration.
      int i=0;
      for ( ; i<maxit && sigma>reduction*reduction*sigma0; i++) {

        // std::cout << "d=P(g-Bu)=" << at_c<0>(d.data);

        opBt.apply(d,r);                       // r = B'd
        rSave = r;
        solA.apply(h,r,resA);                  // Ah = r
        r = rSave;

        double alpha = sigma / spX.dot(h,r);   // alpha = sigma / h'r
        lambda.axpy(-alpha,d);                  // lambda = lambda-alpha d
        u.axpy(alpha,h);                      // u = u+alpha h
        opB.applyscaleadd(-alpha,h,g);          // g = g-alpha Bh

        applyPreconditioner(precB,tmp,g);        // p = Pg
        applyPreconditioner(precB,p,tmp); // WARNING

        double sigmaOld = sigma;
        sigma = spY.dot(g,p);
        double beta = sigma / sigmaOld;

        d *= beta; d.axpy(1,p);               // d = beta d + p

        if (verbose>1) {
          if (i%30==0)
            std::printf("%5s %14s %14s\n","Iter","Preco. Res.","Rate");
          std::printf("%5d %14.4e %14.4e\n",i,std::sqrt(sigma),std::sqrt(beta));
        }
      }

      // Fill statistics
      res.clear();
      res.iterations = i;
      res.reduction = std::sqrt(sigma/sigma0);
      res.converged = res.reduction<=reduction;
      res.conv_rate = std::pow(res.reduction,1.0/i);
      res.elapsed = watch.elapsed();

      if (verbose>0)                 // final print
        printf("=== rate=%g, time=%g, time/it=%g, iter=%d\n",res.conv_rate,res.elapsed,res.elapsed/i,i);
    }

    virtual void apply (Domain& x, Range& b, double reduction_, Dune::InverseOperatorResult& res)
    {
      reduction = reduction_;
      this->apply(x,b,res);
    }

    void apply(Domain& x, Range& b)
    {
      Dune::InverseOperatorResult tmpResult;
      apply(x,b,tmpResult);
    }

  private:
    Dune::SeqScalarProduct<X> spX;
    Dune::SeqScalarProduct<Y> spY;

    Dune::LinearOperator<X,X>& opA;
    Dune::InverseOperator<X,X>& solA;

    Dune::LinearOperator<X,Y>& opB;
    Dune::LinearOperator<Y,X>& opBt;
    Dune::Preconditioner<Y,Y>& precB;

    double reduction;
    int    maxit;
    int    verbose;
  };
}

#endif


#include "languageModels.h"
#include <boost/range/adaptor/map.hpp>
#include <vector>
#include <regex>
#include <map>
#include <tuple>
#include "platformTools.h"
#include <boost/filesystem/fstream.hpp>
#include "appInfo.h"
#include <cmath>
#include <gsl_util.h>

using std::string;
using std::u32string;
using std::vector;
using std::regex;
using std::map;
using std::tuple;
using std::make_tuple;
using std::get;
using std::endl;
using boost::filesystem::path;

using unigram_t = string;
using bigram_t = tuple<string, string>;
using trigram_t = tuple<string, string, string>;

map<unigram_t, int> getUnigramCounts(const vector<string>& words) {
	map<unigram_t, int> unigramCounts;
	for (const unigram_t& unigram : words) {
		++unigramCounts[unigram];
	}
	return unigramCounts;
}

map<bigram_t, int> getBigramCounts(const vector<string>& words) {
	map<bigram_t, int> bigramCounts;
	for (auto it = words.begin(); it < words.end() - 1; ++it) {
		++bigramCounts[bigram_t(*it, *(it + 1))];
	}
	return bigramCounts;
}

map<trigram_t, int> getTrigramCounts(const vector<string>& words) {
	map<trigram_t, int> trigramCounts;
	if (words.size() >= 3) {
		for (auto it = words.begin(); it < words.end() - 2; ++it) {
			++trigramCounts[trigram_t(*it, *(it + 1), *(it + 2))];
		}
	}
	return trigramCounts;
}

map<unigram_t, double> getUnigramProbabilities(const vector<string>& words, const map<unigram_t, int>& unigramCounts, const double deflator) {
	map<unigram_t, double> unigramProbabilities;
	for (const auto& pair : unigramCounts) {
		unigram_t unigram = get<0>(pair);
		int unigramCount = get<1>(pair);
		unigramProbabilities[unigram] = double(unigramCount) / words.size() * deflator;
	}
	return unigramProbabilities;
}

map<bigram_t, double> getBigramProbabilities(const map<unigram_t, int>& unigramCounts, const map<bigram_t, int>& bigramCounts, const double deflator) {
	map<bigram_t, double> bigramProbabilities;
	for (const auto& pair : bigramCounts) {
		bigram_t bigram = get<0>(pair);
		int bigramCount = get<1>(pair);
		int unigramPrefixCount = unigramCounts.at(get<0>(bigram));
		bigramProbabilities[bigram] = double(bigramCount) / unigramPrefixCount * deflator;
	}
	return bigramProbabilities;
}

map<trigram_t, double> getTrigramProbabilities(const map<bigram_t, int>& bigramCounts, const map<trigram_t, int>& trigramCounts, const double deflator) {
	map<trigram_t, double> trigramProbabilities;
	for (const auto& pair : trigramCounts) {
		trigram_t trigram = get<0>(pair);
		int trigramCount = get<1>(pair);
		int bigramPrefixCount = bigramCounts.at(bigram_t(get<0>(trigram), get<1>(trigram)));
		trigramProbabilities[trigram] = double(trigramCount) / bigramPrefixCount * deflator;
	}
	return trigramProbabilities;
}

map<unigram_t, double> getUnigramBackoffWeights(
	const map<unigram_t, int>& unigramCounts,
	const map<unigram_t, double>& unigramProbabilities,
	const map<bigram_t, int>& bigramCounts,
	const double discountMass)
{
	map<unigram_t, double> unigramBackoffWeights;
	for (const unigram_t& unigram : unigramCounts | boost::adaptors::map_keys) {
		double denominator = 1;
		for (const bigram_t& bigram : bigramCounts | boost::adaptors::map_keys) {
			if (get<0>(bigram) == unigram) {
				denominator -= unigramProbabilities.at(get<1>(bigram));
			}
		}
		unigramBackoffWeights[unigram] = discountMass / denominator;
	}
	return unigramBackoffWeights;
}

map<bigram_t, double> getBigramBackoffWeights(
	const map<bigram_t, int>& bigramCounts,
	const map<bigram_t, double>& bigramProbabilities,
	const map<trigram_t, int>& trigramCounts,
	const double discountMass)
{
	map<bigram_t, double> bigramBackoffWeights;
	for (const bigram_t& bigram : bigramCounts | boost::adaptors::map_keys) {
		double denominator = 1;
		for (const trigram_t& trigram : trigramCounts | boost::adaptors::map_keys) {
			if (bigram_t(get<0>(trigram), get<1>(trigram)) == bigram) {
				denominator -= bigramProbabilities.at(bigram_t(get<1>(trigram), get<2>(trigram)));
			}
		}
		bigramBackoffWeights[bigram] = discountMass / denominator;
	}
	return bigramBackoffWeights;
}

void createLanguageModelFile(const vector<string>& words, path filePath) {
	const double discountMass = 0.5;
	const double deflator = 1.0 - discountMass;

	map<unigram_t, int> unigramCounts = getUnigramCounts(words);
	map<bigram_t, int> bigramCounts = getBigramCounts(words);
	map<trigram_t, int> trigramCounts = getTrigramCounts(words);

	map<unigram_t, double> unigramProbabilities = getUnigramProbabilities(words, unigramCounts, deflator);
	map<bigram_t, double> bigramProbabilities = getBigramProbabilities(unigramCounts, bigramCounts, deflator);
	map<trigram_t, double> trigramProbabilities = getTrigramProbabilities(bigramCounts, trigramCounts, deflator);

	map<unigram_t, double> unigramBackoffWeights = getUnigramBackoffWeights(unigramCounts, unigramProbabilities, bigramCounts, discountMass);
	map<bigram_t, double> bigramBackoffWeights = getBigramBackoffWeights(bigramCounts, bigramProbabilities, trigramCounts, discountMass);

	boost::filesystem::ofstream file(filePath);
	file << "Generated by " << appName << " " << appVersion << endl << endl;

	file << "\\data\\" << endl;
	file << "ngram 1=" << unigramCounts.size() << endl;
	file << "ngram 2=" << bigramCounts.size() << endl;
	file << "ngram 3=" << trigramCounts.size() << endl << endl;

	file.setf(std::ios::fixed, std::ios::floatfield);
	file.precision(4);
	file << "\\1-grams:" << endl;
	for (const unigram_t& unigram : unigramCounts | boost::adaptors::map_keys) {
		file << log10(unigramProbabilities.at(unigram))
			<< " " << unigram
			<< " " << log10(unigramBackoffWeights.at(unigram)) << endl;
	}
	file << endl;

	file << "\\2-grams:" << endl;
	for (const bigram_t& bigram : bigramCounts | boost::adaptors::map_keys) {
		file << log10(bigramProbabilities.at(bigram))
			<< " " << get<0>(bigram) << " " << get<1>(bigram)
			<< " " << log10(bigramBackoffWeights.at(bigram)) << endl;
	}
	file << endl;

	file << "\\3-grams:" << endl;
	for (const trigram_t& trigram : trigramCounts | boost::adaptors::map_keys) {
		file << log10(trigramProbabilities.at(trigram))
			<< " " << get<0>(trigram) << " " << get<1>(trigram) << " " << get<2>(trigram) << endl;
	}
	file << endl;

	file << "\\end\\" << endl;
}

lambda_unique_ptr<ngram_model_t> createLanguageModel(const vector<string>& words, logmath_t& logMath) {
	path tempFilePath = getTempFilePath();
	createLanguageModelFile(words, tempFilePath);
	auto deleteTempFile = gsl::finally([&]() { boost::filesystem::remove(tempFilePath); });

	return lambda_unique_ptr<ngram_model_t>(
		ngram_model_read(nullptr, tempFilePath.string().c_str(), NGRAM_ARPA, &logMath),
		[](ngram_model_t* lm) { ngram_model_free(lm); });
}


//#define EIGEN_USE_BLAS

#include <Eigen/Dense>
#include <math.h>
extern "C" {
  #include <csr.h>
}
#include <iostream>

#include <stdexcept>

#include "chol.h"
#include "linop.h"

using namespace Eigen;
using namespace std;

extern "C" void dsyrk_(char *uplo, char *trans, int *m, int *n, double *alpha, double a[],
            int *lda, double *beta, double c[], int *ldc);
extern "C" void dgemm_(char *transa, char *transb, int *m, int *n, int *k, double *alpha,
            double a[], int *lda, double b[], int *ldb, double *beta, double c[],
            int *ldc);

/*
extern "C" void dsymm_(char *side, char *uplo, int *m, int *n, double *alpha, double a[],
            int *lda, double b[], int *ldb, double *beta, double c[], int *ldc);
*/

// out = bcsr * b (for vectors)
void A_mul_B(Eigen::VectorXd & out, BinaryCSR & csr, Eigen::VectorXd & b) {
  if (csr.nrow != out.size()) {throw std::runtime_error("csr.nrow must equal out.size()");}
  if (csr.ncol != b.size())   {throw std::runtime_error("csr.ncol must equal b.size()");}
  bcsr_A_mul_B( out.data(), & csr, b.data() );
}

// OUT' = bcsr * B' (for matrices)
void A_mul_Bt(Eigen::MatrixXd & out, BinaryCSR & csr, Eigen::MatrixXd & B) {
  if (csr.nrow != out.cols()) {throw std::runtime_error("csr.nrow must equal out.cols()");}
  if (csr.ncol != B.cols())   {throw std::runtime_error("csr.ncol must equal b.cols()");}
  if (out.rows() != B.rows()) {throw std::runtime_error("out.rows() must equal B.rows()");}
  bcsr_A_mul_Bn( out.data(), & csr, B.data(), B.rows() );
}

// out = bcsr * b (for vectors)
void A_mul_B(Eigen::VectorXd & out, CSR & csr, Eigen::VectorXd & b) {
  if (csr.nrow != out.size()) {throw std::runtime_error("csr.nrow must equal out.size()");}
  if (csr.ncol != b.size())   {throw std::runtime_error("csr.ncol must equal b.size()");}
  csr_A_mul_B( out.data(), & csr, b.data() );
}

// OUT' = bcsr * B' (for matrices)
void A_mul_Bt(Eigen::MatrixXd & out, CSR & csr, Eigen::MatrixXd & B) {
  if (csr.nrow != out.cols()) {throw std::runtime_error("csr.nrow must equal out.cols()");}
  if (csr.ncol != B.cols())   {throw std::runtime_error("csr.ncol must equal b.cols()");}
  if (out.rows() != B.rows()) {throw std::runtime_error("out.rows() must equal B.rows()");}
  csr_A_mul_Bn( out.data(), & csr, B.data(), B.rows() );
}

void At_mul_A(Eigen::MatrixXd & out, SparseFeat & A) {
  if (out.cols() != A.cols()) {
    throw std::runtime_error("At_mul_A(SparseFeat): out.cols() must equal A.cols()");
  }
  if (out.cols() != out.rows()) {
    throw std::runtime_error("At_mul_A(SparseFeat): out must be square matrix.)");
  }

  out.setZero();
  const int nfeat = A.M.ncol;

#pragma omp parallel for schedule(dynamic, 8)
  for (int f1 = 0; f1 < nfeat; f1++) {
    int end = A.Mt.row_ptr[f1 + 1];
    out(f1, f1) = end - A.Mt.row_ptr[f1];
    // looping over all non-zero rows of f1
    for (int i = A.Mt.row_ptr[f1]; i < end; i++) {
      int Mrow = A.Mt.cols[i]; /* row in M */
      int end2 = A.M.row_ptr[Mrow + 1];
      for (int j = A.M.row_ptr[Mrow]; j < end2; j++) {
        int f2 = A.M.cols[j];
        if (f1 < f2) {
          out(f2, f1) += 1;
        }
      }
    }
  }
}

void At_mul_A(Eigen::MatrixXd & out, Eigen::MatrixXd & A) {
  // TODO: use blas
  out.triangularView<Eigen::Lower>() = A.transpose() * A;
}

void At_mul_A(Eigen::MatrixXd & out, SparseDoubleFeat & A) {
  if (out.cols() != A.cols()) {
    throw std::runtime_error("At_mul_A(SparseDoubleFeat): out.cols() must equal A.cols()");
  }
  if (out.cols() != out.rows()) {
    throw std::runtime_error("At_mul_A(SparseDoubleFeat): out must be square matrix.)");
  }
  out.setZero();
  const int nfeat = A.M.ncol;

#pragma omp parallel for schedule(dynamic, 8)
  for (int f1 = 0; f1 < nfeat; f1++) {
    // looping over all non-zero rows of f1
    for (int i = A.Mt.row_ptr[f1], end = A.Mt.row_ptr[f1 + 1]; i < end; i++) {
      int Mrow     = A.Mt.cols[i]; /* row in M */
      double val1  = A.Mt.vals[i]; /* value for Mrow */

      for (int j = A.M.row_ptr[Mrow], end2 = A.M.row_ptr[Mrow + 1]; j < end2; j++) {
        int f2 = A.M.cols[j];
        if (f1 <= f2) {
          out(f2, f1) += A.M.vals[j] * val1;
        }
      }
    }
  }
}

void At_mul_A_blas(Eigen::MatrixXd & A, double* AtA) {
  char lower  = 'L';
  char trans  = 'T';
  int  m      = A.cols();
  int  n      = A.rows();
  double one  = 1.0;
  double zero = 0.0;
  dsyrk_(&lower, &trans, &m, &n, &one, A.data(),
         &n, &zero, AtA, &m);
}

void A_mul_At_blas(Eigen::MatrixXd & A, double* AAt) {
  char lower  = 'L';
  char trans  = 'N';
  int  m      = A.cols();
  int  n      = A.rows();
  double one  = 1.0;
  double zero = 0.0;
  dsyrk_(&lower, &trans, &n, &m, &one, A.data(),
         &n, &zero, AAt, &n);
}

void A_mul_B_blas(Eigen::MatrixXd & Y, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  if (Y.rows() != A.rows()) {throw std::runtime_error("A.rows() must equal Y.rows()");}
  if (Y.cols() != B.cols()) {throw std::runtime_error("B.cols() must equal Y.cols()");}
  if (A.cols() != B.rows()) {throw std::runtime_error("B.rows() must equal A.cols()");}
  int m = A.rows();
  int n = B.cols();
  int k = A.cols();
  char transA = 'N';
  char transB = 'N';
  double alpha = 1.0;
  double beta  = 0.0;
  dgemm_(&transA, &transB, &m, &n, &k, &alpha, A.data(), &m, B.data(), &k, &beta, Y.data(), &m);
}

void At_mul_B_blas(Eigen::MatrixXd & Y, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  if (Y.rows() != A.cols()) {throw std::runtime_error("A.rows() must equal Y.rows()");}
  if (Y.cols() != B.cols()) {throw std::runtime_error("B.cols() must equal Y.cols()");}
  if (A.rows() != B.rows()) {throw std::runtime_error("B.rows() must equal A.cols()");}
  int m = A.cols();
  int n = B.cols();
  int k = A.rows();
  char transA = 'T';
  char transB = 'N';
  double alpha = 1.0;
  double beta  = 0.0;
  dgemm_(&transA, &transB, &m, &n, &k, &alpha, A.data(), &k, B.data(), &k, &beta, Y.data(), &m);
}

void A_mul_Bt_blas(Eigen::MatrixXd & Y, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  if (Y.rows() != A.rows()) {throw std::runtime_error("A.rows() must equal Y.rows()");}
  if (Y.cols() != B.rows()) {throw std::runtime_error("B.rows() must equal Y.cols()");}
  if (A.cols() != B.cols()) {throw std::runtime_error("B.cols() must equal A.cols()");}
  int m = A.rows();
  int n = B.rows();
  int k = A.cols();
  char transA = 'N';
  char transB = 'T';
  double alpha = 1.0;
  double beta  = 0.0;
  dgemm_(&transA, &transB, &m, &n, &k, &alpha, A.data(), &m, B.data(), &n, &beta, Y.data(), &m);
}
/*
void Asym_mul_B_left(double beta, Eigen::MatrixXd & Y, double alpha, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  cblas_dsymm(CblasColMajor, CblasLeft, CblasLower, Y.rows(), Y.cols(), alpha, A.data(), A.rows(), B.data(), B.rows(), beta, Y.data(), Y.rows());
}

void Asym_mul_B_right(double beta, Eigen::MatrixXd & Y, double alpha, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  cblas_dsymm(CblasColMajor, CblasRight, CblasLower, Y.rows(), Y.cols(), alpha, A.data(), A.rows(), B.data(), B.rows(), beta, Y.data(), Y.rows());
}*/

template<> void AtA_mul_B(Eigen::MatrixXd & out, SparseFeat & A, double reg, Eigen::MatrixXd & B, Eigen::MatrixXd & tmp) {
  // solution update:
  A_mul_Bt(tmp, A.M, B);
  // move KP += reg * P here with nowait from previous loop
  // http://stackoverflow.com/questions/30496365/parallelize-the-addition-of-a-vector-of-matrices-in-openmp
  A_mul_Bt(out, A.Mt, tmp);

  int ncol = out.cols(), nrow = out.rows();
#pragma omp parallel for schedule(static)
  for (int col = 0; col < ncol; col++) {
    for (int row = 0; row < nrow; row++) {
      out(row, col) += reg * B(row, col);
    }
  }
}

template<>
void AtA_mul_B(Eigen::MatrixXd & out, SparseDoubleFeat & A, double reg, Eigen::MatrixXd & B, Eigen::MatrixXd & tmp) {
  // solution update:
  A_mul_Bt(tmp, A.M, B);
  // move KP += reg * P here with nowait from previous loop
  // http://stackoverflow.com/questions/30496365/parallelize-the-addition-of-a-vector-of-matrices-in-openmp
  A_mul_Bt(out, A.Mt, tmp);
  int ncol = out.cols(), nrow = out.rows();
#pragma omp parallel for schedule(static)
  for (int col = 0; col < ncol; col++) {
    for (int row = 0; row < nrow; row++) {
      out(row, col) += reg * B(row, col);
    }
  }
}

template<>
void AtA_mul_B(Eigen::MatrixXd & out, Eigen::MatrixXd & A, double reg, Eigen::MatrixXd & B, Eigen::MatrixXd & tmp) {
  // solution update:
  A_mul_Bt_blas(tmp, B, A);
  // move KP += reg * P here with nowait from previous loop
  // http://stackoverflow.com/questions/30496365/parallelize-the-addition-of-a-vector-of-matrices-in-openmp
  A_mul_B_blas(out, B, A);

  int ncol = out.cols(), nrow = out.rows();
#pragma omp parallel for schedule(static)
  for (int col = 0; col < ncol; col++) {
    for (int row = 0; row < nrow; row++) {
      out(row, col) += reg * B(row, col);
    }
  }
}

void makeSymmetric(Eigen::MatrixXd & A) {
  A = A.selfadjointView<Eigen::Lower>();
}

/**
 * A is [n x k] matrix
 * returns [n x n] matrix A * A'
 */
Eigen::MatrixXd A_mul_At_combo(Eigen::MatrixXd & A) {
  MatrixXd AAt(A.rows(), A.rows());
  A_mul_At_combo(AAt, A);
  AAt = AAt.selfadjointView<Eigen::Lower>();
  return AAt;
}

/** A   is [n x k] matrix
 *  out is [n x n] matrix
 *  A and out are column-ordered
 *  computes out = A * A'
 *  (storing only lower triangular part)
 */
void A_mul_At_combo(Eigen::MatrixXd & out, Eigen::MatrixXd & A) {
  if (out.rows() >= 128) {
    // using blas for larger matrices
    A_mul_At_blas(A, out.data());
  } else {
    A_mul_At_omp(out, A);
  }
}

void A_mul_At_omp(Eigen::MatrixXd & out, Eigen::MatrixXd & A) {
  const int n = A.rows();
  const int k = A.cols();
  int nthreads = -1;
  double* x = A.data();
  if (A.rows() != out.rows()) {
    throw std::runtime_error("A.rows() must equal out.rows()");
  }

#pragma omp parallel
  {
#pragma omp single
    {
      nthreads = omp_get_num_threads();
    }
  }
  std::vector<MatrixXd> Ys;
  Ys.resize(nthreads, MatrixXd(n, n));

#pragma omp parallel
  {
    const int ithread  = omp_get_thread_num();
    int rows_per_thread = (int) 8 * ceil(k / 8.0 / nthreads);
    int row_start = rows_per_thread * ithread;
    int row_end   = rows_per_thread * (ithread + 1);
    if (row_start >= k) {
      Ys[ithread].setZero();
    } else {
      if (row_end > k) {
        row_end = k;
      }
      double* xi = & x[ row_start * n ];
      int nrows  = row_end - row_start;
      MatrixXd X = Map<MatrixXd>(xi, n, nrows);
      MatrixXd & Y = Ys[ithread];
      Y.triangularView<Eigen::Lower>() = X * X.transpose();
    }
  }
  for (int i = 0; i < n; i++) {
    for (int j = i; j < n; j++) {
      double tmp = 0;
      for (int k = 0; k < nthreads; k++) {
        tmp += Ys[k](j, i);
      }
      out(j, i) = tmp;
    }
  }
}

void A_mul_Bt_omp_sym(Eigen::MatrixXd & out, Eigen::MatrixXd & A, Eigen::MatrixXd & B) {
  const int n = A.rows();
  const int k = A.cols();
  int nthreads = -1;
  double* x  = A.data();
  double* x2 = B.data();
  assert(A.rows() == B.rows());
  assert(A.cols() == B.cols());
  if (A.rows() != out.rows()) {
    throw std::runtime_error("A.rows() must equal out.rows()");
  }

#pragma omp parallel
  {
#pragma omp single
    {
      nthreads = omp_get_num_threads();
    }
  }
  std::vector<MatrixXd> Ys;
  Ys.resize(nthreads, MatrixXd(n, n));

#pragma omp parallel
  {
    const int ithread  = omp_get_thread_num();
    int rows_per_thread = (int) 8 * ceil(k / 8.0 / nthreads);
    int row_start = rows_per_thread * ithread;
    int row_end   = rows_per_thread * (ithread + 1);
    if (row_start >= k) {
      Ys[ithread].setZero();
    } else {
      if (row_end > k) {
        row_end = k;
      }
      double* xi  = & x[ row_start * n ];
      double* x2i = & x2[ row_start * n ];
      int nrows   = row_end - row_start;
      MatrixXd X  = Map<MatrixXd>(xi, n, nrows);
      MatrixXd X2 = Map<MatrixXd>(x2i, n, nrows);
      MatrixXd & Y = Ys[ithread];
      Y.triangularView<Eigen::Lower>() = X * X2.transpose();
    }
  }
  for (int i = 0; i < n; i++) {
    for (int j = i; j < n; j++) {
      double tmp = 0;
      for (int k = 0; k < nthreads; k++) {
        tmp += Ys[k](j, i);
      }
      out(j, i) = tmp;
    }
  }
}

Eigen::VectorXd col_square_sum(SparseFeat & A) {
  const int ncol = A.cols();
  VectorXd out(ncol);
#pragma omp parallel for schedule(static)
  for (int col = 0; col < ncol; col++) {
    out(col) = A.Mt.row_ptr[col + 1] - A.Mt.row_ptr[col];
  }
  return out;
}

Eigen::VectorXd col_square_sum(SparseDoubleFeat & A) {
  const int ncol = A.cols();
  const int* row_ptr = A.Mt.row_ptr;
  const double* vals = A.Mt.vals;
  VectorXd out(ncol);

#pragma omp parallel for schedule(dynamic, 256)
  for (int col = 0; col < ncol; col++) {
    double tmp = 0;
    int i   = row_ptr[col];
    int end = row_ptr[col + 1];
    for (; i < end; i++) {
      tmp += vals[i] * vals[i];
    }
    out(col) = tmp;
  }
  return out;
}