// 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 "Skeleton.h"

#include <fstream>
#include <map>
#include <string>
#include <sstream>

#include "math_helper.h"

Skeleton::Skeleton()
{
	origin.identity();
}

Skeleton::~Skeleton()
{
	if (root)
		delete root;
}

Bone *Skeleton::get_root() const { return root; }
const Mat4 &Skeleton::get_origin() const { return origin; }
Vec3 Skeleton::get_origin_vec() const { return Vec3(origin(0, 3), origin(1, 3), origin(2, 3)); }
void Skeleton::set_origin(const Vec3 &v) { origin = translate(v); }
void Skeleton::set_origin(const Mat4 &m) { origin = m; }

enum ParseState
{
	Ignore,
	Root,
	BoneData,
	Hierarchy,
};

bool Skeleton::fromASFFile(const std::string &filename)
{
	origin.identity();

	std::ifstream fin;
	reset_bounding_box();
	ParseState state = Ignore;

	Bone *current_node = nullptr;

	bones.clear();

	int n_dofs;

	try
	{
#ifdef _WIN32
		std::wstring wfilename = cgv::utils::str2wstr(filename);
		fin.open(wfilename);
#else
		fin.open(filename);
#endif
		if (!fin.good())
			return false;
		while (!fin.eof())
		{
			char buf[CHARS_PER_LINE];
			fin.getline(buf, CHARS_PER_LINE);
			std::string str(buf);
			str = trim(str); //remove whitespaces
			if (str.find('#') == 0)
				continue; //don't handle comments

			if (str.find(':') == 0)
			{
				if (str.find(":version") == 0)
					version = str.substr(9);
				else if (str.find(":name") == 0)
					name = str.substr(6);
				else if (str.find(":root") == 0)
				{
					state = Root;
					current_node = new Bone();
					current_node->set_name("root");
					bones["root"] = current_node;
				}
				else if (str.find(":bonedata") == 0)
					state = BoneData;
				else if (str.find(":hierarchy") == 0)
					state = Hierarchy;
				else
					state = Ignore;
				continue;
			}
			switch (state)
			{
			case Ignore:
				//ignore this line
				break;
			case Root:
				if (str.find("order") == 0)
				{
					std::string dofs = str.substr(6);
					std::string dofstr;
					std::stringstream ss(dofs);
					n_dofs = 0;
					while (!ss.eof())
					{
						ss >> dofstr;
						AtomicTransform *dof;
						if (dofstr.find("RX") == 0)
							dof = new AtomicXRotationTransform();
						else if (dofstr.find("RY") == 0)
							dof = new AtomicYRotationTransform();
						else if (dofstr.find("RZ") == 0)
							dof = new AtomicZRotationTransform();
						else if (dofstr.find("TX") == 0)
							dof = new AtomicXTranslationTransform();
						else if (dofstr.find("TY") == 0)
							dof = new AtomicYTranslationTransform();
						else if (dofstr.find("TZ") == 0)
							dof = new AtomicZTranslationTransform();
						dof->set_value(0);
						current_node->add_dof(dof);
						++n_dofs;
					}
				}
				break;
			case BoneData:
				if (str.find("begin") == 0)
					current_node = new Bone();
				else if (str.find("name") == 0)
				{
					std::string name = str.substr(5);
					current_node->set_name(name);
					bones[name] = current_node;
				}
				else if (str.find("direction") == 0)
				{
					std::string direction = str.substr(10);
					float dx, dy, dz;
					std::stringstream ss(direction);
					ss >> dx >> dy >> dz;
					current_node->set_direction_in_world_space(Vec3(dx, dy, dz));
				}
				else if (str.find("length") == 0)
				{
					std::string length = str.substr(7);
					float l;
					std::stringstream ss(length);
					ss >> l;
					current_node->set_length(l);
				}
				else if (str.find("axis") == 0)
				{
					std::string axis = str.substr(5);
					float a[3];
					std::string order;
					std::stringstream ss(axis);
					ss >> a[0] >> a[1] >> a[2] >> order;
					for (int i = 0; i < 3; ++i)
					{
						AtomicTransform *t;
						if (order.at(i) == 'X')
							t = new AtomicXRotationTransform();
						else if (order.at(i) == 'Y')
							t = new AtomicYRotationTransform();
						else if (order.at(i) == 'Z')
							t = new AtomicZRotationTransform();
						t->set_value(a[i]);
						current_node->add_axis_rotation(t);
					}
				}
				else if (str.find("dof") == 0)
				{
					std::string dofs = str.substr(4);
					std::string dofstr;
					std::stringstream ss(dofs);
					n_dofs = 0;
					while (!ss.eof())
					{
						ss >> dofstr;
						AtomicTransform *dof;
						if (dofstr.find("rx") == 0)
							dof = new AtomicXRotationTransform();
						else if (dofstr.find("ry") == 0)
							dof = new AtomicYRotationTransform();
						else if (dofstr.find("rz") == 0)
							dof = new AtomicZRotationTransform();
						current_node->add_dof(dof);
						++n_dofs;
					}
				}
				else if (str.find("limits") == 0)
				{
					std::string limits = str.substr(7);
					for (int i = 0; i < n_dofs; ++i)
					{
						if (i > 0)
						{
							fin.getline(buf, CHARS_PER_LINE);
							limits = std::string(buf);
							limits = trim(limits);
						}

						limits = limits.substr(limits.find('(') + 1);
						std::stringstream ss(limits);
						float lower, upper;
						ss >> lower >> upper;
						current_node->get_dof(n_dofs - i - 1)->set_limits(lower, upper);
					}
				}
				break;
			case Hierarchy:
				if (str.find("begin") == std::string::npos && str.find("end") == std::string::npos)
				{
					std::stringstream ss(str);
					std::string parentstr;
					ss >> parentstr;
					Bone *parent = bones[parentstr];
					std::string childstr;
					while (!ss.eof())
					{
						ss >> childstr;
						Bone *child = bones[childstr];
						parent->add_child(child);
					}
				}
			}
		}
		fin.close();
		root = bones["root"];

		//root->revert_dofs();

		postprocess(root, Vec3(0, 0, 0));
		return true;
	}
	catch (...)
	{
		fin.close();
		return false;
	}
}

void Skeleton::postprocess(Bone *node, const Vec3 &global_position)
{
	//For display adaptation
	auto bone_offset_in_global_system = node->get_direction_in_world_space() * node->get_length();
	auto bone_tip_in_global_system = global_position + bone_offset_in_global_system;
	add_point(bone_tip_in_global_system);

	node->calculate_matrices();
	int n = node->childCount();
	for (int i = 0; i < n; ++i)
	{
		postprocess(node->child_at(i), bone_tip_in_global_system);
	}
}

Vec3 into_vec3(const Vec4 &v)
{
	return Vec3(v.x(), v.y(), v.z());
}

void print_vec3(const Vec3 &v)
{
	std::cout << "(" << v.x() << ", " << v.y() << ", " << v.z() << ")" << std::endl;
}

Bone *find_root(Bone *bone)
{
	float len = length(bone->get_bone_local_root_position() - bone->get_bone_local_tip_position());

	if (len > 0.0f)
		return bone;

	for (int i = 0; i < bone->childCount(); i++)
	{
		Bone *new_root = find_root(bone->child_at(i));

		if (new_root != nullptr)
			return new_root;
	}

	return nullptr;
}

void Skeleton::get_bones(Bone *bone, std::vector<Bone *> &bone_list)
{
	float len = length(bone->get_bone_local_root_position() - bone->get_bone_local_tip_position());

	if (len > 0.0f)
		bone_list.emplace_back(bone);

	for (int i = 0; i < bone->childCount(); i++)
	{
		get_bones(bone->child_at(i), bone_list);
	}
}

int Skeleton::find_index(Bone *bone, std::vector<Bone *> &bone_list)
{
	for (int i = 0; i < bone_list.size(); i++)
	{
		if (bone == bone_list[i])
		{
			return i;
		}
	}

	return -1;
}

void calc_points(Bone *bone, std::vector<Vec3> &points, Vec3 parent_pos)
{
	Vec3 new_pos = parent_pos + into_vec3(bone->get_bone_local_tip_position());

	points.emplace_back(new_pos);

	for (int i = 0; i < bone->childCount(); i++)
	{
		calc_points(bone->child_at(i), points, new_pos);
	}
}

void calc_points(Bone *bone, std::vector<Vec3> &points, Vec3 parent_pos, float correction)
{
	Vec3 new_pos = parent_pos + (into_vec3(bone->get_bone_local_tip_position()) / correction);

	points.emplace_back(new_pos);

	for (int i = 0; i < bone->childCount(); i++)
	{
		calc_points(bone->child_at(i), points, new_pos, correction);
	}
}

struct BonePoint
{
	Bone *bone;
	Vec3 point;
};

void calc_points_off_set(Bone *bone, std::vector<BonePoint> &points, Vec3 parent_pos, float correction)
{
	Vec3 new_pos = parent_pos + (into_vec3(bone->get_bone_local_tip_position()) / correction);

	BonePoint bp = {bone, new_pos};

	points.emplace_back(bp);

	for (int i = 0; i < bone->childCount(); i++)
	{
		calc_points_off_set(bone->child_at(i), points, new_pos, correction);
	}
}

float abs(float &v)
{
	if (v < 0.0f)
	{
		return -v;
	}

	return v;
}

struct Differences
{
	Vec3 min;
	Vec3 max;
	Vec3 diff;
	float max_diff;
};

Differences calc_diffs(std::vector<Vec3> &points)
{
	Vec3 min(0.0f, 0.0f, 0.0f);
	Vec3 max(0.0f, 0.0f, 0.0f);

	for (auto &point : points)
	{
		//print_vec3(point);

		if (point.x() < min.x())
		{
			min.x() = point.x();
		}

		if (point.y() < min.y())
		{
			min.y() = point.y();
		}

		if (point.z() < min.z())
		{
			min.z() = point.z();
		}

		if (point.x() > max.x())
		{
			max.x() = point.x();
		}

		if (point.y() > max.y())
		{
			max.y() = point.y();
		}

		if (point.z() > max.z())
		{
			max.z() = point.z();
		}
	}

	float x_diff = abs(min.x()) + abs(max.x());
	float y_diff = abs(min.y()) + abs(max.y());
	float z_diff = abs(min.z()) + abs(max.z());

	float max_diff = std::max(x_diff, std::max(y_diff, z_diff));

	Differences diffs = {
		min, max, Vec3(x_diff, y_diff, z_diff), max_diff};

	return diffs;
}

float nip(float v)
{
	if (v < 0.0f)
		return -v;

	return 0.0f;
}

Vec3 nip(Vec3 v)
{
	return Vec3(nip(v.x()), nip(v.y()), nip(v.z()));
}

float correct_fract(float &f)
{
	float almost_zero = 0.0000001f;

	if (abs(f) < almost_zero)
		return 0.0f;

	return f;
}

Vec3 correct_fract(Vec3 &v)
{
	return Vec3(correct_fract(v.x()), correct_fract(v.y()), correct_fract(v.z()));
}

Vec3 find_point(Bone *bone, std::vector<BonePoint> &bps)
{
	for (BonePoint &bp : bps)
	{
		if (bp.bone == bone)
			return correct_fract(bp.point);
	}

	// bone not found, something went wrong
	abort();
}

void Skeleton::write_pinocchio_file(const std::string &filename)
{
	std::ofstream o;
#ifdef _WIN32
	std::wstring wfilename = cgv::utils::str2wstr(filename);
	o.open(wfilename, std::ios::out);
#else
	o.open(filename, std::ios::out);
#endif

	if (o)
	{
		/*Task 4.1: Write Pinocchio file into o */

		std::vector<Bone *> bone_list;
		get_bones(root, bone_list);

		// gather all points
		std::vector<Vec3> points;
		calc_points(root, points, Vec3(0.0f, 0.0f, 0.0f));

		Differences diffs = calc_diffs(points);

		std::vector<Vec3> points2;
		calc_points(root, points2, Vec3(0.0f, 0.0f, 0.0f), diffs.max_diff);

		Differences diffs2 = calc_diffs(points2);

		Vec3 offset(nip(diffs2.min));

		std::cout << "offset: ";
		print_vec3(offset);

		std::vector<BonePoint> points3;
		calc_points_off_set(root, points3, offset, diffs.max_diff);

		//std::cout << "max_diff: " << max_diff << std::endl;

		Bone *new_root = find_root(root);

		// ---------------------------- info ----------------------------
		int i = 0;
		for (Bone *bone : bone_list)
		{
			//float len = length(bone->get_bone_local_tip_position());

			int parent_index = find_index(bone->get_parent(), bone_list);

			if (parent_index == -1)
			{
				if (bone != new_root)
				{
					parent_index = find_index(new_root, bone_list);
				}
			}

			//std::cout << i++ << ": " << bone->get_name() << ": " << len << " parent: " << parent_index << std::endl;

			Vec3 c = find_point(bone, points3);

			o << i++ << " " << c.x() << " " << c.y() << " " << c.z() << " " << parent_index << std::endl;
		}
		//std::cout << std::endl
		//		  << "---------------------------" << std::endl
		//		  << std::endl;
	}
	o.close();
}

Bone *get_bone(int i, std::vector<Bone *> &bone_list)
{
	if (i < 0)
		return nullptr;

	if (i >= bone_list.size())
		return nullptr;

	return bone_list[i];
}

void Skeleton::read_pinocchio_file(std::string filename)
{
	std::ifstream o;
#ifdef _WIN32
	std::wstring wfilename = cgv::utils::str2wstr(filename);
	o.open(wfilename, std::ios::in);
#else
	o.open(filename, std::ios::in);
#endif
	if (o)
	{
		std::vector<Bone *> bone_list;
		get_bones(root, bone_list);

		std::string line;

		while (std::getline(o, line))
		{
			std::istringstream iss(line);

			int bone_id, parent_id;
			float x, y, z;

			iss >> bone_id >> x >> y >> z >> parent_id;

			//std::cout << bone_id << " " << x << " " << y << " " << z << " " << parent_id << std::endl;

			Vec3 position(x, y, z);
			Bone *current_bone = get_bone(bone_id, bone_list);

			if (current_bone == nullptr)
				abort();

			current_bone->set_length(position.normalize());
			current_bone->set_direction_in_world_space(position);

			// TODO: add new orientation

			current_bone->calculate_matrices();
		}

		/*Task 4.1: Read Pinocchio file */
	}

	o.close();

	postprocess(root, get_origin_vec());
}

void Skeleton::get_skinning_matrices(std::vector<Mat4> &matrices)
{
	/*Task 4.6: Calculate skinning matrices */
}

Bone *Skeleton::find_bone(const std::string &name) const
{
	auto it = bones.find(name);
	if (it == bones.end())
		return nullptr;
	else
		return it->second;
}