#ifndef STAN__PROB__DISTRIBUTIONS__MULTIVARIATE__CONTINUOUS__MULTI_NORMAL_SUFFICIENT_HPP
#define STAN__PROB__DISTRIBUTIONS__MULTIVARIATE__CONTINUOUS__MULTI_NORMAL_SUFFICIENT_HPP

#include <boost/random/normal_distribution.hpp>
#include <boost/random/variate_generator.hpp>
#include <stan/math/prim/mat/err/check_ldlt_factor.hpp>
#include <stan/math/prim/mat/err/check_symmetric.hpp>
#include <stan/math/prim/scal/err/check_size_match.hpp>
#include <stan/math/prim/scal/err/check_finite.hpp>
#include <stan/math/prim/scal/err/check_not_nan.hpp>
#include <stan/math/prim/scal/err/check_positive.hpp>
#include <stan/math/prim/mat/fun/trace_inv_quad_form_ldlt.hpp>
#include <stan/math/prim/mat/fun/log_determinant_ldlt.hpp>
#include <stan/math/prim/scal/meta/return_type.hpp>
#include <stan/math/prim/scal/meta/VectorViewMvt.hpp>
#include <stan/math/prim/scal/meta/max_size_mvt.hpp>
#include <stan/math/prim/scal/meta/include_summand.hpp>

namespace stan {

  namespace prob {

    template <typename T_sample, typename T_loc, typename T_covar>
    typename boost::math::tools::promote_args<T_sample, typename scalar_type<T_loc>::type, T_covar>::type
    multi_normal_sufficient_log(const int sampleSize,
				const Eigen::Matrix<T_sample,Eigen::Dynamic,1>& sampleMu,
				const Eigen::Matrix<T_sample,Eigen::Dynamic,Eigen::Dynamic>& sampleSigma,
				const T_loc& mu,
				const Eigen::Matrix<T_covar,Eigen::Dynamic,Eigen::Dynamic>& Sigma) {
      static const char *function("stan::prob::multi_normal_sufficient_log");
      typedef typename boost::math::tools::promote_args<T_sample, typename scalar_type<T_loc>::type, T_covar>::type param_t;
      typedef param_t lp_type;
      lp_type lp(0.0);

      using stan::math::check_size_match;
      using stan::math::check_finite;
      using stan::math::check_not_nan;
      using stan::math::check_positive;
      using stan::math::check_symmetric;
      using stan::math::check_ldlt_factor;

      check_size_match(function,
                       "Rows of covariance parameter", sampleSigma.rows(),
                       "columns of covariance parameter", sampleSigma.cols());
      check_positive(function, "Covariance matrix rows", sampleSigma.rows());
      check_symmetric(function, "Covariance matrix", sampleSigma);

      check_size_match(function,
                       "Rows of covariance parameter", Sigma.rows(),
                       "columns of covariance parameter", Sigma.cols());
      check_positive(function, "Covariance matrix rows", Sigma.rows());
      check_symmetric(function, "Covariance matrix", Sigma);

      check_size_match(function,
                       "Size of data location", sampleMu.size(),
                       "size of model location", mu.size());
      check_size_match(function,
                       "Size of data covariance", sampleSigma.rows(),
                       "size of model covariance", Sigma.rows());

      stan::math::LDLT_factor<param_t,Eigen::Dynamic,Eigen::Dynamic> ldlt_Sigma(Sigma);
      check_ldlt_factor(function, "LDLT_Factor of covariance parameter", ldlt_Sigma);

      Eigen::Matrix<param_t, Eigen::Dynamic, Eigen::Dynamic> ss;
      ss = mdivide_left_ldlt(ldlt_Sigma, sampleSigma);

      lp += (ss.diagonal().sum() + log_determinant_ldlt(ldlt_Sigma)) * (sampleSize - 1);

      lp_type lp_location(0.0);
      {
	Eigen::Matrix<param_t, Eigen::Dynamic, 1> y_minus_mu(mu.size());

	for (int j = 0; j < mu.size(); j++)
	  y_minus_mu(j) = mu(j) - sampleMu(j);

	lp_location = trace_inv_quad_form_ldlt(ldlt_Sigma, y_minus_mu) * sampleSize;
	// Could avoid re-solving Sigma
	// lp_location = quad_form(ss, y_minus_mu).diagonal().sum() * sampleSize;
      }
      return (lp + lp_location) * -0.5;
    }
  }
}

#endif


#pragma once

#include <Core/CoreMacros.hpp>
#include <Core/Math/Math.hpp>
#include <Eigen/Core>
#include <Eigen/Geometry> //homogeneous
#include <random>

namespace Ra {
namespace Core {
namespace Utils {

/*!
 * Colors are defined as vector4, i.e. 4 Scalars in RGBA order.
 * displayable colors should have all their coordinates between 0 and 1.
 *
 * \warning Vector arithmetics can be used to add, substract, multiply or divide colors with
 * scalars and other colors. In that case, alpha is treated as any other component. It is fine for
 * most cases, however adding two colors (e.g. Blue()+Red()) may lead to unconsistent
 * alpha: Blue()+Red() = {1, 0, 1, 2}.
 * In that case, only the rgb values need to be added:
 * \code
 * Color result = Color::fromRGB( Blue().rgb()+Red().rgb() ); // result = {1, 0, 1, 1}
 * \endcode
 * The validity of alpha can be checked using hasValidAlpha().
 */
template <typename _Scalar>
class ColorBase : public Eigen::Matrix<_Scalar, 4, 1>
{
  public:
    using VectorType = Eigen::Matrix<_Scalar, 4, 1>;

    explicit inline ColorBase() : ColorBase( _Scalar( 1. ), _Scalar( 1. ), _Scalar( 1. ) ) {}

    template <typename S2>
    inline ColorBase( S2 r, S2 g, S2 b, S2 alpha = S2( 1 ) ) :
        VectorType( _Scalar( r ), _Scalar( g ), _Scalar( b ), _Scalar( alpha ) ) {}

    /// Copy constructor
    template <typename S2>
    inline ColorBase( const ColorBase<S2>& other ) : VectorType( other.template cast<_Scalar>() ) {}

    /// Copy constructor from Eigen expressions: `ColorBase<_Scalar> c( (*this) * 255 );`
    template <typename Derived>
    inline ColorBase( const Eigen::MatrixBase<Derived>& v ) :
        VectorType( v.template cast<_Scalar>() ) {}

    /// cast operator, mandatory to use Vector arithmetic
    operator VectorType() { return *this; }

    /// convert the color expressed in sRGB color space to linear RGB
    static inline ColorBase sRGBToLinearRGB( const ColorBase& srgb ) {
        ColorBase<_Scalar> c( srgb );
        for ( auto& u : c.rgb() )
        {
            if ( u < 0.04045_ra ) { u /= 12.92_ra; }
            else
            { u = std::pow( ( u + 0.055_ra ) / 1.055_ra, 2.4_ra ); }
        }
        return c;
    }

    /// convert the color expressed in linear RGB color space to sRGB
    static inline ColorBase linearRGBTosRGB( const ColorBase& lrgb ) {
        ColorBase<_Scalar> c( lrgb );
        for ( auto& u : c.rgb() )
        {
            if ( u < 0.0031308_ra ) { u *= 12.92_ra; }
            else
            { u = 1.055_ra * std::pow( u, 1_ra / 2.4_ra ) - 0.055_ra; }
        }
        return c;
    }

    template <typename Derived>
    static inline ColorBase fromRGB( const Eigen::MatrixBase<Derived>& rgb,
                                     Scalar alpha = Scalar( 1. ) ) {
        ColorBase c( rgb.template cast<_Scalar>().homogeneous() );
        c.alpha() = alpha;
        return c;
    }

    Eigen::Block<VectorType, 3, 1> rgb() { return ( *this ).template head<3>(); }
    const Eigen::Block<VectorType, 3, 1> rgb() const { return ( *this ).template head<3>(); }

    Scalar alpha() const { return ( *this )( 3 ); }
    Scalar& alpha() { return ( *this )( 3 ); }
    bool hasValidAlpha() const { return Math::checkRange( alpha(), 0_ra, 1_ra ); }

    static inline ColorBase<_Scalar> Alpha() {
        return ColorBase<_Scalar>( _Scalar( 0. ), _Scalar( 0. ), _Scalar( 0. ), _Scalar( 0. ) );
    }

    static inline ColorBase<_Scalar> Black() {
        return ColorBase<_Scalar>( _Scalar( 0. ), _Scalar( 0. ), _Scalar( 0. ) );
    }

    static inline ColorBase<_Scalar> Red() {
        return ColorBase<_Scalar>( _Scalar( 1. ), _Scalar( 0. ), _Scalar( 0. ) );
    }

    static inline ColorBase<_Scalar> Green() {
        return ColorBase<_Scalar>( _Scalar( 0. ), _Scalar( 1. ), _Scalar( 0. ) );
    }

    static inline ColorBase<_Scalar> Blue() {
        return ColorBase<_Scalar>( _Scalar( 0. ), _Scalar( 0. ), _Scalar( 1. ) );
    }

    static inline ColorBase<_Scalar> Yellow() {
        return ColorBase<_Scalar>( _Scalar( 1. ), _Scalar( 1. ), _Scalar( 0. ) );
    }

    static inline ColorBase<_Scalar> Magenta() {
        return ColorBase<_Scalar>( _Scalar( 1. ), _Scalar( 0. ), _Scalar( 1. ) );
    }

    static inline ColorBase<_Scalar> Cyan() {
        return ColorBase<_Scalar>( _Scalar( 0. ), _Scalar( 1. ), _Scalar( 1. ) );
    }

    static inline ColorBase<_Scalar> White() {
        return ColorBase<_Scalar>( _Scalar( 1. ), _Scalar( 1. ), _Scalar( 1. ) );
    }

    static inline ColorBase<_Scalar> Grey( _Scalar f = _Scalar( 0.5 ), _Scalar a = _Scalar( 1. ) ) {
        return ColorBase<_Scalar>( f, f, f, a );
    }

    static inline ColorBase<_Scalar> Skin() {
        return ColorBase<_Scalar>( _Scalar( 1.0 ), _Scalar( 0.87 ), _Scalar( 0.74 ) );
    }
    // Convert to/from various int formats

    static inline ColorBase<_Scalar> fromChars( uchar r, uchar g, uchar b, uchar a = 0xff ) {
        return ColorBase<_Scalar>( _Scalar( r ) / 255.0f,
                                   _Scalar( g ) / 255.0f,
                                   _Scalar( b ) / 255.0f,
                                   _Scalar( a ) / 255.0f );
    }

    static inline ColorBase<_Scalar> fromRGBA32( uint32_t rgba ) {
        uchar r = uchar( ( rgba >> 24 ) & 0xff );
        uchar g = uchar( ( rgba >> 16 ) & 0xff );
        uchar b = uchar( ( rgba >> 8 ) & 0xff );
        uchar a = uchar( ( rgba >> 0 ) & 0xff );
        return fromChars( r, g, b, a );
    }

    static inline ColorBase<_Scalar> fromARGB32( uint32_t argb ) {
        uchar a = uchar( ( argb >> 24 ) & 0xff );
        uchar r = uchar( ( argb >> 16 ) & 0xff );
        uchar g = uchar( ( argb >> 8 ) & 0xff );
        uchar b = uchar( ( argb >> 0 ) & 0xff );
        return fromChars( r, g, b, a );
    }

    static inline ColorBase<_Scalar> fromHSV( const _Scalar hue,
                                              const _Scalar saturation = 1.0,
                                              const _Scalar value      = 1.0,
                                              const _Scalar alpha      = 1.0 ) {
        ColorBase<_Scalar> c;

        if ( saturation == 0.0f )
        {
            c[0] = c[1] = c[2] = value;
            c[3]               = alpha;
            return c;
        }
        _Scalar h  = ( ( hue == 1.0f ) ? 0.0f : hue ) * 6.0f;
        int i      = int( std::floor( h ) );
        _Scalar v1 = value * ( 1.0f - saturation );
        _Scalar v2 = value * ( 1.0f - ( saturation * ( h - i ) ) );
        _Scalar v3 = value * ( 1.0f - ( saturation * ( 1.0f - h - i ) ) );
        switch ( i )
        {
        case 0: {
            c[0] = value;
            c[1] = v3;
            c[2] = v1;
        }
        break;
        case 1: {
            c[0] = v2;
            c[1] = value;
            c[2] = v1;
        }
        break;
        case 2: {
            c[0] = v1;
            c[1] = value;
            c[2] = v3;
        }
        break;
        case 3: {
            c[0] = v1;
            c[1] = v2;
            c[2] = value;
        }
        break;
        case 4: {
            c[0] = v3;
            c[1] = v1;
            c[2] = value;
        }
        break;
        default: {
            c[0] = value;
            c[1] = v1;
            c[2] = v2;
        }
        break;
        }
        c[3] = alpha;
        return c;
    }

    inline uint32_t toRGBA32() const {
        ColorBase<_Scalar> c( ( *this ) * 255 );
        Eigen::Matrix<int, 4, 1> scaled( c.x(), c.y(), c.z(), c.w() );
        return ( uint32_t( scaled( 0 ) ) << 24 ) | ( uint32_t( scaled( 1 ) ) << 16 ) |
               ( uint32_t( scaled( 2 ) ) << 8 ) | ( uint32_t( scaled( 3 ) ) << 0 );
    }

    inline uint32_t toARGB32() const {
        ColorBase<_Scalar> c( ( *this ) * 255 );
        Eigen::Matrix<int, 4, 1> scaled( c.x(), c.y(), c.z(), c.w() );
        return ( uint32_t( scaled( 3 ) ) << 24 ) | ( uint32_t( scaled( 0 ) ) << 16 ) |
               ( uint32_t( scaled( 1 ) ) << 8 ) | ( uint32_t( scaled( 2 ) ) << 0 );
    }

    static inline std::vector<ColorBase<_Scalar>> scatter( const uint size, const _Scalar gamma ) {
        std::vector<ColorBase<_Scalar>> color( size );
        if ( size > 1 )
            for ( uint i = 0; i < size; ++i )
            {
                color[i] = fromHSV( ( _Scalar( i ) / _Scalar( size - 1 ) ) * 0.777 );
                color[i] = ( color[i] + ColorBase<_Scalar>::Constant( gamma ) ) * 0.5;
            }
        else
        { color[0] = Red(); }
        std::shuffle( color.begin(), color.end(), std::mt19937( std::random_device()() ) );
        return color;
    }
};

using Color  = ColorBase<Scalar>;
using Colorf = ColorBase<float>;
using Colord = ColorBase<double>;

} // namespace Utils
} // namespace Core
} // namespace Ra


/*=============================================================================

PHAS0100ASSIGNMENT2: PHAS0100 Assignment 2 Gravitational N-body Simulation

Copyright (c) University College London (UCL). All rights reserved.

This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.

See LICENSE.txt in the top level directory for details.

=============================================================================*/

#include <CLI11.hpp>
#include <nbsimBasicTypes.h>
#include <nbsimMyFunctions.h>
#include <nbsimExceptionMacro.h>
#include <iostream>
#include <string.h>
#include <stdio.h>
#include <ctime>
#include <chrono>
#include <omp.h>
#include <iomanip>
#include "nbsimSolarSystemData.ipp"

// Example, header-only library, included in project for simplicity's sake.
#include <Eigen/Dense>
#include <unsupported/Eigen/NonLinearOptimization>

#define YEAR2SEC 365*24*60*60
/**
 * \brief Demo file to check that includes and library linkage is correct.
 */
int main(int argc, char** argv)
{

	int returnStatus = EXIT_FAILURE;
	int ibody = 0;
	Eigen::Vector3d pos, vel, r, rcom, p;
	double mass, mtotal;
	std::shared_ptr<nbsim::MassiveParticle> planet[9];
	double ts = 0.0000274 * YEAR2SEC, dur = 1* YEAR2SEC;

	std::string names[9];
	CLI::App app{"Solar system simulator"};
	//app.require_subcommand(1);

	CLI11_PARSE(app, argc, argv);

	/* Init planets */
	for (int ibody=0; ibody <9; ibody++) {
		//std::cout << nbsim::solarSystemData[ibody].name << std::endl;
		names[ibody] = nbsim::solarSystemData[ibody].name;
		mass = nbsim::solarSystemData[ibody].mu / GRAV;
		pos = nbsim::solarSystemData[ibody].position;
		vel = nbsim::solarSystemData[ibody].velocity;
		planet[ibody] = std::make_shared<nbsim::MassiveParticle>(pos, vel, mass);
	}
	/* Add mutual attraction */
	for (int i=0; i <9; i++) {
		for (int j=0; j<9; j++) {
			//printf("Adding %d to %d\n", j, i);
			if (i != j) planet[i]->addAttractor(planet[j]);
		}
	}

	/* start timing */
	std::clock_t c_start = std::clock();
	auto t_start = std::chrono::high_resolution_clock::now();
	/* simulate */
	for (double t=0; t<dur; t+=ts) {
#ifdef DEBUG
		{
			/* calculate rcom and ptotal */
			rcom << 0,0,0;
			p << 0,0,0;
			mtotal = 0;
			for (int i=0; i<9; i++) {
				mass = planet[i]->getMass();
				rcom += mass * planet[i]->getPosition();
				p += mass * planet[i]->getVelocity();
				mtotal += mass;
			}
			rcom = rcom / mtotal;
			std::cout << "rcom = " << std::endl << rcom << std::endl
				<< "ptotal = " << std::endl << p <<std::endl
				<< "mtotal = " << mtotal << std::endl;
		}
#endif
#pragma omp parallel for
		for (int i=0; i<9; i++) {
			planet[i]->calculateAcceleration();
		}
#pragma omp parallel for
		for (int i=0; i<9; i++) {
			planet[i]->integrateTimestep(ts);
		}
	}
	/* stop timing */
    std::clock_t c_end = std::clock();
    auto t_end = std::chrono::high_resolution_clock::now();

	/* output res */
	for (int i=0; i<9; i++) {
		std::cout << names[i] << std::endl;
		std::cout << planet[i]->getPosition() << std::endl;
	}
		{
			/* calculate rcom and ptotal */
			rcom << 0,0,0;
			p << 0,0,0;
			mtotal = 0;
			for (int i=0; i<9; i++) {
				mass = planet[i]->getMass();
				rcom += mass * planet[i]->getPosition();
				p += mass * planet[i]->getVelocity();
				mtotal += mass;
			}
			rcom = rcom / mtotal;
			std::cout << "rcom = " << std::endl << rcom << std::endl
				<< "ptotal = " << std::endl << p <<std::endl
				<< "mtotal = " << mtotal << std::endl;
		}

    std::cout << std::fixed << std::setprecision(2) << "CPU time used: "
              << 1000.0 * (c_end - c_start) / CLOCKS_PER_SEC / 1000 << " s\n"
              << "Wall clock time passed: "
              << std::chrono::duration<double, std::milli>(t_end-t_start).count() / 1000
              << " s\n";
	return 0;
}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2007-2015 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2015 Bruno Lalande, Paris, France.
// Copyright (c) 2009-2015 Mateusz Loskot, London, UK.

// This file was modified by Oracle on 2015.
// Modifications copyright (c) 2015 Oracle and/or its affiliates.

// Contributed and/or modified by Menelaos Karavelas, on behalf of Oracle

// Parts of Boost.Geometry are redesigned from Geodan's Geographic Library
// (geolib/GGL), copyright (c) 1995-2010 Geodan, Amsterdam, the Netherlands.

// Use, modification and distribution is subject to the Boost Software License,
// Version 1.0. (See accompanying file LICENSE_1_0.txt or copy at
// http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_GEOMETRY_STRATEGIES_TRANSFORM_MATRIX_TRANSFORMERS_HPP
#define BOOST_GEOMETRY_STRATEGIES_TRANSFORM_MATRIX_TRANSFORMERS_HPP

#include <cstddef>

#include <boost/qvm/mat.hpp>
#include <boost/qvm/vec.hpp>
#include <boost/qvm/mat_access.hpp>
#include <boost/qvm/vec_access.hpp>
#include <boost/qvm/mat_operations.hpp>
#include <boost/qvm/vec_mat_operations.hpp>
#include <boost/qvm/map_mat_mat.hpp>
#include <boost/qvm/map_mat_vec.hpp>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/core/coordinate_dimension.hpp>
#include <boost/geometry/core/coordinate_promotion.hpp>
#include <boost/geometry/core/cs.hpp>
#include <boost/geometry/util/math.hpp>
#include <boost/geometry/util/select_coordinate_type.hpp>
#include <boost/geometry/util/select_most_precise.hpp>

namespace boost { namespace geometry
{

namespace strategy { namespace transform
{

namespace detail { namespace matrix_transformer
{

template
<
    typename Point,
    std::size_t Dimension = 0,
    std::size_t DimensionCount = geometry::dimension<Point>::value
>
struct set_point_from_vec
{
    template <typename Vector>
    static inline void apply(Point & p, Vector const& v)
    {
        typedef typename geometry::coordinate_type<Point>::type coord_t;
        set<Dimension>(p, boost::numeric_cast<coord_t>(qvm::A<Dimension>(v)));
        set_point_from_vec<Point, Dimension + 1, DimensionCount>::apply(p, v);
    }
};

template
<
    typename Point,
    std::size_t DimensionCount
>
struct set_point_from_vec<Point, DimensionCount, DimensionCount>
{
    template <typename Vector>
    static inline void apply(Point &, Vector const&) {}
};

template
<
    typename Point,
    std::size_t Dimension = 0,
    std::size_t DimensionCount = geometry::dimension<Point>::value
>
struct set_vec_from_point
{
    template <typename Vector>
    static inline void apply(Point const& p, Vector & v)
    {
        qvm::A<Dimension>(v) = get<Dimension>(p);
        set_vec_from_point<Point, Dimension + 1, DimensionCount>::apply(p, v);
    }
};

template
<
    typename Point,
    std::size_t DimensionCount
>
struct set_vec_from_point<Point, DimensionCount, DimensionCount>
{
    template <typename Vector>
    static inline void apply(Point const&, Vector &) {}
};

template
<
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class matrix_transformer
{
protected :
    typedef CalculationType ct;
    typedef boost::qvm::mat<ct, Dimension2 + 1, Dimension1 + 1> matrix_type;
    matrix_type m_matrix;
public :
    matrix_type const& matrix() const { return m_matrix; }
    template <typename P1, typename P2>
    inline bool apply(P1 const& p1, P2& p2) const
    {
        assert_dimension_greater_equal<P1,Dimension1>();
        assert_dimension_greater_equal<P2,Dimension2>();
        qvm::vec<ct,Dimension1 + 1> p1temp;
        qvm::A<Dimension1>(p1temp) = 1;
        qvm::vec<ct,Dimension2 + 1> p2temp;
        set_vec_from_point<P1, 0, Dimension1>::apply(p1, p1temp);
        p2temp = m_matrix * p1temp;
        set_point_from_vec<P2, 0, Dimension2>::apply(p2, p2temp);
        return true;
    }

};

}} // namespace detail::matrix_transform

/*!
\brief Affine transformation strategy in Cartesian system.
\details The strategy serves as a generic definition of an affine transformation
         matrix and procedure for applying it to a given point.
\see http://en.wikipedia.org/wiki/Affine_transformation
     and http://www.devmaster.net/wiki/Transformation_matrices
\ingroup strategies
\tparam Dimension1 number of dimensions to transform from
\tparam Dimension2 number of dimensions to transform to
 */
template
<
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class matrix_transformer : public detail::matrix_transformer::matrix_transformer<CalculationType, Dimension1, Dimension2>
{
public:
    template<typename Matrix>
    inline matrix_transformer(Matrix const& matrix)
    {
        qvm::assign(this->m_matrix, matrix);
    }
    inline matrix_transformer() {}
};

template <typename CalculationType>
class matrix_transformer<CalculationType, 2, 2> : public detail::matrix_transformer::matrix_transformer<CalculationType, 2, 2>
{
    typedef CalculationType ct;
public :
    template<typename Matrix>
    inline matrix_transformer(Matrix const& matrix)
    {
        qvm::assign(this->m_matrix, matrix);
    }

    inline matrix_transformer() {}

    inline matrix_transformer(
                ct const& m_0_0, ct const& m_0_1, ct const& m_0_2,
                ct const& m_1_0, ct const& m_1_1, ct const& m_1_2,
                ct const& m_2_0, ct const& m_2_1, ct const& m_2_2)
    {
        qvm::A<0,0>(this->m_matrix) = m_0_0;   qvm::A<0,1>(this->m_matrix) = m_0_1;   qvm::A<0,2>(this->m_matrix) = m_0_2;
        qvm::A<1,0>(this->m_matrix) = m_1_0;   qvm::A<1,1>(this->m_matrix) = m_1_1;   qvm::A<1,2>(this->m_matrix) = m_1_2;
        qvm::A<2,0>(this->m_matrix) = m_2_0;   qvm::A<2,1>(this->m_matrix) = m_2_1;   qvm::A<2,2>(this->m_matrix) = m_2_2;
    }

    template <typename P1, typename P2>
    inline bool apply(P1 const& p1, P2& p2) const
    {
        assert_dimension_greater_equal<P1, 2>();
        assert_dimension_greater_equal<P2, 2>();

        ct const& c1 = get<0>(p1);
        ct const& c2 = get<1>(p1);

        typedef typename geometry::coordinate_type<P2>::type ct2;
        set<0>(p2, boost::numeric_cast<ct2>(c1 * qvm::A<0,0>(this->m_matrix) + c2 * qvm::A<0,1>(this->m_matrix) + qvm::A<0,2>(this->m_matrix)));
        set<1>(p2, boost::numeric_cast<ct2>(c1 * qvm::A<1,0>(this->m_matrix) + c2 * qvm::A<1,1>(this->m_matrix) + qvm::A<1,2>(this->m_matrix)));

        return true;
    }
};

// It IS possible to go from 3 to 2 coordinates
template <typename CalculationType>
class matrix_transformer<CalculationType, 3, 2> : public detail::matrix_transformer::matrix_transformer<CalculationType, 3, 2>
{
    typedef CalculationType ct;
public :
    template<typename Matrix>
    inline matrix_transformer(Matrix const& matrix)
    {
        qvm::assign(this->m_matrix, matrix);
    }

    inline matrix_transformer() {}

    inline matrix_transformer(
                ct const& m_0_0, ct const& m_0_1, ct const& m_0_2,
                ct const& m_1_0, ct const& m_1_1, ct const& m_1_2,
                ct const& m_2_0, ct const& m_2_1, ct const& m_2_2)
    {
        qvm::A<0,0>(this->m_matrix) = m_0_0;   qvm::A<0,1>(this->m_matrix) = m_0_1;   qvm::A<0,2>(this->m_matrix) = 0;   qvm::A<0,3>(this->m_matrix) = m_0_2;
        qvm::A<1,0>(this->m_matrix) = m_1_0;   qvm::A<1,1>(this->m_matrix) = m_1_1;   qvm::A<1,2>(this->m_matrix) = 0;   qvm::A<1,3>(this->m_matrix) = m_1_2;
        qvm::A<2,0>(this->m_matrix) = m_2_0;   qvm::A<2,1>(this->m_matrix) = m_2_1;   qvm::A<2,2>(this->m_matrix) = 0;   qvm::A<2,3>(this->m_matrix) = m_2_2;
    }

    template <typename P1, typename P2>
    inline bool apply(P1 const& p1, P2& p2) const
    {
        assert_dimension_greater_equal<P1, 3>();
        assert_dimension_greater_equal<P2, 2>();

        ct const& c1 = get<0>(p1);
        ct const& c2 = get<1>(p1);
        ct const& c3 = get<2>(p1);

        typedef typename geometry::coordinate_type<P2>::type ct2;

        set<0>(p2, boost::numeric_cast<ct2>(
            c1 * qvm::A<0,0>(this->m_matrix) + c2 * qvm::A<0,1>(this->m_matrix) + c3 * qvm::A<0,2>(this->m_matrix) + qvm::A<0,3>(this->m_matrix)));
        set<1>(p2, boost::numeric_cast<ct2>(
            c1 * qvm::A<1,0>(this->m_matrix) + c2 * qvm::A<1,1>(this->m_matrix) + c3 * qvm::A<1,2>(this->m_matrix) + qvm::A<1,3>(this->m_matrix)));

        return true;
    }

};

template <typename CalculationType>
class matrix_transformer<CalculationType, 3, 3> : public detail::matrix_transformer::matrix_transformer<CalculationType, 3, 3>
{
    typedef CalculationType ct;
public :
    template<typename Matrix>
    inline matrix_transformer(Matrix const& matrix)
    {
        qvm::assign(this->m_matrix, matrix);
    }

    inline matrix_transformer() {}

    inline matrix_transformer(
                ct const& m_0_0, ct const& m_0_1, ct const& m_0_2, ct const& m_0_3,
                ct const& m_1_0, ct const& m_1_1, ct const& m_1_2, ct const& m_1_3,
                ct const& m_2_0, ct const& m_2_1, ct const& m_2_2, ct const& m_2_3,
                ct const& m_3_0, ct const& m_3_1, ct const& m_3_2, ct const& m_3_3
                )
    {
        qvm::A<0,0>(this->m_matrix) = m_0_0; qvm::A<0,1>(this->m_matrix) = m_0_1; qvm::A<0,2>(this->m_matrix) = m_0_2; qvm::A<0,3>(this->m_matrix) = m_0_3;
        qvm::A<1,0>(this->m_matrix) = m_1_0; qvm::A<1,1>(this->m_matrix) = m_1_1; qvm::A<1,2>(this->m_matrix) = m_1_2; qvm::A<1,3>(this->m_matrix) = m_1_3;
        qvm::A<2,0>(this->m_matrix) = m_2_0; qvm::A<2,1>(this->m_matrix) = m_2_1; qvm::A<2,2>(this->m_matrix) = m_2_2; qvm::A<2,3>(this->m_matrix) = m_2_3;
        qvm::A<3,0>(this->m_matrix) = m_3_0; qvm::A<3,1>(this->m_matrix) = m_3_1; qvm::A<3,2>(this->m_matrix) = m_3_2; qvm::A<3,3>(this->m_matrix) = m_3_3;
    }

    template <typename P1, typename P2>
    inline bool apply(P1 const& p1, P2& p2) const
    {
        assert_dimension_greater_equal<P1, 3>();
        assert_dimension_greater_equal<P2, 3>();

        ct const& c1 = get<0>(p1);
        ct const& c2 = get<1>(p1);
        ct const& c3 = get<2>(p1);

        typedef typename geometry::coordinate_type<P2>::type ct2;

        set<0>(p2, boost::numeric_cast<ct2>(
            c1 * qvm::A<0,0>(this->m_matrix) + c2 * qvm::A<0,1>(this->m_matrix) + c3 * qvm::A<0,2>(this->m_matrix) + qvm::A<0,3>(this->m_matrix)));
        set<1>(p2, boost::numeric_cast<ct2>(
            c1 * qvm::A<1,0>(this->m_matrix) + c2 * qvm::A<1,1>(this->m_matrix) + c3 * qvm::A<1,2>(this->m_matrix) + qvm::A<1,3>(this->m_matrix)));
        set<2>(p2, boost::numeric_cast<ct2>(
            c1 * qvm::A<2,0>(this->m_matrix) + c2 * qvm::A<2,1>(this->m_matrix) + c3 * qvm::A<2,2>(this->m_matrix) + qvm::A<2,3>(this->m_matrix)));

        return true;
    }
};

/*!
\brief Strategy of translate transformation in Cartesian system.
\details Translate moves a geometry a fixed distance in 2 or 3 dimensions.
\see http://en.wikipedia.org/wiki/Translation_%28geometry%29
\ingroup strategies
\tparam Dimension1 number of dimensions to transform from
\tparam Dimension2 number of dimensions to transform to
 */
template
<
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class translate_transformer
{
};

template<typename CalculationType>
class translate_transformer<CalculationType, 2, 2> : public matrix_transformer<CalculationType, 2, 2>
{
public :
    // To have translate transformers compatible for 2/3 dimensions, the
    // constructor takes an optional third argument doing nothing.
    inline translate_transformer(CalculationType const& translate_x,
                CalculationType const& translate_y,
                CalculationType const& = 0)
        : matrix_transformer<CalculationType, 2, 2>(
                1, 0, translate_x,
                0, 1, translate_y,
                0, 0, 1)
    {}
};

template <typename CalculationType>
class translate_transformer<CalculationType, 3, 3> : public matrix_transformer<CalculationType, 3, 3>
{
public :
    inline translate_transformer(CalculationType const& translate_x,
                CalculationType const& translate_y,
                CalculationType const& translate_z)
        : matrix_transformer<CalculationType, 3, 3>(
                1, 0, 0, translate_x,
                0, 1, 0, translate_y,
                0, 0, 1, translate_z,
                0, 0, 0, 1)
    {}

};

/*!
\brief Strategy of scale transformation in Cartesian system.
\details Scale scales a geometry up or down in all its dimensions.
\see http://en.wikipedia.org/wiki/Scaling_%28geometry%29
\ingroup strategies
\tparam Dimension1 number of dimensions to transform from
\tparam Dimension2 number of dimensions to transform to
*/
template
<
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class scale_transformer
{
};

template
<
    typename CalculationType,
    std::size_t Dimension1
>
class scale_transformer<CalculationType, Dimension1, Dimension1> : public matrix_transformer<CalculationType, Dimension1, Dimension1>
{
public:
    inline scale_transformer(CalculationType const& scale)
    {
        boost::qvm::set_identity(this->m_matrix);
        this->m_matrix*=scale;
        qvm::A<Dimension1,Dimension1>(this->m_matrix) = 1;
    }
};

template <typename CalculationType>
class scale_transformer<CalculationType, 2, 2> : public matrix_transformer<CalculationType, 2, 2>
{

public :
    inline scale_transformer(CalculationType const& scale_x,
                CalculationType const& scale_y,
                CalculationType const& = 0)
        : matrix_transformer<CalculationType, 2, 2>(
                scale_x, 0,       0,
                0,       scale_y, 0,
                0,       0,       1)
    {}

    inline scale_transformer(CalculationType const& scale)
        : matrix_transformer<CalculationType, 2, 2>(
                scale, 0,     0,
                0,     scale, 0,
                0,     0,     1)
    {}
};

template <typename CalculationType>
class scale_transformer<CalculationType, 3, 3> : public matrix_transformer<CalculationType, 3, 3>
{
public :
    inline scale_transformer(CalculationType const& scale_x,
                CalculationType const& scale_y,
                CalculationType const& scale_z)
        : matrix_transformer<CalculationType, 3, 3>(
                scale_x, 0,       0,       0,
                0,       scale_y, 0,       0,
                0,       0,       scale_z, 0,
                0,       0,       0,       1)
    {}

    inline scale_transformer(CalculationType const& scale)
        : matrix_transformer<CalculationType, 3, 3>(
                scale, 0,     0,     0,
                0,     scale, 0,     0,
                0,     0,     scale, 0,
                0,     0,     0,     1)
    {}
};

#ifndef DOXYGEN_NO_DETAIL
namespace detail
{

template <typename DegreeOrRadian>
struct as_radian
{};

template <>
struct as_radian<radian>
{
    template <typename T>
    static inline T get(T const& value)
    {
        return value;
    }
};

template <>
struct as_radian<degree>
{
    template <typename T>
    static inline T get(T const& value)
    {
        typedef typename promote_floating_point<T>::type promoted_type;
        return value * math::d2r<promoted_type>();
    }

};

template
<
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class rad_rotate_transformer
    : public transform::matrix_transformer<CalculationType, Dimension1, Dimension2>
{
public :
    inline rad_rotate_transformer(CalculationType const& angle)
        : transform::matrix_transformer<CalculationType, Dimension1, Dimension2>(
                 cos(angle), sin(angle), 0,
                -sin(angle), cos(angle), 0,
                 0,          0,          1)
    {}
};

} // namespace detail
#endif // DOXYGEN_NO_DETAIL

/*!
\brief Strategy for rotate transformation in Cartesian coordinate system.
\details Rotate rotates a geometry by a specified angle about a fixed point (e.g. origin).
\see http://en.wikipedia.org/wiki/Rotation_%28mathematics%29
\ingroup strategies
\tparam DegreeOrRadian degree/or/radian, type of rotation angle specification
\note A single angle is needed to specify a rotation in 2D.
      Not yet in 3D, the 3D version requires special things to allow
      for rotation around X, Y, Z or arbitrary axis.
\todo The 3D version will not compile.
 */
template
<
    typename DegreeOrRadian,
    typename CalculationType,
    std::size_t Dimension1,
    std::size_t Dimension2
>
class rotate_transformer : public detail::rad_rotate_transformer<CalculationType, Dimension1, Dimension2>
{

public :
    inline rotate_transformer(CalculationType const& angle)
        : detail::rad_rotate_transformer
            <
                CalculationType, Dimension1, Dimension2
            >(detail::as_radian<DegreeOrRadian>::get(angle))
    {}
};

}} // namespace strategy::transform

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_STRATEGIES_TRANSFORM_MATRIX_TRANSFORMERS_HPP


#include "hyperopt.h"

#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/max.hpp>
#include <boost/accumulators/statistics/mean.hpp>
#include <boost/accumulators/statistics/min.hpp>
#include <boost/accumulators/statistics/variance.hpp>
#include <boost/math/distributions/students_t.hpp>
#include <cmath>
#include <tbb/mutex.h>
#include <tbb/parallel_for_each.h>
#include <tbb/parallel_for.h>

#include "export.h"
#include "ndarray_ops.h"
#include "predict.h"
#include "train.h"
#include "utils.h"

namespace curfil {

bool continueSearching(const std::vector<double>& currentBestLosses,
        const std::vector<double>& currentRunLosses) {

    static const size_t MIN_SAMPLES = 3;

    if (currentBestLosses.empty()) {
        CURFIL_INFO("current best is empty. continue searching");
        return true;
    }

    // we want a minimum of three samples
    if (currentRunLosses.size() < MIN_SAMPLES) {
        CURFIL_INFO("too few elements in current run: " << currentRunLosses.size()
                << ". continue searching");
        return true;
    }

    boost::accumulators::accumulator_set<double,
            boost::accumulators::features<boost::accumulators::tag::max,
                    boost::accumulators::tag::mean,
                    boost::accumulators::tag::variance> > acc1;

    boost::accumulators::accumulator_set<double,
            boost::accumulators::features<boost::accumulators::tag::min,
                    boost::accumulators::tag::mean,
                    boost::accumulators::tag::variance> > acc2;

    acc1 = std::for_each(currentBestLosses.begin(), currentBestLosses.end(), acc1);
    acc2 = std::for_each(currentRunLosses.begin(), currentRunLosses.end(), acc2);

    double worstCurrentBest = boost::accumulators::max(acc1);
    double bestOfCurrentLosses = boost::accumulators::min(acc2);

    // continue if any loss is at least as good as the worst loss of the currently best
    if (bestOfCurrentLosses <= worstCurrentBest) {
        CURFIL_INFO("best run of current parameters is at least as good as the worst run with the best parameters: "
                << bestOfCurrentLosses << " <= " << worstCurrentBest
                << ". continue searching");
        return true;
    }

    // http://www.boost.org/doc/libs/1_53_0/libs/math/doc/sf_and_dist/html/math_toolkit/dist/stat_tut/weg/st_eg/two_sample_students_t.html

    double Sm1 = boost::accumulators::mean(acc1);                 // Sm1 = Sample 1 Mean.
    double variance1 = boost::accumulators::variance(acc1);
    double Sd1 = std::sqrt(variance1);                            // Sd1 = Sample 1 Standard Deviation.
    unsigned Sn1 = currentBestLosses.size();                      // Sn1 = Sample 1 Size.

    double Sm2 = boost::accumulators::mean(acc2);                 // Sm2 = Sample 2 Mean.
    double variance2 = boost::accumulators::variance(acc2);
    double Sd2 = std::sqrt(variance2);                            // Sd2 = Sample 2 Standard Deviation.
    unsigned Sn2 = currentRunLosses.size();                       // Sn2 = Sample 2 Size.
    double alpha = 0.005;                                         // alpha = Significance Level.

    // Degrees of freedom:
    double v = Sn1 + Sn2 - 2;
    // Pooled variance:
    double sp = sqrt(((Sn1 - 1) * Sd1 * Sd1 + (Sn2 - 1) * Sd2 * Sd2) / v);
    // t-statistic:
    double t_stat = (Sm1 - Sm2) / (sp * sqrt(1.0 / Sn1 + 1.0 / Sn2));

    boost::math::students_t dist(v);
    double q = boost::math::cdf(boost::math::complement(dist, fabs(t_stat)));

    CURFIL_INFO("Sm1: " << Sm1 << ", Sd1: " << Sd1 << ", Sn1: " << Sn1);
    CURFIL_INFO("Sm2: " << Sm2 << ", Sd2: " << Sd2 << ", Sn2: " << Sn2);
    CURFIL_INFO("q: " << q << ", alpha: " << alpha);

    // continue if we reject hypothesis that sample 1 mean loss is less than sample 2 mean loss. then we do not know and need to continue searching
    return (q > alpha);
}

double Result::getLoss() const {
    double accuracy;
    switch (lossFunctionType) {
        case LossFunctionType::CLASS_ACCURACY:
            accuracy = getClassAccuracy();
            break;
        case LossFunctionType::CLASS_ACCURACY_WITHOUT_VOID:
            accuracy = getClassAccuracyWithoutVoid();
            break;
        case LossFunctionType::PIXEL_ACCURACY:
            accuracy = getPixelAccuracy();
            break;
        case LossFunctionType::PIXEL_ACCURACY_WITHOUT_VOID:
            accuracy = getPixelAccuracyWithoutVoid();
            break;
        default:
            throw std::runtime_error("unknown type of loss function");
    }
    assertProbability(accuracy);
    return 1.0 - accuracy;
}

mongo::BSONObj Result::toBSON() const {

    mongo::BSONObjBuilder o(256);
    o << "randomSeed" << getRandomSeed();
    o << "loss" << getLoss();
    o << "classAccuracy" << getClassAccuracy();
    o << "classAccuracyWithoutVoid" << getClassAccuracyWithoutVoid();
    o << "pixelAccuracy" << getPixelAccuracy();
    o << "pixelAccuracyWithoutVoid" << getPixelAccuracyWithoutVoid();

    mongo::BSONArrayBuilder confusionMatrixArray;
    for (size_t label = 0; label < confusionMatrix.getNumClasses(); label++) {
        mongo::BSONArrayBuilder confusionMatrixRow;
        for (size_t prediction = 0; prediction < confusionMatrix.getNumClasses(); prediction++) {
            double probability = confusionMatrix(label, prediction);
            confusionMatrixRow.append(probability);
        }
        confusionMatrixArray.append(confusionMatrixRow.arr());
    }
    o << "confusionMatrix" << confusionMatrixArray.arr();

    return o.obj();
}

LossFunctionType HyperoptClient::parseLossFunction(const std::string& lossFunctionString) {
    if (lossFunctionString == "classAccuracy") {
        return LossFunctionType::CLASS_ACCURACY;
    } else if (lossFunctionString == "classAccuracyWithoutVoid") {
        return LossFunctionType::CLASS_ACCURACY_WITHOUT_VOID;
    }
    else if (lossFunctionString == "pixelAccuracy") {
        return LossFunctionType::PIXEL_ACCURACY;
    }
    else if (lossFunctionString == "pixelAccuracyWithoutVoid") {
        return LossFunctionType::PIXEL_ACCURACY_WITHOUT_VOID;
    } else {
        throw std::runtime_error(std::string("unknown type of loss function: ") + lossFunctionString);
    }
}

const Result HyperoptClient::test(const RandomForestImage& randomForest,
        const std::vector<LabeledRGBDImage>& testImages) {

    tbb::mutex totalMutex;
    utils::Average averageAccuracy;
    utils::Average averageAccuracyWithoutVoid;

    std::vector<int> indices(testImages.size(), 0);
    for (size_t i = 0; i < testImages.size(); i++) {
        indices[i] = i;
    }

    const LabelType numClasses = randomForest.getNumClasses();

    CURFIL_INFO("testing " << testImages.size() << " images with " << static_cast<int>(numClasses) << " classes");

    ConfusionMatrix totalConfusionMatrix(numClasses);

    tbb::parallel_for_each(indices.begin(), indices.end(), [&](const int& i) {
        const RGBDImage& image = testImages[i].getRGBDImage();
        const LabelImage& groundTruth = testImages[i].getLabelImage();

        LabelImage prediction = randomForest.predict(image);

        tbb::mutex::scoped_lock lock(totalMutex);

        ConfusionMatrix confusionMatrix;
        double accuracy = calculatePixelAccuracy(prediction, groundTruth, true, &confusionMatrix);
        double accuracyWithoutVoid = calculatePixelAccuracy(prediction, groundTruth, false);

        totalConfusionMatrix += confusionMatrix;

        averageAccuracy.addValue(accuracy);
        averageAccuracyWithoutVoid.addValue(accuracyWithoutVoid);
    });

    tbb::mutex::scoped_lock lock(totalMutex);
    double accuracy = averageAccuracy.getAverage();
    double accuracyWithoutVoid = averageAccuracyWithoutVoid.getAverage();

    CURFIL_INFO("accuracy (no void): " << accuracy << " (" << accuracyWithoutVoid << ")");

    return Result(totalConfusionMatrix, accuracy, accuracyWithoutVoid, lossFunction);
}

RandomForestImage HyperoptClient::train(size_t trees,
        const TrainingConfiguration& configuration,
        const std::vector<LabeledRGBDImage>& trainImages) {

    CURFIL_INFO("trees: " << trees);
    CURFIL_INFO(configuration);

    // Train

    RandomForestImage randomForest(trees, configuration);

    // parallel training is not thoroughly tested yet
    static const bool trainTreesSequentially = true;

    utils::Timer trainTimer;
    randomForest.train(trainImages, trainTreesSequentially);
    trainTimer.stop();

    CURFIL_INFO("training took " << trainTimer.format(2) <<
            " (" << std::setprecision(3) << trainTimer.getSeconds() / 60.0 << " min)");

    mongo::BSONObjBuilder featureBuilder(64);
    for (const auto& featureCount : randomForest.countFeatures()) {
        featureBuilder.append(featureCount.first, static_cast<int>(featureCount.second));
    }

    mongo::BSONObjBuilder builder(64);
    builder.append("training_time_millis", trainTimer.getMilliseconds());
    builder.append("featureCounts", featureBuilder.obj());

    log(1, builder.obj());

    return randomForest;
}

double HyperoptClient::measureTrueLoss(unsigned int numTrees, TrainingConfiguration configuration,
        const double histogramBias, double& variance) {

    boost::accumulators::accumulator_set<double,
            boost::accumulators::features<boost::accumulators::tag::mean,
                    boost::accumulators::tag::variance> > acc;

    static const size_t TRUE_LOSS_RUNS = 2;

    CURFIL_INFO("measuring true loss");

    Sampler sampler(randomSeed, 1, 100000);

    for (size_t run = 0; run < TRUE_LOSS_RUNS; run++) {

        CURFIL_INFO("true loss run " << (run + 1) << "/" << TRUE_LOSS_RUNS);

        const int seedOfRun = sampler.getNext();

        configuration.setRandomSeed(seedOfRun);

        RandomForestImage forest = train(numTrees, configuration, allRGBDImages);
        forest.normalizeHistograms(histogramBias);
        Result result = test(forest, allTestImages);

        result.setRandomSeed(seedOfRun);

        CURFIL_INFO(result.getConfusionMatrix());

        log(2, BSON("trueLossRun" << static_cast<int>(run)
                << "result" << result.toBSON()));

        acc(result.getLoss());
    }

    variance = boost::accumulators::variance(acc);

    double trueLoss = boost::accumulators::mean(acc);

    CURFIL_INFO("true loss: " << trueLoss);

    return trueLoss;
}

double HyperoptClient::getAverageLossAndVariance(const std::vector<Result>& results, double& variance) {

    boost::accumulators::accumulator_set<double,
            boost::accumulators::features<boost::accumulators::tag::mean,
                    boost::accumulators::tag::variance> > acc;

    for (const Result& result : results) {
        acc(result.getLoss());
    }

    variance = boost::accumulators::variance(acc);

    return boost::accumulators::mean(acc);
}

void HyperoptClient::randomSplit(const int randomSeed, const double testRatio,
        std::vector<LabeledRGBDImage>& trainImages,
        std::vector<LabeledRGBDImage>& testImages) {

    boost::mt19937 rng(randomSeed);
    boost::uniform_real<> dist(0.0, 1.0);

    trainImages.clear();
    testImages.clear();

    while (testImages.empty() || trainImages.empty()) {
        for (const auto& image : allRGBDImages) {
            double random = dist(rng);
            if (random < testRatio) {
                testImages.push_back(image);
            } else {
                trainImages.push_back(image);
            }
        }
    }

    CURFIL_INFO("random split of " << testRatio <<
            ". train images: " << trainImages.size() <<
            ", test images: " << testImages.size());
}

void HyperoptClient::handle_task(const mongo::BSONObj& task) {
    try {
        CURFIL_INFO("got task object: " << task.toString());

        const unsigned int numTrees = getParameterDouble(task, "numTrees");
        const unsigned int samplesPerImage = getParameterDouble(task, "samplesPerImage");
        const unsigned int featureCount = getParameterDouble(task, "featureCount");
        const unsigned int minSampleCount = getParameterDouble(task, "minSampleCount");
        const int maxDepth = getParameterDouble(task, "maxDepth");
        const uint16_t boxRadius = getParameterDouble(task, "boxRadius");
        const uint16_t regionSize = getParameterDouble(task, "regionSize");
        const uint16_t thresholds = getParameterDouble(task, "thresholds");
        const double histogramBias = getParameterDouble(task, "histogramBias");
        const AccelerationMode accelerationMode = AccelerationMode::GPU_ONLY;

        std::vector<Result> results;

        Sampler sampler(randomSeed, 1, 100000);

        static const int RUNS = 5;
        const double testRatio = 1.0 / RUNS;

        std::vector<double> currentRunLosses;

        for (int run = 0; run < RUNS; run++) {

            CURFIL_INFO("starting run " << (run + 1) << "/" << RUNS);

            const int seedOfRun = sampler.getNext();

            std::vector<LabeledRGBDImage> trainImages;
            std::vector<LabeledRGBDImage> testImages;

            randomSplit(seedOfRun, testRatio, trainImages, testImages);

            unsigned int imageCacheSize = 0;
            unsigned int maxSamplesPerBatch = 0;

            determineImageCacheSizeAndSamplesPerBatch(trainImages, deviceIds, featureCount, thresholds,
                    imageCacheSizeMB, imageCacheSize, maxSamplesPerBatch);

            TrainingConfiguration configuration(seedOfRun, samplesPerImage, featureCount, minSampleCount, maxDepth,
                    boxRadius, regionSize, thresholds, numThreads, maxImages, imageCacheSize, maxSamplesPerBatch,
                    accelerationMode, useCIELab, useDepthFilling, deviceIds, subsamplingType, ignoredColors);

            mongo::BSONObj msg = BSON("run" << run
                    << "randomSeed" << seedOfRun
                    << "numTrainImages" << static_cast<int>(trainImages.size())
                    << "numTestImages" << static_cast<int>(testImages.size()));

            log(3, msg);
            checkpoint();

            RandomForestImage forest = train(numTrees, configuration, trainImages);
            forest.normalizeHistograms(histogramBias);
            Result result = test(forest, testImages);
            result.setRandomSeed(seedOfRun);

            results.push_back(result);

            currentRunLosses.push_back(result.getLoss());

            log(3, result.toBSON());
            checkpoint();

            mongo::BSONObj bestTask;
            if (get_best_task(bestTask)) {

                CURFIL_INFO("best task so far: " << bestTask.getObjectField("result").toString());

                std::vector<double> currentBestLosses;

                mongo::BSONElementSet values;
                bestTask.getFieldsDotted("result.results.loss", values);

                for (const mongo::BSONElement& loss : values) {
                    currentBestLosses.push_back(loss.Double());
                }

                if (!continueSearching(currentBestLosses, currentRunLosses)) {
                    log(1, BSON("stopSearching" << true
                            << "run" << run
                            << "currentBestLosses" << currentBestLosses
                            << "currentRunLosses" << currentRunLosses ));
                    CURFIL_INFO("stop searching");
                    break;
                } else {
                    CURFIL_INFO("continue searching");
                }
            } else {
                CURFIL_INFO("no finished task so far");
            }

        }

        double lossVariance;
        double loss = getAverageLossAndVariance(results, lossVariance);

        unsigned int imageCacheSize = 0;
        unsigned int maxSamplesPerBatch = 0;

        determineImageCacheSizeAndSamplesPerBatch(allRGBDImages, deviceIds, featureCount, thresholds,
                imageCacheSizeMB, imageCacheSize, maxSamplesPerBatch);

        TrainingConfiguration configuration(randomSeed, samplesPerImage, featureCount, minSampleCount, maxDepth,
                boxRadius, regionSize, thresholds, numThreads, maxImages, imageCacheSize, maxSamplesPerBatch,
                accelerationMode, useCIELab, useDepthFilling, deviceIds, subsamplingType, ignoredColors);

        double trueLossVariance;
        double trueLoss = measureTrueLoss(numTrees, configuration, histogramBias, trueLossVariance);

        mongo::BSONObjBuilder builder(64);

        builder << "status" << "ok";
        builder << "loss" << loss << "loss_variance" << lossVariance;
        builder << "true_loss" << trueLoss << "true_loss_variance" << trueLossVariance;

        std::vector<mongo::BSONObj> resultsDetails;
        for (size_t i = 0; i < results.size(); i++) {
            resultsDetails.push_back(results[i].toBSON());
        }

        builder.append("results", resultsDetails);

        finish(builder.obj(), true);

    } catch (const std::runtime_error& e) {
        CURFIL_ERROR(e.what());
        finish(BSON("status" << "fail" << "why" << e.what()), false);
        throw e;
    }

}

double HyperoptClient::getParameterDouble(const mongo::BSONObj& task, const std::string& field) {
    const mongo::BSONObj vals = task.getObjectField("vals");
    if (!vals.hasField(field.c_str())) {
        throw std::runtime_error(task.toString() + " has no value '" + field + "'");
    }
    const mongo::BSONElement value = vals.getField(field);
    const std::vector<mongo::BSONElement> a = value.Array();
    return a.at(0).Double();
}

void HyperoptClient::run() {

    mongo::BSONObj result;
    if (get_best_task(result)) {
        CURFIL_INFO("best result so far: " << result["result"].toString());
    } else {
        CURFIL_INFO("no best result so far");
    }

    CURFIL_INFO("running client");
    bool log_when_waiting = true;
    while (true) {
        mongo::BSONObj task;
        if (!get_next_task(task)) {
            if (log_when_waiting)
                CURFIL_INFO("no task in queue. waiting...");
            log_when_waiting = false;
            sleep(1);
        } else {
            handle_task(task);
            log_when_waiting = true;
        }
    }
}

}


#ifndef MATHTOOLBOX_ACQUISITION_FUNCTIONS_HPP
#define MATHTOOLBOX_ACQUISITION_FUNCTIONS_HPP

#include <Eigen/Core>
#include <functional>

namespace mathtoolbox
{
    double GetExpectedImprovement(const Eigen::VectorXd&                               x,
                                  const std::function<double(const Eigen::VectorXd&)>& mu,
                                  const std::function<double(const Eigen::VectorXd&)>& sigma,
                                  const Eigen::VectorXd&                               x_best);

    Eigen::VectorXd
    GetExpectedImprovementDerivative(const Eigen::VectorXd&                                        x,
                                     const std::function<double(const Eigen::VectorXd&)>&          mu,
                                     const std::function<double(const Eigen::VectorXd&)>&          sigma,
                                     const Eigen::VectorXd&                                        x_best,
                                     const std::function<Eigen::VectorXd(const Eigen::VectorXd&)>& mu_derivative,
                                     const std::function<Eigen::VectorXd(const Eigen::VectorXd&)>& sigma_derivative);

    /// \param hyperparam The hyperparameter that controls the trade-off of exploration and exploitation. Specifically,
    /// this hyperparameter corresponds to the square root of the beta in [Srinivas et al. ICML '10]. Setting this to
    /// zero means pure exploitation, and setting to a very large value means (almost) pure exploration. This value
    /// needs to be non-negative.
    double GetGaussianProcessUpperConfidenceBound(const Eigen::VectorXd&                               x,
                                                  const std::function<double(const Eigen::VectorXd&)>& mu,
                                                  const std::function<double(const Eigen::VectorXd&)>& sigma,
                                                  const double                                         hyperparam);

    /// \param hyperparam The hyperparameter that controls the trade-off of exploration and exploitation. Specifically,
    /// this hyperparameter corresponds to the square root of the beta in [Srinivas et al. ICML '10]. Setting this to
    /// zero means pure exploitation, and setting to a very large value means (almost) pure exploration. This value
    /// needs to be non-negative.
    Eigen::VectorXd GetGaussianProcessUpperConfidenceBoundDerivative(
        const Eigen::VectorXd&                                        x,
        const std::function<double(const Eigen::VectorXd&)>&          mu,
        const std::function<double(const Eigen::VectorXd&)>&          sigma,
        const double                                                  hyperparam,
        const std::function<Eigen::VectorXd(const Eigen::VectorXd&)>& mu_derivative,
        const std::function<Eigen::VectorXd(const Eigen::VectorXd&)>& sigma_derivative);
} // namespace mathtoolbox

#endif // MATHTOOLBOX_ACQUISITION_FUNCTIONS_HPP


/*

PICCANTE
The hottest HDR imaging library!
http://vcg.isti.cnr.it/piccante

Copyright (C) 2014
Visual Computing Laboratory - ISTI CNR
http://vcg.isti.cnr.it
First author: Francesco Banterle

This Source Code Form is subject to the terms of the Mozilla Public
License, v. 2.0. If a copy of the MPL was not distributed with this
file, You can obtain one at http://mozilla.org/MPL/2.0/.

*/

#ifndef PIC_COMPUTER_VISION_FUNDAMENTAL_HPP
#define PIC_COMPUTER_VISION_FUNDAMENTAL_HPP

#include <vector>
#include <random>
#include <stdlib.h>

#include "../base.hpp"

#include "../image.hpp"

#include "../filtering/filter_luminance.hpp"
#include "../filtering/filter_gaussian_2d.hpp"

#include "../util/math.hpp"
#include "../util/eigen_util.hpp"

#include "../features_matching/orb_descriptor.hpp"
#include "../features_matching/feature_matcher.hpp"
#include "../features_matching/binary_feature_lsh_matcher.hpp"
#include "../computer_vision/nelder_mead_opt_fundamental.hpp"

#ifndef PIC_DISABLE_EIGEN

#ifndef PIC_EIGEN_NOT_BUNDLED
    #include "../externals/Eigen/Dense"
    #include "../externals/Eigen/SVD"
    #include "../externals/Eigen/Geometry"
#else
    #include <Eigen/Dense>
    #include <Eigen/SVD>
    #include <Eigen/Geometry>
#endif

#endif

namespace pic {

#ifndef PIC_DISABLE_EIGEN

/**
 * @brief estimateFundamental estimates the foundamental matrix between image 1 to image 2
 * @param points0 is an array of points computed from image 1.
 * @param points1 is an array of points computed from image 2.
 * @return It returns the fundamental matrix, F_{1,2}.
 */
PIC_INLINE Eigen::Matrix3d estimateFundamental(std::vector< Eigen::Vector2f > &points0,
                                    std::vector< Eigen::Vector2f > &points1)
{
    Eigen::Matrix3d F;

    if((points0.size() != points1.size()) || (points0.size() < 8)) {
        F.setZero();
        return F;
    }

    //shift and scale points for numerical stability
    Eigen::Vector3f transform_0 = ComputeNormalizationTransform(points0);
    Eigen::Vector3f transform_1 = ComputeNormalizationTransform(points1);

    Eigen::Matrix3d mat_0 = getShiftScaleMatrix(transform_0);
    Eigen::Matrix3d mat_1 = getShiftScaleMatrix(transform_1);

    Eigen::MatrixXd A(points0.size(), 9);

    //set up the linear system
    for(unsigned int i = 0; i < points0.size(); i++) {

        //transform coordinates for increasing stability of the system
        Eigen::Vector2f p0 = points0[i];
        Eigen::Vector2f p1 = points1[i];

        p0[0] = (p0[0] - transform_0[0]) / transform_0[2];
        p0[1] = (p0[1] - transform_0[1]) / transform_0[2];

        p1[0] = (p1[0] - transform_1[0]) / transform_1[2];
        p1[1] = (p1[1] - transform_1[1]) / transform_1[2];

        A(i, 0) = p0[0] * p1[0];
        A(i, 1) = p0[0] * p1[1];
        A(i, 2) = p0[0];
        A(i, 3) = p0[1] * p1[0];
        A(i, 4) = p0[1] * p1[1];
        A(i, 5) = p0[1];
        A(i, 6) = p1[0];
        A(i, 7) = p1[1];
        A(i, 8) = 1.0;
    }

    //solve the linear system
    Eigen::JacobiSVD< Eigen::MatrixXd > svd(A, Eigen::ComputeFullV);
    Eigen::MatrixXd V = svd.matrixV();

    int n = int(V.cols()) - 1;

    F(0, 0) = V(0, n);
    F(1, 0) = V(1, n);
    F(2, 0) = V(2, n);

    F(0, 1) = V(3, n);
    F(1, 1) = V(4, n);
    F(2, 1) = V(5, n);

    F(0, 2) = V(6, n);
    F(1, 2) = V(7, n);
    F(2, 2) = V(8, n);

    //compute the final F matrix
    Eigen::Matrix3d mat_1_t = Eigen::Transpose< Eigen::Matrix3d >(mat_1);
    F = mat_1_t * F * mat_0;

    //enforce singularity
    Eigen::JacobiSVD< Eigen::MatrixXd > svdF(F, Eigen::ComputeThinU | Eigen::ComputeThinV);
    Eigen::Matrix3d Uf = svdF.matrixU();
    Eigen::Matrix3d Vf = svdF.matrixV();
    Eigen::Vector3d Df = svdF.singularValues();
    Df[2] = 0.0;

    Eigen::Matrix3d F_new = Uf * DiagonalMatrix(Df) * Eigen::Transpose< Eigen::Matrix3d >(Vf);

    double norm = MAX(Df[0], Df[1]);
    return F_new / norm;
}

/**
 * @brief estimateFundamentalRansac
 * @param points0
 * @param points1
 * @param inliers
 * @param maxIterations
 * @return
 */
PIC_INLINE Eigen::Matrix3d estimateFundamentalRansac(std::vector< Eigen::Vector2f > &points0,
                                          std::vector< Eigen::Vector2f > &points1,
                                          std::vector< unsigned int > &inliers,
                                          unsigned int maxIterations = 100,
                                          double threshold = 0.01,
                                          unsigned int seed = 1)
{
    if(points0.size() < 9) {
        return estimateFundamental(points0, points1);
    }

    Eigen::Matrix3d F;
    int nSubSet = 8;

    std::mt19937 m(seed);

    unsigned int n = int(points0.size());

    unsigned int *subSet = new unsigned int [nSubSet];

    inliers.clear();

    for(unsigned int i = 0; i < maxIterations; i++) {
        getRandomPermutation(m, subSet, nSubSet, n);

        std::vector< Eigen::Vector2f > sub_points0;
        std::vector< Eigen::Vector2f > sub_points1;

        for(int j = 0; j < nSubSet; j++) {
            unsigned int k = subSet[j];
            sub_points0.push_back(points0[k]);
            sub_points1.push_back(points1[k]);
        }

        Eigen::Matrix3d tmpF = estimateFundamental(sub_points0, sub_points1);

        //is it a good one?
        std::vector< unsigned int > tmp_inliers;

        for(unsigned int j = 0; j < n; j++) {
            Eigen::Vector3d p0 = Eigen::Vector3d(points0[j][0], points0[j][1], 1.0);
            Eigen::Vector3d p1 = Eigen::Vector3d(points1[j][0], points1[j][1], 1.0);

            Eigen::Vector3d tmpF_p0 = tmpF * p0;
            double n0 = sqrt(tmpF_p0[0] * tmpF_p0[0] + tmpF_p0[1] * tmpF_p0[1]);
            if(n0 >  0.0) {
                tmpF_p0 /= n0;
            }

            double err = fabs(tmpF_p0.dot(p1));

            if(err < threshold){
                tmp_inliers.push_back(j);
            }
        }

        //get the inliers
        if(tmp_inliers.size() > inliers.size()) {
            F = tmpF;
            inliers.clear();
            inliers.assign(tmp_inliers.begin(), tmp_inliers.end());
        }
    }

    //improve estimate with inliers only
    if(inliers.size() > 7) {

        #ifdef PIC_DEBUG
            printf("Better estimate using inliers only.\n");
        #endif

        std::vector< Eigen::Vector2f > sub_points0;
        std::vector< Eigen::Vector2f > sub_points1;

        for(unsigned int i = 0; i < inliers.size(); i++) {
            sub_points0.push_back(points0[inliers[i]]);
            sub_points1.push_back(points1[inliers[i]]);
        }

        F = estimateFundamental(sub_points0, sub_points1);
    }

    return F;
}

/**
 * @brief estimateFundamentalWithNonLinearRefinement
 * @param F
 * @return
 */
PIC_INLINE Eigen::Matrix3d estimateFundamentalWithNonLinearRefinement(std::vector< Eigen::Vector2f > &points0,
                                                           std::vector< Eigen::Vector2f > &points1,
                                                           std::vector< unsigned int >    &inliers,
                                                           unsigned int maxIterationsRansac = 100,
                                                           double thresholdRansac = 0.01,
                                                           unsigned int seed = 1,
                                                           unsigned int maxIterationsNonLinear = 10000,
                                                           float thresholdNonLinear = 1e-4f
                                                           )
{
    Eigen::Matrix3d F = estimateFundamentalRansac(points0, points1, inliers, maxIterationsRansac, thresholdRansac, seed);

    //non-linear refinement using Nelder-Mead
    NelderMeadOptFundamental nmf(points0, points1, inliers);

    float F_data_opt[9];
    nmf.run(getLinearArrayFromMatrix(F), 9, thresholdNonLinear, maxIterationsNonLinear, &F_data_opt[0]);
    F = getMatrixdFromLinearArray(F_data_opt, 3, 3);

    return F;
}

/**
 * @brief noramalizeFundamentalMatrix
 * @param F
 * @return
 */
PIC_INLINE Eigen::Matrix3d noramalizeFundamentalMatrix(Eigen::Matrix3d F)
{
    Eigen::JacobiSVD< Eigen::Matrix3d > svdF(F, Eigen::ComputeThinU | Eigen::ComputeThinV);
    Eigen::Matrix3d Uf = svdF.matrixU();
    Eigen::Matrix3d Vf = svdF.matrixV();
    Eigen::Vector3d Df = svdF.singularValues();
    Df[2] = 0.0;

    Eigen::Matrix3d F_new = Uf * DiagonalMatrix(Df) * Eigen::Transpose< Eigen::Matrix3d >(Vf);

    double norm = MAX(Df[0], Df[1]);
    return F_new / norm;
}

/**
 * @brief extractFundamentalMatrix
 * @param M0
 * @param M1
 * @param e0
 * @param e1
 * @return
 */
PIC_INLINE Eigen::Matrix3d extractFundamentalMatrix(Eigen::Matrix34d &M0, Eigen::Matrix34d &M1, Eigen::VectorXd &e0, Eigen::VectorXd &e1) {

    Eigen::Matrix3d M0_3 = getSquareMatrix(M0);
    Eigen::Matrix3d M1_3 = getSquareMatrix(M1);

    Eigen::Matrix3d M0_inv = M0_3.inverse();
    Eigen::Vector3d c0 = - M0_inv * getLastColumn(M0);
    e1 = M1 * addOne(c0);

    Eigen::Matrix3d M1_inv = M1_3.inverse();
    Eigen::Vector3d c1 = - M1_inv * getLastColumn(M1);
    e0 = M0 * addOne(c1);

    Eigen::Matrix3d F;

    F(0, 0) =  0.0;
    F(0, 1) = -e1(2);
    F(0, 2) =  e1(1);

    F(1, 0) =  e1(2);
    F(1, 1) =  0.0;
    F(1, 2) = -e1(0);

    F(2, 0) = -e1(1);
    F(2, 1) =  e1(0);
    F(2, 2) =  0.0;

    F = F * M1_3 * M0_inv;

    Eigen::JacobiSVD< Eigen::Matrix3d > svdF(F, Eigen::ComputeThinU | Eigen::ComputeThinV);
    Eigen::Vector3d Df = svdF.singularValues();

    double norm = MAX(Df[0], MAX(Df[1], Df[2]));
    return F / norm;
}

/**
 * @brief estimateFundamentalFromImages
 * @param img0
 * @param img1
 * @return
 */
PIC_INLINE  Eigen::Matrix3d estimateFundamentalFromImages(Image *img0,
                                                          Image *img1,
                                                          std::vector< Eigen::Vector2f > &m0,
                                                          std::vector< Eigen::Vector2f > &m1,
                                                          std::vector< unsigned int > &inliers)
{
    Eigen::Matrix3d F;
    if(img0 == NULL || img1 == NULL) {
        return F;
    }

    m0.clear();
    m1.clear();
    inliers.clear();

    //corners
    std::vector< Eigen::Vector2f > corners_from_img0;
    std::vector< Eigen::Vector2f > corners_from_img1;

    //compute the luminance images
    Image *L0 = FilterLuminance::execute(img0, NULL, LT_CIE_LUMINANCE);
    Image *L1 = FilterLuminance::execute(img1, NULL, LT_CIE_LUMINANCE);

    //extract corners
    HarrisCornerDetector hcd(2.5f, 5);
    hcd.execute(L0, &corners_from_img0);
    hcd.execute(L1, &corners_from_img1);

    //compute ORB descriptors for each corner and image

    //apply a gaussian filter to luminance images
    Image *L0_flt = FilterGaussian2D::execute(L0, NULL, 2.5f);
    Image *L1_flt = FilterGaussian2D::execute(L1, NULL, 2.5f);

    //compute ORB descriptor
    ORBDescriptor b_desc(31, 512);

    std::vector< unsigned int *> descs0;
    b_desc.getAll(L0_flt, corners_from_img0, descs0);

    std::vector< unsigned int *> descs1;
    b_desc.getAll(L1_flt, corners_from_img1, descs1);

    //match ORB descriptors
    std::vector< Eigen::Vector3i > matches;
    int n = b_desc.getDescriptorSize();

    //BinaryFeatureBruteForceMatcher bffm_bin(&descs1, n);
    BinaryFeatureLSHMatcher bffm_bin(&descs1, n, 64);
    bffm_bin.getAllMatches(descs0, matches);

    //get matches
    FeatureMatcher<unsigned int>::filterMatches(corners_from_img0, corners_from_img1, matches, m0, m1);

    //estimate the fundamental matrix
    F = estimateFundamentalWithNonLinearRefinement(m0, m1, inliers, 1000, 0.5, 1, 1000, 1e-4f);

    delete L0;
    delete L1;
    delete L0_flt;
    delete L1_flt;

    return F;
}

#endif // PIC_DISABLE_EIGEN

} // end namespace pic

#endif // PIC_COMPUTER_VISION_FUNDAMENTAL_HPP


// Software License Agreement (BSD-3-Clause)
//
// Copyright 2018 The University of North Carolina at Chapel Hill
//
// 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 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.

//! @author Jeff Ichnowski

#ifndef LINEAR_H
#define LINEAR_H

#include <Eigen/Dense>
#include <Eigen/Geometry>

namespace nao_cup {
    template <typename S>
    using Transform = Eigen::Transform<S, 3, Eigen::Isometry>;

    template <typename S>
    using Vec2 = Eigen::Matrix<S, 2, 1>;
    template <typename S>
    using Vec3 = Eigen::Matrix<S, 3, 1>;
    template <typename S>
    using Vec4 = Eigen::Matrix<S, 4, 1>;

    template <typename S>
    S v3_len(const Vec3<S> *v) { return v->norm(); }

    template <typename S>
    S v3_dot(const Vec3<S> *a, const Vec3<S> *b) {
        return a->dot(*b);
    }

    template <typename S>
    S v3_dist(const Vec3<S> *a, const Vec3<S> *b) {
        return (*a - *b).norm();
    }

    template <typename S>
    void v3_sub(Vec3<S>* r, const Vec3<S>* a, const Vec3<S>* b) { *r = *a - *b; }

    template <typename S>
    void v3_add(Vec3<S>* r, const Vec3<S>* a, const Vec3<S>* b) { *r = *a + *b; }

    template <typename S>
    void v3_scale(Vec3<S>* r, const Vec3<S> *v, S s) { *r = *v * s; }

    template <typename S>
    void m4_mul(
        Transform<S>* r,
        const Transform<S> *a,
        const Transform<S> *b)
    {
        *r = *a * *b;
    }

    template <typename S>
    void m4_rotate(
        Transform<S> *m,
        const Transform<S> *t,
        S a, S x, S y, S z)
    {
        *m = *t * Eigen::AngleAxis<S>(a, Vec3<S>(x, y, z));
    }

    template <typename S>
    void m4_translate(Transform<S> *m, const Transform<S> *t, S x, S y, S z) {
        *m = *t * Eigen::Translation<S, 3>(x, y, z);
    }

    template <typename S>
    void m4_transform_i3(Vec3<S> *r, const Transform<S> *m, S x, S y, S z) {
        *r = *m * Vec3<S>(x, y, z);
    }

    template <typename S>
    void m4_extract_translation(
        Vec3<S>* r,
        const Transform<S> *m)
    {
        *r = m->translation().template head<3>();
    }

    template <typename S>
    void m4_transform_i(Vec4<S> *r, const Transform<S> *m, S x, S y, S z, S w) {
        *r = *m * Vec4<S>(x, y, z, w);
    }

    template <typename S>
    void m4_transform_i(Vec3<S>* r, const Transform<S> *m, S x, S y, S z) {
        *r = *m * Vec3<S>(x, y, z);
    }

    template <typename S>
    void v3_norm(Vec3<S> *r, const Vec3<S> *a) {
        *r = a->normalized();
    }

    template <typename S>
    S v3_dist_segment_point(const Vec3<S> *s0, const Vec3<S> *s1, const Vec3<S> *pt) {
        Vec3<S> v, w;
        S c1, c2, b;

        v3_sub(&v, s1, s0);
        v3_sub(&w, pt, s0);

        c1 = v3_dot(&w, &v);
        if (c1 <= 0.0) {
            return v3_len(&w);
        }

        c2 = v3_dot(&v, &v);
        if (c2 <= c1) {
            return v3_dist(pt, s1);
        }

        b = c1 / c2;

        v3_scale(&v, &v, b);
        v3_add(&v, s0, &v);

        return v3_dist(&v, pt);
    }
}

#endif /* LINEAR_H */


#include <iostream>
#include <optional>
#include <sstream>
#include <stdexcept>

#include <ceres/ceres.h>
#include <gflags/gflags.h>
#include <glog/logging.h>
#include <Eigen/Dense>
#include <boost/asio.hpp>

#include "farm_ng/calibration/align_sensor_rig.pb.h"
#include "farm_ng/calibration/calibrate_multi_view_apriltag_rig.pb.h"
#include "farm_ng/calibration/calibrate_multi_view_lidar.pb.h"
#include "farm_ng/calibration/local_parameterization.h"
#include "farm_ng/calibration/multi_view_apriltag_rig_calibrator.h"
#include "farm_ng/calibration/multi_view_lidar_model.pb.h"

#include "farm_ng/perception/apriltag.h"
#include "farm_ng/perception/point_cloud.h"

#include "farm_ng/core/blobstore.h"
#include "farm_ng/core/event_log_reader.h"
#include "farm_ng/core/init.h"
#include "farm_ng/core/ipc.h"

#include "farm_ng/perception/pose_utils.h"
#include "farm_ng/perception/tensor.h"

#include "farm_ng/perception/time_series.h"

#include <Eigen/Dense>

DEFINE_string(result, "aligned_sensor_rig.json", "Output path.");
DEFINE_string(output_config, "",
              "output a config json file to the given path.");
DEFINE_string(config, "", "Load config from ajson file.");

DEFINE_string(multi_view_lidar_model, "",
              "Path to the multi_view_lidar_model to align");

DEFINE_int32(floor_tag, 313, "A tag on the floor.");
DEFINE_double(floor_distance_threshold, 0.1,
              "Distance from floor to consider tags belonging to the floor.");

namespace fs = boost::filesystem;

namespace farm_ng::calibration {
Eigen::Matrix3Xd ToMatrix(const std::vector<Eigen::Vector3d>& points) {
  return Eigen::Map<const Eigen::Matrix3Xd>(points[0].data(), 3, points.size());
}
Eigen::Hyperplane<double, 3> FitPlaneToPoints(const Eigen::Matrix3Xd& points) {
  Eigen::Vector3d mean = points.rowwise().mean();
  Eigen::Matrix3Xd points_centered = points.colwise() - mean;
  auto svd =
      points_centered.jacobiSvd(Eigen::ComputeFullU | Eigen::ComputeThinV);
  Eigen::Vector3d normal = svd.matrixU().col(2);
  LOG(INFO) << "p=" << mean.transpose() << " n=" << normal.transpose();
  return Eigen::Hyperplane<double, 3>(normal, mean);
}

// based on https://github.com/nghiaho12/rigid_transform_3D/blob/master/rigid_transform_3D.py
// http://nghiaho.com/?page_id=671
Sophus::SE3d FitAPoseB(const Eigen::Matrix3Xd& points_a,
                       const Eigen::Matrix3Xd& points_b) {
  Eigen::Vector3d mean_a = points_a.rowwise().mean();
  Eigen::Vector3d mean_b = points_b.rowwise().mean();
  Eigen::Matrix3Xd points_a_centered = points_a.colwise() - mean_a;
  Eigen::Matrix3Xd points_b_centered = points_b.colwise() - mean_b;
  Eigen::MatrixXd a_H_b = points_a_centered * points_b_centered.transpose();
  auto svd = a_H_b.jacobiSvd(Eigen::ComputeFullU | Eigen::ComputeFullV);
  Eigen::MatrixXd V = svd.matrixV();
  Eigen::MatrixXd U = svd.matrixU();
  Eigen::Matrix3d b_R_a = V * U.transpose();
  if (b_R_a.determinant() < 0) {
    V.col(2) *= -1;
    b_R_a = V * U.transpose();
    CHECK_GE(b_R_a.determinant(), 0);
  }
  Eigen::Matrix4d b_pose_a = Eigen::Matrix4d::Identity();
  b_pose_a.block<3, 3>(0, 0) = b_R_a;
  b_pose_a.block<3, 1>(0, 3) = mean_b - b_R_a * mean_a;
  auto a_pose_b = Sophus::SE3d(b_pose_a).inverse();
  return a_pose_b;
}

int align_sensor_rig(AlignSensorRigConfiguration config) {
  CHECK_GT(config.floor_tag_ids_size(), 0);
  MultiViewLidarModel lidar_model =
      core::ReadProtobufFromResource<MultiViewLidarModel>(config.model());

  perception::PoseGraph pose_graph;
  for (const perception::ApriltagRig::Node& node :
       lidar_model.apriltag_rig().nodes()) {
    pose_graph.AddPose(node.pose());
  }
  std::string floor_tag_name = perception::FrameRigTag(
      lidar_model.apriltag_rig().name(), config.floor_tag_ids(0));
  perception::PoseGraph floor_tag_graph =
      pose_graph.AveragePoseGraph(floor_tag_name);
  double floor_distance_thresh = config.floor_distance_threshold();
  std::vector<Eigen::Vector3d> floor_points_tag_rig;

  for (const perception::ApriltagRig::Node& node :
       lidar_model.apriltag_rig().nodes()) {
    auto floor_tag_pose_tag =
        floor_tag_graph.CheckAverageAPoseB(floor_tag_name, node.frame_name());
    double err = 0;
    auto points_tag = perception::PointsTag(node);
    for (auto point_tag : points_tag) {
      // LOG(INFO) << "point_tag: " << point_tag.transpose();
      err = std::max(err, std::abs((floor_tag_pose_tag * point_tag).z()));
    }
    if (err < floor_distance_thresh) {
      LOG(INFO) << node.frame_name() << " on floor err: " << err;
      Sophus::SE3d tag_rig_pose_tag = ProtoToSophus(
          node.pose(), lidar_model.apriltag_rig().name(), node.frame_name());
      for (auto point_tag : points_tag) {
        floor_points_tag_rig.push_back(tag_rig_pose_tag * point_tag);
      }
    } else {
      // LOG(INFO) << node.frame_name() << " err : " << err;
    }
  }

  std::vector<Eigen::Vector3d> points_camera_rig;

  for (int i = 0; i < lidar_model.measurements_size(); ++i) {
    const auto& m_i0 = lidar_model.measurements(i);
    Sophus::SE3d camera_rig_pose_apriltag_rig = ProtoToSophus(
        m_i0.camera_rig_pose_apriltag_rig(), lidar_model.camera_rig().name(),
        lidar_model.apriltag_rig().name());
    for (auto floor_point_tag_rig : floor_points_tag_rig) {
      points_camera_rig.push_back(camera_rig_pose_apriltag_rig *
                                  floor_point_tag_rig);
    }
  }
  perception::PoseGraph sensor_rig_graph;
  sensor_rig_graph.AddPoses(lidar_model.lidar_poses());
  auto plane_camera_rig = FitPlaneToPoints(ToMatrix(points_camera_rig));
  for (auto pose : lidar_model.lidar_poses()) {
    auto sensor_name = pose.frame_b();
    auto camera_rig_pose_sensor = sensor_rig_graph.CheckAverageAPoseB(
        lidar_model.camera_rig().name(), sensor_name);
    LOG(INFO) << "sensor_name: " << sensor_name << " distance to ground: "
              << plane_camera_rig.signedDistance(
                     camera_rig_pose_sensor.translation());
  }

  std::vector<Eigen::Vector3d> points_base;
  std::vector<Eigen::Vector3d> points_rig;
  Eigen::Hyperplane<double, 3> plane_base(Eigen::Vector3d(0, 0, 1.0),
                                          Eigen::Vector3d(0, 0, 0));
  for (auto pose : config.base_pose_sensor_measured()) {
    Sophus::SE3d base_pose_sensor =
        perception::ProtoToSophus(pose, config.base_frame(), pose.frame_b());
    points_base.push_back(base_pose_sensor.translation());
    points_base.push_back(
        plane_base.projection(base_pose_sensor.translation()));

    Sophus::SE3d rig_pose_sensor = sensor_rig_graph.CheckAverageAPoseB(
        lidar_model.camera_rig().name(), pose.frame_b());
    points_rig.push_back(rig_pose_sensor.translation());
    points_rig.push_back(
        plane_camera_rig.projection(rig_pose_sensor.translation()));
  }
  Sophus::SE3d base_pose_rig =
      FitAPoseB(ToMatrix(points_base), ToMatrix(points_rig));
  AlignSensorRigResult result;
  result.set_base_frame(config.base_frame());
  for (auto pose : lidar_model.lidar_poses()) {
    auto sensor_name = pose.frame_b();
    auto camera_rig_pose_sensor = sensor_rig_graph.CheckAverageAPoseB(
        lidar_model.camera_rig().name(), sensor_name);
    auto base_pose_sensor = base_pose_rig * camera_rig_pose_sensor;
    perception::SophusToProto(base_pose_sensor, config.base_frame(),
                              sensor_name, result.add_base_pose_sensor());
    LOG(INFO) << result.base_pose_sensor(result.base_pose_sensor_size() - 1)
                     .ShortDebugString();
  }

  farm_ng::core::WriteProtobufToJsonFile(FLAGS_result, result);

  return 0;
}

}  // namespace farm_ng::calibration

int main(int argc, char* argv[]) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  FLAGS_logtostderr = 1;
  std::string filename = boost::filesystem::path(argv[0]).filename().string();
  google::InitGoogleLogging(filename.c_str());
  farm_ng::calibration::AlignSensorRigConfiguration config;
  if (!FLAGS_config.empty()) {
    config = farm_ng::core::ReadProtobufFromJsonFile<
        farm_ng::calibration::AlignSensorRigConfiguration>(FLAGS_config);
  } else {
    config.set_floor_distance_threshold(FLAGS_floor_distance_threshold);
    config.add_floor_tag_ids(FLAGS_floor_tag);
    config.mutable_model()->set_path(FLAGS_multi_view_lidar_model);
    config.mutable_model()->set_content_type(
        farm_ng::core::ContentTypeProtobufBinary<
            farm_ng::calibration::MultiViewLidarModel>());
  }
  if (!FLAGS_output_config.empty()) {
    farm_ng::core::WriteProtobufToJsonFile(FLAGS_output_config, config);
    return 0;
  }
  return farm_ng::calibration::align_sensor_rig(config);
}


#include <algorithm>
//#include <boost/iterator/counting_iterator.hpp>
#include <ibs>
#include <iterator>
#include <math.h>
#include <numeric>
#include <vector>
// gets the particle times and uses the phase acceptance and number of bins
// to generate an integer which is than binned/histogram by
// ParticleTimesToHistogram
struct ParticleTimesToInteger {
  /*
  0 -> tauahat = half the phase acceptance
  1 -> nbins = numbers of bins to use in binning
  2 -> t synchronous to shift the center of the distribution to zero for the
  binning
  */
  int operator()(std::vector<double> &data, std::vector<double> &params) const {
    double out;
    double dtsamp2 = 2 * params[0] / params[1];
    // std::printf("%-30s %12.8e\n", "dtsamp2", dtsamp2);
    out = (data[4] - params[2] + params[0]) / dtsamp2;
    // std::printf("%-30s %12.8e\n", "data[4]", data[4]);
    // std::printf("%-30s %12.8e\n", "out doub", out);
    out = (int)(out + 0.5);
    // std::printf("%-30s %12.8e\n", "out int", out);
    return out;
  }
};

void denseHistogram(std::vector<int> data, std::vector<int> &histogram,
                    int nbins, double tauhat) {
  std::vector<int> input = data;
  std::sort(input.begin(), input.end());

  int numBins = nbins + 1;
  histogram.resize(numBins);
  for (int i = 0; i < numBins; i++) {
    auto upper = std::upper_bound(input.begin(), input.end(), i);
    histogram[i] = std::distance(input.begin(), upper);
  }
  /*
    for (std::vector<int>::iterator i = histogram.begin(); i != histogram.end();
         i++) {
      std::printf("%12i\n", *i);
    }
  */
  std::adjacent_difference(histogram.begin(), histogram.end(),
                           histogram.begin());
}

std::vector<int>
ParticlesTimesToHistogram(std::vector<std::vector<double>> &data, int nbins,
                          double tauhat, double ts) {

  int n = data.size();
  std::vector<int> timecomponent(n);
  std::vector<int> histogram;
  std::vector<double> param;

  param.push_back(tauhat);
  param.push_back(nbins);
  param.push_back(ts);
  std::vector<std::vector<double>> p(n);
  std::fill(p.begin(), p.end(), param);
  /*
  std::printf("%-30s %12.8e\n", "tauhat", tauhat);
  std::printf("%-30s %12i\n", "nbins", nbins);
  std::printf("%-30s %12.8e\n", "ts", ts);

  for (std::vector<double>::iterator i = param.begin(); i != param.end(); i++) {
    std::printf("%12.8e\n", *i);
  }

  // print out for debug
  std::for_each(p.begin(), p.end(), [](std::vector<double> &particle) {
    std::printf("%12.8e %12.8e %12.8e \n", particle[0], particle[1],
                particle[2]);
  });

  std::printf("\n");
*/
  std::transform(data.begin(), data.end(), p.begin(), timecomponent.begin(),
                 ParticleTimesToInteger());
  /*
    for (std::vector<int>::iterator i = timecomponent.begin();
         i != timecomponent.end(); i++) {
      std::printf("%12i\n", *i);
    }
    */
  denseHistogram(timecomponent, histogram, nbins, tauhat);
  /*
  for (std::vector<int>::iterator i = histogram.begin(); i != histogram.end();
       i++) {
    std::printf("%3i", *i);
  };
  std::printf("\n");
  */
  return histogram;
}

/*
 ********************************************************************************
 ********************************************************************************
 * REF:
 * https://stackoverflow.com/questions/14924912/computing-column-sums-of-matrix-vectorvectordouble-with-iterators
 ********************************************************************************
 */
std::vector<double> getColumnMeans(std::vector<std::vector<double>> &dist) {
  std::vector<double> colsums(dist[0].size());

  std::for_each(dist.begin(), dist.end(), [&](const std::vector<double> &row) {
    std::transform(
        row.begin(), row.end(), colsums.begin(), colsums.begin(),
        [&](double d1, double d2) { return (d1 + d2) / dist.size(); });
  });

  return colsums;
}

std::vector<double> vectorMultiply(std::vector<double> x,
                                   std::vector<double> y) {
  std::transform(x.begin(), x.end(), y.begin(), x.begin(),
                 std::multiplies<double>());
  return x;
}

std::vector<double> vectorAdd(std::vector<double> x, std::vector<double> y) {
  std::transform(x.begin(), x.end(), y.begin(), x.begin(), std::plus<double>());
  return x;
}

std::vector<double> vectorSub(std::vector<double> x, std::vector<double> y) {
  std::transform(x.begin(), x.end(), y.begin(), x.begin(),
                 std::minus<double>());
  return x;
}
std::vector<double> CalcRMS(std::vector<std::vector<double>> &dist) {
  std::vector<double> avg = getColumnMeans(dist);

  // for (std::vector<double>::const_iterator i = avg.begin(); i != avg.end();
  // ++i)
  //  std::printf("%-30s %12.8e\n", "avg", *i);

  std::vector<std::vector<double>> avgarr(dist.size());
  std::fill(avgarr.begin(), avgarr.end(), avg);

  // std::for_each(avgarr.begin(), avgarr.end(), [](std::vector<double> &a) {
  //  std::printf("%12.8e %12.8e %12.8e %12.8e %12.8e %12.8e\n", a[0], a[1],
  //  a[2],
  //              a[3], a[4], a[5]);
  //});
  std::vector<std::vector<double>> distcopy = dist;
  std::transform(dist.begin(), dist.end(), avgarr.begin(), distcopy.begin(),
                 vectorSub);
  std::transform(distcopy.begin(), distcopy.end(), distcopy.begin(),
                 distcopy.begin(), vectorMultiply);

  std::vector<double> MS = getColumnMeans(distcopy);
  std::transform(MS.begin(), MS.end(), MS.begin(), (double (*)(double))sqrt);

  return MS;
}

std::vector<double> HistogramToSQRTofCumul(std::vector<int> inputHistogram,
                                           double coeff) {

  /* The name is a badly chose here as it was originally used to calculate
   * cumulated distributions which      */
  /* turned out to be not necessary - the function takes a histogram as input,
   * multiplies it with a constant  */
  /* vector before taking the sqrt of each element. This produces a vector
   * that is used in the IBSNew routine */
  /* to multiply with particle momenta representing the IBS contribution */
  int n = inputHistogram.size();

  std::vector<double> vcoeff(n);
  std::fill(vcoeff.begin(), vcoeff.end(), coeff);
  /*
  for (std::vector<double>::const_iterator i = vcoeff.begin();
       i != vcoeff.end(); ++i)
    std::printf("%-30s %12.8e\n", "avg", *i);
    */
  // fill constant vector

  // multiply with constant
  std::transform(inputHistogram.begin(), inputHistogram.end(), vcoeff.begin(),
                 vcoeff.begin(), std::multiplies<double>());

  // take sqrt
  std::transform(vcoeff.begin(), vcoeff.end(), vcoeff.begin(),
                 (double (*)(double))sqrt);

  return vcoeff;
}

std::map<std::string, double> readInput(std::string filename) {
  std::vector<std::string> ALLOWEDKEYS{
      "bucket",      "atomNumber", "charge",    "nMacro",
      "nReal",       "nbins",      "sigs",      "seed",
      "ex",          "ey",         "timeRatio", "fracibstot",
      "ibsCoupling", "model",      "nturns",    "nwrite"};

  std::map<std::string, double> out;
  std::string line;
  std::ifstream file(filename);

  std::getline(file, line);

  // check if file is open
  if (file.is_open()) {
    std::string key;
    double value;
    std::istringstream iss(line);
    iss >> key >> value;

    vector<string>::iterator it =
        find(ALLOWEDKEYS.begin(), ALLOWEDKEYS.end(), key);
    // cout << key << " " << value << " " << endl;
    if (it != ALLOWEDKEYS.end()) {
      out[key] = value;
    }
    while (!file.eof()) {
      std::getline(file, line);
      std::istringstream iss(line);
      iss >> key >> value;

      vector<string>::iterator it =
          find(ALLOWEDKEYS.begin(), ALLOWEDKEYS.end(), key);
      // cout << key << " " << value << " " << endl;
      if (it != ALLOWEDKEYS.end()) {
        out[key] = value;
      }
    }

    file.close();
  }
  return out;
}

bool isInLong(std::vector<double> particle, double tauhat, double synctime) {
  return !(abs(particle[4] - synctime) < tauhat);
}


#include "drake/solvers/fbstab/components/riccati_linear_solver.h"

#include <cmath>
#include <iostream>
#include <stdexcept>

#include <Eigen/Dense>

#include "drake/solvers/fbstab/components/mpc_data.h"
#include "drake/solvers/fbstab/components/mpc_residual.h"
#include "drake/solvers/fbstab/components/mpc_variable.h"

namespace drake {
namespace solvers {
namespace fbstab {

using VectorXd = Eigen::VectorXd;
using MatrixXd = Eigen::MatrixXd;

RiccatiLinearSolver::RiccatiLinearSolver(int N, int nx, int nu, int nc) {
  if (N <= 0 || nx <= 0 || nu <= 0 || nc <= 0) {
    throw std::runtime_error(
        "In RiccatiLinearSolver::RiccatiLinearSolver: all inputs must be "
        "positive.");
  }
  N_ = N;
  nx_ = nx;
  nu_ = nu;
  nc_ = nc;
  nz_ = (N_ + 1) * (nx_ + nu_);
  nl_ = (N_ + 1) * (nx_);
  nv_ = (N_ + 1) * (nc_);

  Q_.resize(N + 1);
  S_.resize(N + 1);
  R_.resize(N + 1);

  P_.resize(N + 1);
  SG_.resize(N + 1);
  M_.resize(N + 1);
  L_.resize(N + 1);
  SM_.resize(N + 1);
  AM_.resize(N + 1);

  h_.resize(N + 1);
  th_.resize(N + 1);

  for (int i = 0; i < N + 1; i++) {
    Q_[i].resize(nx, nx);
    S_[i].resize(nu, nx);
    R_[i].resize(nu, nu);

    P_[i].resize(nx, nu);
    SG_[i].resize(nu, nu);
    M_[i].resize(nx, nx);
    L_[i].resize(nx, nx);
    SM_[i].resize(nu, nx);
    AM_[i].resize(nx, nx);

    h_[i].resize(nx);
    th_[i].resize(nx);
  }

  gamma_.resize(nv_);
  mus_.resize(nv_);
  Gamma_.resize(nc, N + 1);

  Etemp_.resize(nc, nx);
  Ltemp_.resize(nc, nu);
  Linv_.resize(nx, nx);

  tx_.resize(nx);
  tu_.resize(nu);
  tl_.resize(nx);
  r1_.resize(nz_);
  r2_.resize(nl_);
  r3_.resize(nv_);
}

bool RiccatiLinearSolver::Initialize(const MpcVariable& x,
                                     const MpcVariable& xbar, double sigma) {
  const MpcData* const data = x.data();
  if (xbar.data_ != data) {
    throw std::runtime_error(
        "In RiccatiLinearSolver::Initialize: x and xbar have mismatched "
        "problem data.");
  }
  if (!MpcVariable::SameSize(x, xbar)) {
    throw std::runtime_error(
        "In RiccatiLinearSolver::Initialize: x and xbar are not the same "
        "size.");
  }
  if (sigma <= 0) {
    throw std::runtime_error(
        "In RiccatiLinearSolver::Initialize: sigma must be positive.");
  }

  Eigen::Vector2d temp;
  for (int i = 0; i < nv_; i++) {
    const double ys = x.y()(i) + sigma * (x.v()(i) - xbar.v()(i));
    temp = PFBGradient(ys, x.v()(i));

    gamma_(i) = temp(0);
    mus_(i) = temp(1) + sigma * temp(0);
    Gamma_(i) = gamma_(i) / mus_(i);
  }
  // Compute the barrier augmented Hessian.
  for (int i = 0; i < N_ + 1; i++) {
    const MatrixXd& Ei = (*data->E_)[i];
    const MatrixXd& Li = (*data->L_)[i];

    Q_[i].triangularView<Eigen::Lower>() =
        (*data->Q_)[i] + sigma * MatrixXd::Identity(nx_, nx_);
    R_[i].triangularView<Eigen::Lower>() =
        (*data->R_)[i] + sigma * MatrixXd::Identity(nu_, nu_);
    S_[i] = (*data->S_)[i];

    // Add barriers associated with E(i)x(i) + L(i)u(i) + d(i) <=0
    // Q(i) += E(i)'*diag(Gamma(i))*E(i)
    Etemp_.noalias() = Gamma_.col(i).asDiagonal() * Ei;
    Q_[i].triangularView<Eigen::Lower>() += Ei.transpose() * Etemp_;

    // R(i) += L(i)'*diag(Gamma(i))*L(i)
    Ltemp_.noalias() = Gamma_.col(i).asDiagonal() * Li;
    R_[i].triangularView<Eigen::Lower>() += Li.transpose() * Ltemp_;

    // S(i) += L(i) ' * diag(Gamma(i)) * E(i)
    S_[i].noalias() += Li.transpose() * Etemp_;
  }

  // Begin the matrix potion of the Riccati recursion.
  // Base case: L(0) = chol(sigma*I).
  L_[0] = sqrt(sigma) * MatrixXd::Identity(nx_, nx_);

#define FBSTAB_LLT_CHECK(llt)           \
  {                                     \
    if (llt.info() != Eigen::Success) { \
      return false;                     \
    }                                   \
  }

  for (int i = 0; i < N_; i++) {
    // Compute inv(L(i)) then
    // compute QQ = Q+inv(L*L') = Q + inv(L)'*inv(L)
    // and factor M = chol(QQ) in place.
    Linv_ = MatrixXd::Identity(nx_, nx_);
    L_[i].triangularView<Eigen::Lower>().solveInPlace(Linv_);
    L_[i].triangularView<Eigen::Lower>().transpose().solveInPlace(Linv_);
    M_[i].triangularView<Eigen::Lower>() = Q_[i] + Linv_;
    Eigen::LLT<Eigen::Ref<MatrixXd>> llt1(M_[i]);
    FBSTAB_LLT_CHECK(llt1);

    // Compute AM = A*inv(M)'.
    AM_[i] = (*data->A_)[i];
    M_[i]
        .triangularView<Eigen::Lower>()
        .transpose()
        .solveInPlace<Eigen::OnTheRight>(AM_[i]);

    // Compute SM = S*inv(M)'.
    SM_[i] = S_[i];
    M_[i]
        .triangularView<Eigen::Lower>()
        .transpose()
        .solveInPlace<Eigen::OnTheRight>(SM_[i]);

    // Factor SG = chol(R - SM*SM') in place.
    SG_[i].noalias() = R_[i] - SM_[i] * SM_[i].transpose();
    Eigen::LLT<Eigen::Ref<MatrixXd>> llt2(SG_[i]);
    FBSTAB_LLT_CHECK(llt2);

    // Compute P = (A*inv(QQ)S' - B)*inv(SG)',
    //           = (AM*SM' - B)*inv(SG)'.
    P_[i].noalias() = AM_[i] * SM_[i].transpose();
    P_[i] -= (*data->B_)[i];
    SG_[i]
        .triangularView<Eigen::Lower>()
        .transpose()
        .solveInPlace<Eigen::OnTheRight>(P_[i]);

    // Compute L(i+1) = chol(Pi)
    // where Pi = P*P' + AM*AM' + sigma I.
    L_[i + 1] = sigma * MatrixXd::Identity(nx_, nx_);
    L_[i + 1].noalias() += P_[i] * P_[i].transpose();
    L_[i + 1].noalias() += AM_[i] * AM_[i].transpose();
    Eigen::LLT<Eigen::Ref<MatrixXd>> llt3(L_[i + 1]);
    FBSTAB_LLT_CHECK(llt3);
  }

  // Finish the recursion, i.e., perform the i = N step.
  Linv_ = MatrixXd::Identity(nx_, nx_);
  L_[N_].triangularView<Eigen::Lower>().solveInPlace(Linv_);
  L_[N_].triangularView<Eigen::Lower>().transpose().solveInPlace(Linv_);

  // Compute M = chol(Q + inv(L*L')).
  M_[N_].triangularView<Eigen::Lower>() = Q_[N_] + Linv_;
  Eigen::LLT<Eigen::Ref<MatrixXd>> llt4(M_[N_]);
  FBSTAB_LLT_CHECK(llt4);

  // Compute SM = S*inv(M)'.
  SM_[N_] = S_[N_];
  M_[N_]
      .triangularView<Eigen::Lower>()
      .transpose()
      .solveInPlace<Eigen::OnTheRight>(SM_[N_]);

  // Compute SG = chol(R - SM*SM').
  SG_[N_].noalias() = R_[N_] - SM_[N_] * SM_[N_].transpose();
  Eigen::LLT<Eigen::Ref<MatrixXd>> llt5(SG_[N_]);
  FBSTAB_LLT_CHECK(llt5);

#undef FBSTAB_LLT_CHECK
  return true;
}

bool RiccatiLinearSolver::Solve(const MpcResidual& r, MpcVariable* dx) const {
  const MpcData* const data = dx->data();
  if (r.nz_ != dx->nz_ || r.nl_ != dx->nl_ || r.nv_ != dx->nv_) {
    throw std::runtime_error(
        "In RiccatiLinearSolver::Solve: r and dx size mismatch.");
  }
  // Compute the post-elimination residual,
  // r1 = rz - A'*(rv./mus) and r2 = -rl.
  r1_ = r.z_;
  r3_ = r.v_.cwiseQuotient(mus_);  // r3_ is used as a temp here
  data->gemvAT(r3_, -1.0, 1.0, &r1_);
  r2_ = -r.l_;
  // Get reshaped aliases for r1 and r2.
  Eigen::Map<MatrixXd> r1(r1_.data(), nx_ + nu_, N_ + 1);
  Eigen::Map<MatrixXd> r2(r2_.data(), nx_, N_ + 1);

  // Begin the vector portion of the Riccati recursion.
  // Base case: theta(0) = -rl(0), h(0) = inv(L*L')*theta(0) - rx(0).
  th_[0] = r2.col(0);
  h_[0] = th_[0];
  L_[0].triangularView<Eigen::Lower>().solveInPlace(h_[0]);
  L_[0].triangularView<Eigen::Lower>().transpose().solveInPlace(h_[0]);
  h_[0].noalias() -= r1.block(0, 0, nx_, 1);  // r1(0) = [rx(0);ru(0)]

  // Main loop:
  for (int i = 0; i < N_; i++) {
    // Compute theta(i+1).
    // tx = inv(M)*h
    tx_ = h_[i];
    M_[i].triangularView<Eigen::Lower>().solveInPlace(tx_);

    // tu = inv(SG)*(SM*tx + ru)
    // r1(i) = [rx(i);ru(i)], block extracts ru(i)
    tu_.noalias() = SM_[i] * tx_;
    tu_.noalias() += r1.block(nx_, i, nu_, 1);
    SG_[i].triangularView<Eigen::Lower>().solveInPlace(tu_);

    const auto rlp = r2.col(i + 1);
    th_[i + 1].noalias() = P_[i] * tu_ + AM_[i] * tx_;
    th_[i + 1].noalias() += rlp;

    // Compute h(i+1).
    const auto rxp = r1.block(0, i + 1, nx_, 1);
    h_[i + 1] = th_[i + 1];
    L_[i + 1].triangularView<Eigen::Lower>().solveInPlace(h_[i + 1]);
    L_[i + 1].triangularView<Eigen::Lower>().transpose().solveInPlace(
        h_[i + 1]);
    h_[i + 1].noalias() -= rxp;
  }

  // Begin the backwards recursion for the solution
  // by computing xN,uN, and lN.
  // u(N) = inv(SG*SG')*(SM*inv(M)*h + ru)
  tx_ = h_[N_];
  M_[N_].triangularView<Eigen::Lower>().solveInPlace(tx_);
  tu_.noalias() = SM_[N_] * tx_;
  tu_ += r1.block(nx_, N_, nu_, 1);
  SG_[N_].triangularView<Eigen::Lower>().solveInPlace(tu_);
  SG_[N_].triangularView<Eigen::Lower>().transpose().solveInPlace(tu_);

  // x(N) = -inv(M')*(inv(M)*h + SM'*u(N))
  tx_ = h_[N_];
  M_[N_].triangularView<Eigen::Lower>().solveInPlace(tx_);
  tx_.noalias() += SM_[N_].transpose() * tu_;
  M_[N_].triangularView<Eigen::Lower>().transpose().solveInPlace(tx_);
  tx_ *= -1.0;

  // l(N) = -inv(L*L')* (xN + theta)
  tl_ = tx_ + th_[N_];
  L_[N_].triangularView<Eigen::Lower>().solveInPlace(tl_);
  L_[N_].triangularView<Eigen::Lower>().transpose().solveInPlace(tl_);
  tl_ *= -1.0;

  // Copy these into the solution vector.
  // Using reshaped aliases to make indexing easier.
  Eigen::Map<MatrixXd> dz(dx->z_->data(), nx_ + nu_, N_ + 1);
  Eigen::Map<MatrixXd> dl(dx->l_->data(), nx_, N_ + 1);

  dz.block(0, N_, nx_, 1) = tx_;
  dz.block(nx_, N_, nu_, 1) = tu_;
  dl.col(N_) = tl_;

  // The main backwards recursion loop.
  for (int i = N_ - 1; i >= 0; i--) {
    // Solve SG'*u(i) = inv(SG)*(SM*inv(M)*h + ru) + P'*l(i+1)
    tx_ = h_[i];
    M_[i].triangularView<Eigen::Lower>().solveInPlace(tx_);

    // This is an alias.
    auto ui = dz.block(nx_, i, nu_, 1);  // dz(i) = [xi;ui], extract ui
    ui.noalias() = SM_[i] * tx_;
    ui += r1.block(nx_, i, nu_, 1);  // SM*tx + ru
    SG_[i].triangularView<Eigen::Lower>().solveInPlace(ui);
    ui.noalias() += P_[i].transpose() * dl.col(i + 1);
    SG_[i].triangularView<Eigen::Lower>().transpose().solveInPlace(ui);

    // Solve -M'*x(i) = inv(M)*h + SM'*u(i) + AM'*l(i+1)
    // This is an alias.
    auto xi = dz.block(0, i, nx_, 1);

    xi = h_[i];
    M_[i].triangularView<Eigen::Lower>().solveInPlace(xi);
    xi.noalias() += SM_[i].transpose() * ui;
    xi.noalias() += AM_[i].transpose() * dl.col(i + 1);
    M_[i].triangularView<Eigen::Lower>().transpose().solveInPlace(xi);
    xi *= -1.0;

    // Solve -L*L' * l(i) = theta(i) + x(i).
    auto li = dl.col(i);
    li = th_[i] + xi;
    L_[i].triangularView<Eigen::Lower>().solveInPlace(li);
    L_[i].triangularView<Eigen::Lower>().transpose().solveInPlace(li);
    li *= -1.0;
  }

  // Recover the inequality duals by solving
  // diag(mus)* dv = (rv + diag(gamma)*A*dz).
  VectorXd& dv = dx->v();
  dv = r.v_;
  // r3_ = A*dz, r3_ is being used as a temp.
  data->gemvA(dx->z(), 1.0, 0.0, &r3_);
  dv += gamma_.asDiagonal() * r3_;
  dv = dv.cwiseQuotient(mus_);

  // Compute dy = b - A*dz.
  VectorXd& dy = dx->y();
  data->gemvA(dx->z(), -1.0, 0.0, &dy);
  data->axpyb(1.0, &dy);

  return true;
}

Eigen::Vector2d RiccatiLinearSolver::PFBGradient(double a, double b) const {
  const double r = sqrt(a * a + b * b);
  const double d = 1.0 / sqrt(2.0);

  Eigen::Vector2d v;
  if (r < zero_tolerance_) {
    v(0) = alpha_ * (1.0 - d);
    v(1) = alpha_ * (1.0 - d);

  } else if ((a > 0) && (b > 0)) {
    v(0) = alpha_ * (1.0 - a / r) + (1.0 - alpha_) * b;
    v(1) = alpha_ * (1.0 - b / r) + (1.0 - alpha_) * a;

  } else {
    v(0) = alpha_ * (1.0 - a / r);
    v(1) = alpha_ * (1.0 - b / r);
  }

  return v;
}

}  // namespace fbstab
}  // namespace solvers
}  // namespace drake