#include <memory>

// TnzExt includes
#include "ext/plasticskeleton.h"
#include "ext/plasticskeletondeformation.h"

// STD includes
#include <limits>
#include <map>
#include <algorithm>

// Boost includes
#include <boost/multi_index_container.hpp>
#include <boost/multi_index/ordered_index.hpp>
#include <boost/multi_index/member.hpp>

// Qt includes
#include <QMutex>
#include <QMutexLocker>

#include "ext/plasticdeformerstorage.h"

//***********************************************************************************************
//    Storage multi-index map  definition
//***********************************************************************************************

namespace
{

typedef PlasticDeformerDataGroup DataGroup;

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

typedef std::pair<const SkD *, int> DeformedSkeleton;

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

struct Key {
	const TMeshImage *m_mi;
	DeformedSkeleton m_ds;

	std::shared_ptr<DataGroup> m_dataGroup;

public:
	Key(const TMeshImage *mi, const SkD *sd, int skelId)
		: m_mi(mi), m_ds(sd, skelId), m_dataGroup() {}

	bool operator<(const Key &other) const
	{
		return (m_mi < other.m_mi) || ((!(other.m_mi < m_mi)) && (m_ds < other.m_ds));
	}
};

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

using namespace boost::multi_index;

typedef boost::multi_index_container<Key, indexed_by<

											  ordered_unique<identity<Key>>,
											  ordered_non_unique<
												  tag<TMeshImage>, member<Key, const TMeshImage *, &Key::m_mi>>,
											  ordered_non_unique<
												  tag<DeformedSkeleton>, member<Key, DeformedSkeleton, &Key::m_ds>>

											  >> DeformersSet;

typedef DeformersSet::nth_index<0>::type DeformersByKey;
typedef DeformersSet::index<TMeshImage>::type DeformersByMeshImage;
typedef DeformersSet::index<DeformedSkeleton>::type DeformersByDeformedSkeleton;

} // namespace

//***********************************************************************************************
//    Initialization stage  functions
//***********************************************************************************************

namespace
{

void initializeSO(PlasticDeformerData &data, const TTextureMeshP &mesh)
{
	data.m_so.reset(new double[mesh->facesCount()]);
}

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

void initializeDeformerData(PlasticDeformerData &data, const TTextureMeshP &mesh)
{
	initializeSO(data, mesh); // Allocates SO data

	// Also, allocate suitable input-output arrays for the deformation
	data.m_output.reset(new double[2 * mesh->verticesCount()]);
}

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

void initializeDeformersData(DataGroup *group, const TMeshImage *meshImage)
{
	group->m_datas.reset(new PlasticDeformerData[meshImage->meshes().size()]);

	// Push a PlasticDeformer for each mesh in the image
	const std::vector<TTextureMeshP> &meshes = meshImage->meshes();
	int fTotal = 0; // Also count total # of faces

	int m, mCount = meshes.size();
	for (m = 0; m != mCount; ++m) {
		fTotal += meshes[m]->facesCount();
		initializeDeformerData(group->m_datas[m], meshes[m]);
	}

	// Initialize the vector of sorted faces
	std::vector<std::pair<int, int>> &sortedFaces = group->m_sortedFaces;

	sortedFaces.reserve(fTotal);
	for (m = 0; m != mCount; ++m) {
		const TTextureMesh &mesh = *meshes[m];

		int f, fCount = mesh.facesCount();
		for (f = 0; f != fCount; ++f)
			sortedFaces.push_back(std::make_pair(f, m));
	}
}

} // namespace

//***********************************************************************************************
//    Handle processing  functions
//***********************************************************************************************

namespace
{

void transformHandles(std::vector<PlasticHandle> &handles, const TAffine &aff)
{
	// Transforms handles through deformAff AND applies mi's dpi scale inverse
	std::vector<PlasticHandle>::size_type h, hCount = handles.size();
	for (h = 0; h != hCount; ++h)
		handles[h].m_pos = aff * handles[h].m_pos;
}

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

void transformHandles(std::vector<TPointD> &handles, const TAffine &aff)
{
	// Transforms handles through deformAff AND applies mi's dpi scale inverse
	std::vector<PlasticHandle>::size_type h, hCount = handles.size();
	for (h = 0; h != hCount; ++h)
		handles[h] = aff * handles[h];
}

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

void processHandles(DataGroup *group, double frame, const TMeshImage *meshImage,
					const SkD *sd, int skelId, const TAffine &deformationAffine)
{
	assert(sd);

	const PlasticSkeletonP &skeleton = sd->skeleton(skelId);

	if (!skeleton || skeleton->verticesCount() == 0) {
		group->m_handles.clear();
		group->m_dstHandles.clear();

		group->m_compiled |= PlasticDeformerStorage::HANDLES;
		group->m_upToDate |= PlasticDeformerStorage::HANDLES;

		return;
	}

	int mCount = meshImage->meshes().size();

	if (!(group->m_upToDate & PlasticDeformerStorage::HANDLES)) {
		// Compile handles if necessary
		if (!(group->m_compiled & PlasticDeformerStorage::HANDLES)) {
			// Build and transform handles
			group->m_handles = skeleton->verticesToHandles();
			::transformHandles(group->m_handles, deformationAffine);

			// Prepare a vector for handles' face hints
			for (int m = 0; m != mCount; ++m)
				group->m_datas[m].m_faceHints.resize(group->m_handles.size(), -1);

			group->m_compiled |= PlasticDeformerStorage::HANDLES;
		}

		// Then, build destination handles
		PlasticSkeleton deformedSkeleton; // NOTE: Could this be moved to the group as well?
		sd->storeDeformedSkeleton(skelId, frame, deformedSkeleton);

		// Copy deformed skeleton data into input deformation parameters
		group->m_dstHandles = std::vector<TPointD>(deformedSkeleton.vertices().begin(), deformedSkeleton.vertices().end());
		::transformHandles(group->m_dstHandles, deformationAffine);

		group->m_upToDate |= PlasticDeformerStorage::HANDLES;
	}
}

} // namespace

//***********************************************************************************************
//    Stacking Order processing  functions
//***********************************************************************************************

namespace
{

bool updateHandlesSO(DataGroup *group, const SkD *sd, int skelId, double frame)
{
	assert(sd);

	const PlasticSkeletonP &skeleton = sd->skeleton(skelId);

	if (!skeleton || skeleton->verticesCount() == 0) {
		group->m_soMin = group->m_soMax = 0.0;
		return false;
	}

	// Copy SO values to data's handles
	// Return whether values changed with respect to previous ones
	bool changed = false;

	assert(group->m_handles.size() == skeleton->verticesCount());

	int h, hCount = group->m_handles.size();
	{
		tcg::list<PlasticSkeletonVertex>::iterator vt = skeleton->vertices().begin();

		for (h = 0; h != hCount; ++h, ++vt) {
			const SkVD *vd = sd->vertexDeformation(vt->name());
			if (!vd)
				continue;

			double so = vd->m_params[SkVD::SO]->getValue(frame);

			PlasticHandle &handle = group->m_handles[h];
			if (handle.m_so != so) {
				group->m_handles[h].m_so = so;
				changed = true;
			}
		}
	}

	if (changed) {
		// Rebuild SO minmax
		group->m_soMax = -(group->m_soMin = (std::numeric_limits<double>::max)());

		for (h = 0; h != hCount; ++h) {
			const double &so = group->m_handles[h].m_so;

			group->m_soMin = std::min(group->m_soMin, so);
			group->m_soMax = std::max(group->m_soMax, so);
		}
	}

	return changed;
}

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

void interpolateSO(DataGroup *group, const TMeshImage *meshImage)
{
	int m, mCount = meshImage->meshes().size();

	if (group->m_handles.size() == 0) {
		// No handles case, fill in with 0s

		for (m = 0; m != mCount; ++m) {
			const TTextureMesh &mesh = *meshImage->meshes()[m];
			PlasticDeformerData &data = group->m_datas[m];

			std::fill(data.m_so.get(), data.m_so.get() + mesh.facesCount(), 0.0);
		}

		return;
	}

	// Apply handles' SO values to each mesh
	for (m = 0; m != mCount; ++m) {
		const TTextureMesh &mesh = *meshImage->meshes()[m];
		PlasticDeformerData &data = group->m_datas[m];

		// Interpolate so values
		std::unique_ptr<double[]> verticesSO(new double[mesh.verticesCount()]);

		::buildSO(verticesSO.get(), mesh, group->m_handles, &data.m_faceHints.front());

		// Make the mean of each face's vertex values and store that
		int f, fCount = mesh.facesCount();
		for (f = 0; f != fCount; ++f) {
			int v0, v1, v2;
			mesh.faceVertices(f, v0, v1, v2);

			data.m_so[f] = (verticesSO[v0] + verticesSO[v1] + verticesSO[v2]) / 3.0;
		}
	}
}

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

struct FaceLess {
	const PlasticDeformerDataGroup *m_group;

public:
	FaceLess(const PlasticDeformerDataGroup *group) : m_group(group) {}

	bool operator()(const std::pair<int, int> &a, const std::pair<int, int> &b)
	{
		return (m_group->m_datas[a.second].m_so[a.first] < m_group->m_datas[b.second].m_so[b.first]);
	}
};

// Must be invoked after updateSO
void updateSortedFaces(PlasticDeformerDataGroup *group)
{
	FaceLess comp(group);
	std::sort(group->m_sortedFaces.begin(), group->m_sortedFaces.end(), comp);
}

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

void processSO(DataGroup *group, double frame, const TMeshImage *meshImage,
			   const SkD *sd, int skelId, const TAffine &deformationAffine)
{
	// SO re-interpolate values along the mesh if either:
	//  1. Recompilation was requested (ie some vertex may have been added/removed)
	//  2. OR the value of one of the handle has changed

	bool interpolate = !(group->m_compiled & PlasticDeformerStorage::SO);

	if (!(group->m_upToDate & PlasticDeformerStorage::SO)) // implied by (interpolate == true)
	{
		interpolate = updateHandlesSO(group, sd, skelId, frame) || interpolate; // Order is IMPORTANT

		if (interpolate) {
			interpolateSO(group, meshImage);
			updateSortedFaces(group);
		}

		group->m_compiled |= PlasticDeformerStorage::SO;
		group->m_upToDate |= PlasticDeformerStorage::SO;
	}
}

} // namespace

//***********************************************************************************************
//    Mesh Deform processing  functions
//***********************************************************************************************

namespace
{

void processMesh(DataGroup *group, double frame, const TMeshImage *meshImage,
				 const SkD *sd, int skelId, const TAffine &deformationAffine)
{
	if (!(group->m_upToDate & PlasticDeformerStorage::MESH)) {
		int m, mCount = meshImage->meshes().size();

		if (!(group->m_compiled & PlasticDeformerStorage::MESH)) {
			for (m = 0; m != mCount; ++m) {
				const TTextureMeshP &mesh = meshImage->meshes()[m];
				PlasticDeformerData &data = group->m_datas[m];

				data.m_deformer.initialize(mesh);
				data.m_deformer.compile(group->m_handles, data.m_faceHints.empty() ? 0 : &data.m_faceHints.front());
				data.m_deformer.releaseInitializedData();
			}

			group->m_compiled |= PlasticDeformerStorage::MESH;
		}

		const TPointD *dstHandlePos = group->m_dstHandles.empty() ? 0 : &group->m_dstHandles.front();

		for (m = 0; m != mCount; ++m) {
			PlasticDeformerData &data = group->m_datas[m];
			data.m_deformer.deform(dstHandlePos, data.m_output.get());
		}

		group->m_upToDate |= PlasticDeformerStorage::MESH;
	}
}

} // namespace

//***********************************************************************************************
//    PlasticDeformerData  implementation
//***********************************************************************************************

PlasticDeformerData::PlasticDeformerData()
{
}

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

PlasticDeformerData::~PlasticDeformerData()
{
}

//***********************************************************************************************
//    PlasticDeformerDataGroup  implementation
//***********************************************************************************************

PlasticDeformerDataGroup::PlasticDeformerDataGroup()
	: m_datas()
	, m_compiled(PlasticDeformerStorage::NONE)
	, m_upToDate(PlasticDeformerStorage::NONE)
	, m_outputFrame((std::numeric_limits<double>::max)())
	, m_soMin()
	, m_soMax()
{
}

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

PlasticDeformerDataGroup::~PlasticDeformerDataGroup()
{
}

//***********************************************************************************************
//    PlasticDeformerStorage::Imp  definition
//***********************************************************************************************

class PlasticDeformerStorage::Imp
{
public:
	QMutex m_mutex;			  //!< Access mutex - needed for thread-safety
	DeformersSet m_deformers; //!< Set of deformers, ordered by mesh image, deformation, and affine.

public:
	Imp() : m_mutex(QMutex::Recursive) {}
};

//***********************************************************************************************
//    PlasticDeformerStorage  implementation
//***********************************************************************************************

PlasticDeformerStorage::PlasticDeformerStorage()
	: m_imp(new Imp)
{
}

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

PlasticDeformerStorage::~PlasticDeformerStorage()
{
}

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

PlasticDeformerStorage *PlasticDeformerStorage::instance()
{
	static PlasticDeformerStorage theInstance;
	return &theInstance;
}

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

PlasticDeformerDataGroup *PlasticDeformerStorage::deformerData(
	const TMeshImage *meshImage, const PlasticSkeletonDeformation *deformation, int skelId)
{
	QMutexLocker locker(&m_imp->m_mutex);

	// Search for the corresponding deformation in the storage
	Key key(meshImage, deformation, skelId);

	DeformersByKey::iterator dt = m_imp->m_deformers.find(key);
	if (dt == m_imp->m_deformers.end()) {
		// No deformer was found. Allocate it.
		key.m_dataGroup = std::make_shared<PlasticDeformerDataGroup>();
		initializeDeformersData(key.m_dataGroup.get(), meshImage);

		dt = m_imp->m_deformers.insert(key).first;
	}

	return dt->m_dataGroup.get();
}

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

const PlasticDeformerDataGroup *PlasticDeformerStorage::process(
	double frame, const TMeshImage *meshImage,
	const PlasticSkeletonDeformation *deformation, int skelId,
	const TAffine &skeletonAffine, DataType dataType)
{
	QMutexLocker locker(&m_imp->m_mutex);

	PlasticDeformerDataGroup *group = deformerData(meshImage, deformation, skelId);

	// On-the-fly checks for data invalidation
	if (group->m_skeletonAffine != skeletonAffine) {
		group->m_upToDate = NONE;
		group->m_compiled = NONE;
		group->m_skeletonAffine = skeletonAffine;
	}

	if (group->m_outputFrame != frame) {
		group->m_upToDate = NONE;
		group->m_outputFrame = frame;
	}

	bool doMesh = (dataType & MESH);
	bool doSO = (dataType & SO) || doMesh;
	bool doHandles = (bool)dataType;

	// Process data
	if (doHandles)
		processHandles(group, frame, meshImage, deformation, skelId, skeletonAffine);

	if (doSO)
		processSO(group, frame, meshImage, deformation, skelId, skeletonAffine);

	if (doMesh)
		processMesh(group, frame, meshImage, deformation, skelId, skeletonAffine);

	return group;
}

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

const PlasticDeformerDataGroup *PlasticDeformerStorage::processOnce(
	double frame,
	const TMeshImage *meshImage,
	const PlasticSkeletonDeformation *deformation, int skelId,
	const TAffine &skeletonAffine,
	DataType dataType)
{
	PlasticDeformerDataGroup *group = new PlasticDeformerDataGroup;
	initializeDeformersData(group, meshImage);

	bool doMesh = (dataType & MESH);
	bool doSO = (dataType & SO) || doMesh;
	bool doHandles = (bool)dataType;

	// Process data
	if (doHandles)
		processHandles(group, frame, meshImage, deformation, skelId, skeletonAffine);

	if (doSO)
		processSO(group, frame, meshImage, deformation, skelId, skeletonAffine);

	if (doMesh)
		processMesh(group, frame, meshImage, deformation, skelId, skeletonAffine);

	return group;
}

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

void PlasticDeformerStorage::invalidateMeshImage(
	const TMeshImage *meshImage, int recompiledData)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformersByMeshImage &deformers = m_imp->m_deformers.get<TMeshImage>();

	DeformersByMeshImage::iterator dBegin(deformers.lower_bound(meshImage));
	if (dBegin == deformers.end())
		return;

	DeformersByMeshImage::iterator dt, dEnd(deformers.upper_bound(meshImage));
	for (dt = dBegin; dt != dEnd; ++dt) {
		dt->m_dataGroup->m_outputFrame = (std::numeric_limits<double>::max)(); // Schedule for redeformation
		if (recompiledData)
			dt->m_dataGroup->m_compiled &= ~recompiledData; // Schedule for recompilation, too
	}
}

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

void PlasticDeformerStorage::invalidateSkeleton(
	const PlasticSkeletonDeformation *deformation, int skelId, int recompiledData)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformedSkeleton ds(deformation, skelId);

	DeformersByDeformedSkeleton &deformers = m_imp->m_deformers.get<DeformedSkeleton>();

	DeformersByDeformedSkeleton::iterator dBegin(deformers.lower_bound(ds));
	if (dBegin == deformers.end())
		return;

	DeformersByDeformedSkeleton::iterator dt, dEnd(deformers.upper_bound(ds));
	for (dt = dBegin; dt != dEnd; ++dt) {
		dt->m_dataGroup->m_outputFrame = (std::numeric_limits<double>::max)(); // Schedule for redeformation
		if (recompiledData)
			dt->m_dataGroup->m_compiled &= ~recompiledData; // Schedule for recompilation, too
	}
}

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

void PlasticDeformerStorage::invalidateDeformation(
	const PlasticSkeletonDeformation *deformation, int recompiledData)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformersByDeformedSkeleton &deformers = m_imp->m_deformers.get<DeformedSkeleton>();

	DeformedSkeleton dsBegin(deformation, -(std::numeric_limits<int>::max)()),
		dsEnd(deformation, (std::numeric_limits<int>::max)());

	DeformersByDeformedSkeleton::iterator dBegin(deformers.lower_bound(dsBegin));
	DeformersByDeformedSkeleton::iterator dEnd(deformers.upper_bound(dsEnd));

	if (dBegin == dEnd)
		return;

	for (DeformersByDeformedSkeleton::iterator dt = dBegin; dt != dEnd; ++dt) {
		dt->m_dataGroup->m_outputFrame = (std::numeric_limits<double>::max)(); // Schedule for redeformation
		if (recompiledData)
			dt->m_dataGroup->m_compiled &= ~recompiledData; // Schedule for recompilation, too
	}
}

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

void PlasticDeformerStorage::releaseMeshData(const TMeshImage *meshImage)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformersByMeshImage &deformers = m_imp->m_deformers.get<TMeshImage>();

	DeformersByMeshImage::iterator dBegin(deformers.lower_bound(meshImage));
	if (dBegin == deformers.end())
		return;

	deformers.erase(dBegin, deformers.upper_bound(meshImage));
}

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

void PlasticDeformerStorage::releaseSkeletonData(const SkD *deformation, int skelId)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformedSkeleton ds(deformation, skelId);

	DeformersByDeformedSkeleton &deformers = m_imp->m_deformers.get<DeformedSkeleton>();

	DeformersByDeformedSkeleton::iterator dBegin(deformers.lower_bound(ds));
	if (dBegin == deformers.end())
		return;

	deformers.erase(dBegin, deformers.upper_bound(ds));
}

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

void PlasticDeformerStorage::releaseDeformationData(const SkD *deformation)
{
	QMutexLocker locker(&m_imp->m_mutex);

	DeformersByDeformedSkeleton &deformers = m_imp->m_deformers.get<DeformedSkeleton>();

	DeformedSkeleton dsBegin(deformation, -(std::numeric_limits<int>::max)()),
		dsEnd(deformation, (std::numeric_limits<int>::max)());

	DeformersByDeformedSkeleton::iterator dBegin(deformers.lower_bound(dsBegin));
	DeformersByDeformedSkeleton::iterator dEnd(deformers.upper_bound(dsEnd));

	if (dBegin == dEnd)
		return;

	deformers.erase(dBegin, dEnd);
}

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

void PlasticDeformerStorage::clear()
{
	QMutexLocker locker(&m_imp->m_mutex);

	m_imp->m_deformers.clear();
}