#include "precompile.h"
#include <iostream>

#include "nis_solver.h"

#ifdef ENABLE_MATLAB_DEBUGGING
#include "../matlab_debugger.h"
extern MatlabDebugger* g_debugger;
#endif

NISSolver::NISSolver()
{
	m_initialized = false;
}

NISSolver::~NISSolver()
{
}

void NISSolver::Init(const VectorX& x, const EigenMatrixi& elements_indices, std::vector<Handle>& handles, bool mode)
{
	m_mode = mode;

	for (std::vector<Constraint*>::iterator it = m_constraints.begin(); it != m_constraints.end(); ++it)
	{
		delete (*it);
	}
	m_constraints.clear();

	m_f_size = 0;
	m_constraint_num = 0;
	m_attachment_constraint_num = 0;
	m_element_num = elements_indices.rows();
	// setup element constraints
	if (mode == MODE_2D) // triangle constraints
	{
		for (unsigned int i = 0; i < elements_indices.rows(); ++i)
		{
			TriangleElementConstraint *tric = new TriangleElementConstraint(i, &m_mu, &m_lambda, &m_enable_non_inversion, &m_alpha, &m_beta, &m_slackness, elements_indices(i, 0), elements_indices(i, 1), elements_indices(i, 2), x);
			m_constraints.push_back(tric);
		}
		m_f_size += 4*elements_indices.rows();
		m_constraint_num += 2*elements_indices.rows();
	}
	else // tet constraints
	{
		for (unsigned int i = 0; i < elements_indices.rows(); ++i)
		{
			TetElementConstraint *tetc = new TetElementConstraint(i, &m_mu, &m_lambda, &m_enable_non_inversion, &m_alpha, &m_beta, &m_gamma, &m_slackness, elements_indices(i, 0), elements_indices(i, 1), elements_indices(i, 2), elements_indices(i, 3), x);
			m_constraints.push_back(tetc);
		}
		m_f_size += 9*elements_indices.rows();
		m_constraint_num += 3*elements_indices.rows();
	}

	// setup attachment constraints
	unsigned int dim = (mode == MODE_2D)?2:3;
	unsigned int size = x.size() / dim;

	for (std::vector<Handle>::iterator it = handles.begin(); it != handles.end(); ++it)
	{
		std::vector<unsigned int>& indices = it->Indices();
		for (unsigned int i = 0; i != indices.size(); ++i)
		{
			AttachmentConstraint *ac = new AttachmentConstraint(&m_stiffness_attachment, indices[i], it->ID(), i, dim);
			m_constraints.push_back(ac);
			m_f_size += dim;
			m_attachment_constraint_num ++;
		}
	}
	m_vertices_num = size;

	// allocate spaces
	m_handle_vertices.resize(handles.size());
	m_per_element_rotation_2d.resize(elements_indices.rows());
	m_per_element_rotation_3d.resize(elements_indices.rows());

	m_initialized = true;
}

void NISSolver::PrepareHandleInfo(const VectorX& x, std::vector<Handle>& handles, EigenMatrix& per_element_weights)
{
	if (m_initialized == false)
		return;

	// handle vertices:
	for (unsigned int i = 0; i != m_handle_vertices.size(); ++i)
	{
		m_handle_vertices[i] = handles[i].Vertices();
	}

	// rotation 2d
	for (unsigned int i = 0; i != m_per_element_rotation_2d.size(); ++i)
	{
		m_per_element_rotation_2d[i] = 0;
		for (unsigned int j = 0; j != per_element_weights.cols(); ++j)
		{
			m_per_element_rotation_2d[i] += per_element_weights(i, j) * handles[j].RotationAngle2d();
		}
	}

	// rotation 3d
	for (unsigned int i = 0; i != m_per_element_rotation_3d.size(); ++i)
	{
		EigenQuaternion q;
		// Q-SLERP
		q.vec().setZero();
		q.w() = 0;
		for (unsigned int j = 0; j != per_element_weights.cols(); ++j)
		{
			q.vec() += per_element_weights(i, j) * handles[j].Rotation3d().vec();
			q.w() += per_element_weights(i, j) * handles[j].Rotation3d().w();
		}
		q.normalize();
		m_per_element_rotation_3d[i] = q;
	}

	// update to constraints
	UpdateHandleInfo();
}

//void NISSolver::PrepareHandleInfoDisconnectedMesh(const VectorX& disconnected_x, std::vector<std::vector<unsigned int>>& parent_to_children_map, std::vector<Handle>& handles, EigenMatrix& per_element_weights)
//{
//	if (m_initialized == false)
//		return;
//
//	// handle vertices:
//	for (unsigned int i = 0; i != m_handle_vertices.size(); ++i)
//	{
//		m_handle_vertices[i] = handles[i].Vertices();
//	}
//
//	// rotation 2d
//	for (unsigned int i = 0; i != m_per_element_rotation_2d.size(); ++i)
//	{
//		m_per_element_rotation_2d[i] = 0;
//		for (unsigned int j = 0; j != per_element_weights.cols(); ++j)
//		{
//			m_per_element_rotation_2d[i] += per_element_weights(i, j) * handles[j].RotationAngle2d();
//		}
//	}
//
//	// rotation 3d
//	for (unsigned int i = 0; i != m_per_element_rotation_3d.size(); ++i)
//	{
//		EigenQuaternion q;
//		// Q-SLERP
//		q.vec().setZero();
//		q.w() = 0;
//		for (unsigned int j = 0; j != per_element_weights.cols(); ++j)
//		{
//			q.vec() += per_element_weights(i, j) * handles[j].Rotation3d().vec();
//			q.w() += per_element_weights(i, j) * handles[j].Rotation3d().w();
//		}
//		q.normalize();
//		m_per_element_rotation_3d[i] = q;
//	}
//
//	// update to constraints
//	UpdateHandleInfo();
//}


// GetDetFInfo
void NISSolver::GetDetF(const VectorX& x, std::vector<float>& list_detf)
{
	if (m_initialized == false)
		return;

	TriangleElementConstraint* tri;
	TetElementConstraint* tet;
	unsigned int i = 0;
	for (std::vector<Constraint*>::iterator it = m_constraints.begin(); it != m_constraints.end(); ++it)
	{
		if ((*it)->Type() == CONSTRAINT_TYPE_TRIANGLE_ELEMENT)
		{
			tri = dynamic_cast<TriangleElementConstraint*>(*it);
			const EigenMatrix2& F = tri->GetF(x);
			ScalarType value = F.determinant();
			list_detf[3*i] = value;
			list_detf[3*i+1] = value;
			list_detf[3*i+2] = value;
			i++;
		}
		if ((*it)->Type() == CONSTRAINT_TYPE_TET_ELEMENT)
		{
			tet = dynamic_cast<TetElementConstraint*>(*it);
			const EigenMatrix3& F = tet->GetF(x);
			ScalarType value = F.determinant();
			list_detf[4*i] = value;
			list_detf[4*i+1] = value;
			list_detf[4*i+2] = value;
			list_detf[4*i+3] = value;
			i++;
		}
	}
}

// GetW11Info
void NISSolver::GetW11(const VectorX& x, std::vector<float>& list_w11)
{
	if (m_initialized == false)
		return;

	TriangleElementConstraint* tri;
	TetElementConstraint* tet;
	unsigned int i = 0;
	for (std::vector<Constraint*>::iterator it = m_constraints.begin(); it != m_constraints.end(); ++it)
	{
		if ((*it)->Type() == CONSTRAINT_TYPE_TRIANGLE_ELEMENT)
		{
			tri = dynamic_cast<TriangleElementConstraint*>(*it);
			const EigenMatrix2& W = tri->GetW(x);
			ScalarType value = W(0, 0);
			list_w11[3*i] = value;
			list_w11[3*i+1] = value;
			list_w11[3*i+2] = value;
			i++;
		}
		if ((*it)->Type() == CONSTRAINT_TYPE_TET_ELEMENT)
		{
			tet = dynamic_cast<TetElementConstraint*>(*it);
			const EigenMatrix3& W = tet->GetW(x);
			ScalarType value = W(0, 0);
			list_w11[4*i] = value;
			list_w11[4*i+1] = value;
			list_w11[4*i+2] = value;
			list_w11[4*i+3] = value;
			i++;
		}
	}
}
// Get Det(W) Info
void NISSolver::GetDetW(const VectorX& x, std::vector<float>& list_detw)
{
	if (m_initialized == false)
		return;

	TriangleElementConstraint* tri;
	TetElementConstraint* tet;
	unsigned int i = 0;
	for (std::vector<Constraint*>::iterator it = m_constraints.begin(); it != m_constraints.end(); ++it)
	{
		if ((*it)->Type() == CONSTRAINT_TYPE_TRIANGLE_ELEMENT)
		{
			tri = dynamic_cast<TriangleElementConstraint*>(*it);
			const EigenMatrix2& W = tri->GetW(x);
			ScalarType value = W.determinant();
			list_detw[3*i] = value;
			list_detw[3*i+1] = value;
			list_detw[3*i+2] = value;
			i++;
		}
		if ((*it)->Type() == CONSTRAINT_TYPE_TET_ELEMENT)
		{
			tet = dynamic_cast<TetElementConstraint*>(*it);
			const EigenMatrix3& W = tet->GetW(x);
			ScalarType value = W.determinant();
			list_detw[3*i] = value;
			list_detw[3*i+1] = value;
			list_detw[3*i+2] = value;
			i++;
		}
	}
}

ScalarType NISSolver::GetSmallestDetW(const VectorX& x)
{
	if (m_initialized == false)
		return 1;

	TriangleElementConstraint* tri;
	TetElementConstraint* tet;
	unsigned int i = 0;
	ScalarType min = INFINITY;
	for (std::vector<Constraint*>::iterator it = m_constraints.begin(); it != m_constraints.end(); ++it)
	{
		if ((*it)->Type() == CONSTRAINT_TYPE_TRIANGLE_ELEMENT)
		{
			tri = dynamic_cast<TriangleElementConstraint*>(*it);
			const EigenMatrix2& W = tri->GetW(x);
			ScalarType value =W.determinant();
			if (min > value)
			{
				min = value;
			}
			i++;
		}
		if ((*it)->Type() == CONSTRAINT_TYPE_TET_ELEMENT)
		{
			tet = dynamic_cast<TetElementConstraint*>(*it);
			const EigenMatrix3& W = tet->GetW(x);
			ScalarType value = W.determinant();
			if (min > value)
			{
				min = value;
			}
			i++;
		}
	}
}