#include <ai.h>

#include <Alembic/Util/All.h>
#include <Alembic/AbcCoreAbstract/All.h>
#include <Alembic/AbcCoreOgawa/All.h>
#ifdef SIMBA_ENABLE_ALEMBIC_HDF5
#include <Alembic/AbcCoreHDF5/All.h>
#endif // SIMBA_ENABLE_ALEMBIC_HDF5
#include <Alembic/Abc/All.h>
#include <Alembic/AbcCoreFactory/All.h>
#include <Alembic/AbcGeom/All.h>
#include <Alembic/AbcCollection/All.h>
#include <Alembic/AbcMaterial/All.h>

#include <vector>
#include <map>
#include <iostream>
#include <fstream>
#include <cassert>

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

#include <OpenEXR/ImathVec.h>

template <typename T> void interpolate(const Imath::Vec3<T>& P1,
									   const Imath::Vec3<T>& T1,
									   const Imath::Vec3<T>& P2,
									   const Imath::Vec3<T>& T2,
									   T                     s,
									   Imath::Vec3<T>&       P)
{
	T h1 =  2.0*std::pow(s,3.0) - 3.0*std::pow(s,2.0) + 1.0; // calculate basis function 1
	T h2 = -2.0*std::pow(s,3.0) + 3.0*std::pow(s,2.0);       // calculate basis function 2
	T h3 =      std::pow(s,3.0) - 2.0*std::pow(s,2.0) + s;   // calculate basis function 3
	T h4 =      std::pow(s,3.0) -     std::pow(s,2.0);       // calculate basis function 4

	P = h1*P1 + // multiply and sum all functions
		h2*P2 + // together to build the interpolated
		h3*T1 + // point along the curve.
		h4*T2;
}

typedef boost::multi_array<Imath::V3f, 2> V3fSamplingArray2D;
typedef std::vector<AtUInt32> AtUInt32Container;
typedef std::vector<uint64_t> UInt64Container;
typedef std::vector<std::string> StringContainer;
typedef std::vector<float> FloatContainer;
struct ArnoldMeshData
{
	V3fSamplingArray2D _vlist_data_array;
	// Topological stable being assumed
	AtUInt32Container _nsides_data;
	AtUInt32Container _vidxs_data;
};
struct ArnoldPointsData
{
	V3fSamplingArray2D _points_data_array;
	// Topological stable being assumed
	UInt64Container _ids_data;
	FloatContainer _radius_data;
};
struct AlembicPointsData
{
	AlembicPointsData(const Imath::V3f& i_position, const Imath::V3f& i_velocity)
	: _position(i_position)
	, _velocity(i_velocity)
	{}
	Imath::V3f _position;
	Imath::V3f _velocity;
};
typedef std::map<uint64_t,AlembicPointsData> AlembicPointsDataIndexedMap;

void make_arnold_points(const std::string& 		name,
						const ArnoldPointsData& i_arnold_points_data,
						float 					i_shutter_open,
						float 					i_shutter_close)
{
	AtByte nkeys = i_arnold_points_data._points_data_array.size();
	if (nkeys<1)
		return;
	// Assumes topological stability
	AtUInt32 nelements = i_arnold_points_data._points_data_array[0].size();
	// name
	AtNode *points = AiNode("points");
	AiNodeSetStr(points, "name", name.c_str());

	std::cout << boost::format("nelements=%1% nkeys=%2%") % nelements % int(nkeys)<< std::endl;

	// points
	AiNodeSetArray(points, "points", AiArrayConvert(nelements,nkeys,AI_TYPE_POINT,i_arnold_points_data._points_data_array.data()));

	// radius
	AiNodeSetArray(points, "radius", AiArrayConvert(i_arnold_points_data._radius_data.size(),1,AI_TYPE_FLOAT,i_arnold_points_data._radius_data.data()));

	// deform_time_samples
	float deform_time_samples[2] = {i_shutter_open,i_shutter_close};
	AiNodeSetArray(points, "deform_time_samples", AiArrayConvert(2,1,AI_TYPE_FLOAT,deform_time_samples));
}

void make_arnold_polymesh(const std::string& name,
						  const ArnoldMeshData& i_arnold_mesh_data,
						  float i_shutter_open,
						  float i_shutter_close)
{
	AtByte nkeys = i_arnold_mesh_data._vlist_data_array.size();
	if (nkeys<1)
		return;
	// Assumes topological stability
	AtUInt32 nelements = i_arnold_mesh_data._vlist_data_array[0].size();
	// name
	AtNode *polymesh = AiNode("polymesh");
	AiNodeSetStr(polymesh, "name", name.c_str());

	std::cout << boost::format("nelements=%1% nkeys=%2%") % nelements % int(nkeys)<< std::endl;

	// vlist
	AiNodeSetArray(polymesh, "vlist", AiArrayConvert(nelements,nkeys,AI_TYPE_POINT,i_arnold_mesh_data._vlist_data_array.data()));

	// nsides
	AiNodeSetArray(polymesh, "nsides", AiArrayConvert(i_arnold_mesh_data._nsides_data.size(),1,AI_TYPE_UINT,i_arnold_mesh_data._nsides_data.data()));

	// vidxs
	AiNodeSetArray(polymesh, "vidxs", AiArrayConvert(i_arnold_mesh_data._vidxs_data.size(),1,AI_TYPE_UINT,i_arnold_mesh_data._vidxs_data.data()));

	// deform_time_samples
	float deform_time_samples[2] = {i_shutter_open,i_shutter_close};
	AiNodeSetArray(polymesh, "deform_time_samples", AiArrayConvert(2,1,AI_TYPE_FLOAT,deform_time_samples));
}

void write_arnold_points_data_to_csv_file(const ArnoldPointsData&   i_arnold_points_data,
										  V3fSamplingArray2D::index i_sample_index,
										  const std::string&        i_csv_filename)
{
	std::ofstream csv_file;
	csv_file.open (i_csv_filename.c_str());

    size_t per_sample_point_count = i_arnold_points_data._points_data_array[i_sample_index].size();
    size_t radius_count = i_arnold_points_data._radius_data.size();

    assert ( per_sample_point_count == radius_count );

    csv_file << "x,y,z,width" << std::endl;

    for (size_t index=0;index<per_sample_point_count;index++)
    {
    	float x = i_arnold_points_data._points_data_array[i_sample_index][index].x;
    	float y = i_arnold_points_data._points_data_array[i_sample_index][index].y;
    	float z = i_arnold_points_data._points_data_array[i_sample_index][index].z;
    	float w = i_arnold_points_data._radius_data[index] * 2.0f;
        csv_file << boost::format("%1%, %2%, %3%, %4%") % x % y % z % w << std::endl;
    }

    csv_file.close();
}

void write_arnold_points_data_to_csv_sequence(const ArnoldPointsData& i_arnold_points_data,
												  const std::string&    i_base_filename)
{
	std::cout << "NICHOLAS write_arnold_points_data_to_csv_sequence() XXXXXXXXXXXXXXXXX" << std::endl;
	V3fSamplingArray2D::index num_samples = i_arnold_points_data._points_data_array.size();
	V3fSamplingArray2D::index num_elements_per_sample = i_arnold_points_data._points_data_array.shape()[1];
	size_t num_radii = i_arnold_points_data._radius_data.size();
	size_t num_ids = i_arnold_points_data._ids_data.size();
	if ( num_radii != num_elements_per_sample )
	{
		std::cerr << boost::format("num_radii = %1% != num_elements_per_sample = %2%") % num_radii % num_elements_per_sample << std::endl;
		return;
	}
	if ( num_ids != num_elements_per_sample )
	{
		std::cerr << boost::format("num_ids = %1% != num_elements_per_sample = %2%") % num_ids % num_elements_per_sample << std::endl;
		return;
	}
	// std::cout << boost::format("num_elements_per_sample = %1%") % num_elements_per_sample << std::endl;
	for (V3fSamplingArray2D::index sample_index=0;sample_index<num_samples;++sample_index)
	{
		std::string numbered_output_filename = (boost::format(i_base_filename) % (sample_index+1)).str();
		// std::cout << boost::format("numbered_output_filename = '%1%'") % numbered_output_filename << std::endl;
		write_arnold_points_data_to_csv_file(i_arnold_points_data,sample_index,numbered_output_filename);
	}
	// std::cout << boost::format("_vlist_data_array num_elements = %1%") % i_arnold_mesh_data._vlist_data_array.size() << std::endl;
}

void write_arnold_mesh_data_to_wavefront_file(const ArnoldMeshData&     i_arnold_mesh_data,
											  V3fSamplingArray2D::index i_sample_index,
											  const std::string&        i_wavefront_filename)
{

	std::ofstream wavefront_file;
	wavefront_file.open (i_wavefront_filename.c_str());

	wavefront_file << "# File exported by Nicholas Yue\n";

    size_t P_size = i_arnold_mesh_data._vlist_data_array[i_sample_index].size();
    size_t indices_size = i_arnold_mesh_data._vidxs_data.size();
    size_t counts_size = i_arnold_mesh_data._nsides_data.size();

    wavefront_file << boost::format("# %1% points") % P_size << std::endl;
    wavefront_file << boost::format("# %1% faces") % counts_size << std::endl;

    wavefront_file << "g" << std::endl;
    for (size_t index=0;index<P_size;index++)
    {
    	float x = i_arnold_mesh_data._vlist_data_array[i_sample_index][index].x;
    	float y = i_arnold_mesh_data._vlist_data_array[i_sample_index][index].y;
    	float z = i_arnold_mesh_data._vlist_data_array[i_sample_index][index].z;
        wavefront_file << boost::format("v %1% %2% %3%") % x % y % z << std::endl;
    }

    wavefront_file << "g" << std::endl;
    size_t indices_index = 0;
    for (size_t index=0;index<counts_size;index++)
    {
        wavefront_file << "f";
    	size_t face_vert_count = i_arnold_mesh_data._nsides_data[index];
    	for (size_t face_vert_index=0;face_vert_index < face_vert_count; ++face_vert_index)
    	{
    		size_t face_vert_index_value = i_arnold_mesh_data._vidxs_data[indices_index] + 1; // OBJ array is not zero based!
    		wavefront_file << boost::format(" %1%") % face_vert_index_value;
    		indices_index++;
    	}
		wavefront_file << std::endl;
    }
    wavefront_file.close();
}

void write_arnold_mesh_data_to_wavefront_sequence(const ArnoldMeshData& i_arnold_mesh_data,
												  const std::string&    i_base_filename)
{

	V3fSamplingArray2D::index num_samples = i_arnold_mesh_data._vlist_data_array.size();
	V3fSamplingArray2D::index num_elements_per_sample = i_arnold_mesh_data._vlist_data_array.shape()[1];
	std::cout << boost::format("num_elements_per_sample = %1%") % num_elements_per_sample << std::endl;
	for (V3fSamplingArray2D::index sample_index=0;sample_index<num_samples;++sample_index)
	{
		std::string numbered_output_filename = (boost::format(i_base_filename) % (sample_index+1)).str();
		std::cout << boost::format("numbered_output_filename = '%1%'") % numbered_output_filename << std::endl;
		write_arnold_mesh_data_to_wavefront_file(i_arnold_mesh_data,sample_index,numbered_output_filename);
	}
	std::cout << boost::format("_vlist_data_array num_elements = %1%") % i_arnold_mesh_data._vlist_data_array.size() << std::endl;
}

void write_arnold_points_data_to_file(const ArnoldPointsData& i_arnold_points_data,
									  const std::string&      i_arnold_filename,
									  float 				  i_shutter_open,
									  float 				  i_shutter_close)
{
	// start an Arnold session
	AiBegin();
//	AiMsgSetLogFileName("pm.log");

	make_arnold_points("test",i_arnold_points_data,i_shutter_open,i_shutter_close);
#ifdef FULL_ARNOLD_SCENE
	// create a sphere geometric primitive
	AtNode *sph = AiNode("sphere");
	AiNodeSetStr(sph, "name", "mysphere");
	AiNodeSetFlt(sph, "radius", 5.0f);

	// create a lambert shader
	AtNode *shader = AiNode("lambert");
	AiNodeSetStr(shader, "name", "mylambert");
	AiNodeSetRGB(shader, "Kd_color", 1.0f, 0.0f, 0.0f);

	// assign the sphere's shader
	AiNodeSetPtr(sph, "shader", shader);

	// create a perspective camera
	AtNode *camera = AiNode("persp_camera");
	AiNodeSetStr(camera, "name", "mycamera");
	// position the camera (alternatively you can set 'matrix')
	AiNodeSetPnt(camera, "position", 0.f, 0.f, 60.f);

	// create a point light source
	AtNode *light = AiNode("point_light");
	AiNodeSetStr(light, "name", "mylight");
	// position the light (alternatively use 'matrix')
	AiNodeSetPnt(light, "position", 0.f, 10.f, 10.f);

	// get the global options node and set some options
	AtNode *options = AiUniverseGetOptions();
	AiNodeSetInt(options, "AA_samples", 1);
	AiNodeSetInt(options, "xres", 320);
	AiNodeSetInt(options, "yres", 240);
	// set the active camera (optional, since there is only one camera)
	AiNodeSetPtr(options, "camera", camera);

	// create an output driver node
	AtNode *driver = AiNode("driver_tiff");
	AiNodeSetStr(driver, "name", "mydriver");
	AiNodeSetStr(driver, "filename", "pm.tiff");

	// create a gaussian filter node
	AtNode *filter = AiNode("gaussian_filter");
	AiNodeSetStr(filter, "name", "myfilter");

	// assign the driver and filter to the main (beauty) AOV, which is called "RGB"
	AtArray *outputs_array = AiArrayAllocate(1, 1, AI_TYPE_STRING);
	AiArraySetStr(outputs_array, 0, "RGBA RGBA myfilter mydriver");
	AiNodeSetArray(options, "outputs", outputs_array);

	// finally, render the image
	// AiRender(AI_RENDER_MODE_CAMERA);
#endif // FULL_ARNOLD_SCENE

	// ... or you can write out an .ass file instead
	AiASSWrite(i_arnold_filename.c_str(), AI_NODE_ALL, false, false);

	// at this point we can shut down Arnold
	AiEnd();

}

void write_arnold_mesh_data_to_file(const ArnoldMeshData& i_arnold_mesh_data,
									const std::string&    i_arnold_filename,
									float 				  i_shutter_open,
									float 				  i_shutter_close)
{
	// start an Arnold session
	AiBegin();
	AiMsgSetLogFileName("pm.log");

	make_arnold_polymesh("test",i_arnold_mesh_data,i_shutter_open,i_shutter_close);

	// create a sphere geometric primitive
	AtNode *sph = AiNode("sphere");
	AiNodeSetStr(sph, "name", "mysphere");
	AiNodeSetFlt(sph, "radius", 5.0f);

	// create a lambert shader
	AtNode *shader = AiNode("lambert");
	AiNodeSetStr(shader, "name", "mylambert");
	AiNodeSetRGB(shader, "Kd_color", 1.0f, 0.0f, 0.0f);

	// assign the sphere's shader
	AiNodeSetPtr(sph, "shader", shader);

	// create a perspective camera
	AtNode *camera = AiNode("persp_camera");
	AiNodeSetStr(camera, "name", "mycamera");
	// position the camera (alternatively you can set 'matrix')
	AiNodeSetPnt(camera, "position", 0.f, 0.f, 60.f);

	// create a point light source
	AtNode *light = AiNode("point_light");
	AiNodeSetStr(light, "name", "mylight");
	// position the light (alternatively use 'matrix')
	AiNodeSetPnt(light, "position", 0.f, 10.f, 10.f);

	// get the global options node and set some options
	AtNode *options = AiUniverseGetOptions();
	AiNodeSetInt(options, "AA_samples", 1);
	AiNodeSetInt(options, "xres", 320);
	AiNodeSetInt(options, "yres", 240);
	// set the active camera (optional, since there is only one camera)
	AiNodeSetPtr(options, "camera", camera);

	// create an output driver node
	AtNode *driver = AiNode("driver_tiff");
	AiNodeSetStr(driver, "name", "mydriver");
	AiNodeSetStr(driver, "filename", "pm.tiff");

	// create a gaussian filter node
	AtNode *filter = AiNode("gaussian_filter");
	AiNodeSetStr(filter, "name", "myfilter");

	// assign the driver and filter to the main (beauty) AOV, which is called "RGB"
	AtArray *outputs_array = AiArrayAllocate(1, 1, AI_TYPE_STRING);
	AiArraySetStr(outputs_array, 0, "RGBA RGBA myfilter mydriver");
	AiNodeSetArray(options, "outputs", outputs_array);

	// finally, render the image
	// AiRender(AI_RENDER_MODE_CAMERA);

	// ... or you can write out an .ass file instead
	AiASSWrite(i_arnold_filename.c_str(), AI_NODE_ALL, false, false);

	// at this point we can shut down Arnold
	AiEnd();

}

bool build_even_motion_relative_time_samples(float i_relative_shutter_open,
											 float i_relative_shutter_close,
											 AtByte i_motion_sample_count,
											 FloatContainer& o_earlier_sampling_time_vector,
											 FloatContainer& o_later_sampling_time_vector)
{
	if (i_motion_sample_count<2)
		return false;
	float shutter_delta = (i_relative_shutter_close - i_relative_shutter_open)/(i_motion_sample_count-1);
	if (shutter_delta < FLT_EPSILON)
		return false;
	o_earlier_sampling_time_vector.clear();
	o_later_sampling_time_vector.clear();
	for (AtByte sample_index=0;sample_index<i_motion_sample_count;sample_index++)
	{
		float time_sample  = i_relative_shutter_open + sample_index * shutter_delta;
		std::cout << boost::format("time_sample = %1%") % time_sample << std::endl;
		if (time_sample>=0.0f)
			o_later_sampling_time_vector.push_back(time_sample);
		else
			o_earlier_sampling_time_vector.push_back(time_sample);
	}
	return true;
}

/*!
 * \param i_relative_shutter_open This can be negative (relative to the current frame)
 * \param i_relative_shutter_close This can be negative (relative to the current frame)
 */
void build_interim_points_for_arnold_ass(const Alembic::AbcGeom::IPointsSchema::Sample* i_sample,
										 AlembicPointsDataIndexedMap&                   o_interim_points)
{
	Alembic::AbcGeom::P3fArraySamplePtr positions = i_sample->getPositions();
	Alembic::AbcGeom::V3fArraySamplePtr velocities = i_sample->getVelocities();
	Alembic::AbcGeom::UInt64ArraySamplePtr ids = i_sample->getIds();
	size_t num_positions = positions->size();
	size_t num_velocities = velocities->size();
	size_t num_ids = ids->size();
	assert ( num_positions == num_ids );
	assert ( num_velocities == num_ids );
	std::pair<AlembicPointsDataIndexedMap::iterator,bool> ret;
	for (size_t positions_index = 0;positions_index < num_positions; positions_index++)
	{
		ret = o_interim_points.insert(AlembicPointsDataIndexedMap::value_type(
				ids->get()[positions_index],
				AlembicPointsData(Imath::V3f(positions->get()[positions_index].x,positions->get()[positions_index].y,positions->get()[positions_index].z),
						Imath::V3f(velocities->get()[positions_index].x,velocities->get()[positions_index].y,velocities->get()[positions_index].z))));
		if (!ret.second)
		{
			std::cerr << "Failed to insert into map" << std::endl;
			return;
		}
	}
	// std::cout << boost::format("XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX o_interim_points = %1%") % o_interim_points.size() << std::endl;
}

// Arnold point data layout in the case of multiple points with multiple motion blur segments
//
//points
//{
// name test
// points 2 3 POINT
// # 2 points
// # 3 motion blur samples per point
// # First sample
// -10 0 0
// 10 0 0
// # Second sample
// -10 5 0
// 10 15 0
// # Third sample
// -15 5 0
// 15 10 0
// radius 1
// deform_time_samples 2 1 FLOAT
//-0.25 0.25
//}
void build_points_for_arnold_ass_from_interim_points(const AlembicPointsDataIndexedMap* i_previous_interim_points,
													 const AlembicPointsDataIndexedMap* i_current_interim_points,
													 const AlembicPointsDataIndexedMap* i_next_interim_points,
													 float								i_relative_shutter_open,
													 float								i_relative_shutter_close,
													 AtByte								i_motion_sample_count,
													 ArnoldPointsData&					o_arnold_points)
{

	// Use the current interim point's id as lookup for the previous and next information
//	V3fSamplingArray2D::index num_test_samples = i_motion_sample_count;
//	std::cout << boost::format("YYYYYYYYYYYYYYYYYYYYYY num_test_samples = %1%") % num_test_samples << std::endl;
//	num_test_samples = 1;

	FloatContainer earlier_sampling_time_vector;
	FloatContainer later_sampling_time_vector;

	if (build_even_motion_relative_time_samples(i_relative_shutter_open,
												i_relative_shutter_close,
												i_motion_sample_count,
												earlier_sampling_time_vector,
												later_sampling_time_vector))
	{
		{
			size_t num_earlier_samples = earlier_sampling_time_vector.size();
			for (size_t index=0;index<num_earlier_samples;++index)
			{
				std::cout << boost::format("earlier_sampling_time_vector[%1%] = %2%") % index % earlier_sampling_time_vector[index] << std::endl;
			}
		}
		{
			size_t num_later_samples = later_sampling_time_vector.size();
			for (size_t index=0;index<num_later_samples;++index)
			{
				std::cout << boost::format("later_sampling_time_vector[%1%] = %2%") % index % later_sampling_time_vector[index] << std::endl;
			}
		}
	}
	o_arnold_points._points_data_array.resize(boost::extents[i_motion_sample_count][i_current_interim_points->size()]);
	V3fSamplingArray2D::index point_index = 0;
	for(AlembicPointsDataIndexedMap::const_iterator iter = i_current_interim_points->begin(); iter != i_current_interim_points->end(); iter++,point_index++)
	{
		uint64_t search_id = iter->first;
		o_arnold_points._ids_data.push_back(search_id);
		o_arnold_points._radius_data.push_back(0.01f);
		bool previous_point_exists = false;
		bool next_point_exists = false;
		// Generate samples for this particular point
		typedef std::vector<Imath::V3f> V3fContainer;

		V3fContainer interpolated_P(i_motion_sample_count);
		AlembicPointsDataIndexedMap::const_iterator find_previous_point_result;
		AlembicPointsDataIndexedMap::const_iterator find_next_point_result;
		if (i_previous_interim_points)
		{
			find_previous_point_result = i_previous_interim_points->find(search_id);
	        if (find_previous_point_result != i_previous_interim_points->end())
	        {
	        	previous_point_exists = true;
	        }
		}

		if (i_next_interim_points)
		{
			find_next_point_result = i_next_interim_points->find(search_id);
	        if (find_next_point_result != i_next_interim_points->end())
	        {
	        	next_point_exists = true;
	        }
		}

		// The are 3 possibilities - start, in-between and end
		if (!previous_point_exists && next_point_exists)
		{
			// START frame
			std::cout << "START frame : points processing" << std::endl;
		}
		else if (previous_point_exists && !next_point_exists)
		{
			// END frame
			std::cout << "END frame : points processing" << std::endl;
		}
		else if (previous_point_exists && next_point_exists)
		{
			// INBETWEEN frame
			std::cout << "INBETWEEN frame : points processing" << std::endl;
			size_t interpolated_P_index = 0;
			size_t num_earlier_sampling_time = earlier_sampling_time_vector.size();
			size_t num_later_sampling_time = later_sampling_time_vector.size();

			Imath::V3f P1;
			Imath::V3f T1;
			Imath::V3f P2;
			Imath::V3f T2;

			// Previous
			P1 = find_previous_point_result->second._position;
			T1 = find_previous_point_result->second._velocity;
			P2 = iter->second._position;
			T2 = iter->second._velocity;
			for (size_t earlier_sample_index=0;earlier_sample_index<num_earlier_sampling_time;++earlier_sample_index)
			{
				float s = 1+earlier_sampling_time_vector[earlier_sample_index];
				interpolate<float>(P1,T1,P2,T2,s,interpolated_P[interpolated_P_index]);
				interpolated_P_index++;
			}

			// Next
			P1 = iter->second._position;
			T1 = iter->second._velocity;
			P2 = find_next_point_result->second._position;
			T2 = find_next_point_result->second._velocity;
			for (size_t next_sample_index=0;next_sample_index<num_later_sampling_time;++next_sample_index)
			{
				float s = later_sampling_time_vector[next_sample_index];
				interpolate<float>(P1,T1,P2,T2,s,interpolated_P[interpolated_P_index]);
				interpolated_P_index++;
			}

			assert ( interpolated_P_index == i_motion_sample_count );
		}
		for (V3fSamplingArray2D::index motion_index = 0;motion_index<i_motion_sample_count;++motion_index)
		{
			o_arnold_points._points_data_array[motion_index][point_index] = interpolated_P[motion_index];
		}
	}
}

///*!
// * \param i_relative_shutter_open This can be negative (relative to the current frame)
// * \param i_relative_shutter_close This can be negative (relative to the current frame)
// */
//void build_points_for_arnold_ass(const Alembic::AbcGeom::IPointsSchema::Sample* i_previous_sample,
//								 const Alembic::AbcGeom::IPointsSchema::Sample* i_current_sample,
//								 const Alembic::AbcGeom::IPointsSchema::Sample* i_next_sample,
//								 float											i_relative_shutter_open,
//								 float											i_relative_shutter_close,
//								 AtByte											i_motion_sample_count,
//								 ArnoldPointsData&                              o_arnold_points)
//{
//
//}
/*!
 * \param i_relative_shutter_open This can be negative (relative to the current frame)
 * \param i_relative_shutter_close This can be negative (relative to the current frame)
 */
void build_polymesh_for_arnold_ass(const Alembic::AbcGeom::IPolyMeshSchema::Sample* i_previous_sample,
								   const Alembic::AbcGeom::IPolyMeshSchema::Sample* i_current_sample,
								   const Alembic::AbcGeom::IPolyMeshSchema::Sample* i_next_sample,
								   float											i_relative_shutter_open,
								   float											i_relative_shutter_close,
								   AtByte											i_motion_sample_count,
								   ArnoldMeshData&                                  o_arnold_mesh)
{
	std::cout << "INTERPOLATE 0000" << std::endl;
	if (!i_current_sample)
		return;

	std::cout << "INTERPOLATE 0100" << std::endl;
	// Assumes topologically stable
	Alembic::AbcGeom::Int32ArraySamplePtr indices = i_current_sample->getFaceIndices();
	Alembic::AbcGeom::Int32ArraySamplePtr counts = i_current_sample->getFaceCounts();

	// Current
	Alembic::AbcGeom::P3fArraySamplePtr current_P = i_current_sample->getPositions();
	Alembic::AbcGeom::V3fArraySamplePtr current_v = i_current_sample->getVelocities();
	size_t current_num_P = current_P->size();
	size_t current_num_v = current_v->size();

	// Topologically constant
	size_t num_nsides = counts->size();
	size_t num_indices = indices->size();
	o_arnold_mesh._nsides_data.resize(num_nsides);
	o_arnold_mesh._vidxs_data.resize(num_indices);
	for (size_t index=0;index<num_nsides;index++)
	{
		o_arnold_mesh._nsides_data[index] = counts->get()[index];
	}
	for (size_t index=0;index<num_indices;index++)
	{
		o_arnold_mesh._vidxs_data[index] = indices->get()[index];
	}

	if ( i_motion_sample_count == 1)
	{
		// Special case, return single time step from current sample
		std::cout << boost::format("num_nsides=%1% num_indices=%2% num_P=%3%") % num_nsides % num_indices % current_num_P << std::endl;
		o_arnold_mesh._vlist_data_array.resize(boost::extents[i_motion_sample_count][current_num_P]);
		for (size_t index=0;index<current_num_P;index++)
		{
			o_arnold_mesh._vlist_data_array[0][index].x = current_P->get()[index].x;
			o_arnold_mesh._vlist_data_array[0][index].y = current_P->get()[index].y;
			o_arnold_mesh._vlist_data_array[0][index].z = current_P->get()[index].z;
		}

		return;
	}

	if (i_previous_sample && i_next_sample)
	{
		std::cout << "INTERPOLATE INBETWEEN" << std::endl;

		Alembic::AbcGeom::P3fArraySamplePtr previous_P = i_previous_sample->getPositions();
		Alembic::AbcGeom::V3fArraySamplePtr previous_v = i_previous_sample->getVelocities();
		size_t previous_num_P = previous_P->size();
		size_t previous_num_v = previous_v->size();

		Alembic::AbcGeom::P3fArraySamplePtr next_P = i_next_sample->getPositions();
		Alembic::AbcGeom::V3fArraySamplePtr next_v = i_next_sample->getVelocities();
		size_t next_num_P = next_P->size();
		size_t next_num_v = next_v->size();

		if (	(previous_num_P==current_num_P)
				&&
				(next_num_P==current_num_P)
				&&
				(previous_num_v==current_num_v)
				&&
				(next_num_v==current_num_v)
				)
		{

			FloatContainer earlier_sampling_time_vector;
			FloatContainer later_sampling_time_vector;
			build_even_motion_relative_time_samples(i_relative_shutter_open,
													i_relative_shutter_close,
													i_motion_sample_count,
													earlier_sampling_time_vector,
													later_sampling_time_vector);

			o_arnold_mesh._vlist_data_array.resize(boost::extents[i_motion_sample_count][current_num_P]);
			assert(i_motion_sample_count==(earlier_sampling_time_vector.size()+later_sampling_time_vector.size()));
			FloatContainer::const_iterator tsIter;
			FloatContainer::const_iterator tsEIter;
			size_t sampling_index = 0; // shared over both earlier_sampling_time_vector and later_sampling_time_vector

			// Start with earlier time sample first so it is monotonically increasing

			// EARLIER
			tsIter = earlier_sampling_time_vector.begin();
			tsEIter = earlier_sampling_time_vector.end();
			for (;tsIter!=tsEIter;++tsIter)
			{

				for (size_t index=0;index<current_num_P;index++)
				{
					Imath::Vec3<float> P1(previous_P->get()[index].x,previous_P->get()[index].y,previous_P->get()[index].z);
					Imath::Vec3<float> T1(previous_v->get()[index].x,previous_v->get()[index].y,previous_v->get()[index].z);
					Imath::Vec3<float> P2(current_P->get()[index].x,current_P->get()[index].y,current_P->get()[index].z);
					Imath::Vec3<float> T2(current_v->get()[index].x,current_v->get()[index].y,current_v->get()[index].z);
					Imath::Vec3<float> P;
					float s = 1+(*tsIter);
					std::cout << boost::format("EARLIER s = %1%") % s << std::endl;
					interpolate<float>(P1,T1,P2,T2,s,P);
					std::cout << boost::format("Earlier [sampling_index = %1%] P1 = %2%, T1 = %3%, P2 = %4%, T2 = %5% [s = %6%] {P =%7%}") % sampling_index % P1 % T1 % P2 % T2 % s % P << std::endl;

					o_arnold_mesh._vlist_data_array[sampling_index][index].x = P.x;
					o_arnold_mesh._vlist_data_array[sampling_index][index].y = P.y;
					o_arnold_mesh._vlist_data_array[sampling_index][index].z = P.z;
				}

				sampling_index++;

			}

			std::cout << "Earlier/Later =================================" << std::endl;

			// LATER
			tsIter = later_sampling_time_vector.begin();
			tsEIter = later_sampling_time_vector.end();
			for (;tsIter!=tsEIter;++tsIter)
			{
				for (size_t index=0;index<current_num_P;index++)
				{

					Imath::Vec3<float> P1(current_P->get()[index].x,current_P->get()[index].y,current_P->get()[index].z);
					Imath::Vec3<float> T1(current_v->get()[index].x,current_v->get()[index].y,current_v->get()[index].z);
					Imath::Vec3<float> P2(next_P->get()[index].x,next_P->get()[index].y,next_P->get()[index].z);
					Imath::Vec3<float> T2(next_v->get()[index].x,next_v->get()[index].y,next_v->get()[index].z);
					Imath::Vec3<float> P;
					float s = *tsIter;
					// std::cout << boost::format("s = %1%") % s << std::endl;
					interpolate<float>(P1,T1,P2,T2,s,P);
					std::cout << boost::format("Later [sampling_index = %1%] P1 = %2%, T1 = %3%, P2 = %4%, T2 = %5% [s = %6%] {P =%7%}") % sampling_index % P1 % T1 % P2 % T2 % s % P << std::endl;

					o_arnold_mesh._vlist_data_array[sampling_index][index].x = P.x;
					o_arnold_mesh._vlist_data_array[sampling_index][index].y = P.y;
					o_arnold_mesh._vlist_data_array[sampling_index][index].z = P.z;
				}

				sampling_index++;
			}
		}
	}
	else if ((i_previous_sample == 0) && i_next_sample)
	{
		std::cout << "INTERPOLATE FIRST FRAME" << std::endl;
	}
	else if (i_previous_sample && (i_next_sample == 0))
	{
		std::cout << "INTERPOLATE LAST FRAME" << std::endl;
	}

}

void print_sample_data(const std::string& i_message,const Alembic::AbcGeom::IPolyMeshSchema::Sample& i_sample)
{
	Alembic::AbcGeom::P3fArraySamplePtr sample_P = i_sample.getPositions();
	Alembic::AbcGeom::V3fArraySamplePtr sample_v = i_sample.getVelocities();
	size_t sample_num_P = sample_P->size();
	size_t sample_num_v = sample_v->size();
	assert(sample_num_P==sample_num_v);
	for (size_t index=0;index<sample_num_P;index++)
	{
		std::cout << boost::format("[%1%] P = %2%, v = %3%") % i_message % sample_P->get()[index] % sample_v->get()[index] << std::endl;
	}

}

void export_points_as_arnold_ass(Alembic::AbcGeom::IPoints& points,
								 Alembic::Abc::index_t      i_start_frame_number,
								 Alembic::Abc::index_t      i_requested_frame_number,
								 const std::string&         i_arnold_filename,
								 AtByte						i_motion_samples,
								 float						i_relative_shutter_open,
								 float						i_relative_shutter_close)
{

    size_t num_samples = points.getSchema().getNumSamples();
    size_t last_sample_index = num_samples - 1;
    Alembic::Abc::int64_t requested_index = i_requested_frame_number - i_start_frame_number;
    std::cout << boost::format("num_samples = %1% i_requested_frame_number = %2% requested_index = %3% last_sample_index = %4%") % num_samples % i_requested_frame_number % requested_index % last_sample_index << std::endl;
    if (requested_index < 0 || requested_index > (num_samples-1))
    	return;

	Alembic::Abc::ISampleSelector current_sample_selector(requested_index);
    Alembic::AbcGeom::IPointsSchema::Sample current_sample;
    points.getSchema().get( current_sample, current_sample_selector );
    AlembicPointsDataIndexedMap current_interim_points;

    build_interim_points_for_arnold_ass(&current_sample,current_interim_points);

    if (requested_index == 0)
    {
    	// First frame
    	Alembic::Abc::ISampleSelector next_sample_selector(requested_index+1);

        Alembic::AbcGeom::IPointsSchema::Sample next_sample;
    	points.getSchema().get( next_sample, next_sample_selector );

        AlembicPointsDataIndexedMap next_interim_points;
        build_interim_points_for_arnold_ass(&next_sample,next_interim_points);

    	ArnoldPointsData arnold_points_data;
    	build_points_for_arnold_ass_from_interim_points(0,
    													&current_interim_points,
														&next_interim_points,
														i_relative_shutter_open,
														i_relative_shutter_close,
														i_motion_samples,
														arnold_points_data);
    	write_arnold_points_data_to_file(arnold_points_data,
    									 i_arnold_filename,
										 i_relative_shutter_open,
										 i_relative_shutter_close);

    }
    else if (requested_index == last_sample_index)
    {
    	// Last frame
    	Alembic::Abc::ISampleSelector previous_sample_selector(requested_index-1);

        Alembic::AbcGeom::IPointsSchema::Sample previous_sample;
    	points.getSchema().get( previous_sample, previous_sample_selector );

        AlembicPointsDataIndexedMap previous_interim_points;
        build_interim_points_for_arnold_ass(&previous_sample,previous_interim_points);

    	ArnoldPointsData arnold_points_data;
    	build_points_for_arnold_ass_from_interim_points(&previous_interim_points,
    													&current_interim_points,
														0,
														i_relative_shutter_open,
														i_relative_shutter_close,
														i_motion_samples,
														arnold_points_data);
    	write_arnold_points_data_to_file(arnold_points_data,
    									 i_arnold_filename,
										 i_relative_shutter_open,
										 i_relative_shutter_close);

    }
    else
    {
    	std::cout << "NICHOLAS XXXXXXXXXXXXXXXXXXXXXXXX" << std::endl;
    	Alembic::Abc::ISampleSelector previous_sample_selector(requested_index-1);

        Alembic::AbcGeom::IPointsSchema::Sample previous_sample;
    	points.getSchema().get( previous_sample, previous_sample_selector );

        AlembicPointsDataIndexedMap previous_interim_points;
        build_interim_points_for_arnold_ass(&previous_sample,previous_interim_points);

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

    	Alembic::Abc::ISampleSelector next_sample_selector(requested_index+1);

        Alembic::AbcGeom::IPointsSchema::Sample next_sample;
    	points.getSchema().get( next_sample, next_sample_selector );

        AlembicPointsDataIndexedMap next_interim_points;
        build_interim_points_for_arnold_ass(&next_sample,next_interim_points);

        // ==================
    	ArnoldPointsData arnold_points_data;
    	build_points_for_arnold_ass_from_interim_points(&previous_interim_points,
    													&current_interim_points,
														&next_interim_points,
														i_relative_shutter_open,
														i_relative_shutter_close,
														i_motion_samples,
														arnold_points_data);
    	// write_arnold_points_data_to_csv_sequence(arnold_points_data,"points_per_sample.%04d.csv");
    	write_arnold_points_data_to_file(arnold_points_data,
    									 i_arnold_filename,
										 i_relative_shutter_open,
										 i_relative_shutter_close);

    }

}

void export_polymesh_as_arnold_ass(Alembic::AbcGeom::IPolyMesh& pmesh,
		   	   	   	   	   	   	   Alembic::Abc::index_t        i_start_frame_number,
								   Alembic::Abc::index_t        i_requested_frame_number,
								   const std::string&           i_arnold_filename,
								   AtByte 						i_motion_samples,
								   float          				i_relative_shutter_open,
								   float          				i_relative_shutter_close
								   )
{

    size_t num_samples = pmesh.getSchema().getNumSamples();
    size_t last_sample_index = num_samples - 1;
    Alembic::Abc::int64_t requested_index = i_requested_frame_number - i_start_frame_number - 1;
    std::cout << boost::format("num_samples = %1% i_requested_frame_number = %2% requested_index = %3% last_sample_index = %4%") % num_samples % i_requested_frame_number % requested_index % last_sample_index << std::endl;
    if (requested_index < 0 || requested_index > (num_samples-1))
    	return;

	Alembic::Abc::ISampleSelector current_sample_selector(requested_index);
    Alembic::AbcGeom::IPolyMeshSchema::Sample current_sample;
	pmesh.getSchema().get( current_sample, current_sample_selector );

    if (requested_index == 0)
    {
    	// First frame
    	Alembic::Abc::ISampleSelector next_sample_selector(requested_index+1);

        Alembic::AbcGeom::IPolyMeshSchema::Sample next_sample;
    	pmesh.getSchema().get( next_sample, next_sample_selector );

    	ArnoldMeshData arnold_mesh_data;
    	build_polymesh_for_arnold_ass(0,
    								  &current_sample,
									  &next_sample,
									  i_relative_shutter_open,
									  i_relative_shutter_close,
									  i_motion_samples,
									  arnold_mesh_data);
    	write_arnold_mesh_data_to_file(arnold_mesh_data,i_arnold_filename,
				  i_relative_shutter_open,
				  i_relative_shutter_close);

    }
    else if (requested_index == last_sample_index)
    {
    	// Last frame
    	Alembic::Abc::ISampleSelector previous_sample_selector(requested_index-1);

        Alembic::AbcGeom::IPolyMeshSchema::Sample previous_sample;
    	pmesh.getSchema().get( previous_sample, previous_sample_selector );

    	ArnoldMeshData arnold_mesh_data;
    	build_polymesh_for_arnold_ass(&previous_sample,
    								  &current_sample,
									  0,
									  i_relative_shutter_open,
									  i_relative_shutter_close,
									  i_motion_samples,
									  arnold_mesh_data);
    	write_arnold_mesh_data_to_file(arnold_mesh_data,i_arnold_filename,
				  i_relative_shutter_open,
				  i_relative_shutter_close);

    }
    else
    {
    	// Frame between first and last frame
    	Alembic::Abc::int64_t previous_index = requested_index - 1;
    	Alembic::Abc::int64_t next_index = requested_index + 1;
    	Alembic::Abc::ISampleSelector previous_sample_selector(previous_index);
    	Alembic::Abc::ISampleSelector next_sample_selector(next_index);
    	std::cout << boost::format("previous_index = %1% requested_index = %2% next_index = %3%") % previous_index % requested_index % next_index << std::endl;
        Alembic::AbcGeom::IPolyMeshSchema::Sample next_sample;
    	pmesh.getSchema().get( next_sample, next_sample_selector );

        Alembic::AbcGeom::IPolyMeshSchema::Sample previous_sample;
    	pmesh.getSchema().get( previous_sample, previous_sample_selector );

    	{
    		print_sample_data("previous",previous_sample);
    		print_sample_data("current",current_sample);
    		print_sample_data("next",next_sample);
    	}
    	ArnoldMeshData arnold_mesh_data;
    	build_polymesh_for_arnold_ass(&previous_sample,
    								  &current_sample,
									  &next_sample,
									  i_relative_shutter_open,
									  i_relative_shutter_close,
									  i_motion_samples,
									  arnold_mesh_data);
    	write_arnold_mesh_data_to_file(arnold_mesh_data,i_arnold_filename,
    								   i_relative_shutter_open,
									   i_relative_shutter_close);
    	write_arnold_mesh_data_to_wavefront_sequence(arnold_mesh_data,"mesh_per_sample.%04d.obj");
    }
}

void flatten_string_array(const StringContainer& i_string_array,
						  const std::string&     i_separator,
						  std::string&           o_flatten_string)
{
	size_t num_strings = i_string_array.size();
	for (size_t index=0;index<num_strings;index++)
	{
		o_flatten_string.append(i_string_array[index]);
		if (index<(num_strings-1)) // We want to avoid a trailing separator!
			o_flatten_string.append(i_separator);
	}
}

void locate_geometry_in_hierarchy(const Alembic::Abc::IObject& top,
								  const StringContainer&       i_hierachy_path,
								  size_t 					  i_requested_index,
								  float          			  i_relative_shutter_open,
								  float          			  i_relative_shutter_close,
								  AtByte						  num_motion_samples,
								  size_t                       i_level = 0)
{
	size_t numChildren = top.getNumChildren();
	std::string unique_object_path;

	for (size_t i=0;i<numChildren;i++)
	{
		std::string child_name =top.getChildHeader(i).getName();
		for (size_t indent=0;indent<i_level;indent++)
			std::cout << "  ";
		Alembic::Abc::IObject child(top,child_name);
		const Alembic::Abc::MetaData &child_md = child.getMetaData();
		std::string metadata_string = child_md.serialize();
		std::cout << boost::format("metadata_string='%1%' child_name='%2%'") % metadata_string % child_name;
		if (Alembic::AbcGeom::IPolyMeshSchema::matches(child_md))
		{
			std::cout << "00000" << std::endl;
	        Alembic::AbcGeom::IPolyMesh mesh(top,child_name);
	        Alembic::AbcGeom::IPolyMeshSchema& schema = mesh.getSchema();

	    	Alembic::Abc::IV3fArrayProperty velocities_property = schema.getVelocitiesProperty();
	    	if (velocities_property.valid())
	    	{
				std::cout << "00100" << std::endl;
	    		StringContainer       _concatenated_hierachy_path = i_hierachy_path;
	    		_concatenated_hierachy_path.push_back(child_name);
	    		flatten_string_array(_concatenated_hierachy_path, "_", unique_object_path);

	    		Alembic::AbcGeom::TimeSamplingPtr ts_ptr = mesh.getSchema().getTimeSampling();
	    		Alembic::AbcGeom::TimeSamplingType timeType = ts_ptr->getTimeSamplingType();
	    		Alembic::AbcGeom::chrono_t tpc = timeType.getTimePerCycle();
	    		// Alembic::AbcGeom::chrono_t fps = 1.0/tpc;
    			// std::cout << boost::format("fps = %1%") % fps << std::endl;
	    		if ( timeType.isUniform() )
	    		{
	    			size_t start_frame = ts_ptr->getStoredTimes()[0] / tpc;
	    			// std::cout << boost::format("start_frame = %1%") % start_frame << std::endl;
	    			std::string arnold_filename = (boost::format("%s.%04d.ass") % unique_object_path % i_requested_index).str();
	    			export_polymesh_as_arnold_ass(mesh,start_frame,i_requested_index,arnold_filename,num_motion_samples,i_relative_shutter_open,i_relative_shutter_close);
	    		}
	    	}
		}
		else if (Alembic::AbcGeom::IPointsSchema::matches(child_md))
		{
			std::cout << " of type Points";
	        Alembic::AbcGeom::IPoints points(top,child_name);
	        Alembic::AbcGeom::IPointsSchema& schema = points.getSchema();

	    	Alembic::Abc::IV3fArrayProperty velocities_property = schema.getVelocitiesProperty();
	    	Alembic::Abc::IUInt64ArrayProperty ids_property = schema.getIdsProperty();
	    	if (velocities_property.valid() && ids_property.valid())
	    	{
	    		StringContainer       _concatenated_hierachy_path = i_hierachy_path;
	    		_concatenated_hierachy_path.push_back(child_name);
	    		flatten_string_array(_concatenated_hierachy_path, "_", unique_object_path);
	    		std::cout << boost::format(" of type Points, unique_object_path : '%1%'") % unique_object_path << std::endl;

	    		Alembic::AbcGeom::TimeSamplingPtr ts_ptr = points.getSchema().getTimeSampling();
	    		Alembic::AbcGeom::TimeSamplingType timeType = ts_ptr->getTimeSamplingType();
	    		Alembic::AbcGeom::chrono_t tpc = timeType.getTimePerCycle();
	    		// Alembic::AbcGeom::chrono_t fps = 1.0/tpc;
    			// std::cout << boost::format("fps = %1%") % fps << std::endl;
	    		if ( timeType.isUniform() )
	    		{
		    		std::cout << " timeType.isUniform() == true" << std::endl;
	    			size_t start_frame = ts_ptr->getStoredTimes()[0] / tpc;
	    			// std::cout << boost::format("start_frame = %1%") % start_frame << std::endl;
	    			std::string arnold_filename = (boost::format("%s.%04d.ass") % unique_object_path % i_requested_index).str();
		    		std::cout << boost::format(" arnold_filename : '%1%'") % arnold_filename << std::endl;
	    			export_points_as_arnold_ass(points,start_frame,i_requested_index,arnold_filename,num_motion_samples,i_relative_shutter_open,i_relative_shutter_close);
	    		}
	    		else if (timeType.isCyclic())
	    		{
		    		std::cout << " timeType.isCyclic() == true" << std::endl;
	    		}
	    		else if (timeType.isAcyclic())
	    		{
		    		std::cout << " timeType.isAcyclic() == true" << std::endl;
	    		}
	    	}
		}
		std::cout << std::endl;
		StringContainer       concatenated_hierachy_path = i_hierachy_path;
		concatenated_hierachy_path.push_back(child_name);
		locate_geometry_in_hierarchy(child,
								concatenated_hierachy_path,
								i_requested_index,
								i_relative_shutter_open,
								i_relative_shutter_close,
								num_motion_samples,
								i_level+1);
	}
}

int main(int argc, char** argv)
{
    if (argc != 4)
    {
        std::cerr << "Usage : " << argv[0] << " <Frame (integer)> <Alembic Archive> <Arnold ASS>" << std::endl;
        return 1;
    }
    size_t frame_to_export(atoi(argv[1]));
    std::string alembic_fileName(argv[2]);
    std::string ass_fileName(argv[3]);
    Alembic::AbcGeom::IArchive alembic_archive;
    alembic_archive = Alembic::AbcGeom::IArchive( Alembic::AbcCoreOgawa::ReadArchive(), alembic_fileName );

	StringContainer       hierachy_path;
	float relative_shutter_open = -0.25f;
	float relative_shutter_close = 0.25f;
	AtByte num_motion_samples = 3;

	locate_geometry_in_hierarchy(alembic_archive.getTop(),hierachy_path,frame_to_export,relative_shutter_open,relative_shutter_close,num_motion_samples);

    return 0;
}

// == Emacs ================
// -------------------------
// Local variables:
// tab-width: 4
// indent-tabs-mode: t
// c-basic-offset: 4
// end:
//
// == vi ===================
// -------------------------
// Format block
// ex:ts=4:sw=4:expandtab
// -------------------------


/*
 * Update.hpp
 *
 *  Created on: Feb 9, 2014
 *      Author: Bloeschm
 */

#ifndef LWF_UPDATEMODEL_HPP_
#define LWF_UPDATEMODEL_HPP_

#include "lightweight_filtering/common.hpp"
#include "lightweight_filtering/ModelBase.hpp"
#include "lightweight_filtering/PropertyHandler.hpp"
#include "lightweight_filtering/SigmaPoints.hpp"
#include "lightweight_filtering/OutlierDetection.hpp"
#include <Eigen/StdVector>

namespace LWF{

template<typename Innovation, typename FilterState, typename Meas, typename Noise, typename OutlierDetection = OutlierDetectionDefault, bool isCoupled = false>
class Update: public ModelBase<Update<Innovation,FilterState,Meas,Noise,OutlierDetection,isCoupled>,Innovation,typename FilterState::mtState,Noise>, public PropertyHandler{
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW
  static_assert(!isCoupled || Noise::D_ == FilterState::noiseExtensionDim_,"Noise Size for coupled Update must match noise extension of prediction!");
  typedef ModelBase<Update<Innovation,FilterState,Meas,Noise,OutlierDetection,isCoupled>,Innovation,typename FilterState::mtState,Noise> mtModelBase;
  typedef FilterState mtFilterState;
  typedef typename mtFilterState::mtState mtState;
  typedef typename mtModelBase::mtInputTuple mtInputTuple;
  typedef typename mtFilterState::mtPredictionMeas mtPredictionMeas;
  typedef typename mtFilterState::mtPredictionNoise mtPredictionNoise;
  typedef Innovation mtInnovation;
  typedef Meas mtMeas;
  typedef Noise mtNoise;
  typedef OutlierDetection mtOutlierDetection;
  mtMeas meas_; // TODO change to pointer, or remove
  static const bool coupledToPrediction_ = isCoupled;
  bool useSpecialLinearizationPoint_;
  bool useImprovedJacobian_;
  bool hasConverged_;
  bool successfulUpdate_;
  mutable bool cancelIteration_;
  mutable int candidateCounter_;
  mutable Eigen::MatrixXd H_;
  Eigen::MatrixXd Hlin_;
  Eigen::MatrixXd boxMinusJac_;
  Eigen::MatrixXd Hn_;
  Eigen::MatrixXd updnoiP_;
  Eigen::MatrixXd noiP_;
  Eigen::MatrixXd preupdnoiP_;
  Eigen::MatrixXd C_;
  mtInnovation y_;
  mutable Eigen::MatrixXd Py_;
  Eigen::MatrixXd Pyinv_;
  typename mtInnovation::mtDifVec innVector_;
  mtInnovation yIdentity_;
  typename mtState::mtDifVec updateVec_;
  mtState linState_;
  double updateVecNorm_;
  Eigen::MatrixXd K_;
  Eigen::MatrixXd Pyx_;
  mutable typename mtState::mtDifVec difVecLinInv_;

  SigmaPoints<mtState,2*mtState::D_+1,2*(mtState::D_+mtNoise::D_)+1,0> stateSigmaPoints_;
  SigmaPoints<mtNoise,2*mtNoise::D_+1,2*(mtState::D_+mtNoise::D_)+1,2*mtState::D_> stateSigmaPointsNoi_;
  SigmaPoints<mtInnovation,2*(mtState::D_+mtNoise::D_)+1,2*(mtState::D_+mtNoise::D_)+1,0> innSigmaPoints_;
  SigmaPoints<mtNoise,2*(mtNoise::D_+mtPredictionNoise::D_)+1,2*(mtState::D_+mtNoise::D_+mtPredictionNoise::D_)+1,2*(mtState::D_)> coupledStateSigmaPointsNoi_;
  SigmaPoints<mtInnovation,2*(mtState::D_+mtNoise::D_+mtPredictionNoise::D_)+1,2*(mtState::D_+mtNoise::D_+mtPredictionNoise::D_)+1,0> coupledInnSigmaPoints_;
  SigmaPoints<LWF::VectorElement<mtState::D_>,2*mtState::D_+1,2*mtState::D_+1,0> updateVecSP_;
  SigmaPoints<mtState,2*mtState::D_+1,2*mtState::D_+1,0> posterior_;
  double alpha_;
  double beta_;
  double kappa_;
  double updateVecNormTermination_;
  int maxNumIteration_;
  int iterationNum_;
  mtOutlierDetection outlierDetection_;
  unsigned int numSequences;
  bool disablePreAndPostProcessingWarning_;
  Update(): H_((int)(mtInnovation::D_),(int)(mtState::D_)),
      Hlin_((int)(mtInnovation::D_),(int)(mtState::D_)),
      boxMinusJac_((int)(mtState::D_),(int)(mtState::D_)),
      Hn_((int)(mtInnovation::D_),(int)(mtNoise::D_)),
      updnoiP_((int)(mtNoise::D_),(int)(mtNoise::D_)),
      noiP_((int)(mtNoise::D_),(int)(mtNoise::D_)),
      preupdnoiP_((int)(mtPredictionNoise::D_),(int)(mtNoise::D_)),
      C_((int)(mtState::D_),(int)(mtInnovation::D_)),
      Py_((int)(mtInnovation::D_),(int)(mtInnovation::D_)),
      Pyinv_((int)(mtInnovation::D_),(int)(mtInnovation::D_)),
      K_((int)(mtState::D_),(int)(mtInnovation::D_)),
      Pyx_((int)(mtInnovation::D_),(int)(mtState::D_)){
    alpha_ = 1e-3;
    beta_ = 2.0;
    kappa_ = 0.0;
    updateVecNormTermination_ = 1e-6;
    maxNumIteration_  = 10;
    updnoiP_.setIdentity();
    updnoiP_ *= 0.0001;
    noiP_.setZero();
    preupdnoiP_.setZero();
    useSpecialLinearizationPoint_ = false;
    useImprovedJacobian_ = false;
    yIdentity_.setIdentity();
    updateVec_.setIdentity();
    refreshNoiseSigmaPoints();
    // refreshUKFParameter();
    mtNoise n;
    n.setIdentity();
    n.registerCovarianceToPropertyHandler_(updnoiP_,this,"UpdateNoise.");
    doubleRegister_.registerScalar("alpha",alpha_);
    doubleRegister_.registerScalar("beta",beta_);
    doubleRegister_.registerScalar("kappa",kappa_);
    doubleRegister_.registerScalar("updateVecNormTermination",updateVecNormTermination_);
    intRegister_.registerScalar("maxNumIteration",maxNumIteration_);
    outlierDetection_.setEnabledAll(false);
    numSequences = 1;
    disablePreAndPostProcessingWarning_ = false;
  };
  virtual ~Update(){};
  void refreshNoiseSigmaPoints(){
    if(noiP_ != updnoiP_){
      noiP_ = updnoiP_;
      stateSigmaPointsNoi_.computeFromZeroMeanGaussian(noiP_);
    }
  }
  void refreshUKFParameter(){
    stateSigmaPoints_.computeParameter(alpha_,beta_,kappa_);
    innSigmaPoints_.computeParameter(alpha_,beta_,kappa_);
    coupledInnSigmaPoints_.computeParameter(alpha_,beta_,kappa_);
    updateVecSP_.computeParameter(alpha_,beta_,kappa_);
    posterior_.computeParameter(alpha_,beta_,kappa_);
    stateSigmaPointsNoi_.computeParameter(alpha_,beta_,kappa_);
    stateSigmaPointsNoi_.computeFromZeroMeanGaussian(noiP_);
    coupledStateSigmaPointsNoi_.computeParameter(alpha_,beta_,kappa_);
  }
  void refreshProperties(){
    refreshPropertiesCustom();
    // refreshUKFParameter();
  }
  virtual void refreshPropertiesCustom(){}
  void eval_(mtInnovation& x, const mtInputTuple& inputs, double dt) const{
    evalInnovation(x,std::get<0>(inputs),std::get<1>(inputs));
  }
  template<int i,typename std::enable_if<i==0>::type* = nullptr>
  void jacInput_(Eigen::MatrixXd& F, const mtInputTuple& inputs, double dt) const{
    jacState(F,std::get<0>(inputs));
  }
  template<int i,typename std::enable_if<i==1>::type* = nullptr>
  void jacInput_(Eigen::MatrixXd& F, const mtInputTuple& inputs, double dt) const{
    jacNoise(F,std::get<0>(inputs));
  }
  virtual void evalInnovation(mtInnovation& y, const mtState& state, const mtNoise& noise) const = 0;
  virtual void evalInnovationShort(mtInnovation& y, const mtState& state) const{
    mtNoise n; // TODO get static for Identity()
    n.setIdentity();
    evalInnovation(y,state,n);
  }
  virtual void jacState(Eigen::MatrixXd& F, const mtState& state) const = 0;
  virtual void jacNoise(Eigen::MatrixXd& F, const mtState& state) const = 0;
  virtual void preProcess(mtFilterState& filterState, const mtMeas& meas, bool& isFinished){
    isFinished = false;
    if(!disablePreAndPostProcessingWarning_){
      std::cout << "Warning: update preProcessing is not implemented!" << std::endl;
    }
  }
  virtual bool extraOutlierCheck(const mtState& state) const{
    return hasConverged_;
  }
  virtual bool generateCandidates(const mtFilterState& filterState, mtState& candidate) const{
    candidate = filterState.state_;
    candidateCounter_++;
    if(candidateCounter_<=1)
      return true;
    else
      return false;
  }
  virtual void postProcess(mtFilterState& filterState, const mtMeas& meas, const mtOutlierDetection& outlierDetection, bool& isFinished){
    isFinished = true;
    if(!disablePreAndPostProcessingWarning_){
      std::cout << "Warning: update postProcessing is not implemented!" << std::endl;
    }
  }
  int performUpdate(mtFilterState& filterState, const mtMeas& meas){
    bool isFinished = true;
    int r = 0;
    do {
      preProcess(filterState,meas,isFinished);
      if(!isFinished){
        switch(filterState.mode_){
          case ModeEKF:
            r = performUpdateEKF(filterState,meas);
            break;
          case ModeUKF:
            r = performUpdateUKF(filterState,meas);
            break;
          case ModeIEKF:
            r = performUpdateIEKF(filterState,meas);
            break;
          default:
            r = performUpdateEKF(filterState,meas);
            break;
        }
      }
      postProcess(filterState,meas,outlierDetection_,isFinished);
      filterState.state_.fix();
      enforceSymmetry(filterState.cov_);
    } while (!isFinished);
    return r;
  }
  int performUpdateEKF(mtFilterState& filterState, const mtMeas& meas){
    meas_ = meas;
    if(!useSpecialLinearizationPoint_){
      this->jacState(H_,filterState.state_);
      Hlin_ = H_;
      this->jacNoise(Hn_,filterState.state_);
      this->evalInnovationShort(y_,filterState.state_);
    } else {
      filterState.state_.boxPlus(filterState.difVecLin_,linState_);
      this->jacState(H_,linState_);
      if(useImprovedJacobian_){
        filterState.state_.boxMinusJac(linState_,boxMinusJac_);
        Hlin_ = H_*boxMinusJac_;
      } else {
        Hlin_ = H_;
      }
      this->jacNoise(Hn_,linState_);
      this->evalInnovationShort(y_,linState_);
    }

    if(isCoupled){
      C_ = filterState.G_*preupdnoiP_*Hn_.transpose();
      Py_ = Hlin_*filterState.cov_*Hlin_.transpose() + Hn_*updnoiP_*Hn_.transpose() + Hlin_*C_ + C_.transpose()*Hlin_.transpose();
    } else {
      Py_ = Hlin_*filterState.cov_*Hlin_.transpose() + Hn_*updnoiP_*Hn_.transpose();
    }
    y_.boxMinus(yIdentity_,innVector_);

    // Outlier detection // TODO: adapt for special linearization point
    outlierDetection_.doOutlierDetection(innVector_,Py_,Hlin_);
    Pyinv_.setIdentity();
    Py_.llt().solveInPlace(Pyinv_);

    // Kalman Update
    if(isCoupled){
      K_ = (filterState.cov_*Hlin_.transpose()+C_)*Pyinv_;
    } else {
      K_ = filterState.cov_*Hlin_.transpose()*Pyinv_;
    }
    filterState.cov_ = filterState.cov_ - K_*Py_*K_.transpose();
    if(!useSpecialLinearizationPoint_){
      updateVec_ = -K_*innVector_;
    } else {
      filterState.state_.boxMinus(linState_,difVecLinInv_);
      updateVec_ = -K_*(innVector_+H_*difVecLinInv_); // includes correction for offseted linearization point, dif must be recomputed (a-b != (-(b-a)))
    }
    filterState.state_.boxPlus(updateVec_,filterState.state_);
    return 0;
  }
  int performUpdateIEKF(mtFilterState& filterState, const mtMeas& meas){
    meas_ = meas;
    successfulUpdate_ = false;
    candidateCounter_ = 0;

    std::vector<double> scores;
    std::vector<mtState, Eigen::aligned_allocator<mtState>> states;
    double bestScore = -1.0;
    mtState bestState;
    MXD bestCov;

    while(generateCandidates(filterState,linState_)){
      cancelIteration_ = false;
      hasConverged_ = false;
      for(iterationNum_=0;iterationNum_<maxNumIteration_ && !hasConverged_ && !cancelIteration_;iterationNum_++){
        this->jacState(H_,linState_);
        this->jacNoise(Hn_,linState_);
        this->evalInnovationShort(y_,linState_);

        if(isCoupled){
          C_ = filterState.G_*preupdnoiP_*Hn_.transpose();
          Py_ = H_*filterState.cov_*H_.transpose() + Hn_*updnoiP_*Hn_.transpose() + H_*C_ + C_.transpose()*H_.transpose();
        } else {
          Py_ = H_*filterState.cov_*H_.transpose() + Hn_*updnoiP_*Hn_.transpose();
        }
        y_.boxMinus(yIdentity_,innVector_);

        // Outlier detection
        outlierDetection_.doOutlierDetection(innVector_,Py_,H_);
        Pyinv_.setIdentity();
        Py_.llt().solveInPlace(Pyinv_);

        // Kalman Update
        if(isCoupled){
          K_ = (filterState.cov_*H_.transpose()+C_)*Pyinv_;
        } else {
          K_ = filterState.cov_*H_.transpose()*Pyinv_;
        }
        filterState.state_.boxMinus(linState_,difVecLinInv_);
        updateVec_ = -K_*(innVector_+H_*difVecLinInv_)+difVecLinInv_; // includes correction for offseted linearization point, dif must be recomputed (a-b != (-(b-a)))
        linState_.boxPlus(updateVec_,linState_);
        updateVecNorm_ = updateVec_.norm();
        hasConverged_ = updateVecNorm_<=updateVecNormTermination_;
      }
      if(extraOutlierCheck(linState_)){
        successfulUpdate_ = true;
        double score = (innVector_.transpose()*Pyinv_*innVector_)(0);
        scores.push_back(score);
        states.push_back(linState_);
        if(bestScore == -1.0 || score < bestScore){
          bestScore = score;
          bestState = linState_;
          bestCov = filterState.cov_ - K_*Py_*K_.transpose();
        }
      }
    }

    if(successfulUpdate_){
      if(scores.size() == 1){
        filterState.state_ = bestState;
        filterState.cov_ = bestCov;
      } else {
        bool foundOtherMin = false;
        for(auto it = states.begin();it!=states.end();it++){
          bestState.boxMinus(*it,difVecLinInv_);
          if(difVecLinInv_.norm()>2*updateVecNormTermination_){
            foundOtherMin = true;
            break;
          }
        }
        if(!foundOtherMin){
          filterState.state_ = bestState;
          filterState.cov_ = bestCov;
        } else {
          successfulUpdate_ = false;
        }
      }
    }
    return 0;
  }
  int performUpdateUKF(mtFilterState& filterState, const mtMeas& meas){
    meas_ = meas;
    handleUpdateSigmaPoints<isCoupled>(filterState);
    y_.boxMinus(yIdentity_,innVector_);

    outlierDetection_.doOutlierDetection(innVector_,Py_,Pyx_);
    Pyinv_.setIdentity();
    Py_.llt().solveInPlace(Pyinv_);

    // Kalman Update
    K_ = Pyx_.transpose()*Pyinv_;
    filterState.cov_ = filterState.cov_ - K_*Py_*K_.transpose();
    updateVec_ = -K_*innVector_;

    // Adapt for proper linearization point
    updateVecSP_.computeFromZeroMeanGaussian(filterState.cov_);
    for(unsigned int i=0;i<2*mtState::D_+1;i++){
      filterState.state_.boxPlus(updateVec_+updateVecSP_(i).v_,posterior_(i));
    }
    posterior_.getMean(filterState.state_);
    posterior_.getCovarianceMatrix(filterState.state_,filterState.cov_);
    return 0;
  }
  template<bool IC = isCoupled, typename std::enable_if<(IC)>::type* = nullptr>
  void handleUpdateSigmaPoints(mtFilterState& filterState){
    coupledStateSigmaPointsNoi_.extendZeroMeanGaussian(filterState.stateSigmaPointsNoi_,updnoiP_,preupdnoiP_);
    for(unsigned int i=0;i<coupledInnSigmaPoints_.L_;i++){
      this->evalInnovation(coupledInnSigmaPoints_(i),filterState.stateSigmaPointsPre_(i),coupledStateSigmaPointsNoi_(i));
    }
    coupledInnSigmaPoints_.getMean(y_);
    coupledInnSigmaPoints_.getCovarianceMatrix(y_,Py_);
    coupledInnSigmaPoints_.getCovarianceMatrix(filterState.stateSigmaPointsPre_,Pyx_);
  }
  template<bool IC = isCoupled, typename std::enable_if<(!IC)>::type* = nullptr>
  void handleUpdateSigmaPoints(mtFilterState& filterState){
    refreshNoiseSigmaPoints();
    stateSigmaPoints_.computeFromGaussian(filterState.state_,filterState.cov_);
    for(unsigned int i=0;i<innSigmaPoints_.L_;i++){
      this->evalInnovation(innSigmaPoints_(i),stateSigmaPoints_(i),stateSigmaPointsNoi_(i));
    }
    innSigmaPoints_.getMean(y_);
    innSigmaPoints_.getCovarianceMatrix(y_,Py_);
    innSigmaPoints_.getCovarianceMatrix(stateSigmaPoints_,Pyx_);
  }
  bool testUpdateJacs(double d = 1e-6,double th = 1e-6){
    mtState state;
    mtMeas meas;
    unsigned int s = 1;
    state.setRandom(s);
    meas.setRandom(s);
    return testUpdateJacs(state,meas,d,th);
  }
  bool testUpdateJacs(const mtState& state, const mtMeas& meas, double d = 1e-6,double th = 1e-6){
    mtInputTuple inputs;
    const double dt = 1.0;
    std::get<0>(inputs) = state;
    std::get<1>(inputs).setIdentity(); // Noise is always set to zero for Jacobians
    meas_ = meas;
    return this->testJacs(inputs,d,th,dt);
  }
};

}

#endif /* LWF_UPDATEMODEL_HPP_ */


#include <Eigen/Core>

#include <vector>
#include <memory>

#include <CGAL/Boolean_set_operations_2.h>

// standard includes
#include <iostream>

#include "SDOT/PolygonRasterize.h"
#include "SDOT/RegularGrid.h"
#include "SDOT/LaguerreDiagram.h"

using namespace sdot;

int main(int argc, char* argv[])
{
  int numPts = 4;
  Eigen::VectorXd costs = Eigen::VectorXd::Ones(numPts);

  Eigen::Matrix2Xd pts(2,numPts);
  pts <<  0.680375,  0.566198,  0.823295,  0.329554,//  0.444451, 0.0452059,  0.270431,  0.904459,  0.271423,  0.716795,
          0.211234,   0.59688,  0.604897,  0.536459;//,   0.10794,  0.257742, 0.0268018,   0.83239,  0.434594,  0.213938;

  std::cout << "Points = \n";
  std::cout << "[[" << pts(0,0) << "," << pts(1,0) << "]";
  for(int i=1; i<numPts; ++i){
    std::cout << ", [" << pts(0,i) << "," << pts(1,i) << "]";
  }
  std::cout << "]" << std::endl;

  Eigen::Matrix2Xd domain(2,4);
  domain << 0.0, 1.0, 1.0, 0.0,
            0.0, 0.0, 1.0, 1.0;
  // domain << 0.0, 1.0, 1.0, 0.0,
  //           0.0, 0.0, 1.0, 1.0;

  LaguerreDiagram diag(domain(0,0), domain(0,1), domain(1,0), domain(1,2), pts, costs);

  auto grid = std::make_shared<RegularGrid>(domain(0,0),domain(1,0), domain(0,2), domain(1,2), 10, 10);

  double area = 0.0;
  Eigen::VectorXd localAreas = Eigen::VectorXd::Zero(numPts);
  Eigen::MatrixXd cellAreas = Eigen::MatrixXd::Zero(grid->NumCells(0), grid->NumCells(1));

  for(int polyInd=0; polyInd<numPts; ++polyInd){
    std::cout << "\n\n==================================\n";
    std::cout << "Polygon " << polyInd << std::endl;
    std::shared_ptr<PolygonRasterizeIter::Polygon_2> poly = diag.GetCell(polyInd)->ToCGAL();

    auto vertIt = poly->vertices_begin();
    std::cout << "[[" << vertIt->x() << "," << vertIt->y() << "]";
    vertIt++;
    for(;  vertIt != poly->vertices_end(); ++vertIt){
      std::cout << ", [" << vertIt->x() << "," << vertIt->y() << "]";
    }
    std::cout << "]" << std::endl;

    PolygonRasterizeIter gridIter(grid,poly);

    //
    // unsigned int oldYInd = gridIter.Indices().second;
    // std::cout << "yind = " << oldYInd << std::endl;
    // std::cout << "yval = " << grid->yMin + oldYInd*grid->dy << std::endl;
    // std::cout << "    ";

    while(gridIter.IsValid()){

      // if(gridIter.Indices().second != oldYInd){
      //   oldYInd = gridIter.Indices().second;
      //   std::cout << "\nyind = " << oldYInd << std::endl;
      //   std::cout << "yval = " << grid->yMin + oldYInd*grid->dy << std::endl;
      //   std::cout << "    ";
      // }

      double cellArea;

      if(gridIter.IsBoundary()){

        cellArea = CGAL::to_double( gridIter.OverlapPoly()->area() );
        //std::cout << "  " << gridIter.Indices().first << ", " << gridIter.Indices().second << " -> Overlap poly = " << *gridIter.OverlapPoly() << std::endl;
        // unsigned int indX = gridIter.Indices().first;
        // unsigned int indY = gridIter.Indices().second;
        //
        // PolygonRasterizeIter::Polygon_2 tempPoly;
        // tempPoly.push_back(PolygonRasterizeIter::Point_2(grid->xMin + indX*grid->dx, grid->yMin+indY*grid->dy));
        // tempPoly.push_back(PolygonRasterizeIter::Point_2(grid->xMin + (indX+1)*grid->dx, grid->yMin+indY*grid->dy));
        // tempPoly.push_back(PolygonRasterizeIter::Point_2(grid->xMin + (indX+1)*grid->dx, grid->yMin+(indY+1)*grid->dy));
        // tempPoly.push_back(PolygonRasterizeIter::Point_2(grid->xMin + indX*grid->dx, grid->yMin+(indY+1)*grid->dy));
        //
        // // Compute the intersection of P and Q.
        // std::list<PolygonRasterizeIter::Polygon_with_holes_2> interList;
        // CGAL::intersection(*poly, tempPoly, std::back_inserter(interList));
        //
        // double trueArea = CGAL::to_double( interList.begin()->outer_boundary().area() );
        //
        // //std::cout << *gridIter.OverlapPoly() << std::endl;
        // double error = cellArea-trueArea;
        // if(std::abs(error)>std::numeric_limits<double>::epsilon()){
        //   // std::cout << "  Cell area = " << cellArea << " with error " << error << std::endl;
        // }
      }else{
        cellArea = grid->dx*grid->dy;
      }

      std::cout << "    area(" << gridIter.Indices().first << "," << gridIter.Indices().second << ")= " << cellArea << std::endl;
      cellAreas(gridIter.Indices().first, gridIter.Indices().second) += cellArea;
      localAreas(polyInd) += std::abs(cellArea);
      area += std::abs(cellArea);

      gridIter.Increment();
    }

    std::cout << "  Polygon area = " << localAreas(polyInd) << std::endl;
    std::cout << "  Polygon area error = " << localAreas(polyInd) - CGAL::to_double( poly->area() ) << std::endl;
    // std::cout << "\n\n";
  }

  // Eigen::VectorXd trueAreas(3);
  // trueAreas << 0.15*0.15, (domain(0,2)-0.15)*0.15 + 0.5*(domain(1,2)-0.15)*(domain(0,2)-0.15), 0.15*(domain(1,2)-0.15) + 0.5*(domain(1,2)-0.15)*(domain(0,2)-0.15);

  // List all cells with errors
  for(unsigned int yInd=0; yInd<grid->NumCells(1); ++yInd){
    for(unsigned int xInd=0; xInd<grid->NumCells(0); ++xInd){
      if(std::abs(cellAreas(xInd,yInd) - grid->dx*grid->dy)>1e-15){
        std::cout << "Cell " << xInd << "," << yInd << " error = " << cellAreas(xInd,yInd) - grid->dx*grid->dy << std::endl;
      }
    }
  }

  std::cout << "Local areas = " << localAreas.transpose() << std::endl;

  //std::cout << "dx*dy = " << grid->dx * grid->dy << std::endl;
  std::cout << "True total area = " << grid->dx*grid->dy*grid->NumCells() << std::endl;
  std::cout << "Total Area = " << area << std::endl;
  std::cout << "Error = " << area - grid->dx*grid->dy*grid->NumCells() << std::endl;
  return 0;
}


/*
Brian Staber (brian.staber@gmail.com)
*/

#ifndef SHINOZUKAPP_LAYEREDCOMP_HPP
#define SHINOZUKAPP_LAYEREDCOMP_HPP

#include "meshpp.hpp"
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/uniform_real_distribution.hpp>
#include <boost/math/special_functions/erf.hpp>
#include <boost/random/normal_distribution.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/beta.hpp>

class shinozuka_layeredcomp
{
public:

    shinozuka_layeredcomp();
    shinozuka_layeredcomp(int & nu, double & L1, double & L2, double & L3);
    ~shinozuka_layeredcomp();

    template<typename typearg>
    double tau_beta(typearg & beta);
    double s_tau(double & tau);

    void generator_gauss_points(Epetra_SerialDenseVector & v, mesh & Mesh, std::vector<int> & phase);
    void generator_one_gauss_point(Epetra_SerialDenseVector & v, mesh & Mesh, std::vector<int> & phase, double & xi, double & eta, double & zeta);

    void icdf_gamma(Epetra_Vector & V, Epetra_Vector & G, double & alpha, double & beta);
    void icdf_beta(Epetra_Vector & V, Epetra_Vector & B, double & tau1, double & tau2);

    int order;
    double l1;
    double l2;
    double l3;
    boost::random::mt19937 rng;
    boost::random::uniform_real_distribution<> phi_;
    boost::random::uniform_real_distribution<> psi_;
};
#endif


// nvcc -O3 CholeskyDecomp_t.cu --expt-relaxed-constexpr -gencode arch=compute_61,code=sm_61 --compiler-options="-Ofast -march=native"
// add -DDOPROF to run  nvprof --metrics all

#include <algorithm>
#include <cassert>
#include <chrono>
#include <iomanip>
#include <iostream>
#include <limits>
#include <memory>
#include <random>

#include <Eigen/Core>
#include <Eigen/Eigenvalues>

#include "DataFormats/Math/interface/choleskyInversion.h"

constexpr int stride() { return 5 * 1024; }
template <int DIM>
using MXN = Eigen::Matrix<double, DIM, DIM>;
template <int DIM>
using MapMX = Eigen::Map<MXN<DIM>, 0, Eigen::Stride<DIM * stride(), stride()> >;

// generate matrices
template <class M>
void genMatrix(M& m) {
  using T = typename std::remove_reference<decltype(m(0, 0))>::type;
  int n = M::ColsAtCompileTime;
  std::mt19937 eng;
  // std::mt19937 eng2;
  std::uniform_real_distribution<T> rgen(0., 1.);

  // generate first diagonal elemets
  for (int i = 0; i < n; ++i) {
    double maxVal = i * 10000 / (n - 1) + 1;  // max condition is 10^4
    m(i, i) = maxVal * rgen(eng);
  }
  for (int i = 0; i < n; ++i) {
    for (int j = 0; j < i; ++j) {
      double v = 0.3 * std::sqrt(m(i, i) * m(j, j));  // this makes the matrix pos defined
      m(i, j) = v * rgen(eng);
      m(j, i) = m(i, j);
    }
  }
}

template <int N>
void go(bool soa) {
  constexpr unsigned int DIM = N;
  using MX = MXN<DIM>;
  std::cout << "testing Matrix of dimension " << DIM << " size " << sizeof(MX) << " in " << (soa ? "SOA" : "AOS")
            << " mode" << std::endl;

  auto start = std::chrono::high_resolution_clock::now();
  auto delta = start - start;

  constexpr unsigned int SIZE = 4 * 1024;

  alignas(128) MX mm[stride()];  // just storage in case of SOA
  double* __restrict__ p = (double*)__builtin_assume_aligned(mm, 128);

  if (soa) {
    for (unsigned int i = 0; i < SIZE; ++i) {
      MapMX<N> m(p + i);
      genMatrix(m);
    }
  } else {
    for (auto& m : mm)
      genMatrix(m);
  }

  std::cout << mm[SIZE / 2](1, 1) << std::endl;

  if (soa)
    for (unsigned int i = 0; i < SIZE; ++i) {
      MapMX<N> m(p + i);
      math::cholesky::invert(m, m);
      math::cholesky::invert(m, m);
    }
  else
    for (auto& m : mm) {
      math::cholesky::invert(m, m);
      math::cholesky::invert(m, m);
    }

  std::cout << mm[SIZE / 2](1, 1) << std::endl;

  constexpr int NKK =
#ifdef DOPROF
      2;
#else
      1000;
#endif
  for (int kk = 0; kk < NKK; ++kk) {
    delta -= (std::chrono::high_resolution_clock::now() - start);
    if (soa)
#ifdef USE_VECTORIZATION_PRAGMA
#pragma GCC ivdep
#ifdef __clang__
#pragma clang loop vectorize(enable) interleave(enable)
#endif
#endif
      for (unsigned int i = 0; i < SIZE; ++i) {
        MapMX<N> m(p + i);
        math::cholesky::invert(m, m);
      }
    else
#ifdef USE_VECTORIZATION_PRAGMA
#pragma GCC ivdep
#ifdef __clang__
#pragma clang loop vectorize(enable) interleave(enable)
#endif
#endif
      for (auto& m : mm) {
        math::cholesky::invert(m, m);
      }

    delta += (std::chrono::high_resolution_clock::now() - start);
  }

  std::cout << mm[SIZE / 2](1, 1) << std::endl;

  double DNNK = NKK;
  std::cout << "x86 computation took " << std::chrono::duration_cast<std::chrono::milliseconds>(delta).count() / DNNK
            << ' ' << " ms" << std::endl;
}

int main() {
  go<2>(false);
  go<3>(false);
  go<4>(false);
  go<5>(false);
  go<6>(false);

  go<2>(true);
  go<3>(true);
  go<4>(true);
  go<5>(true);
  go<6>(true);

  // go<10>();
  return 0;
}


#include <boost/accumulators/statistics/skewness.hpp>


// Boost.Geometry (aka GGL, Generic Geometry Library)
//
// Copyright (c) 2007-2012 Barend Gehrels, 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)
//
// GD example

// GD is a well-known and often used graphic library to write GIF (and other formats)

// To build and run this example:
// 1) download GD from http://www.libgd.org (currently gd-2.0.35 is assumed)
// 2) add 11 files
//          gd.c, gd_gd.c, gd_gif_out.c, gd_io*.c, gd_security.c, gd_topal.c, gdhelpers.c, gdtables.c
//    to the project or makefile or jamfile
// 3) for windows, add define NONDLL to indicate GD is not used as a DLL
//    (Note that steps 2 and 3 are done in the MSVC gd_example project file and property sheets)

#include <cmath>
#include <cstdio>
#include <vector>

#include <fstream>
#include <sstream>

#include <boost/foreach.hpp>

#include <boost/geometry.hpp>
#include <boost/geometry/multi/geometries/multi_polygon.hpp>

#include <boost/geometry/extensions/gis/latlong/latlong.hpp>
#include <boost/geometry/extensions/gis/geographic/strategies/area_huiller_earth.hpp>

#include <gd.h>

namespace bg = boost::geometry;

// ----------------------------------------------------------------------------
// Read an ASCII file containing WKT's
// (note this function is shared by various examples)
// ----------------------------------------------------------------------------
template <typename Geometry>
inline void read_wkt(std::string const& filename, std::vector<Geometry>& geometries)
{
    std::ifstream cpp_file(filename.c_str());
    if (cpp_file.is_open())
    {
        while (! cpp_file.eof() )
        {
            std::string line;
            std::getline(cpp_file, line);
            if (! line.empty())
            {
                Geometry geometry;
                bg::read_wkt(line, geometry);
                geometries.push_back(geometry);
            }
        }
    }
}

int main()
{
    // Adapt if necessary
    std::string filename = "../data/world.wkt";

    // The world is measured in latlong (degrees), so latlong is appropriate.
    // We ignore holes in this sample (below)
    typedef bg::model::ll::point<bg::degree> point_type;
    typedef bg::model::polygon<point_type> polygon_type;
    typedef bg::model::multi_polygon<polygon_type> country_type;

    std::vector<country_type> countries;

    // Read (for example) world countries
    read_wkt(filename, countries);

    // Create a GD image.
    // A world map, as world.shp, is usually mapped in latitude-longitude (-180..180 and -90..90)
    // For this example we use a very simple "transformation"
    // mapping to 0..720 and 0..360
    const double factor = 2.0;
    gdImagePtr im = gdImageCreateTrueColor(int(360 * factor), int(180 * factor));

    // Allocate three colors
    int blue = gdImageColorResolve(im, 0, 52, 255);
    int green = gdImageColorResolve(im, 0, 255, 0);
    int black = gdImageColorResolve(im, 0, 0, 0);

    // Paint background in blue
    gdImageFilledRectangle(im, 0, 0, im->sx, im->sy, blue);

    // Paint all countries in green
    BOOST_FOREACH(country_type const& country, countries)
    {
        BOOST_FOREACH(polygon_type const& polygon, country)
        {
            // Ignore holes, so take only exterior ring
            bg::model::ring<point_type> const& ring = bg::exterior_ring(polygon);

            // If wished, suppress too small polygons.
            // (Note that even in latlong, area is calculated in square meters)
            double const a = bg::area(ring);
            if (std::fabs(a) > 5000.0e6)
            {
                int const n = ring.size();
                gdPoint* points = new gdPoint[n];

                for (int i = 0; i < n; i++)
                {
                    // Translate lon/lat or x/y to GD x/y points
                    points[i].x = int(factor * (bg::get<0>(ring[i]) + 180.0));
                    points[i].y = im->sy - int(factor * (bg::get<1>(ring[i]) + 90.0));
                }

                // Draw the polygon...
                gdImageFilledPolygon(im, points, n, green);
                // .. and the outline in black...
                gdImagePolygon(im, points, n, black);

                delete[] points;
            }
        }
    }

    // Use GD to create a GIF file
    std::FILE* out = std::fopen("world.gif", "wb");
    if (out != NULL)
    {
        gdImageGif(im, out);
        std::fclose(out);
    }

    gdImageDestroy(im);

    return 0;
}


#include <boost/math/special_functions/next.hpp>
#include <boost/random.hpp>

#include <limits>

#include "caffe/common.hpp"
#include "caffe/util/math_functions.hpp"
#include "caffe/util/more_math_functions.hpp"
#include "caffe/util/rng.hpp"

namespace caffe {

template <typename Dtype>
Dtype caffe_cpu_max_element(const int N, const Dtype *x) {
  Dtype max_val = (*x);
  for (int i = 1; i < N; ++i) {
    ++x;
    if (max_val < (*x)) {
      max_val = (*x);
    }
  }
  return max_val;
}

template int caffe_cpu_max_element<int>(const int N, const int *x);
template float caffe_cpu_max_element<float>(const int N, const float *x);
template double caffe_cpu_max_element<double>(const int N, const double *x);

template <typename Dtype>
Dtype caffe_cpu_sum(const int N, const Dtype *x) {
  Dtype sum = (*x);
  for (int i = 1; i < N; ++i) {
    ++x;
    sum += (*x);
  }
  return sum;
}

template int caffe_cpu_sum<int>(const int N, const int *x);
template float caffe_cpu_sum<float>(const int N, const float *x);
template double caffe_cpu_sum<double>(const int N, const double *x);

/**
Warning: the diff of cpg_blob will be modified!
 **/
template <typename Dtype>
int caffe_cpu_threshold_bbox(Blob<Dtype> *cpg_blob, Blob<Dtype> *bboxes_blob,
                             const float fg_threshold, const int gt_label) {
  // CPG should be in the cpg_blob->gpu_data()
  // CPG:	1	1	heigh_im_	width_im
  caffe_set(bboxes_blob->count(), static_cast<Dtype>(-1),
            bboxes_blob->mutable_cpu_data());

  const int num = cpg_blob->shape(0);
  const int channels = cpg_blob->shape(1);
  const int height = cpg_blob->shape(2);
  const int width = cpg_blob->shape(3);
  CHECK_EQ(num, 1) << "we assume num equal 1";
  CHECK_EQ(channels, 1) << "we assume num equal 1";

  const int num_bboxes = bboxes_blob->shape(0);

  // caffe_copy(cpg_blob->count(), cpg_blob->gpu_data(),
  // cpg_blob->mutable_gpu_diff());
  memcpy(cpg_blob->mutable_cpu_diff(), cpg_blob->cpu_data(),
         cpg_blob->count() * sizeof(Dtype));

  // get maximum value
  // SYNC
  const Dtype *cpg_cpu = cpg_blob->cpu_diff();
  const Dtype max_value = caffe_cpu_max_element(cpg_blob->count(), cpg_cpu);
  const Dtype threshold = fg_threshold * (max_value);
  CHECK_GT(max_value, 0) << "max_value should be greater than 0. #cpg_cpg: "
                         << cpg_blob->count();
  /*LOG_IF(INFO, debug_info_) << "max_value: " << (max_value)*/
  /*<< " threshold: " << threshold;*/

  // id x1 y1 x2 y2
  int bbox_num = 0;
  bool is_trace = false;
  Dtype *bbox = bboxes_blob->mutable_cpu_data();
  const Dtype separate = std::min(height, width) / 5;

  Dtype *cpg_mcpu = cpg_blob->mutable_cpu_diff();
  while (true) {
    for (int n = 0; n < num; ++n) {
      for (int h = 0; h < height; ++h) {
        for (int w = 0; w < width; ++w) {
          for (int c = 0; c < channels; ++c) {
            int index = cpg_blob->offset(n, c, h, w);
            if (cpg_mcpu[index] < threshold && cpg_mcpu[index] > -threshold)
              continue;
            if (is_trace && (w - bbox[2] > separate || h - bbox[3] > separate ||
                             bbox[0] - w > separate || bbox[1] - h > separate))
              continue;
            cpg_mcpu[index] = 0;
            if (!is_trace) {
              bbox[0] = width;
              bbox[1] = height;
              bbox[2] = 0;
              bbox[3] = 0;
              is_trace = true;
            }
            if (w < bbox[0]) bbox[0] = w;
            if (h < bbox[1]) bbox[1] = h;
            if (w > bbox[2]) bbox[2] = w;
            if (h > bbox[3]) bbox[3] = h;
          }
        }
      }
    }
    if (!is_trace) break;
    bbox_num++;
    if (bbox_num == num_bboxes) break;
    bbox += bboxes_blob->offset(1);
    is_trace = false;
  }

  /*if (debug_info_) {*/
  /*for (int i = 0; i < max_bb_per_cls_; ++i) {*/
  /*bbox = bboxes_->mutable_cpu_data() + bboxes_->offset(i);*/
  /*if (bbox[0] == -1) break;*/
  /*LOG(INFO) << "bbox label x1 y1 x2 y2: " << gt_label << " " << bbox[0]*/
  /*<< " " << bbox[1] << " " << bbox[2] << " " << bbox[3];*/
  /*}*/
  /*}*/

  return bbox_num;
}

template int caffe_cpu_threshold_bbox<float>(Blob<float> *cpg_blob,
                                             Blob<float> *bboxes_blob,
                                             const float fg_threshold,
                                             const int gt_label);
template int caffe_cpu_threshold_bbox<double>(Blob<double> *cpg_blob,
                                              Blob<double> *bboxes_blob,
                                              const float fg_threshold,
                                              const int gt_label);

}  // namespace caffe


#ifndef CALIBRATOR_PROCESSES_GENERAL_VASICEK_HPP
#define CALIBRATOR_PROCESSES_GENERAL_VASICEK_HPP

#include <boost/bind.hpp>

#include <ql/math/randomnumbers/sobolrsg.hpp>
#include <ql/models/shortrate/onefactormodel.hpp>
#include <ql/math/integrals/kronrodintegral.hpp>

#include <calibrator/global.hpp>
#include <calibrator/processes/generalornsteinuhlenbeckprocess.hpp>

namespace HJCALIBRATOR
{
	//! Time-dependent %Hull-White model class
	/*! This class implements the time-dependent Hull-White model defined by
	\f[
	dr(t) = (\theta(t) - a(t)r(t))dt + \sigma(t) dW_t ,
	\f]
	a risk premium \f$ \lambda \f$ can also be specified.

	\ingroup shortrate
	*/
	class GeneralizedHullWhite : public OneFactorAffineModel, public TermStructureConsistentModel
	{
	public:
		GeneralizedHullWhite( const Handle<YieldTermStructure>& termStructure,
							  const IntegrableParameter& a,
							  const Parameter& sigma );

		// OnFactorModel virtual override
		boost::shared_ptr<Lattice> tree( const TimeGrid& grid ) const override
		{
			return boost::shared_ptr<Lattice>();
		};
		virtual boost::shared_ptr<ShortRateDynamics> dynamics() const override;

		virtual Real discountBondOption( Option::Type type,
										 Real strike,
										 Time maturity,
										 Time bondMaturity ) const override;

		virtual Real discountBondOption( Option::Type type, Real strike,
										 Time maturity, Time bondStart,
										 Time bondMaturity ) const override;

		Parameter a() const { return a_; }
		Parameter sigma() const { return sigma_; }

		Real a( Time t ) const { return a_( t ); }
		Real sigma( Time t ) const { return sigma_( t ); }

		/*! Futures convexity bias (i.e., the difference between
		futures implied rate and forward rate) calculated as in
		G. Kirikos, D. Novak, "Convexity Conundrums", Risk
		Magazine, March 1997.
		http://www.powerfinance.com/convexity/

		\note t and T should be expressed in yearfraction using
		deposit day counter, F_quoted is futures' market price.
		*/
		// to be implemented with variable a & \sigma
		/*
		static Rate convexityBias( Real futurePrice,
								   Time t,
								   Time T,
								   Real sigma,
								   Real a );
		*/

		static std::vector<bool> FixedReversion() {
			std::vector<bool> c( 2 );
			c[0] = true; c[1] = false;
			return c;
		}

	protected:
		void generateArguments();

		virtual Real A( Time t, Time T ) const;
		virtual Real B( Time t, Time T ) const;

		Parameter& a_;
		Parameter& sigma_;

	private:
		class Dynamics;
		class FittingParameter;

		boost::shared_ptr<FittingParameter> alpha_;
	};

	class GeneralizedHullWhite::FittingParameter : public TermStructureFittingParameter
	{
	private:
		class Impl : public Parameter::Impl
		{
		public:
			Impl( const Handle<YieldTermStructure>& termStructure,
				  const IntegrableParameter& a,
				  const Parameter& sigma );
			~Impl()
			{
				gsl_integration_workspace_free( int_wrkspcs_ );
			}

			Real value( const Array&, Time t ) const override;

			Real B( Time t, Time T ) const;
			Real variance( Time s, Time t ) const;

		private:
			Real E( Time t0, Time t1, Real multiplier = 1. ) const;
			Real integrand( Time u, Time t ) const;
			Real integrandVr( Time t ) const;
			Real value_integral( Time t ) const;

			gsl_integration_workspace* int_wrkspcs_;
			//gsl_function integrand_Vr_;
			//gsl_function OneOverE_;

			GaussKronrodAdaptive integrator_;
			//SimpsonIntegral integrator_;
			//boost::function<Real( Real )> OneOverEintegrand_; // Eq. 31 integrand
			//boost::function<Real( Real )> Vrintegrand_; // Eq. 37 integrand

			Handle<YieldTermStructure> termStructure_;
			IntegrableParameter a_;
			Parameter sigma_;
		};
	public:
		FittingParameter( const Handle<YieldTermStructure>& termStructure,
						  const IntegrableParameter& a,
						  const Parameter& sigma )
			: TermStructureFittingParameter( boost::shared_ptr<Parameter::Impl>( new FittingParameter::Impl( termStructure, a, sigma ) ) )
		{}

		Real B( Time t, Time T ) const
		{
			return boost::static_pointer_cast<Impl>(implementation())->B( t, T );
		}

		Real variance( Time t, Time T ) const
		{
			return boost::static_pointer_cast<Impl>(implementation())->variance( t, T );
		}
	};

	//! Short-rate dynamics in the time-dependent Hull-White model
	/*! The short-rate follows an time-dependent Hull-White process */
	class GeneralizedHullWhite::Dynamics : public OneFactorModel::ShortRateDynamics {
	public:
		Dynamics( const FittingParameter& fitting,
				  const IntegrableParameter& a,
				  const Parameter& sigma )
			: ShortRateDynamics( boost::shared_ptr<StochasticProcess1D>(
				new GeneralizedOrnsteinUhlenbeckProcess( a, sigma ) ) )
			, fitting_( fitting )
		{}

		virtual Real variable( Time t, Rate r ) const {
			return r - fitting_( t );
		}
		virtual Real shortRate( Time t, Real x ) const {
			return x + fitting_( t );
		}

		FittingParameter fitting_;
	};

	// inline definitions
	inline boost::shared_ptr<OneFactorModel::ShortRateDynamics>
		GeneralizedHullWhite::dynamics() const
	{
		return boost::shared_ptr<ShortRateDynamics>(
			new Dynamics( *alpha_, static_cast<IntegrableParameter>(a()), sigma() ) );
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::value( const Array&, Time t ) const
	{
		/*
		Rate forwardRate = termStructure_->forwardRate( t, t, Continuous, NoFrequency );
		Real Et = E( 0, t );

		boost::function<Real( Real )> I_t;
		I_t = boost::bind( &GeneralizedHullWhite::FittingParameter::Impl::integrand, this, _1, t );
		gsl_function tmpf = convertToGslFunction( I_t );

		Real result, error;
		gsl_integration_qags( &tmpf, 0, t, 0, 1e-7, 1000, int_wrkspcs_, &result, &error );

		return forwardRate;// +result;
		*/

		Rate forwardRate = termStructure_->forwardRate( t, t, Continuous, NoFrequency );

		Real intsum = value_integral( t );

		return forwardRate + intsum;
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::integrand( Time u, Time t ) const
	{
		/* eq.36 */
		Real sigma_u = sigma_( u );

		return E( 0, u ) * sigma_u * sigma_u * B( u, t );
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::integrandVr( Time t ) const
	{
		Real sigma_t = sigma_( t );
		Real E_t = E( 0, t );

		return E_t * E_t * sigma_t * sigma_t;
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::E( Time t0, Time t1, Real multiplier ) const
	{
		/* eq. 30 */
		return exp( multiplier * a_.integral( t0, t1 ) );
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::B( Time t, Time T ) const
	{
		/* eq. 31 */

		auto lambda = [&, t]( Time u )
		{
			return 1 / E( t, u );
		};

		return integrator_( lambda, t, T );
	}

	inline Real GeneralizedHullWhite::FittingParameter::Impl::variance( Time s, Time t ) const
	{
		/* eq. 37 */
		const Parameter& sigma_i = sigma_;

		auto integrand = [&, sigma_i,  t]( Time u )
		{
			Real sigma = sigma_i( u );
			Real Eut = E( u, t );
			return sigma * sigma / Eut / Eut;
		};

		return integrator_( integrand, s, t );
	}
}

#endif // !CALIBRATOR_PROCESSES_GENERAL_VASICEK_HPP


/**
 * @file test_nonlinear.cpp
 * @brief This file defines the test cases for the nonlinear cost model.
 *
 * @copyright Copyright (c) 2021
 *
 */

#include <boost/test/unit_test.hpp>
#include <boost/test/unit_test_suite.hpp>

#include <eigen3/Eigen/Core>
#include <string>
#include <vector>

#include "proseco_planning/agent/cost_model/costNonLinear.h"
#include "proseco_planning/agent/vehicle.h"
#include "proseco_planning/config/configuration.h"
#include "proseco_planning/config/defaultConfiguration.h"
#include "proseco_planning/util/utilities.h"

using namespace proseco_planning;

struct NonLinearModelFixture {
  NonLinearModelFixture() {
    inputVector(0) = 1.0;
    inputVector(1) = -2.0;
    inputVector(2) = 3.0;
    inputVector(3) = -0.5;
    inputVector(4) = 1.5;
    inputVector(5) = 2.5;
    inputVector(6) = -27.0;
    inputVector(7) = 12.0;
    inputVector(8) = 0.0;
    inputVector(9) = 439.0;

    reluOutputVector(0) = 1.0;
    reluOutputVector(1) = 0.0;
    reluOutputVector(2) = 3.0;
    reluOutputVector(3) = 0.0;
    reluOutputVector(4) = 1.5;
    reluOutputVector(5) = 2.5;
    reluOutputVector(6) = 0.0;
    reluOutputVector(7) = 12.0;
    reluOutputVector(8) = 0.0;
    reluOutputVector(9) = 439.0;

    leakyReluOutputVector(0) = 1.0;
    leakyReluOutputVector(1) = -0.02;
    leakyReluOutputVector(2) = 3.0;
    leakyReluOutputVector(3) = -0.005;
    leakyReluOutputVector(4) = 1.5;
    leakyReluOutputVector(5) = 2.5;
    leakyReluOutputVector(6) = -0.27;
    leakyReluOutputVector(7) = 12.0;
    leakyReluOutputVector(8) = 0.0;
    leakyReluOutputVector(9) = 439.0;

    hiddenLayerOutputVector(0) = 429.5;
    hiddenLayerOutputVector(1) = 429.5;
    hiddenLayerOutputVector(2) = 429.5;
    hiddenLayerOutputVector(3) = 429.5;
    hiddenLayerOutputVector(4) = 429.5;

    Config::create(config::scenarioSimple, config::optionsSimple);

    std::vector<float> w1Vector(50, 1.0);
    const auto w1 = config::CostModel::convertVectorToEigenMatrix(w1Vector, 10, 5);
    std::vector<float> w2Vector(5, 1.0);
    const auto w2 = config::CostModel::convertVectorToEigenMatrix(w2Vector, 5, 1);
    // Set cost Parameter and generate cost model
    const auto configCostModel = config::CostModel("costNonLinear", -50, 20, 100, 30, -1, 0, -500,
                                                   -500, -500, -10, 0, 1, 1, 1, 1, 1, 1, 1, w1, w2);
    costModel =
        std::dynamic_pointer_cast<CostNonLinear>(CostModel::createCostModel(configCostModel));
  }

  ~NonLinearModelFixture() { Config::get()->reset(); }

  Eigen::VectorXd inputVector{Eigen::VectorXd(10)};
  Eigen::VectorXd reluOutputVector{Eigen::VectorXd(10)};
  Eigen::VectorXd leakyReluOutputVector{Eigen::VectorXd(10)};
  Eigen::VectorXd hiddenLayerOutputVector{Eigen::VectorXd(5)};
  std::shared_ptr<CostNonLinear> costModel;
};

BOOST_FIXTURE_TEST_SUITE(non_linear, NonLinearModelFixture)

BOOST_AUTO_TEST_CASE(ReLU) {
  BOOST_REQUIRE(reluOutputVector.isApprox(costModel->ReLU(inputVector)));
}

BOOST_AUTO_TEST_CASE(LeakyReLU) {
  BOOST_REQUIRE(leakyReluOutputVector.isApprox(costModel->LeakyReLU(inputVector)));
}

BOOST_AUTO_TEST_CASE(forward_pass) {
  BOOST_REQUIRE_SMALL(costModel->forwardPass(inputVector) - 2147.5f, 0.00001f);
}

BOOST_AUTO_TEST_SUITE_END()


/**
 * @file expressions.hpp
 *
 */

#pragma once

#include "parameter.hpp"
#include "variable.hpp"

#include <Eigen/Sparse>

namespace Eigen
{

    template <>
    struct NumTraits<cvx::Scalar>
        : NumTraits<double>
    {
        using Real = cvx::Scalar;
        using NonInteger = cvx::Scalar;
        using Nested = cvx::Scalar;

        enum
        {
            IsComplex = 0,
            IsInteger = 0,
            IsSigned = 1,
            RequireInitialization = 1,
            ReadCost = 10,
            AddCost = 200,
            MulCost = 200,
        };
    };

} // namespace Eigen

namespace cvx
{
    class Constraint;
    class OptimizationProblem;
    class Scalar;

    namespace internal
    {
        class Affine;
        class SOCPWrapperBase;
        class QPWrapperBase;

        class Term
        {
        public:
            Term();

            Parameter parameter;
            Variable variable;

            bool operator==(const Term &other) const;
            Term &operator*=(const Parameter &param);
            Term &operator/=(const Parameter &param);

            operator Affine() const;

            friend std::ostream &operator<<(std::ostream &os, const Term &term);
            double evaluate() const;
        };

        class Affine
        {
        public:
            bool operator==(const Affine &other) const;

            Parameter constant = Parameter(0.);
            std::vector<Term> terms;

            friend std::ostream &operator<<(std::ostream &os, const Affine &affine);
            double evaluate() const;
            Affine &operator+=(const Affine &other);
            Affine &operator-=(const Affine &other);
            Affine &operator*=(const Parameter &param);
            Affine &operator/=(const Parameter &param);
            // Affine operator+(const Affine &other) const;
            Affine operator-(const Affine &other) const;
            Affine operator-() const;

            void cleanUp();

            bool isZero() const;
            bool isConstant() const;
            bool isFirstOrder() const;
        };

        class Product
        {
        public:
            explicit Product(const Affine &term);
            Product(const Affine &lhs, const Affine &rhs);
            Affine &firstTerm();
            Affine &secondTerm();
            const Affine &firstTerm() const;
            const Affine &secondTerm() const;
            void toSquaredTerm();
            double evaluate() const;
            bool isSquare() const;

            bool operator==(const Product &other) const;

            friend std::ostream &operator<<(std::ostream &os, const Product &product);

        private:
            std::vector<Affine> factors;
        };

    } // namespace internal

    class Scalar
    {
    public:
        Scalar() = default;
        explicit Scalar(int x);
        Scalar(double x);
        explicit Scalar(double *x);

        Scalar &operator+=(const Scalar &other);
        Scalar &operator-=(const Scalar &other);
        Scalar &operator*=(const Scalar &other);
        Scalar &operator/=(const Scalar &other);
        Scalar operator-() const;
        friend Scalar operator+(const Scalar &lhs, const Scalar &rhs);
        friend Scalar operator-(const Scalar &lhs, const Scalar &rhs);
        friend Scalar operator*(const Scalar &lhs, const Scalar &rhs);
        friend Scalar operator/(const Scalar &lhs, const Scalar &rhs);

        bool operator==(const cvx::Scalar &other) const;

        double evaluate() const;
        size_t getOrder() const;
        bool isNorm() const;

        friend OptimizationProblem;
        friend internal::SOCPWrapperBase;
        friend internal::QPWrapperBase;

        explicit operator double() const;

        friend Scalar sqrt(const Scalar &scalar);
        friend Scalar square(const Scalar &scalar);
        friend Scalar abs2(const Scalar &scalar);

        // friend internal::Parameter::operator Scalar() const;
        friend internal::Variable::operator Scalar() const;

        friend Constraint equalTo(const Scalar &lhs, const Scalar &rhs);
        friend Constraint lessThan(const Scalar &lhs, const Scalar &rhs);
        friend Constraint greaterThan(const Scalar &lhs, const Scalar &rhs);
        friend Constraint box(const Scalar &lower, const Scalar &middle, const Scalar &upper);

    private:
        internal::Affine affine;
        std::vector<internal::Product> products;
        bool norm = false;

        friend std::ostream &operator<<(std::ostream &os, const Scalar &scalar);
    };

    using MatrixX = Eigen::Matrix<cvx::Scalar, Eigen::Dynamic, Eigen::Dynamic>;
    using VectorX = Eigen::Matrix<cvx::Scalar, Eigen::Dynamic, 1>;

    /**
     * @brief Creates a constant parameter.
     *
     * @param p The value of the parameter
     * @return Scalar The constant parameter
     */
    Scalar par(double p);

    /**
     * @brief Creates a dynamic parameter.
     *
     * @details Internally stores pointers to the original values.
     *
     * @warning Do not delete the source before the parameter is no longer required.
     *
     * @param p The value of the parameter
     * @return Scalar The constant parameter
     */
    Scalar dynpar(double &p);

    /**
     * @brief Creates a constant parameter from a dense Eigen type.
     *
     * @tparam Derived
     * @param m A dense Eigen type containing problem parameters
     * @return auto A dense Eigen type with cvx::Scalar as scalar type
     */
    template <typename Derived>
    inline auto par(const Eigen::MatrixBase<Derived> &m)
    {
        return m.template cast<Scalar>().eval();
    }

    /**
     * @brief Creates a constant parameter from a sparse Eigen type.
     *
     * @tparam Derived
     * @param m A sparse Eigen type containing problem parameters
     * @return auto A sparse Eigen type with cvx::Scalar as scalar type
     */
    template <typename Derived>
    inline auto par(const Eigen::SparseMatrixBase<Derived> &m)
    {
        return m.template cast<Scalar>().eval();
    }

    /**
     * @brief Creates a dynamic parameter from a dense Eigen type.
     *
     * @details Internally stores pointers to the original values.
     *
     * @warning Do not delete the source before the parameter is no longer required.
     *
     * @tparam Derived
     * @param m A dense Eigen type containing problem parameters
     * @return auto A dense Eigen type with cvx::Scalar as scalar type
     */
    template <typename Derived>
    auto dynpar(Eigen::MatrixBase<Derived> &m)
    {
        auto result = m.template cast<Scalar>().eval();

        for (int row = 0; row < m.rows(); row++)
        {
            for (int col = 0; col < m.cols(); col++)
            {
                result.coeffRef(row, col) = dynpar(m.coeffRef(row, col));
            }
        }

        return result;
    }

    /**
     * @brief Creates a dynamic parameter from a sparse Eigen type.
     *
     * @details Internally stores pointers to the original values.
     *
     * @warning Do not delete the source before the parameter is no longer required.
     *
     * @tparam T
     * @param m A sparse Eigen type containing problem parameters
     * @return auto A sparse Eigen type with cvx::Scalar as scalar type
     */
    template <typename T>
    auto dynpar(Eigen::SparseMatrix<T> &m)
    {
        auto result = m.template cast<Scalar>().eval();

        for (int k = 0; k < result.nonZeros(); k++)
        {
            result.valuePtr()[k] = dynpar(m.valuePtr()[k]);
        }

        return result;
    }

    /**
     * @brief Evaluates the scalar.
     *
     * @param s The scalar to be evaluated
     * @return double The value of the scalar
     */
    double eval(const Scalar &s);

    /**
     * @brief Evaluates a dense Eigen type
     *
     * @tparam Derived Has to be cvx::Scalar
     * @param m A dense Eigen type to be evaluated
     * @return auto The evaluated dense Eigen type
     */
    template <typename Derived>
    inline auto eval(const Eigen::MatrixBase<Derived> &m)
    {
        return m.template cast<double>();
    }

    /**
     * @brief Evaluates a sparse Eigen type
     *
     * @tparam Derived Has to be cvx::Scalar
     * @param m A sparse Eigen type to be evaluated
     * @return auto The evaluated sparse Eigen type
     */
    template <typename Derived>
    inline auto eval(const Eigen::SparseMatrixBase<Derived> &m)
    {
        return m.template cast<double>();
    }

    inline const Scalar &conj(const Scalar &x) { return x; }
    inline const Scalar &real(const Scalar &x) { return x; }
    inline Scalar imag(const Scalar &) { return Scalar(0.); }
    inline Scalar square(const Scalar &x)
    {
        Scalar new_scalar;
        new_scalar.products = {internal::Product(x.affine)};
        return new_scalar;
    }
    inline Scalar abs2(const Scalar &x)
    {
        return square(x);
    }

} // namespace cvx