/****************************************************************************
 * Copyright (c) 2021 by the Picasso authors                                *
 * All rights reserved.                                                     *
 *                                                                          *
 * This file is part of the Picasso library. Picasso is distributed under a *
 * BSD 3-clause license. For the licensing terms see the LICENSE file in    *
 * the top-level directory.                                                 *
 *                                                                          *
 * SPDX-License-Identifier: BSD-3-Clause                                    *
 ****************************************************************************/

#ifndef PICASSO_FACETGEOMETRY_HPP
#define PICASSO_FACETGEOMETRY_HPP

#include <Picasso_BatchedLinearAlgebra.hpp>

#include <Kokkos_Core.hpp>
#include <Kokkos_Random.hpp>

#include <boost/property_tree/ptree.hpp>

#include <cfloat>
#include <cmath>
#include <fstream>
#include <unordered_map>
#include <vector>

namespace Picasso
{
//---------------------------------------------------------------------------//
template <class MemorySpace>
struct FacetGeometryData
{
    // Get the number of volumes.
    KOKKOS_FUNCTION
    int numVolume() const { return volume_offsets.extent( 0 ); }

    // Given a local volume id get the facets associated with the volume.
    KOKKOS_FUNCTION
    Kokkos::View<float* [4][3], MemorySpace>
    volumeFacets( const int volume_id ) const
    {
        Kokkos::pair<int, int> facet_bounds(
            ( 0 == volume_id ) ? 0 : volume_offsets( volume_id - 1 ),
            volume_offsets( volume_id ) );
        return Kokkos::subview( volume_facets, facet_bounds, Kokkos::ALL(),
                                Kokkos::ALL() );
    }

    // Get the number of surfaces.
    KOKKOS_FUNCTION
    int numSurface() const { return surface_offsets.extent( 0 ); }

    // Given a local surface id get the facets associated with the surface.
    KOKKOS_FUNCTION
    Kokkos::View<float* [4][3], MemorySpace>
    surfaceFacets( const int surface_id ) const
    {
        Kokkos::pair<int, int> facet_bounds(
            ( 0 == surface_id ) ? 0 : surface_offsets( surface_id - 1 ),
            surface_offsets( surface_id ) );
        return Kokkos::subview( surface_facets, facet_bounds, Kokkos::ALL(),
                                Kokkos::ALL() );
    }

    // Volume facets. Ordered as (facet,vector,dim) where vector=0,1,2 are the
    // vertices and vector=3 is the unit normal facing outward from the
    // volume.
    Kokkos::View<float* [4][3], MemorySpace> volume_facets;

    // Volume face offsets. Inclusive scan of volume counts giving the offset
    // into the facet array for each volume.
    Kokkos::View<int*, MemorySpace> volume_offsets;

    // Surface facets. Ordered as (facet,vector,dim) where vector=0,1,2 are the
    // vertices and vector=3 is the unit normal facing outward from the
    // surface.
    Kokkos::View<float* [4][3], MemorySpace> surface_facets;

    // Surface face offsets. Inclusive scan of surface counts giving the offset
    // into the facet array for each surface.
    Kokkos::View<int*, MemorySpace> surface_offsets;

    // Local id of the volume the representes the global axis-aligned bounding
    // box.
    int global_bounding_volume_id;

    // Axis aligned bounding box for all volumes.
    Kokkos::View<float* [6], MemorySpace> volume_bounding_boxes;
};

//---------------------------------------------------------------------------//
template <class MemorySpace>
class FacetGeometry
{
  public:
    // Compute the axis-aligned bounding box of a volume.
    struct BoundingBoxReduce
    {
        typedef float value_type[];
        typedef Kokkos::View<float* [4][3], MemorySpace> facet_view_type;
        typedef Kokkos::View<int*, MemorySpace> offset_view_type;
        typedef typename facet_view_type::size_type size_type;
        size_type value_count;

        facet_view_type facets;
        offset_view_type offsets;
        int volume_id;

        BoundingBoxReduce( const facet_view_type& f, const offset_view_type& o,
                           int v )
            : value_count( 6 )
            , facets( f )
            , offsets( o )
            , volume_id( v )
        {
        }

        KOKKOS_FUNCTION
        void operator()( const size_type facet_id, value_type result ) const
        {
            auto f_start = ( volume_id == 0 ) ? 0 : offsets( volume_id - 1 );
            auto f = f_start + facet_id;
            result[0] =
                fmin( result[0],
                      fmin( facets( f, 0, 0 ),
                            fmin( facets( f, 1, 0 ), facets( f, 2, 0 ) ) ) );
            result[1] =
                fmin( result[1],
                      fmin( facets( f, 0, 1 ),
                            fmin( facets( f, 1, 1 ), facets( f, 2, 1 ) ) ) );
            result[2] =
                fmin( result[2],
                      fmin( facets( f, 0, 2 ),
                            fmin( facets( f, 1, 2 ), facets( f, 2, 2 ) ) ) );
            result[3] =
                fmax( result[3],
                      fmax( facets( f, 0, 0 ),
                            fmax( facets( f, 1, 0 ), facets( f, 2, 0 ) ) ) );
            result[4] =
                fmax( result[4],
                      fmax( facets( f, 0, 1 ),
                            fmax( facets( f, 1, 1 ), facets( f, 2, 1 ) ) ) );
            result[5] =
                fmax( result[5],
                      fmax( facets( f, 0, 2 ),
                            fmax( facets( f, 1, 2 ), facets( f, 2, 2 ) ) ) );
        }

        KOKKOS_FUNCTION
        void join( volatile value_type dst,
                   const volatile value_type src ) const
        {
            dst[0] = fmin( dst[0], src[0] );
            dst[1] = fmin( dst[1], src[1] );
            dst[2] = fmin( dst[2], src[2] );
            dst[3] = fmax( dst[3], src[3] );
            dst[4] = fmax( dst[4], src[4] );
            dst[5] = fmax( dst[5], src[5] );
        }

        KOKKOS_FUNCTION
        void init( value_type v ) const
        {
            v[0] = FLT_MAX;
            v[1] = FLT_MAX;
            v[2] = FLT_MAX;
            v[3] = -FLT_MAX;
            v[4] = -FLT_MAX;
            v[5] = -FLT_MAX;
        }
    };

  public:
    using memory_space = MemorySpace;

    // Default constructor.
    FacetGeometry() = default;

    // Create the geometry from an ASCII STL file.
    template <class ExecutionSpace>
    FacetGeometry( const boost::property_tree::ptree& ptree,
                   const ExecutionSpace& exec_space )
    {
        // Get the geometry parameters.
        const auto& params = ptree.get_child( "geometry" );

        // Read the stl file and create the facet geometry.
        auto stl_ascii_filename = params.get<std::string>( "stl_file" );

        // Containers.
        std::vector<int> volume_ids;
        std::vector<int> surface_ids;
        std::vector<float> volume_facets;
        std::vector<float> surface_facets;

        // Load the file.
        std::ifstream file( stl_ascii_filename );
        if ( !file.is_open() )
            throw std::runtime_error( "Unable to open STL ASCII file" );

        // Read the file.
        std::string buffer;
        std::vector<std::string> tokens;
        tokens.reserve( 10 );
        bool read_volume = false;
        bool read_surface = false;
        while ( !file.eof() )
        {
            // Get the current line.
            std::getline( file, buffer );

            // Break the line up into tokens.
            std::istringstream line( buffer );
            tokens.resize( 0 );
            while ( !line.eof() )
            {
                tokens.push_back( std::string() );
                line >> tokens.back();
            }

            // Parse the line.
            if ( tokens.size() > 0 )
            {
                // New solid.
                if ( tokens[0].compare( "solid" ) == 0 )
                {
                    if ( tokens.size() != 3 )
                        throw std::runtime_error(
                            "STL READER: Expected 3 solid line entries" );

                    // New volume.
                    if ( tokens[1].compare( "Volume" ) == 0 ||
                         tokens[1].compare( "Body" ) == 0 )
                    {
                        volume_ids.push_back( std::atoi( tokens[2].c_str() ) );
                        _volume_facet_count.push_back( 0 );
                        read_volume = true;
                    }

                    // New surface.
                    else if ( tokens[1].compare( "Surface" ) == 0 )
                    {
                        surface_ids.push_back( std::atoi( tokens[2].c_str() ) );
                        _surface_facet_count.push_back( 0 );
                        read_surface = true;
                    }

                    else
                    {
                        throw std::runtime_error(
                            "STL READER: Solids execpted to be Volume/Body or "
                            "Surface" );
                    }
                }

                // Read facet.
                else if ( tokens[0].compare( "facet" ) == 0 )
                {
                    if ( tokens.size() != 5 )
                        throw std::runtime_error(
                            "STL READER: Expected 5 facet line entries" );

                    // Volume facet.
                    if ( read_volume )
                    {
                        ++_volume_facet_count.back();
                    }

                    // Surface facet.
                    else if ( read_surface )
                    {
                        ++_surface_facet_count.back();
                    }
                }

                // Read vertex.
                else if ( tokens[0].compare( "vertex" ) == 0 )
                {
                    if ( tokens.size() != 4 )
                        throw std::runtime_error(
                            "STL READER: Expected 4 vertex line entries" );

                    // Volume vertex coordinates.
                    if ( read_volume )
                    {
                        volume_facets.push_back(
                            std::atof( tokens[1].c_str() ) );
                        volume_facets.push_back(
                            std::atof( tokens[2].c_str() ) );
                        volume_facets.push_back(
                            std::atof( tokens[3].c_str() ) );
                    }

                    // Surface vertex coordinates.
                    if ( read_surface )
                    {
                        surface_facets.push_back(
                            std::atof( tokens[1].c_str() ) );
                        surface_facets.push_back(
                            std::atof( tokens[2].c_str() ) );
                        surface_facets.push_back(
                            std::atof( tokens[3].c_str() ) );
                    }
                }

                // Finish reading a solid.
                else if ( tokens[0].compare( "endsolid" ) == 0 )
                {
                    read_volume = false;
                    read_surface = false;
                }
            }
        }

        // Close the file.
        file.close();

        // Put volume data on device.
        putFileDataOnDevice( volume_ids, _volume_facet_count, volume_facets,
                             _volume_ids, _data.volume_facets,
                             _data.volume_offsets );

        // Put surface data on device.
        putFileDataOnDevice( surface_ids, _surface_facet_count, surface_facets,
                             _surface_ids, _data.surface_facets,
                             _data.surface_offsets );

        // Get the volume id of the global bounding box. The user is required
        // to make an axis-aligned bounding box of their geometry that defines
        // the global bounds of the problem. The user input is the global id
        // of this volume.
        _data.global_bounding_volume_id =
            localVolumeId( params.get<int>( "global_bounding_volume_id" ) );

        // Compute the bounding boxes of all the volumes.
        _data.volume_bounding_boxes = Kokkos::View<float* [6], MemorySpace>(
            Kokkos::ViewAllocateWithoutInitializing( "volume_bounding_boxes" ),
            volume_ids.size() );
        auto host_boxes = Kokkos::create_mirror_view(
            Kokkos::HostSpace(), _data.volume_bounding_boxes );
        for ( int v = 0; v < _data.numVolume(); ++v )
        {
            BoundingBoxReduce reducer( _data.volume_facets,
                                       _data.volume_offsets, v );
            float box[6];
            Kokkos::parallel_reduce(
                "Picasso::FacetGeometry::VolumeBoundingBox",
                Kokkos::RangePolicy<ExecutionSpace>( exec_space, 0,
                                                     _volume_facet_count[v] ),
                reducer, box );
            for ( int i = 0; i < 6; ++i )
                host_boxes( v, i ) = box[i];
        }
        Kokkos::deep_copy( _data.volume_bounding_boxes, host_boxes );

        // Extract the global bounding box to the host.
        auto global_box = Kokkos::subview(
            host_boxes, _data.global_bounding_volume_id, Kokkos::ALL() );
        for ( int i = 0; i < 6; ++i )
            _global_bounding_box[i] = global_box( i );
    }

    // Given a global volume id get the local volume id.
    int localVolumeId( const int global_id ) const
    {
        return _volume_ids.find( global_id )->second;
    }

    // Given a global surface id get the local surface id.
    int localSurfaceId( const int global_id ) const
    {
        return _surface_ids.find( global_id )->second;
    }

    // Given a local volume id get the number of facets that compose the
    // volume.
    int numVolumeFacet( const int local_id ) const
    {
        return _volume_facet_count[local_id];
    }

    // Given a local surface id get the number of facets that compose the
    // surface.
    int numSurfaceFacet( const int local_id ) const
    {
        return _surface_facet_count[local_id];
    }

    // Get the global bounding box.
    const Kokkos::Array<double, 6>& globalBoundingBox() const
    {
        return _global_bounding_box;
    }

    // Get the geometry data.
    const FacetGeometryData<MemorySpace>& data() const { return _data; }

  private:
    // Put file data on device.
    void putFileDataOnDevice(
        const std::vector<int>& solid_ids,
        const std::vector<int>& solid_facet_counts,
        const std::vector<float>& solid_facets,
        std::unordered_map<int, int>& id_map,
        Kokkos::View<float* [4][3], MemorySpace>& device_facets,
        Kokkos::View<int*, MemorySpace>& device_offsets )
    {
        // Allocate solid data.
        int num_solid = solid_ids.size();
        device_offsets = Kokkos::View<int*, MemorySpace>(
            Kokkos::ViewAllocateWithoutInitializing( "offsets" ), num_solid );

        int num_solid_facet = solid_facets.size() / 9;
        device_facets = Kokkos::View<float* [4][3], MemorySpace>(
            Kokkos::ViewAllocateWithoutInitializing( "facets" ),
            num_solid_facet );

        // Compute the offsets for the facet-to-solid mapping.
        auto host_offsets =
            Kokkos::create_mirror_view( Kokkos::HostSpace(), device_offsets );
        for ( int i = 0; i < num_solid; ++i )
        {
            // Map global solid ids to local ids.
            id_map.emplace( solid_ids[i], i );

            // Compute the offset via inclusive scan.
            host_offsets( i ) =
                ( 0 == i ) ? solid_facet_counts[i]
                           : solid_facet_counts[i] + host_offsets( i - 1 );
        }

        // Build the facets.
        auto host_facets =
            Kokkos::create_mirror_view( Kokkos::HostSpace(), device_facets );
        for ( int f = 0; f < num_solid_facet; ++f )
        {
            // Extract the vertices.
            for ( int v = 0; v < 3; ++v )
                for ( int d = 0; d < 3; ++d )
                    host_facets( f, v, d ) = solid_facets[9 * f + 3 * v + d];

            // Compute the normals.
            float v10[3];
            float v20[3];
            for ( int d = 0; d < 3; ++d )
            {
                v10[d] = host_facets( f, 1, d ) - host_facets( f, 0, d );
                v20[d] = host_facets( f, 2, d ) - host_facets( f, 0, d );
            }
            host_facets( f, 3, 0 ) = v10[1] * v20[2] - v10[2] * v20[1];
            host_facets( f, 3, 1 ) = v10[2] * v20[0] - v10[0] * v20[2];
            host_facets( f, 3, 2 ) = v10[0] * v20[1] - v10[1] * v20[0];

            // Scale them to make it a unit normal.
            float nmag =
                std::sqrt( host_facets( f, 3, 0 ) * host_facets( f, 3, 0 ) +
                           host_facets( f, 3, 1 ) * host_facets( f, 3, 1 ) +
                           host_facets( f, 3, 2 ) * host_facets( f, 3, 2 ) );
            for ( int d = 0; d < 3; ++d )
                host_facets( f, 3, d ) /= nmag;
        }

        // Copy to device.
        Kokkos::deep_copy( device_facets, host_facets );
        Kokkos::deep_copy( device_offsets, host_offsets );
    }

  public:
    // Volume ids - global-to-local mapping.
    std::unordered_map<int, int> _volume_ids;

    // Surface ids - global-to-local mapping.
    std::unordered_map<int, int> _surface_ids;

    // Volume facet counts.
    std::vector<int> _volume_facet_count;

    // Surface facet counts.
    std::vector<int> _surface_facet_count;

    // Global bounding box.
    Kokkos::Array<double, 6> _global_bounding_box;

    // Data.
    FacetGeometryData<MemorySpace> _data;
};

//---------------------------------------------------------------------------//
// Facet Geometry operations.
//---------------------------------------------------------------------------//
namespace FacetGeometryOps
{
//---------------------------------------------------------------------------//
// Project a point x along the given direction, r, and determine if it
// projects to the facet. If returns true, the projection solution is
// valid.
//
// y[0] = projection barycentric coordinate 1
// y[1] = projection barycentric coordinate 2
// y[2] = distance to triangle
template <class FacetView>
KOKKOS_FUNCTION bool pointFacetProjection( const float x[3], const float r[3],
                                           const FacetView& facets, const int f,
                                           float y[3] )
{
    // Build the system of equations to solve for intersection. Fire the ray
    // in the Y direction - this choice is arbitary.
    Mat3<float> A;
    Vec3<float> b;
    for ( int i = 0; i < 3; ++i )
    {
        A( i, 0 ) = facets( f, 1, i ) - facets( f, 0, i );
        A( i, 1 ) = facets( f, 2, i ) - facets( f, 0, i );
        A( i, 2 ) = -r[i];
        b( i ) = x[i] - facets( f, 0, i );
    }

    // Check the determinant of the matrix. If zero then the ray is parallel
    // so no intersection.
    auto det_A = !A;
    if ( 0.0 == det_A )
        return false;

    // Solve the system.
    VecView3<float> y_view( y, 1 );
    y_view = A ^ b;

    // Check the solution for inclusion in the triangle.
    return ( y[0] >= 0.0 && y[1] >= 0.0 && y[0] + y[1] <= 1.0 );
}

//---------------------------------------------------------------------------//
// Fire a ray from point x along the given direction, r, and determine if it
// intersects the facet.
template <class FacetView>
KOKKOS_FUNCTION bool rayFacetIntersect( const float x[3], const float r[3],
                                        const FacetView& facets, const int f )
{
    // Project the point and check the distance.
    float y[3];
    auto projects = pointFacetProjection( x, r, facets, f, y );
    return projects && ( y[2] > 0.0 );
}

//---------------------------------------------------------------------------//
// Compute the signed distance from a point to the plane defined by a facet.
template <class FacetView>
KOKKOS_FUNCTION float
distanceToFacetPlane( const float x[3], const FacetView& facets, const int f )
{
    return ( x[0] - facets( f, 0, 0 ) ) * facets( f, 3, 0 ) +
           ( x[1] - facets( f, 0, 1 ) ) * facets( f, 3, 1 ) +
           ( x[2] - facets( f, 0, 2 ) ) * facets( f, 3, 2 );
}

//---------------------------------------------------------------------------//
// Determine if a point is in a volume represented by a view of facets.
template <class FacetView>
KOKKOS_FUNCTION bool pointInVolume( const float x[3],
                                    const FacetView& volume_facets )
{
    // The choice of ray direction is arbitrary so generate a random one. This
    // could potentially help with robustness as floating point noise may
    // avoid edge and vertex intersections which could lead to multiple
    // positive intersections and therefore an incorrect point-in-volume
    // determination. The facet geometry has no connectivity so it is
    // difficult to resolve multiple intersections without accumulating the
    // intersection points and checking for duplicates within some tolerance.
    //
    // Note: Duan has indicated that doing tests with 3 different random rays
    // has been enough to be robust as one is likely to pass out of the 3 if
    // is a true intersection.
    using rand_type =
        Kokkos::Random_XorShift64<typename FacetView::device_type>;
    rand_type rng( 0 );
    float r[3] = { Kokkos::rand<rand_type, float>::draw( rng ),
                   Kokkos::rand<rand_type, float>::draw( rng ),
                   Kokkos::rand<rand_type, float>::draw( rng ) };
    float r_mag_inv = 1.0 / sqrt( r[0] * r[0] + r[1] * r[1] + r[2] * r[2] );
    for ( int d = 0; d < 3; ++d )
        r[d] *= r_mag_inv;

    // Fire rays through each facet and count intersections. If an
    // odd number of intersections, the point is in the volume. This works for
    // convex and non-convex volumes.
    int count = 0;
    for ( std::size_t f = 0; f < volume_facets.extent( 0 ); ++f )
        if ( rayFacetIntersect( x, r, volume_facets, f ) )
            ++count;
    return ( 1 == count % 2 );
}

//---------------------------------------------------------------------------//
// Given a point determine the volume in the given facet geometry in which it
// is located.If it is in the implicit complement, return -1. If it is outside
// of the entire domain, return -2;
template <class MemorySpace>
KOKKOS_FUNCTION int locatePoint( const float x[3],
                                 const FacetGeometryData<MemorySpace>& geom )
{
    // Get the global bounding volume id.
    int gbv = geom.global_bounding_volume_id;

    // Start by checking that the point is in the global bounding volume.
    if ( geom.volume_bounding_boxes( gbv, 0 ) <= x[0] &&
         geom.volume_bounding_boxes( gbv, 1 ) <= x[1] &&
         geom.volume_bounding_boxes( gbv, 2 ) <= x[2] &&
         geom.volume_bounding_boxes( gbv, 3 ) >= x[0] &&
         geom.volume_bounding_boxes( gbv, 4 ) >= x[1] &&
         geom.volume_bounding_boxes( gbv, 5 ) >= x[2] )
    {
        // Check each volume except the global bounding volume for point
        // inclusion.
        for ( int v = 0; v < geom.numVolume(); ++v )
        {
            if ( v != gbv )
            {
                // First check if the point is in the axis-aligned
                // bounding box of the volume.
                if ( geom.volume_bounding_boxes( v, 0 ) <= x[0] &&
                     geom.volume_bounding_boxes( v, 1 ) <= x[1] &&
                     geom.volume_bounding_boxes( v, 2 ) <= x[2] &&
                     geom.volume_bounding_boxes( v, 3 ) >= x[0] &&
                     geom.volume_bounding_boxes( v, 4 ) >= x[1] &&
                     geom.volume_bounding_boxes( v, 5 ) >= x[2] )
                {
                    // If in the bounding box, check against each volume
                    // facet for point inclusion.
                    if ( pointInVolume( x, geom.volumeFacets( v ) ) )
                    {
                        return v;
                    }
                }
            }
        }

        // If the point was not in any volume it is in the implicit
        // complement so return -1.
        return -1;
    }

    // Otherwise point is outside of the global domain including the
    // implicit complement so return -2.
    return -2;
}

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

} // end namespace FacetGeometryOps

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

} // end namespace Picasso

#endif // end PICASSO_FACETGEOMETRY_HPP


#include <boost/hana/fwd/concept/euclidean_ring.hpp>


/*!
  \file gpp_python_expected_improvement.cpp
  \rst
  This file has the logic to invoke C++ functions pertaining to expected improvement from Python.
  The data flow follows the basic 4 step from gpp_python_common.hpp.

  .. NoteL: several internal functions of this source file are only called from ``Export*()`` functions,
  so their description, inputs, outputs, etc. comments have been moved. These comments exist in
  ``Export*()`` as Python docstrings, so we saw no need to repeat ourselves.
\endrst*/
// This include violates the Google Style Guide by placing an "other" system header ahead of C and C++ system headers.  However,
// it needs to be at the top, otherwise compilation fails on some systems with some versions of python: OS X, python 2.7.3.
// Putting this include first prevents pyport from doing something illegal in C++; reference: http://bugs.python.org/issue10910
#include "Python.h"  // NOLINT(build/include)

#include "gpp_python_expected_improvement.hpp"

// NOLINT-ing the C, C++ header includes as well; otherwise cpplint gets confused
#include <string>  // NOLINT(build/include_order)
#include <vector>  // NOLINT(build/include_order)

#include <boost/python/bases.hpp>  // NOLINT(build/include_order)
#include <boost/python/class.hpp>  // NOLINT(build/include_order)
#include <boost/python/def.hpp>  // NOLINT(build/include_order)
#include <boost/python/dict.hpp>  // NOLINT(build/include_order)
#include <boost/python/extract.hpp>  // NOLINT(build/include_order)
#include <boost/python/list.hpp>  // NOLINT(build/include_order)
#include <boost/python/object.hpp>  // NOLINT(build/include_order)

#include "gpp_common.hpp"
#include "gpp_domain.hpp"
#include "gpp_exception.hpp"
#include "gpp_expected_improvement_gpu.hpp"
#include "gpp_geometry.hpp"
#include "gpp_heuristic_expected_improvement_optimization.hpp"
#include "gpp_math.hpp"
#include "gpp_optimization.hpp"
#include "gpp_optimizer_parameters.hpp"
#include "gpp_python_common.hpp"

namespace optimal_learning {

namespace {

double ComputeExpectedImprovementWrapper(const GaussianProcess& gaussian_process,
                                         const boost::python::list& points_to_sample,
                                         const boost::python::list& points_being_sampled,
                                         int num_to_sample, int num_being_sampled,
                                         int max_int_steps, double best_so_far,
                                         bool force_monte_carlo,
                                         RandomnessSourceContainer& randomness_source) {
  PythonInterfaceInputContainer input_container(points_to_sample, points_being_sampled,
                                                gaussian_process.dim(), num_to_sample, num_being_sampled);

  bool configure_for_gradients = false;
  if ((num_to_sample == 1) && (num_being_sampled == 0) && (force_monte_carlo == false)) {
    OnePotentialSampleExpectedImprovementEvaluator ei_evaluator(gaussian_process, best_so_far);
    OnePotentialSampleExpectedImprovementEvaluator::StateType ei_state(ei_evaluator,
                                                                       input_container.points_to_sample.data(),
                                                                       configure_for_gradients);
    return ei_evaluator.ComputeExpectedImprovement(&ei_state);
  } else {
    ExpectedImprovementEvaluator ei_evaluator(gaussian_process, max_int_steps, best_so_far);
    ExpectedImprovementEvaluator::StateType ei_state(ei_evaluator, input_container.points_to_sample.data(),
                                                     input_container.points_being_sampled.data(),
                                                     input_container.num_to_sample,
                                                     input_container.num_being_sampled,
                                                     configure_for_gradients,
                                                     randomness_source.normal_rng_vec.data());
    return ei_evaluator.ComputeExpectedImprovement(&ei_state);
  }
}

boost::python::list ComputeGradExpectedImprovementWrapper(const GaussianProcess& gaussian_process,
                                                          const boost::python::list& points_to_sample,
                                                          const boost::python::list& points_being_sampled,
                                                          int num_to_sample, int num_being_sampled,
                                                          int max_int_steps, double best_so_far,
                                                          bool force_monte_carlo,
                                                          RandomnessSourceContainer& randomness_source) {
  PythonInterfaceInputContainer input_container(points_to_sample, points_being_sampled, gaussian_process.dim(),
                                                num_to_sample, num_being_sampled);

  std::vector<double> grad_EI(num_to_sample*input_container.dim);
  bool configure_for_gradients = true;
  if ((num_to_sample == 1) && (num_being_sampled == 0) && (force_monte_carlo == false)) {
    OnePotentialSampleExpectedImprovementEvaluator ei_evaluator(gaussian_process, best_so_far);
    OnePotentialSampleExpectedImprovementEvaluator::StateType ei_state(ei_evaluator,
                                                                       input_container.points_to_sample.data(),
                                                                       configure_for_gradients);
    ei_evaluator.ComputeGradExpectedImprovement(&ei_state, grad_EI.data());
  } else {
    ExpectedImprovementEvaluator ei_evaluator(gaussian_process, max_int_steps, best_so_far);
    ExpectedImprovementEvaluator::StateType ei_state(ei_evaluator, input_container.points_to_sample.data(),
                                                     input_container.points_being_sampled.data(),
                                                     input_container.num_to_sample,
                                                     input_container.num_being_sampled,
                                                     configure_for_gradients,
                                                     randomness_source.normal_rng_vec.data());
    ei_evaluator.ComputeGradExpectedImprovement(&ei_state, grad_EI.data());
  }

  return VectorToPylist(grad_EI);
}

/*!\rst
  Utility that dispatches EI optimization based on optimizer type and num_to_sample.
  This is just used to reduce copy-pasted code.

  \param
    :optimizer_parameters: python/cpp_wrappers/optimization._CppOptimizerParameters
      Python object containing the DomainTypes domain_type and OptimizerTypes optimzer_type to use as well as
      appropriate parameter structs e.g., NewtonParameters for type kNewton).
      See comments on the python interface for multistart_expected_improvement_optimization_wrapper
    :gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities) that describes the
      underlying GP
    :input_container: PythonInterfaceInputContainer object containing data about points_being_sampled
    :domain: object specifying the domain to optimize over (see gpp_domain.hpp)
    :optimizer_type: type of optimization to use (e.g., null, gradient descent)
    :num_to_sample: how many simultaneous experiments you would like to run (i.e., the q in q,p-EI)
    :best_so_far: value of the best sample so far (must be min(points_sampled_value))
    :max_int_steps: maximum number of MC iterations
    :max_num_threads: maximum number of threads for use by OpenMP (generally should be <= # cores)
    :randomness_source: object containing randomness sources (sufficient for multithreading) used in EI computation
    :status: pydict object; cannot be None
  \output
    :randomness_source: PRNG internal states modified
    :status: modified on exit to describe whether convergence occurred
    :best_points_to_sample[num_to_sample][dim]: next set of points to evaluate
\endrst*/
template <typename DomainType>
void DispatchExpectedImprovementOptimization(const boost::python::object& optimizer_parameters,
                                             const GaussianProcess& gaussian_process,
                                             const PythonInterfaceInputContainer& input_container,
                                             const DomainType& domain,
                                             OptimizerTypes optimizer_type,
                                             int num_to_sample, double best_so_far,
                                             int max_int_steps, int max_num_threads,
                                             bool use_gpu, int which_gpu,
                                             RandomnessSourceContainer& randomness_source,
                                             boost::python::dict& status,
                                             double * restrict best_points_to_sample) {
#ifndef OL_GPU_ENABLED
  (void) which_gpu;  // quiet the compiler warning (unused variable)
#endif

  bool found_flag = false;
  switch (optimizer_type) {
    case OptimizerTypes::kNull: {
      ThreadSchedule thread_schedule(max_num_threads, omp_sched_static);
      // optimizer_parameters must contain an int num_random_samples field, extract it
      int num_random_samples = boost::python::extract<int>(optimizer_parameters.attr("num_random_samples"));

      if (use_gpu == true) {
#ifdef OL_GPU_ENABLED
        CudaComputeOptimalPointsToSampleViaLatinHypercubeSearch(gaussian_process, domain, thread_schedule,
                                                                input_container.points_being_sampled.data(),
                                                                num_random_samples, num_to_sample,
                                                                input_container.num_being_sampled,
                                                                best_so_far, max_int_steps, which_gpu, &found_flag,
                                                                &randomness_source.uniform_generator,
                                                                best_points_to_sample);
#else
        OL_THROW_EXCEPTION(OptimalLearningException, "GPU is not installed or enabled!");
#endif
      } else {
        ComputeOptimalPointsToSampleViaLatinHypercubeSearch(gaussian_process, domain, thread_schedule,
                                                            input_container.points_being_sampled.data(),
                                                            num_random_samples, num_to_sample,
                                                            input_container.num_being_sampled,
                                                            best_so_far, max_int_steps,
                                                            &found_flag, &randomness_source.uniform_generator,
                                                            randomness_source.normal_rng_vec.data(),
                                                            best_points_to_sample);
      }
      status[std::string("lhc_") + domain.kName + "_domain_found_update"] = found_flag;
      break;
    }  // end case kNull optimizer_type
    case OptimizerTypes::kGradientDescent: {
      // optimizer_parameters must contain a optimizer_parameters field
      // of type GradientDescentParameters. extract it
      const GradientDescentParameters& gradient_descent_parameters = boost::python::extract<GradientDescentParameters&>(optimizer_parameters.attr("optimizer_parameters"));
      ThreadSchedule thread_schedule(max_num_threads, omp_sched_dynamic);
      int num_random_samples = boost::python::extract<int>(optimizer_parameters.attr("num_random_samples"));

      bool random_search_only = false;
      if (use_gpu == true) {
#ifdef OL_GPU_ENABLED
        CudaComputeOptimalPointsToSample(gaussian_process, gradient_descent_parameters, domain, thread_schedule,
                                         input_container.points_being_sampled.data(), num_to_sample,
                                         input_container.num_being_sampled, best_so_far, max_int_steps,
                                         random_search_only, num_random_samples, which_gpu, &found_flag,
                                         &randomness_source.uniform_generator, best_points_to_sample);
#else
        OL_THROW_EXCEPTION(OptimalLearningException, "GPU is not installed or enabled!");
#endif
      } else {
        ComputeOptimalPointsToSample(gaussian_process, gradient_descent_parameters, domain, thread_schedule,
                                     input_container.points_being_sampled.data(), num_to_sample,
                                     input_container.num_being_sampled, best_so_far, max_int_steps,
                                     random_search_only, num_random_samples, &found_flag,
                                     &randomness_source.uniform_generator,
                                     randomness_source.normal_rng_vec.data(), best_points_to_sample);
      }
      status[std::string("gradient_descent_") + domain.kName + "_domain_found_update"] = found_flag;
      break;
    }  // end case kGradientDescent optimizer_type
    default: {
      std::fill(best_points_to_sample, best_points_to_sample + input_container.dim*num_to_sample, 0.0);
      OL_THROW_EXCEPTION(OptimalLearningException, "ERROR: invalid optimizer choice. Setting all coordinates to 0.0.");
      break;
    }
  }  // end switch over optimizer_type
}

boost::python::list MultistartExpectedImprovementOptimizationWrapper(const boost::python::object& optimizer_parameters,
                                                                     const GaussianProcess& gaussian_process,
                                                                     const boost::python::list& domain_bounds,
                                                                     const boost::python::list& points_being_sampled,
                                                                     int num_to_sample, int num_being_sampled,
                                                                     double best_so_far, int max_int_steps,
                                                                     int max_num_threads, bool use_gpu, int which_gpu,
                                                                     RandomnessSourceContainer& randomness_source,
                                                                     boost::python::dict& status) {
  // TODO(GH-131): make domain objects constructible from python; and pass them in through
  // the optimizer_parameters python object

  // abort if we do not have enough sources of randomness to run with max_num_threads
  if (unlikely(max_num_threads > static_cast<int>(randomness_source.normal_rng_vec.size()))) {
    OL_THROW_EXCEPTION(LowerBoundException<int>, "Fewer randomness_sources than max_num_threads.", randomness_source.normal_rng_vec.size(), max_num_threads);
  }

  int num_to_sample_input = 0;  // No points to sample; we are generating these via EI optimization
  const boost::python::list points_to_sample_dummy;
  PythonInterfaceInputContainer input_container(points_to_sample_dummy, points_being_sampled, gaussian_process.dim(), num_to_sample_input, num_being_sampled);
  std::vector<ClosedInterval> domain_bounds_C(input_container.dim);
  CopyPylistToClosedIntervalVector(domain_bounds, input_container.dim, domain_bounds_C);

  std::vector<double> best_points_to_sample_C(input_container.dim*num_to_sample);

  DomainTypes domain_type = boost::python::extract<DomainTypes>(optimizer_parameters.attr("domain_type"));
  OptimizerTypes optimizer_type = boost::python::extract<OptimizerTypes>(optimizer_parameters.attr("optimizer_type"));
  switch (domain_type) {
    case DomainTypes::kTensorProduct: {
      TensorProductDomain domain(domain_bounds_C.data(), input_container.dim);

      DispatchExpectedImprovementOptimization(optimizer_parameters, gaussian_process, input_container,
                                              domain, optimizer_type, num_to_sample, best_so_far,
                                              max_int_steps, max_num_threads, use_gpu, which_gpu,
                                              randomness_source,
                                              status, best_points_to_sample_C.data());
      break;
    }  // end case OptimizerTypes::kTensorProduct
    case DomainTypes::kSimplex: {
      SimplexIntersectTensorProductDomain domain(domain_bounds_C.data(), input_container.dim);

      DispatchExpectedImprovementOptimization(optimizer_parameters, gaussian_process, input_container,
                                              domain, optimizer_type, num_to_sample, best_so_far,
                                              max_int_steps, max_num_threads, use_gpu, which_gpu,
                                              randomness_source,
                                              status, best_points_to_sample_C.data());
      break;
    }  // end case OptimizerTypes::kSimplex
    default: {
      std::fill(best_points_to_sample_C.begin(), best_points_to_sample_C.end(), 0.0);
      OL_THROW_EXCEPTION(OptimalLearningException, "ERROR: invalid domain choice. Setting all coordinates to 0.0.");
      break;
    }
  }  // end switch over domain_type

  return VectorToPylist(best_points_to_sample_C);
}

/*!\rst
  Utility that dispatches heuristic EI optimization (solving q,0-EI) based on optimizer type and num_to_sample.
  This is just used to reduce copy-pasted code.

  \param
    :optimizer_parameters: python/cpp_wrappers/optimization._CppOptimizerParameters
      Python object containing the DomainTypes domain_type and OptimizerTypes optimzer_type to use as well as
      appropriate parameter structs e.g., NewtonParameters for type kNewton).
      See comments on the python interface for multistart_expected_improvement_optimization_wrapper
    :gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities) that describes the
      underlying GP
    :domain: object specifying the domain to optimize over (see gpp_domain.hpp)
    :optimizer_type: type of optimization to use (e.g., null, gradient descent)
    :estimation_policy: the policy to use to produce (heuristic) objective function estimates during multi-points EI optimization
    :num_to_sample: how many simultaneous experiments you would like to run (i.e., the q in q,0-EI)
    :best_so_far: value of the best sample so far (must be min(points_sampled_value))
    :max_num_threads: maximum number of threads for use by OpenMP (generally should be <= # cores)
    :randomness_source: object containing randomness sources (sufficient for multithreading) used in EI computation
    :status: pydict object; cannot be None
  \output
    :randomness_source: PRNG internal states modified
    :status: modified on exit to describe whether convergence occurred
    :best_points_to_sample[num_to_sample][dim]: next set of points to evaluate
\endrst*/
template <typename DomainType>
void DispatchHeuristicExpectedImprovementOptimization(const boost::python::object& optimizer_parameters,
                                                      const GaussianProcess& gaussian_process,
                                                      const DomainType& domain,
                                                      OptimizerTypes optimizer_type,
                                                      const ObjectiveEstimationPolicyInterface& estimation_policy,
                                                      int num_to_sample, double best_so_far, int max_num_threads,
                                                      RandomnessSourceContainer& randomness_source,
                                                      boost::python::dict& status,
                                                      double * restrict best_points_to_sample) {
  ThreadSchedule thread_schedule(max_num_threads, omp_sched_dynamic);
  bool found_flag = false;
  switch (optimizer_type) {
    case OptimizerTypes::kNull: {
      // optimizer_parameters must contain an int num_multistarts field, extract it
      int num_random_samples = boost::python::extract<int>(optimizer_parameters.attr("num_random_samples"));

      bool random_search_only = true;
      GradientDescentParameters gradient_descent_parameters(0, 0, 0, 0, 1.0, 1.0, 1.0, 0.0);  // dummy struct; we aren't using gradient descent
      ComputeHeuristicPointsToSample(gaussian_process, gradient_descent_parameters, domain,
                                     estimation_policy, thread_schedule, best_so_far,
                                     random_search_only, num_random_samples, num_to_sample,
                                     &found_flag, &randomness_source.uniform_generator,
                                     best_points_to_sample);

      status[std::string("lhc_") + domain.kName + "_domain_found_update"] = found_flag;
      break;
    }  // end case kNull optimizer_type
    case OptimizerTypes::kGradientDescent: {
      // optimizer_parameters must contain a optimizer_parameters field
      // of type GradientDescentParameters. extract it
      const GradientDescentParameters& gradient_descent_parameters = boost::python::extract<GradientDescentParameters&>(optimizer_parameters.attr("optimizer_parameters"));
      int num_random_samples = boost::python::extract<int>(optimizer_parameters.attr("num_random_samples"));

      bool random_search_only = false;
      ComputeHeuristicPointsToSample(gaussian_process, gradient_descent_parameters, domain,
                                     estimation_policy, thread_schedule, best_so_far,
                                     random_search_only, num_random_samples, num_to_sample,
                                     &found_flag, &randomness_source.uniform_generator,
                                     best_points_to_sample);

      status[std::string("gradient_descent_") + domain.kName + "_domain_found_update"] = found_flag;
      break;
    }  // end case kGradientDescent optimizer_type
    default: {
      std::fill(best_points_to_sample, best_points_to_sample + gaussian_process.dim()*num_to_sample, 0.0);
      OL_THROW_EXCEPTION(OptimalLearningException, "ERROR: invalid optimizer choice. Setting all coordinates to 0.0.");
      break;
    }
  }  // end switch over optimizer_type
}

boost::python::list HeuristicExpectedImprovementOptimizationWrapper(const boost::python::object& optimizer_parameters,
                                                                    const GaussianProcess& gaussian_process,
                                                                    const boost::python::list& domain_bounds,
                                                                    const ObjectiveEstimationPolicyInterface& estimation_policy,
                                                                    int num_to_sample, double best_so_far, int max_num_threads,
                                                                    RandomnessSourceContainer& randomness_source,
                                                                    boost::python::dict& status) {
  // TODO(GH-131): make domain objects constructible from python; and pass them in through
  // the optimizer_parameters python object
  int dim = gaussian_process.dim();
  std::vector<ClosedInterval> domain_bounds_C(dim);
  CopyPylistToClosedIntervalVector(domain_bounds, dim, domain_bounds_C);

  std::vector<double> best_points_to_sample_C(dim*num_to_sample);

  DomainTypes domain_type = boost::python::extract<DomainTypes>(optimizer_parameters.attr("domain_type"));
  OptimizerTypes optimizer_type = boost::python::extract<OptimizerTypes>(optimizer_parameters.attr("optimizer_type"));
  switch (domain_type) {
    case DomainTypes::kTensorProduct: {
      TensorProductDomain domain(domain_bounds_C.data(), dim);

      DispatchHeuristicExpectedImprovementOptimization(optimizer_parameters, gaussian_process, domain,
                                                       optimizer_type, estimation_policy, num_to_sample,
                                                       best_so_far, max_num_threads, randomness_source,
                                                       status, best_points_to_sample_C.data());
      break;
    }  // end case OptimizerTypes::kTensorProduct
    case DomainTypes::kSimplex: {
      SimplexIntersectTensorProductDomain domain(domain_bounds_C.data(), dim);

      DispatchHeuristicExpectedImprovementOptimization(optimizer_parameters, gaussian_process, domain,
                                                       optimizer_type, estimation_policy, num_to_sample,
                                                       best_so_far, max_num_threads, randomness_source,
                                                       status, best_points_to_sample_C.data());
      break;
    }  // end case OptimizerTypes::kSimplex
    default: {
      std::fill(best_points_to_sample_C.begin(), best_points_to_sample_C.end(), 0.0);
      OL_THROW_EXCEPTION(OptimalLearningException, "ERROR: invalid domain choice. Setting all coordinates to 0.0.");
      break;
    }
  }  // end switch over domain_type

  return VectorToPylist(best_points_to_sample_C);
}

boost::python::list EvaluateEIAtPointListWrapper(const GaussianProcess& gaussian_process,
                                                 const boost::python::list& initial_guesses,
                                                 const boost::python::list& points_being_sampled,
                                                 int num_multistarts, int num_to_sample,
                                                 int num_being_sampled, double best_so_far,
                                                 int max_int_steps, int max_num_threads,
                                                 RandomnessSourceContainer& randomness_source,
                                                 boost::python::dict& status) {
  // abort if we do not have enough sources of randomness to run with max_num_threads
  if (unlikely(max_num_threads > static_cast<int>(randomness_source.normal_rng_vec.size()))) {
    OL_THROW_EXCEPTION(LowerBoundException<int>, "Fewer randomness_sources than max_num_threads.", randomness_source.normal_rng_vec.size(), max_num_threads);
  }

  int num_to_sample_input = 0;  // No points to sample; we are generating these via EI optimization
  const boost::python::list points_to_sample_dummy;
  PythonInterfaceInputContainer input_container(points_to_sample_dummy, points_being_sampled, gaussian_process.dim(),
                                                num_to_sample_input, num_being_sampled);
  std::vector<double> result_point_C(input_container.dim);  // not used
  std::vector<double> result_function_values_C(num_multistarts);
  std::vector<double> initial_guesses_C(input_container.dim * num_multistarts);

  CopyPylistToVector(initial_guesses, input_container.dim * num_multistarts, initial_guesses_C);

  ThreadSchedule thread_schedule(max_num_threads, omp_sched_static);
  bool found_flag = false;
  EvaluateEIAtPointList(gaussian_process, thread_schedule, initial_guesses_C.data(),
                        input_container.points_being_sampled.data(), num_multistarts,
                        num_to_sample, input_container.num_being_sampled, best_so_far,
                        max_int_steps, &found_flag, randomness_source.normal_rng_vec.data(),
                        result_function_values_C.data(), result_point_C.data());

  status["evaluate_EI_at_point_list"] = found_flag;

  return VectorToPylist(result_function_values_C);
}

}  // end unnamed namespace

void ExportEstimationPolicies() {
  boost::python::class_<ObjectiveEstimationPolicyInterface, boost::noncopyable>("ObjectiveEstimationPolicyInterface", R"%%(
    Pure abstract (in the C++ sense) base class for objective function estimation (e.g., Constant Liar, Kriging Believer).
    Serves no purpose in Python but is needed by boost to allow pointer casts of derived types to this type.
    )%%", boost::python::no_init);

  boost::python::class_<ConstantLiarEstimationPolicy, boost::python::bases<ObjectiveEstimationPolicyInterface> >("ConstantLiarEstimationPolicy", R"%%(
    Produces objective function estimates at a "point" using the "Constant Liar" heuristic.

    Always outputs function_value = lie_value (the "constant lie") and noise_variance = lie_noise_variance, regardless
    of the value of "point".

    :ivar lie_value: (*float64*) the "constant lie" that this estimator should return
    :ivar lie_noise_variance: (*float64*) the noise_variance to associate to the lie_value (MUST be >= 0.0)
    )%%", boost::python::init<double, double>(R"%%(
    Constructs a ConstantLiarEstimationPolicy object.

    :param lie_value: the "constant lie" that this estimator should return
    :type lie_value: float64 (finite)
    :param lie_noise_variance: the noise_variance to associate to the lie_value (MUST be >= 0.0)
    :type lie_noise_variance: float64 >= 0.0
    )%%"));

  boost::python::class_<KrigingBelieverEstimationPolicy, boost::python::bases<ObjectiveEstimationPolicyInterface> >("KrigingBelieverEstimationPolicy", R"%%(
    Produces objective function estimates at a "point" using the "Kriging Believer" heuristic.

    Requires a valid GaussianProcess (GP) to produce estimates. Computes estimates as:

    * function_value = GP.Mean(point) + std_deviation_coef * sqrt(GP.Variance(point))
    * noise_variance = kriging_noise_variance

    :ivar std_deviation_coef: (*float64*) the relative amount of bias (in units of GP std deviation) to introduce into the GP mean
    :ivar kriging_noise_variance: (*float64*) the noise_variance to associate to each function value estimate (MUST be >= 0.0)
    )%%", boost::python::init<double, double>(R"%%(
    Constructs for KrigingBelieverEstimationPolicy object.

    :param std_deviation_coef: the relative amount of bias (in units of GP std deviation) to introduce into the GP mean
    :type std_deviation_coef: float64 (finite)
    :param kriging_noise_variance: the noise_variance to associate to each function value estimate (MUST be >= 0.0)
    :type kriging_noise_variance: float64 >= 0.0
    )%%"));
}

void ExportExpectedImprovementFunctions() {
  boost::python::def("compute_expected_improvement", ComputeExpectedImprovementWrapper, R"%%(
    Compute expected improvement.
    If ``num_to_sample == 1`` and ``num_being_sampled == 0`` AND ``force_monte_carlo is false``, this will
    use (fast/accurate) analytic evaluation.
    Otherwise monte carlo-based EI computation is used.

    :param gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities)
    :type gaussian_process: GPP.GaussianProcess (boost::python ctor wrapper around optimal_learning::GaussianProcess)
    :param points_to_sample: initial points to load into state (must be a valid point for the problem);
      i.e., points at which to evaluate EI and/or its gradient
    :type points_to_sample: list of float64 with shape (num_to_sample, dim)
    :param points_being_sampled: points that are being sampled in concurrently experiments
    :type points_being_sampled: list of float64 with shape (num_being_sampled, dim)
    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0
    :param num_being_sampled: number of points being sampled concurrently (i.e., the p in q,p-EI)
    :type num_being_sampled: int >= 0
    :param max_int_steps: number of MC integration points in EI
    :type max_int_steps: int >= 0
    :param best_so_far: best known value of objective so far
    :type best_so_far: float64
    :param force_monte_carlo: true to force monte carlo evaluation of EI
    :type force_monte_carlo: bool
    :param randomness_source: object containing randomness sources; only thread 0's source is used
    :type randomness_source: GPP.RandomnessSourceContainer
    :return: computed EI
    :rtype: float64 >= 0.0
    )%%");

  boost::python::def("compute_grad_expected_improvement", ComputeGradExpectedImprovementWrapper, R"%%(
    Compute the gradient of expected improvement evaluated at points_to_sample.
    If num_to_sample = 1 and num_being_sampled = 0 AND force_monte_carlo is false, this will
    use (fast/accurate) analytic evaluation.
    Otherwise monte carlo-based EI computation is used.

    :param gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities)
    :type gaussian_process: GPP.GaussianProcess (boost::python ctor wrapper around optimal_learning::GaussianProcess)
    :param points_to_sample: initial points to load into state (must be a valid point for the problem);
      i.e., points at which to evaluate EI and/or its gradient
    :type points_to_sample: list of float64 with shape (num_to_sample, dim)
    :param points_being_sampled: points that are being sampled in concurrently experiments
    :type points_being_sampled: list of float64 with shape (num_being_sampled, dim)
    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0
    :param num_being_sampled: number of points being sampled concurrently (i.e., the p in q,p-EI)
    :type num_being_sampled: int >= 0
    :param max_int_steps: number of MC integration points in EI
    :type max_int_steps: int >= 0
    :param best_so_far: best known value of objective so far
    :type best_so_far: float64
    :param force_monte_carlo: true to force monte carlo evaluation of EI
    :type force_monte_carlo: bool
    :param randomness_source: object containing randomness sources; only thread 0's source is used
    :type randomness_source: GPP.RandomnessSourceContainer
    :return: gradient of EI (computed at points_to_sample + points_being_sampled, wrt points_to_sample)
    :rtype: list of float64 with shape (num_to_sample, dim)
    )%%");

  boost::python::def("multistart_expected_improvement_optimization", MultistartExpectedImprovementOptimizationWrapper, R"%%(
    Optimize expected improvement (i.e., solve q,p-EI) over the specified domain using the specified optimization method.
    Can optimize for num_to_sample new points to sample (i.e., aka "q", experiments to run) simultaneously.
    Allows the user to specify num_being_sampled (aka "p") ongoing/concurrent experiments.

    The _CppOptimizerParameters object is a python class defined in:
    python/cpp_wrappers/optimization._CppOptimizerParameters
    See that class definition for more details.

    This function expects it to have the fields:

    * domain_type (DomainTypes enum from this file)
    * optimizer_type (OptimizerTypes enum from this file)
    * num_random_samples (int, number of samples to 'dumb' search over, if 'dumb' search is being used.
      e.g., if optimizer = kNull or if to_sample > 1)
    * optimizer_parameters (*Parameters struct (gpp_optimizer_parameters.hpp) where * matches optimizer_type
      unused if optimizer_type == kNull)

    This function also has the option of using GPU to compute general q,p-EI via MC simulation. To enable it,
    make sure you have installed GPU components of MOE, otherwise, it will throw Runtime excpetion.

    .. WARNING:: this function FAILS and returns an EMPTY LIST if the number of random sources < max_num_threads

    :param optimizer_parameters: python object containing the DomainTypes domain_type and
      OptimizerTypes optimzer_type to use as well as
      appropriate parameter structs e.g., NewtonParameters for type kNewton)
    :type optimizer_parameters: _CppOptimizerParameters
    :param gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities)
    :type gaussian_process: GPP.GaussianProcess (boost::python ctor wrapper around optimal_learning::GaussianProcess)
    :param domain: [lower, upper] bound pairs for each dimension
    :type domain: list of float64 with shape (dim, 2)
    :param points_being_sampled: points that are being sampled in concurrently experiments
    :type points_being_sampled: list of float64 with shape (num_being_sampled, dim)
    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0
    :param num_being_sampled: number of points being sampled concurrently (i.e., the p in q,p-EI)
    :type num_being_sampled: int >= 0
    :param best_so_far: best known value of objective so far
    :type best_so_far: float64
    :param max_int_steps: number of MC integration points in EI
    :type max_int_steps: int >= 0
    :param max_num_threads: max number of threads to use during EI optimization
    :type max_num_threads: int >= 1
    :param use_gpu: set to 1 if user wants to use GPU for MC computation
    :type use_gpu: bool
    :param which_gpu: GPU device ID
    :type which_gpu: int >= 0
    :param randomness_source: object containing randomness sources; only thread 0's source is used
    :type randomness_source: GPP.RandomnessSourceContainer
    :param status: pydict object (cannot be None!); modified on exit to describe whether convergence occurred
    :type status: dict
    :return: next set of points to eval
    :rtype: list of float64 with shape (num_to_sample, dim)
    )%%");

  boost::python::def("heuristic_expected_improvement_optimization", HeuristicExpectedImprovementOptimizationWrapper, R"%%(
    Compute a heuristic approximation to the result of multistart_expected_improvement_optimization(). That is, it
    optimizes an approximation to q,0-EI over the specified domain using the specified optimization method.
    Can optimize for num_to_sample (aka "q") new points to sample (i.e., experiments to run) simultaneously.

    Computing q,p-EI for q > 1 or p > 1 is expensive. To avoid that cost, this method "solves" q,0-EI by repeatedly
    optimizing 1,0-EI. We do the following (in C++)::

      for i in xrange(num_to_sample):
        new_point = optimize_1_EI(gaussian_process, ...)
        new_function_value, new_noise_variance = estimation_policy.compute_estimate(new_point, gaussian_process, i)
        gaussian_process.add_point(new_point, new_function_value, new_noise_variance)

    So using estimation_policy, we guess what the real-world objective function value would be, evaluated at the result
    of each 1-EI optimization. Then we treat this estimate as *truth* and feed it back to the gaussian process. The
    ConstantLiar and KrigingBelieverEstimationPolicy objects reproduce the heuristics described in Ginsbourger 2008.

    See gpp_heuristic_expected_improvement_optimization.hpp for further details on the algorithm.

    The _CppOptimizerParameters object is a python class defined in:
    ``python/cpp_wrappers/optimization._CppOptimizerParameters``
    See that class definition for more details.

    This function expects it to have the fields:

    * domain_type (DomainTypes enum from this file)
    * optimizer_type (OptimizerTypes enum from this file)
    * num_random_samples (int, number of samples to 'dumb' search over, if 'dumb' search is being used.
      e.g., if optimizer = kNull or if to_sample > 1)
    * optimizer_parameters (*Parameters struct (gpp_optimizer_parameters.hpp) where * matches optimizer_type
      unused if optimizer_type == kNull)

    .. WARNING:: this function FAILS and returns an EMPTY LIST if the number of random sources < max_num_threads

    :param optimizer_parameters: python object containing the DomainTypes domain_type and
      OptimizerTypes optimzer_type to use as well as
      appropriate parameter structs e.g., NewtonParameters for type kNewton)
    :type optimizer_parameters: _CppOptimizerParameters
    :param gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities)
    :type gaussian_process: GPP.GaussianProcess (boost::python ctor wrapper around optimal_learning::GaussianProcess)
    :param domain: [lower, upper] bound pairs for each dimension
    :type domain: list of float64 with shape (dim, 2)
    :param estimation_policy: the policy to use to produce (heuristic) objective function estimates
      during q,0-EI optimization (e.g., ConstantLiar, KrigingBeliever)
    :type estimation_policy: ObjectiveEstimationPolicyInterface
    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0
    :param best_so_far: best known value of objective so far
    :type best_so_far: float64
    :param max_num_threads: max number of threads to use during EI optimization
    :type max_num_threads: int >= 1
    :param randomness_source: object containing randomness sources; only thread 0's source is used
    :type randomness_source: GPP.RandomnessSourceContainer
    :param status: pydict object (cannot be None!); modified on exit to describe whether convergence occurred
    :type status: dict
    :return: next set of points to eval
    :rtype: list of float64 with shape (num_to_sample, dim)
    )%%");

  boost::python::def("evaluate_EI_at_point_list", EvaluateEIAtPointListWrapper, R"%%(
    Evaluates the expected improvement at each point in initial_guesses; can handle q,p-EI.
    Useful for plotting.

    Equivalent to::

      result = []
      for point in initial_guesses:
          result.append(compute_expected_improvement(point, ...))

    But this method is substantially faster (loop in C++ and multithreaded).

    .. WARNING:: this function FAILS and returns an EMPTY LIST if the number of random sources < max_num_threads

    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0

    :param gaussian_process: GaussianProcess object (holds points_sampled, values, noise_variance, derived quantities)
    :type gaussian_process: GPP.GaussianProcess (boost::python ctor wrapper around optimal_learning::GaussianProcess)
    :param initial_guesses: points at which to evaluate EI
    :type initial_guesses: list of flaot64 with shape (num_multistarts, num_to_sample, dim)
    :param points_being_sampled: points that are being sampled in concurrently experiments
    :type points_being_sampled: list of float64 with shape (num_being_sampled, dim)
    :param num_multistarts: number of points at which to evaluate EI
    :type num_multistarts: int > 0
    :param num_to_sample: number of potential future samples; gradients are evaluated wrt these points (i.e., the "q" in q,p-EI)
    :type num_to_sample: int > 0
    :param num_being_sampled: number of points being sampled concurrently (i.e., the p in q,p-EI)
    :type num_being_sampled: int >= 0
    :param best_so_far: best known value of objective so far
    :type best_so_far: float64
    :param max_int_steps: number of MC integration points in EI
    :type max_int_steps: int >= 0
    :param max_num_threads: max number of threads to use during EI optimization
    :type max_num_threads: int >= 1
    :param randomness_source: object containing randomness sources; only thread 0's source is used
    :type randomness_source: GPP.RandomnessSourceContainer
    :param status: pydict object (cannot be None!); modified on exit to describe whether convergence occurred
    :type status: dict
    :return: EI values at each point of the initial_guesses list, in the same order
    :rtype: list of float64 with shape (num_multistarts, )
    )%%");
}

}  // end namespace optimal_learning


#include <ros/ros.h>
#include <ros/rate.h>
#include <std_msgs/Float64.h>
#include <math.h>
#include <string>
#include <boost/scoped_ptr.hpp>
#include <eigen3/Eigen/Eigen>
#include <plugin_manager/QbAdvanceRefEq.h>
#include <std_msgs/Float64MultiArray.h>
#include <sensor_msgs/JointState.h>
#include <plugin_manager/state_info.h>

# define PI 3.14159

int joints_n;
double T_sample;
std::string ns_name;

std::vector<ros::Publisher> pub_ref1, pub_ref2;
ros::Publisher pub_robot_state, pub_motor1_state, pub_motor2_state;

std::vector<ros::Subscriber> sub_joints_state;
std::vector<ros::Subscriber> sub_mot1_state;
std::vector<ros::Subscriber> sub_mot2_state;
ros::Subscriber sub_ref1, sub_ref2;

std::vector<double> q, dq, tau, stiff;
std::vector<double> th_1, th_2, dth_1, dth_2, tau_1, tau_2;

std_msgs::Float64 msg_ref1, msg_ref2;
std_msgs::Float64MultiArray ref_1, ref_2;

std::vector<std::string> joint_name, ref1_topic_name, ref2_topic_name;
std::vector<std::string> topics_ref1, topics_ref2, topics_state;
std::vector<std::string> mot1_state, mot2_state;
std::string topic_state_name;
std::string  topic_mot1state_name, topic_mot2state_name;

sensor_msgs::JointState robot_state, state_1, state_2;

void JointStateCallback(const plugin_manager::state_infoConstPtr &msg, int i)
{
  q[i] = msg->q;
  dq[i] = msg->dq;
  tau[i] = msg->tau;
  stiff[i] = msg->stiff;
}

// Callbacks for the references
void Ref1_Callback(const std_msgs::Float64MultiArray &msg)
{
  ref_1 = msg;
}
void Ref2_Callback(const std_msgs::Float64MultiArray &msg)
{
  ref_2 = msg;
}
// Callbacks for the motors data
void Mot1StateCallback(const sensor_msgs::JointStateConstPtr &msg, int i)
{
/*  msg.position.resize(1);
  msg->velocity.resize(1);
  msg->effort.resize(1);*/
  th_1[i] = msg->position[0];
  dth_1[i] = msg->velocity[0];
  tau_1[i] = msg->effort[0];
}
void Mot2StateCallback(const sensor_msgs::JointStateConstPtr &msg, int i)
{
/*  msg->position.resize(1);
  msg->velocity.resize(1);
  msg->effort.resize(1);*/
  th_2[i] = msg->position[0];
  dth_2[i] = msg->velocity[0];
  tau_2[i] = msg->effort[0];
}

void Initialization(ros::NodeHandle n_)
{
  // Check if namespace is passed and get it
  if (!n_.getParam("namespace", ns_name))
  {
    ROS_ERROR("Specify namespace of robot");
    exit(1);
  }
  // Check if joint number is passed and get it
  if (!n_.getParam("joints_number", joints_n))
  {
    ROS_ERROR("Specify number of joint");
    exit(1);
  }
  // Check if sample time is passed and get it
  if (!n_.getParam("T_sample", T_sample))
  {
    ROS_ERROR("Specify the sample time");
    exit(1);
  }

  // Resize array variables
  pub_ref1.resize(joints_n);
  pub_ref2.resize(joints_n);
  sub_joints_state.resize(joints_n);
  sub_mot1_state.resize(joints_n);
  sub_mot2_state.resize(joints_n);
  joint_name.resize(joints_n);
  ref1_topic_name.resize(joints_n);
  ref2_topic_name.resize(joints_n);
  topics_ref1.resize(joints_n);
  topics_ref2.resize(joints_n);
  topics_state.resize(joints_n);
  mot1_state.resize(joints_n);
  mot2_state.resize(joints_n);
  ref_1.data.resize(joints_n);
  ref_2.data.resize(joints_n);
  q.resize(joints_n);
  dq.resize(joints_n);
  tau.resize(joints_n);
  stiff.resize(joints_n);
  th_1.resize(joints_n);
  dth_1.resize(joints_n);
  tau_1.resize(joints_n);
  th_2.resize(joints_n);
  dth_2.resize(joints_n);
  tau_2.resize(joints_n);
  robot_state.name.resize(joints_n);
  robot_state.position.resize(joints_n);
  robot_state.velocity.resize(joints_n);
  robot_state.effort.resize(joints_n);
  state_1.name.resize(joints_n);
  state_1.position.resize(joints_n);
  state_1.velocity.resize(joints_n);
  state_1.effort.resize(joints_n);
  state_2.name.resize(joints_n);
  state_2.position.resize(joints_n);
  state_2.velocity.resize(joints_n);
  state_2.effort.resize(joints_n);

  // Check if joint name vector is passed and get it
  if (!n_.getParam("joints_name", joint_name))
  {
    ROS_ERROR("Specify the joint name vector");
    exit(1);
  }
  // Subscribe to equilibirum position topics
  if (!n_.getParam("ref1_topic_name", ref1_topic_name) || !n_.getParam("ref2_topic_name", ref2_topic_name))
  {
    ROS_ERROR("Specify the name of the subscriber topic");
    exit(1);
  }

  // Create the subscriber topics' names
  for(int i = 0; i<joints_n; i++)
  {
    topics_ref1[i]  = "/" +  ns_name + "/" + joint_name[i] + ref1_topic_name[i];
    topics_ref2[i]  = "/" +  ns_name + "/" + joint_name[i] + ref2_topic_name[i];
    topics_state[i] = "/" +  ns_name + "/" + joint_name[i] + "/link_state";
    mot1_state[i] = "/" +  ns_name + "/" + joint_name[i] + "/motor_1_state";
    mot2_state[i] = "/" +  ns_name + "/" + joint_name[i] + "/motor_2_state";
  }
  topic_state_name = "/" +  ns_name + "/robot_state";
  topic_mot1state_name = "/" +  ns_name + "/motor_1_state";
  topic_mot2state_name = "/" +  ns_name + "/motor_2_state";

  // Advertise topics' publishers
  for(int i = 0; i<joints_n; i++)
  {
    pub_ref1[i] =  n_.advertise<std_msgs::Float64>(topics_ref1[i], 10);
    pub_ref2[i] =  n_.advertise<std_msgs::Float64>(topics_ref2[i], 10);
  }
  // Advertise robot state and motors state publishers
  pub_robot_state = n_.advertise<sensor_msgs::JointState>(topic_state_name, 10);
  pub_motor1_state = n_.advertise<sensor_msgs::JointState>(topic_mot1state_name, 10);
  pub_motor2_state = n_.advertise<sensor_msgs::JointState>(topic_mot2state_name, 10);

  // Subscribe to the joint state
  for(int i = 0; i<joints_n; i++)
  {
    sub_joints_state[i] = n_.subscribe<plugin_manager::state_info>(topics_state[i], 10, boost::bind(JointStateCallback, _1, i));
    sub_mot1_state[i] = n_.subscribe<sensor_msgs::JointState>(mot1_state[i], 10, boost::bind(Mot1StateCallback, _1, i));
    sub_mot2_state[i] = n_.subscribe<sensor_msgs::JointState>(mot2_state[i], 10, boost::bind(Mot2StateCallback, _1, i));
  }

  std::string sub_ref1_name = ns_name + "/reference_1";
  std::string sub_ref2_name = ns_name + "/reference_2";

  sub_ref1 = n_.subscribe(sub_ref1_name, 10, &Ref1_Callback);
  sub_ref2 = n_.subscribe(sub_ref2_name, 10, &Ref2_Callback);

}

//-----------------------------------------------------
//                                                 main
//-----------------------------------------------------
int main(int argc, char **argv)
{
  ros::init(argc, argv, "plugin_manager_node");

  ros::NodeHandle n_;

  Initialization(n_);

  std::string start_msg = "Plugin manager started on " + ns_name;
  ROS_WARN_STREAM(start_msg);

  ros::Rate r(1/T_sample);

  while(ros::ok())
  {

    for(int i = 0; i<joints_n; i++)
    {
      msg_ref1.data = ref_1.data[i];
      msg_ref2.data = ref_2.data[i];

      pub_ref1[i].publish(msg_ref1);
      pub_ref2[i].publish(msg_ref2);

      // assign stamp time
      ros::Time act_time = ros::Time::now();
      state_1.header.stamp = act_time;
      state_2.header.stamp = act_time;
      robot_state.header.stamp = act_time;

      // Fill data for the robot's link
      robot_state.name[i] = joint_name[i];
      robot_state.position[i] = q[i];
      robot_state.velocity[i] = dq[i];
      robot_state.effort[i] = tau[i];

      // Fill data for the motor 1 variables of the robot
      state_1.name[i] = joint_name[i];
      state_1.position[i] = th_1[i];
      state_1.velocity[i] = dth_1[i];
      state_1.effort[i] = tau_1[i];

      // Fill data for the motor 2 variables of the robot
      state_2.name[i] = joint_name[i];
      state_2.position[i] = th_2[i];
      state_2.velocity[i] = dth_2[i];
      state_2.effort[i] = tau_2[i];
    }

    pub_robot_state.publish(robot_state);
    pub_motor1_state.publish(state_1);
    pub_motor2_state.publish(state_2);

    ros::spinOnce();
    r.sleep();

  }

  return 0;
}


#ifndef _MATRIX_BLOCK_MATRIX_HPP_
#define _MATRIX_BLOCK_MATRIX_HPP_

#include <boost/type.hpp>
#include <boost/multi_array.hpp>

#include "matrix/matrix.hpp"
#include "matrix/matrix_base.hpp"

#include <stdlib.h>
#include <algorithm>
#include <vector>
#include <numeric>

#include "iterator/iterator.hpp"
#include "util/util.hpp"
#include "util/array.hpp"

namespace ublas = boost::numeric::ublas;

namespace matrix {

    using namespace boost;

    template<typename T>
    struct block_matrix_semantics {
	typedef int index_type;
	typedef T* block;
	typedef const T* const_block;
	virtual block b(index_type i, index_type j) = 0;
	virtual const_block b(index_type i, index_type j) const = 0;
    };

    template<typename T>
    struct block_matrix_base : virtual matrix_semantics<T>,
	virtual block_matrix_semantics<T>
    {
	typedef matrix_semantics<T> semantics;
	typedef typename semantics::index_type index_type;

	typedef size_t size_type;
	typedef size_type shape_type[2];
	typedef T value_type;
	typedef T element;
	typedef T* pointer;
	typedef const T* const_pointer;
    };

    template<typename T>
    class block_matrix : public virtual block_matrix_base<T> {
    public:
	typedef block_matrix_base<T> base;

	typedef typename base::index_type index_type;
	typedef typename base::size_type size_type;
	typedef typename base::shape_type shape_type;

	typedef typename base::reference reference;
	typedef typename base::const_reference const_reference;

	typedef typename base::pointer pointer;
	typedef typename base::const_pointer const_pointer;

	typedef typename base::block block;
	typedef typename base::const_block const_block;

	typedef util::Array<size_type,2> size_array;

	void resize(const size_type *matrix, const size_type *block) {
	    set(matrix); set_block(block);
	}

	void resize(size_t size1, size_t size2, size_t block1, size_t block2) {
	    size_t size[2] = { size1, size2 };
	    size_t block[2] = {block1, block2 };
	    resize(size, block);
	    allocate();
	}

	block_matrix() {
	    shape_type shape = { 0, 0 };
	    shape_type block_shape = { 1, 1 };
	    set(shape);
	    set_block(block_shape);
	    this->raw(NULL);
	}

	block_matrix(size_type m, size_type n) {
	    resize(size_array(m, n), size_array(1));
	    allocate();
	}

	block_matrix(const shape_type &shape, const shape_type &block_shape) {
	    resize(size_array(shape), size_array(block_shape));
	    allocate();
	}

	block_matrix(const shape_type &shape, const shape_type &block_shape,
		     pointer p) {
	    resize(size_array(shape), size_array(block_shape));
	    this->raw(p);
	}

	~block_matrix() { this->free(); }

	void allocate() {
	    raw(new T[size()]);
	    allocated_ = true;
	}

	void free() {
	    if (allocated_) delete[] raw_;
	    raw(NULL);
	}

	void raw(pointer raw) {
	    allocated_ = false;
	    raw_ = raw;
	}

	pointer raw() { return raw_; }

	const_pointer raw() const { return raw_; }

	pointer data() { return raw_; }

	const_pointer data() const { return raw_; }

	 const shape_type& shape() const { return shape_; }

	virtual size_type size1() const { return shape()[0]; }

	virtual size_type size2() const { return shape()[1]; }

	 const size_type size() const { return size1()*size2(); }

	size_t block1()const { return block_shape_[0]; }
	size_t block2()const { return block_shape_[1]; }

	 const shape_type& block_shape() const { return block_shape_; }

	size_type block_size() const { return block_size_; }

	void clear() {
	    std::fill(data(), data() + size(), 0);
	}

 	virtual reference operator()(int i, int j) {
	    return raw_[reference_index(i,j)];
	}

 	virtual const_reference operator()(int i, int j) const {
	    return raw_[reference_index(i,j)];
	}

	virtual block b(int a, int b) { return block_(a,b); }

	virtual const_block b(int a, int b) const { return block_(a,b); }

	size_t block_ld() const {return ld_block_; }

	void operator+=(const block_matrix &m) {
	    check_range(m);
	    if (block1() == m.block1() && block2() == m.block2()) {
		for (size_t i = 0; i < size(); ++i) {
		    data()[i] += m.data()[i];
		}
	    }
	    else {
		this->operator+=<block_matrix>(m);
	    }
	}

	template<class  M>
	void operator += (const M & m) {
	    check_range(m);
	    for (size_t j = 0; j < size2(); ++j) {
		for (size_t i = 0; i < size1(); ++i) {
		    (*this)(i,j) += m(i,j);
		}
	    }
	}

    protected:
	friend class meta_matrix<block_matrix>;
	friend class block_meta_matrix<block_matrix>;
	friend class boost::multi_array<block_matrix,2>;

	shape_type shape_, block_shape_, block_num_;
	size_type block_size_, ld_block_;

	pointer raw_;
	bool allocated_;

	template<class M>
	void check_range(const M & m) const {
	    if ( size1() != m.size1() || size2() != m.size2())
		throw std::range_error("invalid dimensions");
	}

	void set(const size_type *shape) {
	    for (int i = 0; i < 2; ++i) { shape_[i] = shape[i]; }
	}

	void set_block(const size_type *block_shape) {
	    block_size_ = 1;
	    for (int i = 0; i < 2; ++i) {
		assert(shape_[i]%block_shape[i] == 0);
		block_shape_[i] = block_shape[i];
		block_num_[i] = shape_[i]/block_shape_[i];
		block_size_ *= block_shape_[i];
	    }
	    ld_block_ = block_num_[0]*block_size_;
	}

	block block_(int a, int b) const {
	    ASSERT_RANGE(a, 0, block_num_[0]);
	    ASSERT_RANGE(b, 0, block_num_[1]);
	    return raw_ + block_index(a,b);
	}

	index_type block_index(index_type i, index_type j) const {
	    return i*block_size() + j*ld_block_;
	}

	index_type reference_index(index_type i, index_type j) const {
	    ASSERT_RANGE(i, 0, shape()[0]);
	    ASSERT_RANGE(j, 0, shape()[1]);
	    index_type a = i/block_shape_[0];
	    index_type b = j/block_shape_[1];
	    i -= a*block_shape_[0];
	    j -= b*block_shape_[1];
	    //printf("%i,%i\n", block_index(a,b) + (i + j*block_shape_[0]), size());
	    assert((block_index(a,b) + (i + j*block_shape_[0])) < size());
	    return block_index(a,b) + (i + j*block_shape_[0]);
	}

	int block_element(int i, int j, int a, int b) const {
	    i -= a*block_shape()[0];
	    j -= b*block_shape()[1];
	    return i + j*block_shape()[0];
	}

    };

}

#endif /* _MATRIX_BLOCK_MATRIX_HPP_ */


#include <algorithm>
#include <random>
#include <unordered_map>
#include <vector>

#include <iostream>

#include <benchmark/benchmark.h>

#include <boost/container/flat_map.hpp>

using value_type = std::pair<int, int>;

std::vector<value_type> random_vector(std::size_t n) {
  std::vector<value_type> ret(n);

  std::mt19937 mersenne_engine(n);
  std::uniform_int_distribution<int> dist;

  std::transform(ret.begin(), ret.end(), ret.begin(),
                 [&dist, &mersenne_engine](value_type) -> value_type {
                   return {dist(mersenne_engine), dist(mersenne_engine)};
                 });

  return ret;
}

template <typename Container>
static void BM_fill_random(benchmark::State &state) {
  auto values = random_vector(state.range(0));

  while (state.KeepRunning()) {
    benchmark::DoNotOptimize(Container(values.begin(), values.end()));
  }
}

template <typename Container>
static void BM_fill_ordered(benchmark::State &state) {
  auto values = random_vector(state.range(0));
  std::sort(values.begin(), values.end(),
            [](value_type const &a, value_type const &b) {
              return a.first < b.first;
            });

  while (state.KeepRunning()) {
    benchmark::DoNotOptimize(Container(values.begin(), values.end()));
  }
}

template <typename Container>
static void BM_sort_and_fill(benchmark::State &state) {
  auto values = random_vector(state.range(0));
  std::sort(values.begin(), values.end(),
            [](value_type const &a, value_type const &b) {
              return a.first < b.first;
            });

  while (state.KeepRunning()) {
    benchmark::DoNotOptimize(Container(values.begin(), values.end()));
  }
}

template <typename Container>
static void BM_traversal(benchmark::State &state) {
  auto values = random_vector(state.range(0));
  Container map(values.begin(), values.end());

  while (state.KeepRunning()) {
    benchmark::DoNotOptimize(std::accumulate(
        map.begin(), map.end(), 0,
        [](int s, value_type const &e) { return s + e.second; }));
  }
}

template <typename Container>
static void BM_random_access(benchmark::State &state) {
  auto values = random_vector(state.range(0));

  int i = 0;
  for (auto &e : values) {
    e.first = i++;
  }

  Container map(values.begin(), values.end());

  std::mt19937 mersenne_engine(state.range(0));
  std::uniform_int_distribution<int> dist(0, state.range(0) - 1);

  int sum = 0;
  while (state.KeepRunning()) {
    auto const id = dist(mersenne_engine);

    benchmark::DoNotOptimize(sum += map.at(id));
  }
}

BENCHMARK_TEMPLATE(BM_random_access, std::unordered_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_random_access, boost::container::flat_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_fill_random, std::unordered_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_fill_random, boost::container::flat_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_fill_ordered, std::unordered_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_fill_ordered, boost::container::flat_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_traversal, std::unordered_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);
BENCHMARK_TEMPLATE(BM_traversal, boost::container::flat_map<int, int>)
    ->RangeMultiplier(2)
    ->Range(8, 8 << 10);

BENCHMARK_MAIN();