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

#include <ndt/ndt_voxel.hpp>
#include <ndt/utils.hpp>
#include <Eigen/LU>
#include <limits>
#include "common/types.hpp"

using autoware::common::types::bool8_t;

namespace autoware
{
namespace localization
{
namespace ndt
{
constexpr uint32_t DynamicNDTVoxel::NUM_POINT_THRESHOLD;

DynamicNDTVoxel::DynamicNDTVoxel()
{
  // default constructors use uninitialized values.
  m_centroid.setZero();
  m_M2.setZero();
  m_covariance.setZero();
}

void DynamicNDTVoxel::add_observation(const Point & pt)
{
  const auto last_count = m_num_points++;

  const auto last_centroid = m_centroid;
  const auto last_delta = pt - last_centroid;
  m_centroid = last_centroid + (last_delta / m_num_points);

  const auto current_delta = pt - m_centroid;
  m_M2 += last_delta * current_delta.transpose();

  if (usable()) {
    // TODO(yunus.caliskan): Apply numerical stability enhancing steps described in:
    // http://www.diva-portal.org/smash/get/diva2:276162/FULLTEXT02.pdf, pg 60
    m_covariance = m_M2 / last_count;
  }

  // set invertibility to unknown since the covariance has changed
  m_invertible = Invertibility::UNKNOWN;
}

bool8_t DynamicNDTVoxel::try_stabilize()
{
  bool8_t invertible = try_stabilize_covariance(m_covariance);
  if (invertible) {
    m_invertible = Invertibility::INVERTIBLE;
  } else {
    m_invertible = Invertibility::NOT_INVERTIBLE;
  }

  return invertible;
}

bool8_t DynamicNDTVoxel::usable() const noexcept
{
  return m_num_points >= NUM_POINT_THRESHOLD;
}

const Eigen::Matrix3d & DynamicNDTVoxel::covariance() const
{
  if (!usable()) {
    throw std::out_of_range(
            "DynamicNDTVoxel: Cannot get covariance from a "
            "voxel without sufficient number of points");
  }
  return m_covariance;
}

std::experimental::optional<Eigen::Matrix3d> DynamicNDTVoxel::inverse_covariance() const
{
  if (!usable()) {
    throw std::out_of_range(
            "DynamicNDTVoxel: Cannot get covariance from a "
            "voxel without sufficient number of points");
  }

  if (m_invertible == Invertibility::NOT_INVERTIBLE) {
    // if stabilization has been performed and covariance is not invertible
    return {};
  }

  Eigen::Matrix3d inv_covariance;
  bool8_t invertible;
  m_covariance.computeInverseWithCheck(inv_covariance, invertible);
  if (invertible) {
    return inv_covariance;
  } else {
    return {};
  }
}

const Eigen::Vector3d & DynamicNDTVoxel::centroid() const
{
  // Using the overloaded function as the parent function will use the hidden occupancy check
  if (!usable()) {
    throw std::out_of_range("DynamicNDTVoxel: Cannot get centroid from an unoccupied voxel");
  }
  return m_centroid;
}

uint64_t DynamicNDTVoxel::count() const noexcept
{
  return m_num_points;
}

/////////////////////////////////////////////////

StaticNDTVoxel::StaticNDTVoxel()
{
  m_centroid.setZero();
  m_inv_covariance.setZero();
}

StaticNDTVoxel::StaticNDTVoxel(const Point & centroid, const Cov & inv_covariance)
: m_centroid{centroid}, m_inv_covariance{inv_covariance}, m_occupied{true}
{}

Eigen::Matrix3d StaticNDTVoxel::covariance() const
{
  Eigen::Matrix3d covariance;
  if (m_occupied) {
    bool8_t invertible{false};
    m_inv_covariance.computeInverseWithCheck(covariance, invertible);
    if (!invertible) {
      throw std::out_of_range("StaticNDTVoxel: Inverse covariance is not invertible");
    }
  } else {
    throw std::out_of_range("StaticNDTVoxel: Cannot get covariance from an unoccupied voxel");
  }
  return covariance;
}

const Eigen::Vector3d & StaticNDTVoxel::centroid() const
{
  if (!m_occupied) {
    throw std::out_of_range("StaticNDTVoxel: Cannot get centroid from an unoccupied voxel");
  }
  return m_centroid;
}

const Eigen::Matrix3d & StaticNDTVoxel::inverse_covariance() const
{
  if (!m_occupied) {
    throw std::out_of_range(
            "StaticNDTVoxel: Cannot get inverse covariance "
            "from an unoccupied voxel");
  }
  return m_inv_covariance;
}

bool8_t StaticNDTVoxel::usable() const noexcept
{
  return m_occupied;
}
}  // namespace ndt
}  // namespace localization
}  // namespace autoware


/*
* Software License Agreement (BSD License)
*
*  Copyright (c) 2014-2015, Timm Linder, Social Robotics Lab, University of Freiburg
*  Copyright (c) 2006-2012, Matthias Luber, Luciano Spinello and Kai O. Arras, Social Robotics Laboratory and
*    Oscar Martinez, Autonomous Intelligent Systems, University of Freiburg
*  All rights reserved.
*
*  Redistribution and use in source and binary forms, with or without
*  modification, are permitted provided that the following conditions are met:
*
*  * Redistributions of source code must retain the above copyright notice, this
*    list of conditions and the following disclaimer.
*  * Redistributions in binary form must reproduce the above copyright notice,
*    this list of conditions and the following disclaimer in the documentation
*    and/or other materials provided with the distribution.
*  * Neither the name of the copyright holder nor the names of its contributors
*    may be used to endorse or promote products derived from this software
*    without specific prior written permission.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
*  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
*  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
*  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
*  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
*  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
*  SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
*  CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
*  OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
*  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

#include <srl_laser_features/features/feature10.h>
#include <Eigen/QR>
#include <float.h>

namespace srl_laser_features
{

Feature10::Feature10(bool extended) : Feature(), m_extended(extended)
{
}

void Feature10::evaluate(const Segment& segment, Eigen::VectorXd& result) const
{
    result = Eigen::VectorXd::Zero(getNDimensions());

    const size_t numPoints = segment.points.size();
    if (numPoints >= 2) {
        double px[numPoints];
        double py[numPoints];

        // copy
        for (size_t pIndex = 0; pIndex < numPoints; ++pIndex) {
            px[pIndex] = segment.points[pIndex](0);
            py[pIndex] = segment.points[pIndex](1);
        }

        double m, q, diff;
        result(0) = 0.0;
        if (m_extended) {
            result(1) = 0.0;
            result(2) = DBL_MAX;
            result(3) = 0.0;
        }
        fitLine(numPoints, px, py, m, q);

        // check for vertical line
        if (fabs(m) > 10e+10) {
            // residual sum
            double diff;
            for (size_t i = 0; i < numPoints; ++i) {
                diff = segment.points[i](0) - segment.mean(0);
                result(0) += diff * diff;
                if (m_extended) {
                    if (diff < result(2)) {
                        result(2) = diff;
                    }
                    if (diff > result(3)) {
                        result(3) = diff;
                    }
                }
            }
        }
        else {
            // residual sum
            for (size_t i = 0; i < numPoints; ++i) {
                diff = m * segment.points[i](0) + q - segment.points[i](1);
                result(0) += diff * diff;
                if (m_extended) {
                    if (diff < result(2)) {
                        result(2) = diff;
                    }
                    if (diff > result(3)) {
                        result(3) = diff;
                    }
                }
            }
        }

        if (m_extended) {
            result(1) = result(0) / numPoints;
            if (result(3) != 0.0) {
                result(4) = result(2) / result(3);
            }
        }
    }
    else {
        result.setConstant(getNDimensions(), -1.0);
    }
}

void Feature10::fitLine(int n, double* xData, double* yData, double& m, double& q)
{
    Eigen::MatrixXd A(n, 2);
    Eigen::VectorXd B(n);
    Eigen::VectorXd C(2);

    // fill A and B
    for (int i = 0; i < n; i++) {
        A(i, 0) = xData[i];
        A(i, 1) = 1.0;
        B(i) = yData[i];
    }

    // solve
    Eigen::FullPivHouseholderQR<Eigen::MatrixXd> QR(A);
    C = QR.solve(B);
    m = C(0);
    q = C(1);
}

} // end of namespace srl_laser_features


// 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


#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); });
}


// This file is part of LatNet Builder.
//
// Copyright (C) 2012-2021  The LatNet Builder author's, supervised by Pierre L'Ecuyer, Universite de Montreal.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "latbuilder/Util.h"
#include "netbuilder/Helpers/Path.h"
#include <cmath>
#include <cstdlib>
#include <fstream>
#include <boost/filesystem.hpp>

#include <NTL/ZZ.h>

namespace LatBuilder
{

//===============================================================================
Polynomial PolynomialFromInt(uInteger x)
{
   Polynomial P;
   P.SetLength(LENGTH_UINTEGER);
   uInteger y=x;
   long i=0;
    do{
       P[i++] = y%2;
   } while(y >>= 1);

    P.normalize();
    return P;
}

//================================================================================

uInteger IndexOfPolynomial(Polynomial P)
{
   uInteger x=0;
   uInteger q=1;

   for (int i = 0; i <= deg(P); i++)
    {
        if(IsOne(coeff(P,i))){
            x+= q;
       }
       q <<= 1;
    }
    return x;
}

//================================================================================

std::vector<uInteger> primeFactors(uInteger n, bool raise)
{
  NTL::PrimeSeq s;
   unsigned long p;
   p = s.next();
   std::vector<uInteger> factors;
   while (p <= n) {
      ldiv_t qr = std::ldiv(n, p);
      uInteger factor = 1;
      while (qr.rem == 0) {
         n = qr.quot;
         factor *= p;
         if (!raise)
            break;
         qr = std::ldiv(n, p);
      }
      if (factor > 1)
         factors.push_back(factor);
      p = s.next();
   }
   return factors;
}

//================================================================================

std::map<uInteger, uInteger> primeFactorsMap(uInteger n)
{
  NTL::PrimeSeq s;
   unsigned long factor;
   factor = s.next();
   std::map<uInteger, uInteger> factors;
   while (factor <= n) {
      ldiv_t qr = std::ldiv(n, factor);
      while (qr.rem == 0) {
         n = qr.quot;
         factors[factor]++;
         qr = std::ldiv(n, factor);
      }
      factor = s.next();
   }
   return factors;
}

//================================================================================

std::pair<long long, long long> egcd(uInteger a, uInteger b)
{
   typedef std::pair<long long, long long> result_type;
   result_type cur{0, 1};
   result_type last{1, 0};
   while (b != 0) {
      ldiv_t qr = std::ldiv(a, b);
      a = b;
      b = qr.rem;
      std::swap(cur, last);
      cur.first = cur.first - qr.quot * last.first;
      cur.second = cur.second - qr.quot * last.second;
   }
   return last;
}

//================================================================================

uInteger Vm(const Polynomial& h, const Polynomial& P)
{

   long m = deg(P);
   NTL::vector<NTL::GF2> w;
   w.resize(m);
   uInteger res = 0;
   for(long i=0; i<m; i++){
      w[i] = (m-i-1>deg(h)) ? NTL::GF2(0) : coeff(h,m-i-1);
      for(long j=0; j<i; j++){
         w[i] += w[j]* coeff(P,m-(i-j));
      }
   }
   uInteger q = 1;
   for(long i=m-1; i>=0; i--){
      res += IsOne(w[i])*q;
      q *= 2;
   }

   return res;

}

uInteger log2Int(unsigned int n){
    unsigned int res=0;
    while (n>1){
    n = n >> 1;
    res ++;
    }
    return res;
}

//================================================================================

const char* ws = " \t\n\r\f\v";

// trim from end (right)
inline std::string& rtrim(std::string& s, const char* t = ws)
{
      s.erase(s.find_last_not_of(t) + 1);
      return s;
}

// trim from beginning (left)
inline std::string& ltrim(std::string& s, const char* t = ws)
{
      s.erase(0, s.find_first_not_of(t));
      return s;
}

// trim from both ends (left & right)
inline std::string& trim(std::string& s, const char* t = ws)
{
      return ltrim(rtrim(s, t), t);
}

std::string getDefaultPolynomial(unsigned int degree)
{
    if (degree <= 32)
    {
        std::string path = NetBuilder::PATH_TO_LATNETBUILDER_DIR + "/../share/latnetbuilder/data/default_polys.csv";
        if (boost::filesystem::exists(path)){
            std::ifstream file(path);
            std::string sent;
            do
            {
            getline(file,sent);
            trim(sent);
            }
            while (sent != "###");

            getline(file,sent);

            for(unsigned int d = 0; d <= degree; ++d)
            {
                getline(file,sent);
            }
            return sent;
        }
        else{
            throw std::runtime_error("Unable to locate data folder. The value of PATH_TO_LATNETBUILDER_DIR is probably incorrect. See netbuilder/Path.h.");
        }
    }
    return "";
}

std::string to_string(LatticeType LT){

    static const char * LatticeTypeStrings[] = { "Ordinary", "Polynomial", "Digital" };

    return LatticeTypeStrings[(int) LT];
}

} // namespace LatBuilder


/***********************************************************************************************************************
*  OpenStudio(R), Copyright (c) 2008-2021, Alliance for Sustainable Energy, LLC, and other contributors. All rights reserved.
*
*  Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
*  following conditions are met:
*
*  (1) Redistributions of source code must retain the above copyright notice, this list of conditions and the following
*  disclaimer.
*
*  (2) Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
*  disclaimer in the documentation and/or other materials provided with the distribution.
*
*  (3) Neither the name of the copyright holder nor the names of any contributors may be used to endorse or promote products
*  derived from this software without specific prior written permission from the respective party.
*
*  (4) Other than as required in clauses (1) and (2), distributions in any form of modifications or other derivative works
*  may not use the "OpenStudio" trademark, "OS", "os", or any other confusingly similar designation without specific prior
*  written permission from Alliance for Sustainable Energy, LLC.
*
*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER(S) AND ANY CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
*  INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
*  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER(S), ANY CONTRIBUTORS, THE UNITED STATES GOVERNMENT, OR THE UNITED
*  STATES DEPARTMENT OF ENERGY, NOR ANY OF THEIR EMPLOYEES, BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
*  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
*  USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
*  STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
*  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
***********************************************************************************************************************/

#ifndef UTILITIES_DATA_MATRIX_HPP
#define UTILITIES_DATA_MATRIX_HPP

#include "Vector.hpp"
#include "../UtilitiesAPI.hpp"
#include "../core/Logger.hpp"

#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/lu.hpp>

namespace openstudio {

/// Matrix
typedef boost::numeric::ublas::matrix<double> Matrix;

//////////////////////////////////////////////////////////////////////////
// Begin SWIG'able, copy and paste into Matrix.i
//////////////////////////////////////////////////////////////////////////

/// new operators

UTILITIES_API bool operator==(const Matrix& lhs, const Matrix& rhs);
UTILITIES_API bool operator!=(const Matrix& lhs, const Matrix& rhs);

/// common methods

/// linear interpolation of the function v = f(x, y) at point xi, yi
/// assumes that x and y are strictly increasing
UTILITIES_API double interp(const Vector& x, const Vector& y, const Matrix& v, double xi, double yi, InterpMethod interpMethod = LinearInterp,
                            ExtrapMethod extrapMethod = NoneExtrap);

/// linear interpolation of the function v = f(x, y) at points xi, yi
/// assumes that x and y are strictly increasing
UTILITIES_API Vector interp(const Vector& x, const Vector& y, const Matrix& v, const Vector& xi, double yi, InterpMethod interpMethod = LinearInterp,
                            ExtrapMethod extrapMethod = NoneExtrap);

/// linear interpolation of the function v = f(x, y) at points xi, yi
/// assumes that x and y are strictly increasing
UTILITIES_API Vector interp(const Vector& x, const Vector& y, const Matrix& v, double xi, const Vector& yi, InterpMethod interpMethod = LinearInterp,
                            ExtrapMethod extrapMethod = NoneExtrap);

/// linear interpolation of the function v = f(x, y) at points xi, yi
/// assumes that x and y are strictly increasing
UTILITIES_API Matrix interp(const Vector& x, const Vector& y, const Matrix& v, const Vector& xi, const Vector& yi,
                            InterpMethod interpMethod = LinearInterp, ExtrapMethod extrapMethod = NoneExtrap);

/// matrix product
UTILITIES_API Matrix prod(const Matrix& lop, const Matrix& rop);

/// vector product
UTILITIES_API Vector prod(const Matrix& m, const Vector& v);

/// outer product
UTILITIES_API Matrix outerProd(const Vector& lhs, const Vector& rhs);

/// take the natural logarithm of Matrix elements, componentwise
UTILITIES_API Matrix log(const Matrix& v);

/// take the logarithm of Matrix elements with respect to base, componentwise
UTILITIES_API Matrix log(const Matrix& v, double base);

/// generates a M x N Matrix whose elements come from the uniform distribution on [a,b].
UTILITIES_API Matrix randMatrix(double a, double b, unsigned M, unsigned N);

/// sum of all elements
UTILITIES_API double sum(const Matrix& matrix);

/// maximum of all elements
UTILITIES_API double maximum(const Matrix& matrix);

/// minimum of all elements
UTILITIES_API double minimum(const Matrix& matrix);

/// mean of all elements
UTILITIES_API double mean(const Matrix& matrix);

/// get the connected components from an NxN adjacency matrix (1.0 for i-j connected, 0.0 for i-j not connected)
UTILITIES_API std::vector<std::vector<unsigned>> findConnectedComponents(const Matrix& matrix);

// from the boost vault:
// The following code inverts the matrix input using LU-decomposition with backsubstitution of unit vectors. Reference: Numerical Recipes in C, 2nd ed., by Press, Teukolsky, Vetterling & Flannery.
/// Matrix inversion routine, using lu_factorize and lu_substitute in uBLAS to invert a matrix */
template <class T>
bool invert(const boost::numeric::ublas::matrix<T>& input, boost::numeric::ublas::matrix<T>& inverse) {

  // create a working copy of the input
  boost::numeric::ublas::matrix<T> A(input);

  // create a permutation matrix for the LU-factorization
  boost::numeric::ublas::permutation_matrix<std::size_t> pm(A.size1());

  // perform LU-factorization
  typename boost::numeric::ublas::matrix<T>::size_type res = boost::numeric::ublas::lu_factorize(A, pm);
  if (res != 0) {
    LOG_FREE(Info, "boost.ublas",
             "boost::numeric::ublas::lu_factorize returned res = " << res << ", A = " << A << ", pm = " << pm << " for input = " << input);
    return false;
  }

  // create identity matrix of "inverse"
  inverse.assign(boost::numeric::ublas::identity_matrix<T>(A.size1()));

  // backsubstitute to get the inverse
  try {
    boost::numeric::ublas::lu_substitute(A, pm, inverse);
  } catch (std::exception& e) {
    LOG_FREE(Info, "boost.ublas", "boost::numeric::ublas::lu_substitute threw exception '" << e.what() << "' for A = " << A << ", pm = " << pm);
    return false;
  }

  return true;
}

}  // namespace openstudio

#endif  //UTILITIES_DATA_MATRIX_HPP


/*
 * DivRightHandSide.cc
 *
 *  Created on: 29.06.2017
 *      Author: thies
 */

#include <deal.II/base/exceptions.h>
#include <deal.II/base/types.h>
#include <deal.II/base/work_stream.h>
#include <deal.II/fe/fe.h>
#include <deal.II/fe/fe_update_flags.h>
#include <deal.II/fe/fe_values.h>

#include <base/DiscretizedFunction.h>
#include <forward/DivRightHandSide.h>

#include <functional>
#include <vector>

namespace wavepi {
namespace forward {
using namespace dealii;
using namespace wavepi::base;

template<int dim>
DivRightHandSide<dim>::DivRightHandSide(std::shared_ptr<DiscretizedFunction<dim>> a,
      std::shared_ptr<DiscretizedFunction<dim>> b, std::shared_ptr<DiscretizedFunction<dim>> u)
      : a(a), b(b), u(u) {
}

template<int dim>
DivRightHandSide<dim>::AssemblyScratchData::AssemblyScratchData(const FiniteElement<dim> &fe,
      const Quadrature<dim> &quad)
      : fe_values(fe, quad, update_values | update_gradients | update_quadrature_points | update_JxW_values) {
}

template<int dim>
DivRightHandSide<dim>::AssemblyScratchData::AssemblyScratchData(const AssemblyScratchData &scratch_data)
      :
            fe_values(scratch_data.fe_values.get_fe(), scratch_data.fe_values.get_quadrature(),
                  update_values | update_gradients | update_quadrature_points | update_JxW_values) {
}

template<int dim>
void DivRightHandSide<dim>::copy_local_to_global(Vector<double> &result, const AssemblyCopyData &copy_data) {
   for (unsigned int i = 0; i < copy_data.local_dof_indices.size(); ++i)
      result(copy_data.local_dof_indices[i]) += copy_data.cell_rhs(i);
}

template<int dim>
void DivRightHandSide<dim>::local_assemble(const Vector<double> &a, const Vector<double> &b, const Vector<double> &u,
      const typename DoFHandler<dim>::active_cell_iterator &cell, AssemblyScratchData &scratch_data,
      AssemblyCopyData &copy_data) {
   const unsigned int dofs_per_cell = scratch_data.fe_values.get_fe().dofs_per_cell;
   const unsigned int n_q_points = scratch_data.fe_values.get_quadrature().size();

   copy_data.cell_rhs.reinit(dofs_per_cell);
   copy_data.local_dof_indices.resize(dofs_per_cell);
   scratch_data.fe_values.reinit(cell);

   cell->get_dof_indices(copy_data.local_dof_indices);

   for (unsigned int q_point = 0; q_point < n_q_points; ++q_point)
      for (unsigned int i = 0; i < dofs_per_cell; ++i)
         for (unsigned int ka = 0; ka < dofs_per_cell; ++ka) {
            for (unsigned int ku = 0; ku < dofs_per_cell; ++ku)
               copy_data.cell_rhs(i) -= a[copy_data.local_dof_indices[ka]]
                     * scratch_data.fe_values.shape_value(ka, q_point) * u[copy_data.local_dof_indices[ku]]
                     * scratch_data.fe_values.shape_grad(ku, q_point) * scratch_data.fe_values.shape_grad(i, q_point)
                     * scratch_data.fe_values.JxW(q_point);

            copy_data.cell_rhs(i) += b[copy_data.local_dof_indices[ka]]
                  * scratch_data.fe_values.shape_value(ka, q_point) * scratch_data.fe_values.shape_value(i, q_point)
                  * scratch_data.fe_values.JxW(q_point);
         }
}

template<int dim>
void DivRightHandSide<dim>::create_right_hand_side(const DoFHandler<dim> &dof, const Quadrature<dim> &quad,
      Vector<double> &rhs) const {
   AssertThrow(a != nullptr, ExcZero());
   AssertThrow(b != nullptr, ExcZero());
   AssertThrow(u != nullptr, ExcZero());

   const Vector<double> &ca = a->get_function_coefficients_by_time(this->get_time());
   const Vector<double> &cb = b->get_function_coefficients_by_time(this->get_time());
   const Vector<double> &cu = u->get_function_coefficients_by_time(this->get_time());

   Assert(ca.size() == dof.n_dofs(), ExcDimensionMismatch(ca.size(), dof.n_dofs()));
   Assert(cb.size() == dof.n_dofs(), ExcDimensionMismatch(cb.size(), dof.n_dofs()));
   Assert(cu.size() == dof.n_dofs(), ExcDimensionMismatch(cu.size(), dof.n_dofs()));

   WorkStream::run(dof.begin_active(), dof.end(),
         std::bind(&DivRightHandSide<dim>::local_assemble, *this, ca, cb, cu, std::placeholders::_1,
               std::placeholders::_2, std::placeholders::_3),
         std::bind(&DivRightHandSide<dim>::copy_local_to_global, *this, std::ref(rhs), std::placeholders::_1),
         AssemblyScratchData(dof.get_fe(), quad), AssemblyCopyData());
}

template class DivRightHandSide<1> ;
template class DivRightHandSide<2> ;
template class DivRightHandSide<3> ;

} /* namespace forward */
} /* namespace wavepi */


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

/**
 * @file cardinality.hpp
 *
 * @brief SAT cardinality constraints
 *
 * @author Mathias Soeken
 * @author Heinz Riener
 *
 * @since  2.2
 */

#ifndef CARDINALITY_HPP
#define CARDINALITY_HPP

#include <numeric>

#include <boost/assign/std/vector.hpp>
#include <boost/range.hpp>
#include <boost/range/adaptors.hpp>
#include <boost/range/algorithm_ext/iota.hpp>
#include <boost/range/algorithm_ext/push_back.hpp>

#include <core/utils/range_utils.hpp>
#include <classical/sat/sat_solver.hpp>
#include <classical/sat/operations/logic.hpp>

using namespace boost::assign;

/**
 * Carsten Sinz's CNF encoding of Boolean cardinality constraints [1].
 * The implementation is mainly based on the Knuth's descriptions [2].
 *
 * [1] Carsten Sinz. Towards an optimal CNF encoding of boolean
 * cardinality constraints, CP, pages 827-831, 2005.
 *
 * [2] Donald E. Knuth, The Art of Computer Programming, Volume 4,
 * Pre-Fascicle 6A, page 8ff.
 */

namespace cirkit
{

template<class S, class C>
void atmost_one( S& solver, const C& x )
{
  for ( auto i = boost::begin( x ); i != boost::end( x ); ++i )
  {
    if ( i + 1 == boost::end( x ) ) continue;
    for ( auto j = i + 1; j != boost::end( x ); ++j )
    {
      add_clause( solver )( {-*i, -*j} );
    }
  }
}

template<class S, class C>
inline void one_hot( S& solver, const C& x )
{
  add_clause( solver )( x );
  atmost_one( solver, x );
}

template < class S >
int less_or_equal_sinz( S& solver, const clause_t& x, unsigned r, int sid, int spol = 1 )
{
  const unsigned n = x.size();

  //std::cerr << "n = " << n << " r = " << r << " sid = " << sid << '\n';
  assert( n > r );
  for ( auto j = 0u; j < n - r - 1; ++j )
  {
    for ( auto k = 0u; k < r; ++k )
    {
      // std::cerr << "(" << j << "," << k << ")" << '\n';
      //const int l1 = spol * -(sid + j * (n - r) + k);
      //const int l2 = spol * (sid + (j + 1u) * (n - r) + k);
      const int l1 = spol * -( sid + j * r + k );
      const int l2 = spol * (sid + ( j + 1u ) * r + k );
      add_clause( solver )( {l1, l2} );
    }
  }

  for ( auto j = 0u; j < (n - r); ++j )
  {
    for ( auto k = 0u; k <= r; ++k )
    {
      // std::cerr << "(" << j << "," << k << ")" << '\n';
      clause_t clause;
      clause += -x[ j + k ];
      //if ( k > 0u ) clause += spol * -(sid + j * ( n - r ) + k - 1u );
      //if ( k < r ) clause += spol * (sid + j * ( n - r ) + k );
      if ( k > 0u ) clause += spol * -(sid + j * r + k - 1u );
      if ( k < r ) clause += spol * (sid + j * r + k );
      add_clause( solver )( clause );
    }
  }

  sid += (n - r)*r/* + 1u*/;
  //std::cerr << "next sid: " << sid << '\n';
  return sid;
}

template < class S >
int greater_or_equal_sinz( S& solver, const clause_t& x, unsigned r, int sid, int spol = 1 )
{
  using boost::adaptors::transformed;

  clause_t xn;
  boost::push_back( xn, x | transformed( []( int v ) { return -v; } ) );
  return less_or_equal_sinz( solver, xn, x.size() - r, sid, spol );
}

template < class S >
int equals_sinz( S& solver, const clause_t& x, unsigned r, int sid )
{
  //std::cout << "call equals_size with x.size = " << x.size() << ", r = " << r << ", and sid = " << sid << std::endl;

  if ( r == 0u )
  {
    for ( const auto& l : x )
    {
      add_clause( solver )( {-l} );
    }
    return sid;
  }

  if ( r == x.size() )
  {
    for ( const auto& l : x )
    {
      add_clause( solver )( {l} );
    }
    return sid;
  }

  sid = less_or_equal_sinz( solver, x, r, sid, 1 );
  return greater_or_equal_sinz( solver, x, r, sid, -1 );
}

/******************************************************************************
 * Binary tree encoding from Bailleux and Boufkhad                            *
 ******************************************************************************/
template<class S>
int less_or_equal_bailleux_boufkhad_assume( S& solver, int assume, const clause_t& x, unsigned r, int sid )
{
  const auto n = x.size();                    /* number of variables */
  std::vector<unsigned> ts_count( 2 * n, 1 ); /* pre-assign t's with 1 (for the leaves), we don't use index 0 */
  std::vector<int> ts( 2 * n, 0 );
  std::vector<int> pol( 2 * n, 1 );

  auto offset = sid;
  for ( auto k = n - 1; k > 1; --k )
  {
    ts_count[k] = std::min( ts_count[k << 1] + ts_count[( k << 1 ) + 1], r );
    ts[k] = offset;
    offset += ts_count[k];
  }
  for ( auto k = n; k < 2 * n; ++k )
  {
    ts[k] = abs( x[k - n] );
    pol[k] = x[k - n] > 0 ? 1 : -1;
  }

  for ( auto k = 2u; k < n; ++k )
  {
    for ( auto i = 0u; i <= ts_count[k << 1]; ++i )
    {
      for ( auto j = 0u; j <= ts_count[( k << 1 ) + 1]; ++j )
      {
        if ( ( i + j < 1 ) || ( i + j > ts_count[k] + 1 ) ) { continue; }

        std::vector<int> clause = { -assume };
        if ( i > 0 ) { clause.push_back( -pol[k << 1] * ( ts[k << 1] + i - 1 ) ); }
        if ( j > 0 ) { clause.push_back( -pol[( k << 1 ) + 1] * ( ts[( k << 1 ) + 1] + j - 1 ) ); }
        if ( i + j <= r ) { clause.push_back( pol[k] * ( ts[k] + i + j - 1 ) ); }

        if ( !clause.empty() )
        {
          add_clause( solver )( clause );
        }
      }
    }
  }

  for ( auto i = 0; i <= static_cast<int>( ts_count[2] ); ++i )
  {
    for ( auto j = 0; j <= static_cast<int>( ts_count[3] ); ++j )
    {
      if ( i + j == static_cast<int>( r ) + 1 )
      {
        assert( i && j );
        add_clause( solver )( { -assume, -pol[2] * ( ts[2] + i - 1 ), -pol[3] * ( ts[3] + j - 1 )} );
      }
    }
  }

  return offset;
}

template<class S>
int less_or_equal_bailleux_boufkhad( S& solver, const clause_t& x, unsigned r, int sid )
{
  const auto n = x.size();                    /* number of variables */
  std::vector<unsigned> ts_count( 2 * n, 1 ); /* pre-assign t's with 1 (for the leaves), we don't use index 0 */
  std::vector<int> ts( 2 * n, 0 );
  std::vector<int> pol( 2 * n, 1 );

  auto offset = sid;
  for ( int k = n - 1; k > 1; --k )
  {
    ts_count[k] = std::min( ts_count[k << 1] + ts_count[( k << 1 ) + 1], r );
    ts[k] = offset;
    offset += ts_count[k];
  }
  for ( auto k = n; k < 2 * n; ++k )
  {
    ts[k] = abs( x[k - n] );
    pol[k] = x[k - n] > 0 ? 1 : -1;
  }

  for ( auto k = 2u; k < n; ++k )
  {
    for ( auto i = 0u; i <= ts_count[k << 1]; ++i )
    {
      for ( auto j = 0u; j <= ts_count[( k << 1 ) + 1]; ++j )
      {
        if ( ( i + j < 1 ) || ( i + j > ts_count[k] + 1 ) ) { continue; }

        std::vector<int> clause;
        if ( i > 0 ) { clause.push_back( -pol[k << 1] * ( ts[k << 1] + i - 1 ) ); }
        if ( j > 0 ) { clause.push_back( -pol[( k << 1 ) + 1] * ( ts[( k << 1 ) + 1] + j - 1 ) ); }
        if ( i + j <= r ) { clause.push_back( pol[k] * ( ts[k] + i + j - 1 ) ); }

        if ( !clause.empty() )
        {
          add_clause( solver )( clause );
        }
      }
    }
  }

  for ( auto i = 0; i <= static_cast<int>( ts_count[2] ); ++i )
  {
    for ( auto j = 0; j <= static_cast<int>( ts_count[3] ); ++j )
    {
      if ( i + j == static_cast<int>( r ) + 1 )
      {
        assert( i && j );
        add_clause( solver )( {-pol[2] * ( ts[2] + i - 1 ), -pol[3] * ( ts[3] + j - 1 )} );
      }
    }
  }

  return offset;
}

template<class S>
int greater_or_equal_bailleux_boufkhad( S& solver, const clause_t& x, unsigned r, int sid )
{
  using boost::adaptors::transformed;

  clause_t xn;
  boost::push_back( xn, x | transformed( []( int v ) { return -v; } ) );
  return less_or_equal_bailleux_boufkhad( solver, xn, x.size() - r, sid );
}

template < class S >
int equals_bailleux_boufkhad( S& solver, const clause_t& x, unsigned r, int sid )
{
  if ( r == 0u )
  {
    for ( const auto& l : x )
    {
      add_clause( solver )( {-l} );
    }
    return sid;
  }

  if ( r == x.size() )
  {
    for ( const auto& l : x )
    {
      add_clause( solver )( {l} );
    }
    return sid;
  }

  sid = less_or_equal_bailleux_boufkhad( solver, x, r, sid );
  return greater_or_equal_bailleux_boufkhad( solver, x, r, sid );
}

/******************************************************************************
 * Pairwise cardinality networks                                              *
 ******************************************************************************/

// Based on [M. Codish and M. Zazon-Ivry, LPAR, 2016, 154-172] and https://bitbucket.org/alanmi/abc/src/f78c2854aa5948fea93533c3934fb0f6a2c3c785/src/aig/gia/giaSatMap.c?at=default&fileviewer=file-view-default

// namespace detail
// {

// template<class S>
// void sat_add_half_sorter( S& solver, int a, int b, int y, int z )
// {
//   add_clause( solver )( {a, -y} );
//   add_clause( solver )( {a, -z} );
//   add_clause( solver )( {b, -y, -z} );
//   // add_clause( solver )( {-a, y} );
//   // add_clause( solver )( {-a, z} );
//   // add_clause( solver )( {-b, y, z} );
// }

// template<class S>
// void sat_pairwise_cardinality_network_add_sorter( S& solver, std::vector<int>& vvars, int i, int k, int *pnvars )
// {
//   const auto ivar1 = (*pnvars)++;
//   const auto ivar2 = (*pnvars)++;

//   sat_add_half_sorter( solver, ivar1, ivar2, vvars[i], vvars[k] );
//   vvars[i] = ivar1;
//   vvars[k] = ivar2;
// }

// template<class S>
// void sat_pairwise_cardinality_network_add_constraint_merge( S& solver, std::vector<int>& vvars, int lo, int hi, int r, int *pnvars )
// {
//   const auto step = r * 2;
//   if ( step >= hi - lo ) return;

//   sat_pairwise_cardinality_network_add_constraint_merge( solver, vvars, lo, hi - r, step, pnvars );
//   sat_pairwise_cardinality_network_add_constraint_merge( solver, vvars, lo + r, hi, step, pnvars );
//   for ( auto i = lo + r; i < hi - r; i += step )
//   {
//     sat_pairwise_cardinality_network_add_sorter( solver, vvars, i, i + r, pnvars );
//   }
// }

// template<class S>
// void sat_pairwise_cardinality_network_add_constraint_range( S& solver, std::vector<int>& vvars, int lo, int hi, int *pnvars )
// {
//   if ( hi - lo < 1 ) return;

//   const auto mid = lo + ( hi - lo ) / 2;
//   for ( auto i = lo; i <= mid; ++i )
//   {
//     sat_pairwise_cardinality_network_add_sorter( solver, vvars, i, i + ( hi - lo + 1 ) / 2, pnvars );
//   }
//   sat_pairwise_cardinality_network_add_constraint_range( solver, vvars, lo, mid, pnvars );
//   sat_pairwise_cardinality_network_add_constraint_range( solver, vvars, mid + 1, hi, pnvars );
//   sat_pairwise_cardinality_network_add_constraint_merge( solver, vvars, lo, hi, 1, pnvars );
// }

// template<class S>
// int sat_pairwise_cardinality_network_add_constraint_pairwise( S& solver, std::vector<int>& vvars )
// {
//   int nvars = vvars.size();
//   sat_pairwise_cardinality_network_add_constraint_range( solver, vvars, 0, nvars - 1, &nvars );
//   return nvars;
// }

// }

// template<class S>
// int sat_pairwise_cardinality_network2( S& solver, unsigned log_n, int sid, std::vector<int>* ovars = nullptr )
// {
//   auto nvars = 1u << log_n;
//   const auto nvars_alloc = nvars + 2u * (nvars * log_n * ( log_n - 1 ) / 4 + nvars - 1);

//   std::vector<int> vvars( nvars );
//   boost::iota( vvars, sid );

//   std::cout << "[i] original vvars: " << any_join( vvars, " " ) << std::endl;

//   const auto nvars_real = detail::sat_pairwise_cardinality_network_add_constraint_pairwise( solver, vvars );
//   assert( nvars_real == nvars_alloc );

//   if ( ovars )
//   {
//     *ovars = vvars;
//   }

//   return sid + nvars_alloc;
// }

/******************************************************************************
 * PW implementation based on paper                                           *
 ******************************************************************************/

namespace detail
{

template<class S>
void sat_pairwise_comparator( S& solver, int a, int b, int c, int d)
{
  //logic_or( solver, a, b, c );
  //logic_and( solver, a, b, d );

  add_clause( solver )( {c, -a} );
  add_clause( solver )( {c, -b} );
  add_clause( solver )( {d, -a, -b} );
}

template<class S>
void sat_pairwise_split( S& solver, const std::vector<int>& as, const std::vector<int>& bs, const std::vector<int>& cs )
{
  assert( bs.size() == cs.size() );
  assert( bs.size() << 1u == as.size() );

  for ( auto i = 0u; i < bs.size(); ++i )
  {
    sat_pairwise_comparator( solver, as[2 * i], as[2 * i + 1], bs[i], cs[i] );
  }
}

template<class S>
int sat_pairwise_merge( S& solver, const std::vector<int>& as, const std::vector<int>& bs, const std::vector<int>& cs, int sid )
{
  const auto n = as.size();

  if ( n == 1u )
  {
    equals( solver, as[0], cs[0] );
    equals( solver, bs[0], cs[1] );
    return sid;
  }

  /* fill ds and es */
  std::vector<int> ds( n ), es( n );

  ds[0u] = cs[0];
  for ( auto i = 1u; i < n; ++i ) { ds[i] = sid++; }

  for ( auto i = 0u; i < n - 1; ++i ) { es[i] = sid++; }
  es[n - 1] = cs[2 * n - 1];

  /* fill a0, b0, a1, b1 */
  std::vector<int> a0( n / 2 ), b0( n / 2 ), a1( n / 2 ), b1( n / 2 );
  for ( auto i = 0u; i < n / 2; ++i )
  {
    a0[i] = as[2 * i];
    a1[i] = as[2 * i + 1];
    b0[i] = bs[2 * i];
    b1[i] = bs[2 * i + 1];
  }

  sid = sat_pairwise_merge( solver, a0, b0, ds, sid );
  sid = sat_pairwise_merge( solver, a1, b1, es, sid );

  for ( auto i = 0u; i < n - 1; ++i )
  {
    sat_pairwise_comparator( solver, es[i], ds[i + 1], cs[2 * i + 1], cs[2 * i + 2] );
  }

  /* sort outputs */
  for ( auto i = 0u; i < cs.size() - 1; ++i )
  {
    add_clause( solver )( {cs[i], -cs[i + 1]} );
  }

  return sid;
}

template<class S>
int sat_pairwise_sort( S& solver, const std::vector<int>& as, const std::vector<int>& ds, int sid )
{
  const auto n = as.size();

  if ( n == 1u )
  {
    equals( solver, as[0], ds[0] );
    return sid;
  }

  std::vector<int> bs( n / 2 ), cs( n / 2 ), bbs( n / 2), ccs( n / 2 );

  boost::iota( bs, sid ); sid += n / 2;
  boost::iota( cs, sid ); sid += n / 2;
  boost::iota( bbs, sid ); sid += n / 2;
  boost::iota( ccs, sid ); sid += n / 2;

  sat_pairwise_split( solver, as, bs, cs );
  sid = sat_pairwise_sort( solver, bs, bbs, sid );
  sid = sat_pairwise_sort( solver, cs, ccs, sid );
  return sat_pairwise_merge( solver, bbs, ccs, ds, sid );
}

}

template<class S>
int sat_pairwise_cardinality_network( S& solver, const std::vector<int>& ivars, int sid, std::vector<int>* ovars = nullptr )
{
  std::vector<int> ds( ivars.size() );
  boost::iota( ds, sid );
  sid += ivars.size();

  if ( ovars )
  {
    *ovars = ds;
  }

  return detail::sat_pairwise_sort( solver, ivars, ds, sid );
}

template<class S>
int sat_pairwise_cardinality_network( S& solver, unsigned log_n, int sid, std::vector<int>* ovars = nullptr )
{
  const auto n = 1 << log_n;
  std::vector<int> as( n );

  boost::iota( as, sid ); sid += n;

  return sat_pairwise_cardinality_network( solver, as, sid, ovars );
}

}

#endif

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