#include "arap.h"
#include "graphics/meshloader.h"

#include <iostream>
#include <set>
#include <map>
#include <vector>

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <Eigen/SVD>

using namespace std;
using namespace Eigen;

ARAP::ARAP() {}

void ARAP::init(Eigen::Vector3f &coeffMin, Eigen::Vector3f &coeffMax)
{
    vector<Vector3f> vertices;
    vector<Vector3i> triangles;

    // If this doesn't work for you, remember to change your working directory
    if (MeshLoader::loadTriMesh("meshes/sphere.obj", vertices, triangles)) {
        m_shape.init(vertices, triangles);
    }

    // Students, please don't touch this code: get min and max for viewport stuff
    MatrixX3f all_vertices = MatrixX3f(vertices.size(), 3);
    int i = 0;
    for (unsigned long i = 0; i < vertices.size(); ++i) {
        all_vertices.row(i) = vertices[i];
    }
    coeffMin = all_vertices.colwise().minCoeff();
    coeffMax = all_vertices.colwise().maxCoeff();
}



// Move an anchored vertex, defined by its index, to targetPosition
void ARAP::move(int vertex, Vector3f targetPosition)
{
    std::vector<Eigen::Vector3f> new_vertices = m_shape.getVertices();
    const std::unordered_set<int>& anchors = m_shape.getAnchors();

	std::cout << "anchors size" << anchors.size() << std::endl;
	std::cout << "targetPosition" << targetPosition << std::endl;

    // TODO: implement ARAP here
    new_vertices[vertex] = targetPosition;

	// make a copy of the vertices into the matrix points
	typedef Matrix<float, 3, Dynamic> PM; // position matrix
	PM p(3, m_shape.getVertices().size());
	for (int i = 0; i < m_shape.getVertices().size(); ++i)
	{
		p.col(i) = m_shape.getVertices()[i];
		// p.col(i) = new_vertices[i];
	}
	PM p_prime(3, m_shape.getVertices().size());
	for (int i = 0; i < m_shape.getVertices().size(); ++i)
	{
		p_prime.col(i) = new_vertices[i];
	}

	// compute the cotan weights
	typedef Triplet<float> Tri;
	vector<Tri> triplets;
	triplets.reserve(3 *  m_shape.getFaces().size());
	for (DenseIndex f = 0; f < m_shape.getFaces().size(); ++f)
	{
		Vector3i face = m_shape.getFaces()[f];
		Vector3f v0 = m_shape.getVertices()[face[0]];
		Vector3f v1 = m_shape.getVertices()[face[1]];
		Vector3f v2 = m_shape.getVertices()[face[2]];

		Vector3f e0 = v1 - v0;
		Vector3f e1 = v2 - v1;
		Vector3f e2 = v0 - v2;

		// protect against degen triangles
		float s0 = e0.squaredNorm();
		s0 = max(s0, 1e-8f);
		float s1 = e1.squaredNorm();
		s1 = max(s1, 1e-8f);
		float s2 = e2.squaredNorm();
		s2 = max(s2, 1e-8f);

		s0 = sqrt(s0);
		s1 = sqrt(s1);
		s2 = sqrt(s2);

		float semi_perimeter = (s0 + s1 + s2) * 0.5f;
		float area = sqrt(semi_perimeter * (semi_perimeter - s0) * (semi_perimeter - s1) * (semi_perimeter - s2));
		area = max(area, 1e-8f);

		float cot0 = (-s0 * s0 + s1 * s1 + s2 * s2) / (4.0f * area);
		cot0 = max(cot0, 1e-10f) * 0.5f;
		float cot1 = (s0 * s0 - s1 * s1 + s2 * s2) / (4.0f * area);
		cot1 = max(cot1, 1e-10f) * 0.5f;
		float cot2 = (s0 * s0 + s1 * s1 - s2 * s2) / (4.0f * area);
		cot2 = max(cot2, 1e-10f) * 0.5f;

		pair<int, int> edge0 =
			face[0] > face[1] ? make_pair(face[1], face[0]) : make_pair(face[0], face[1]);
		pair<int, int> edge1 =
			face[1] > face[2] ? make_pair(face[2], face[1]) : make_pair(face[1], face[2]);
		pair<int, int> edge2 =
			face[2] > face[0] ? make_pair(face[0], face[2]) : make_pair(face[2], face[0]);

		triplets.push_back(Tri(edge0.first, edge0.second, cot0));
		triplets.push_back(Tri(edge1.first, edge1.second, cot1));
		triplets.push_back(Tri(edge2.first, edge2.second, cot2));
		triplets.push_back(Tri(edge0.second, edge0.first, cot0));
		triplets.push_back(Tri(edge1.second, edge1.first, cot1));
		triplets.push_back(Tri(edge2.second, edge2.first, cot2));
	}
	// build the sparse matrix of edge weights
	SparseMatrix<float, RowMajor> edge_weights(m_shape.getVertices().size(), m_shape.getVertices().size());
	edge_weights.reserve(VectorXi::Constant(m_shape.getVertices().size(), 7));
	edge_weights.setZero();
	edge_weights.setFromTriplets(triplets.begin(), triplets.end());

	// initialize the rotation matrices
	typedef Matrix<float, 3, 3> R;
	// TODO: move this to only init once
	vector<R> rotations(m_shape.getVertices().size(), R::Identity(3, 3));
//	rotations.clear();
//	rotations.resize(m_shape.getVertices().size(), R::Identity(3, 3));

	// make a map that reduces the constrained vertices (vtx -> free vtx)
	vector<Index> vtx_to_free_vtx = vector<Index>(m_shape.getVertices().size(), -1);
	// vtx_to_free_vtx.reserve(m_shape.getVertices().size());
	Index count_free = 0;
	for (int i = 0; i < m_shape.getVertices().size(); ++i)
	{
		if (!anchors.contains(i))
		{
			// if were not an anchor, we are free
			vtx_to_free_vtx[i] = count_free;
			count_free++;
		}
	}

	// update pprime with the constraints into the linear system
	for (Index i = 0; i < m_shape.getVertices().size(); ++i)
	{
		if (vtx_to_free_vtx[i] == -1)
		{
			// if we're an anchor, set this position
			p_prime.col(i) = new_vertices[i];
		}
	}

	// setup the linear system we will use to solve
	SparseMatrix<float, RowMajor> L(count_free, count_free);
	// L.resize(count_free, count_free);
	L.reserve(VectorXi::Constant(count_free, 7));
	L.setZero();

	// setup bfixed
	PM b_fixed = PM::Zero(3, count_free);
	// b_fixed.resize(3, count_free);

	// make the L matrix using triplet method
	vector<Tri> triplets_L;
	triplets_L.reserve(7 * count_free);
	for (Index i = 0; i < edge_weights.outerSize(); i++)
	{
		Index free_index = vtx_to_free_vtx[i];
		if (free_index == -1)
		{
			// do not add constraints to the linear system
			continue;
		}

		for (SparseMatrix<float, RowMajor>::InnerIterator it(edge_weights, i); it; ++it)
		{
			Index j = it.col();
			Index free_index_j = vtx_to_free_vtx[j];
			float wij = it.value();
			if (free_index_j == -1)
			{
				// if the neighbor is an anchor, add to b_fixed
				Vector3f location = new_vertices[j];
				b_fixed.col(free_index) += wij * location;
			}
			else
			{
				triplets_L.push_back(Tri(free_index, free_index_j, -wij));
			}
			triplets_L.push_back(Tri(free_index, free_index, wij));
		}
	}
	// create and solve the L matrix
	L.setFromTriplets(triplets_L.begin(), triplets_L.end());
	SimplicialLDLT<SparseMatrix<float>> solver;
	solver.compute(L);

	for (int iter = 0; iter < 5; iter++)
	{
		// estimate the rotations
		rotations.clear();
		rotations.resize(m_shape.getVertices().size(), R::Identity(3, 3));
		typedef Matrix<float, 3, 3> S;
		for (Index i = 0; i < edge_weights.outerSize(); ++i)
		{
			const auto &pi = p.col(i);
			const auto &ppi = p_prime.col(i);
			S cov = S::Zero(3, 3);

			for (SparseMatrix<float, RowMajor>::InnerIterator it(edge_weights, i); it; ++it)
			{
				// get the weight
				float wij = it.value();

				// get the neighbor vertex
				const auto &pj = p.col(it.col());
				const auto &ppj = p_prime.col(it.col());

				Vector3f d = pi - pj;
				Vector3f dp = ppi - ppj;

				cov += wij * d * (dp.transpose());
			}

			JacobiSVD<S> svd(cov, ComputeFullU | ComputeFullV);
			const S &v = svd.matrixV();
			const S u_trans = svd.matrixU().transpose();

			S I = S::Identity(3, 3);
			I(2, 2) = (v * u_trans).determinant();
			rotations[i] = v * I * u_trans;
		}

		// estimate the positions
		// make a copy of b_fixed to modify
		PM b = b_fixed;
		for (Index i = 0; i < edge_weights.outerSize(); ++i)
		{
			// TODO add constraint check here
			Index free_index = vtx_to_free_vtx[i];
			if (free_index == -1)
			{
				// do not consider constraints
				continue;
			}

			for (SparseMatrix<float, RowMajor>::InnerIterator it(edge_weights, i); it; ++it)
			{
				// get the weight
				float wij = it.value();

				// get the rotation matrices for the vertices
				const S &R = rotations[i];
				const S &Rj = rotations[it.col()];

				// calculate the energy
				Eigen::Matrix<float, 3, 1> pt = (R + Rj) * (p.col(i) - p.col(it.col()));
				pt *= wij * 0.5f;
				b.col(free_index) += pt;
			}
		}

		// solve the linear system for each dimension
		Matrix<float, Dynamic, 1> U;
		for (int j = 0; j < 3; ++j)
		{
			U = solver.solve(b.row(j).transpose());

			Index idx = 0;
			for (int k = 0; k < m_shape.getVertices().size(); ++k)
			{
				if (vtx_to_free_vtx[k] != -1)
				{
					p_prime(j, k) = U(idx);
					idx++;
				}
			}
		}
	}


	// update the vertices with pprime
	for (int i = 0; i < m_shape.getVertices().size(); ++i)
	{
		new_vertices[i] = p_prime.col(i);
	}

    // Here are some helpful controls for the application
    //
    // - You start in first-person camera mode
    //   - WASD to move, left-click and drag to rotate
    //   - R and F to move vertically up and down
    //
    // - C to change to orbit camera mode
    //
    // - Right-click (and, optionally, drag) to anchor/unanchor points
    //   - Left-click an anchored point to move it around
    //
    // - Minus and equal keys (click repeatedly) to change the size of the vertices

    m_shape.setVertices(new_vertices);
}