// This source code is property of the Computer Graphics and Visualization
// chair of the TU Dresden. Do not distribute!
// Copyright (C) CGV TU Dresden - All Rights Reserved
//
#include "IKViewer.h"

#include "math_helper.h"
#include "debughelpers.h"

#include <unordered_set>

#include <cgv/gui/mouse_event.h>
#include <cgv/math/inv.h>
#include <cgv/math/mat.h>

IKViewer::IKViewer(DataStore *data)
	: node("IK Viewer"), data(data), modifying(false), target_position(0, 0, 0, 1), max_iterations(20)
{
	connect(data->endeffector_changed, this, &IKViewer::endeffector_changed);
	connect(data->base_changed, this, &IKViewer::base_changed);
}

void IKViewer::endeffector_changed(Bone *b)
{
	std::cout << "endeffector changed!" << std::endl;

	calculate_kinematic_chain(data->get_base(), data->get_endeffector());

	post_redraw();
}

void IKViewer::base_changed(Bone *b)
{
	std::cout << "base changed!" << std::endl;

	calculate_kinematic_chain(data->get_base(), data->get_endeffector());

	post_redraw();
}

int IKViewer::find_element(std::vector<Bone *> &v, Bone *element)
{
	for (int i = 0; i < v.size(); i++)
	{
		if (v[i] == element)
		{
			return i;
		}
	}

	return -1;
}

void IKViewer::calculate_kinematic_chain(Bone *base, Bone *endeffector)
{
	std::cout << "calculate_kinematic_chain" << std::endl;

	if (base)
		std::cout << "Base bone: " << base->get_name() << std::endl;
	else
	{
		std::cout << "no base set" << std::endl;
		std::cout << std::endl;
		return;
	}

	if (endeffector)
		std::cout << "End bone: " << endeffector->get_name() << std::endl;
	else
	{
		std::cout << "no endeffector set" << std::endl;
		std::cout << std::endl;
		return;
	}

	std::cout << std::endl;

	/*Task 3.1: Calculate kinematic chain*/
	kinematic_chain.clear();

	this->base = base;
	this->endeffector = endeffector;

	// get list of bones in order of endeffector to base
	std::vector<Bone *> endeffector_tree;
	std::vector<Bone *> base_tree;

	// gather all bones from endeffector to root
	Bone *current_bone = endeffector;

	while (1)
	{
		endeffector_tree.emplace_back(current_bone);

		// break if we reached base bone
		if (current_bone->get_parent() == nullptr)
			break;

		// set next bone in hierarchy
		current_bone = current_bone->get_parent();
	}

	// gather all bones from base to root
	current_bone = base;

	while (1)
	{
		base_tree.emplace_back(current_bone);

		// break if we reached base bone
		if (current_bone->get_parent() == nullptr)
			break;

		// set next bone in hierarchy
		current_bone = current_bone->get_parent();
	}

	// add all transform from endeffector to common ancestor and to reverse base

	// state == 0 means, iterating through endeffector_tree
	// state == 1 means, reverse iterating through base_tree
	int state = 0;

	int endeffector_index = 0;
	int base_index = 0;

	//TODO: check if common ancestor is visited twice

	while (1)
	{
		if (state == 0)
		{
			if (endeffector_index >= endeffector_tree.size())
				break;

			Bone *bone = endeffector_tree[endeffector_index];

			for (int i = bone->dof_count() - 1; i >= 0; i--)
			{
				kinematic_chain.emplace_front(bone->get_dof(i));
			}

			Transform *tmp = new StaticTransform(bone->calculate_transform_prev_to_current_without_dofs());
			std::shared_ptr<Transform> transform = std::shared_ptr<Transform>(tmp);
			kinematic_chain.emplace_front(transform);

			int index = find_element(base_tree, bone);
			if (index != -1)
			{
				base_index = index;
				state = 1;
				continue;
			}

			endeffector_index++;
		}
		else if (state == 1)
		{
			if (base_index < 0)
				break;

			Bone *bone = base_tree[base_index];

			for (int i = bone->dof_count() - 1; i >= 0; i--)
			{
				Transform *tmp = new InverseTransform(bone->get_dof(i));
				std::shared_ptr<Transform> inverse_dof = std::shared_ptr<Transform>(tmp);
				kinematic_chain.emplace_front(inverse_dof);
			}

			std::shared_ptr<Transform> inverse_static = std::shared_ptr<Transform>(new StaticTransform(inv(bone->calculate_transform_prev_to_current_without_dofs())));
			kinematic_chain.emplace_front(inverse_static);

			base_index--;
		}
		else
		{
			// can't ever happen, but when it does, we will know
			abort();
		}
	}

	current_endeffector_matrix.identity();

	for (auto &transform : kinematic_chain)
	{
		current_endeffector_matrix = current_endeffector_matrix * transform->calculate_matrix();
	}

	target_position = current_endeffector_matrix * endeffector->get_bone_local_tip_position();

	print_vec3(target_position);
	std::cout << std::endl;
}

void IKViewer::optimize()
{
	//used for correct GUI behavior
	data->dof_changed_by_ik = true;

	auto skeleton_size = (data->get_skeleton()->getMax() - data->get_skeleton()->getMin());
	float distance_threshold = 0.0001f * std::max({skeleton_size.x(), skeleton_size.y(), skeleton_size.z()});

	//split the current matrix in:
	//  before_dof -> dof -> after_dof

	/*Task 3.3: Implement CCD	*/
	for (int i = 0; i < max_iterations; i++)
	{
		for (auto it = kinematic_chain.rbegin(); it != kinematic_chain.rend(); i++)
		{
			//it->optimize_value(/*local_vector*/, target_position);
		}

		current_endeffector_matrix.identity();

		for (auto &transform : kinematic_chain)
		{
			current_endeffector_matrix = current_endeffector_matrix * transform->calculate_matrix();
		}

		Vec4 current_endeffector_position = current_endeffector_matrix * endeffector->get_bone_local_tip_position();

		if (distance_threshold >= length(current_endeffector_position - target_position))
			break;
	}

	//used for correct GUI behavior
	data->dof_changed_by_ik = false;
}

void IKViewer::set_target_position_2d(int x, int y)
{
	auto ctx = get_context();
	if (!ctx)
		return;

	//Get the model-view-projection matrix
	GLfloat mv[16];
	GLfloat proj[16];
	glGetFloatv(GL_MODELVIEW_MATRIX, mv);
	glGetFloatv(GL_PROJECTION_MATRIX, proj);

	cgv::math::mat<float> modelview(4, 4, mv);
	cgv::math::mat<float> projm(4, 4, proj);

	cgv::math::mat<float> mvp = projm * modelview;
	cgv::math::mat<float> inv_mvp = cgv::math::inv_44(mvp);

	//Get the z-value of the current target
	auto projected_target = mvp * cgv::math::vec<float>(4, &target_position[0]);
	float z = projected_target.z() / projected_target.w();

	//Unproject the screen position into the scene
	int width = ctx->get_width();
	int height = ctx->get_height();

	float proj_position[4] = {(2.0f * x) / width - 1.0f, (-2.0f * y) / height + 1.0f, z, 1.0f};
	auto unprojected = inv_mvp * cgv::math::vec<float>(4, proj_position);
	unprojected *= 1.0f / unprojected.w();

	target_position.x() = unprojected.x();
	target_position.y() = unprojected.y();
	target_position.z() = unprojected.z();

	if (data->get_endeffector())
		optimize();

	post_redraw();
}

bool IKViewer::handle(event &e)
{
	if (e.get_kind() == EID_MOUSE)
	{
		cgv::gui::mouse_event me = (cgv::gui::mouse_event &)e;

		switch (me.get_action())
		{
		case MA_PRESS:
			if (me.get_button() == MB_LEFT_BUTTON && me.get_modifiers() == EM_CTRL)
			{
				modifying = true;
				set_target_position_2d(me.get_x(), me.get_y());
				return true;
			}
			break;
		case MA_RELEASE:
			modifying = false;
			break;
		case MA_DRAG:
			if (modifying)
			{
				set_target_position_2d(me.get_x(), me.get_y());
				return true;
			}
			break;
		default:
			break;
		}
	}
	return false;
}

void IKViewer::stream_help(std::ostream &os)
{
}

void IKViewer::draw(cgv::render::context &ctx)
{
	if (!data->get_skeleton())
		return;

	auto skeleton_size = (data->get_skeleton()->getMax() - data->get_skeleton()->getMin());
	float scale = 0.2f * std::max({skeleton_size.x(), skeleton_size.y(), skeleton_size.z()});

	glBegin(GL_LINES);

	if (data->get_endeffector())
	{
		glColor3f(1, 1, 1);

		glVertex3f(target_position.x() - scale, target_position.y(), target_position.z());
		glVertex3f(target_position.x() + scale, target_position.y(), target_position.z());

		glVertex3f(target_position.x(), target_position.y() - scale, target_position.z());
		glVertex3f(target_position.x(), target_position.y() + scale, target_position.z());

		glVertex3f(target_position.x(), target_position.y(), target_position.z() - scale);
		glVertex3f(target_position.x(), target_position.y(), target_position.z() + scale);
	}

	if (data->get_base())
	{
		glColor3f(0.5f, 1, 0.5f);

		glVertex3f(current_base_matrix(0, 3) - scale, current_base_matrix(1, 3), current_base_matrix(2, 3));
		glVertex3f(current_base_matrix(0, 3) + scale, current_base_matrix(1, 3), current_base_matrix(2, 3));

		glVertex3f(current_base_matrix(0, 3), current_base_matrix(1, 3) - scale, current_base_matrix(2, 3));
		glVertex3f(current_base_matrix(0, 3), current_base_matrix(1, 3) + scale, current_base_matrix(2, 3));

		glVertex3f(current_base_matrix(0, 3), current_base_matrix(1, 3), current_base_matrix(2, 3) - scale);
		glVertex3f(current_base_matrix(0, 3), current_base_matrix(1, 3), current_base_matrix(2, 3) + scale);
	}

	glEnd();
}

void IKViewer::create_gui()
{
	add_member_control(this, "Max Iterations", max_iterations, "value_slider", "min=1;max=100");
}