#pragma once

#include "stdafx.h"

#include "processedFile.h"
#include "planeExtract.h"

class Alignment {
public:

	static MeshDataf makeNormalMesh(const vec3f& p0, const vec3f& p1, const vec4f& color = vec4f(0.8f, 0.2f, 0.2f, 1.0f))
	{
		TriMeshf normalMesh0 = Shapesf::sphere(0.05f, p0, 10, 10, color);
		TriMeshf normalMesh1 = Shapesf::cylinder(p0, p1, 0.025f, 10, 10, color);
		MeshDataf nMesh = normalMesh0.computeMeshData();
		nMesh.merge(normalMesh1.computeMeshData());
		return nMesh;
	}



	static vec3f upVectorFromViews(const SensorData& sd) {
		vec3f v(0.0f, 0.0f, 0.0f);
		for (size_t i = 0; i < sd.m_frames.size(); i++) {
			const mat4f& t = sd.m_frames[i].getCameraToWorld();
			if (sd.m_frames[i].getCameraToWorld()(0, 0) == -std::numeric_limits<float>::infinity()) continue;

			const vec3f cameraUp(0.0f, -1.0f, 0.0f);
			const vec3f worldUp = (t.getRotation() * cameraUp).getNormalized();
			v += worldUp;
		}
		v /= (float)sd.m_frames.size();
		return v.getNormalized();
	}

	static bool hasAccel(const SensorData& sd, unsigned int numThresh = 10) {
		unsigned int numValidAccel = 0;
		for (const SensorData::IMUFrame& f : sd.m_IMUFrames) {
			if (f.acceleration != vec3d(0.0)) numValidAccel++;
		}

		if (numValidAccel > numThresh) return true;
		else return false;
	}

	static vec3f upVectorFromAccel(const SensorData& sd)
	{
		if (sd.m_IMUFrames.size() == 0) throw MLIB_EXCEPTION("no imu data found");

		vec3f v(0.0f, 0.0f, 0.0f);
		for (size_t i = 0; i < sd.m_frames.size(); i++) {
			const SensorData::IMUFrame& f = sd.findClosestIMUFrame(i);
			
			if (sd.m_frames[i].getCameraToWorld()(0, 0) == -std::numeric_limits<float>::infinity()) continue;

			if (f.acceleration == vec3d::origin) {
				std::cout << "invalid IMU acceleration data entry at " << i << "-th frame" << std::endl;
				continue;
			}

			vec3f cameraUp = -vec3f((float)f.acceleration.x, (float)f.acceleration.y, (float)f.acceleration.z).getNormalized();
			const mat4f& t = sd.m_frames[i].getCameraToWorld();
			const vec3f worldUp = (t.getRotation() * cameraUp).getNormalized();
			v += worldUp;
		}

		return v.getNormalized();
	}

	static bool hasGravity(const SensorData& sd, unsigned int numThresh = 10)
	{
		unsigned int numValidGravity = 0;
		for (const SensorData::IMUFrame& f : sd.m_IMUFrames) {
			if (f.gravity != vec3d(0.0)) numValidGravity++;
		}

		if (numValidGravity > numThresh) return true;
		else return false;
	}

	static vec3f upVectorFromGravity(const SensorData& sd)
	{
		if (sd.m_IMUFrames.size() == 0) throw MLIB_EXCEPTION("no imu data found");

		vec3f v(0.0f, 0.0f, 0.0f);
		for (size_t i = 0; i < sd.m_frames.size(); i++) {
			const SensorData::IMUFrame& f = sd.findClosestIMUFrame(i);
			if (sd.m_frames[i].getCameraToWorld()(0, 0) == -std::numeric_limits<float>::infinity()) continue;

			if (f.gravity == vec3d::origin) {
				std::cout << "invalid IMU gravity data entry at " << i << "-th frame" << std::endl;
				continue;
			}

			vec3f cameraUp = vec3f((float)f.gravity.x, (float)f.gravity.y, (float)f.gravity.z).getNormalized();

			cameraUp = vec3f(cameraUp.y, cameraUp.x, cameraUp.z);

			const mat4f& t = sd.m_frames[i].getCameraToWorld();
			const vec3f worldUp = (t.getRotation() * cameraUp).getNormalized();
			v += worldUp;
		}
		v /= (float)sd.m_frames.size();

		return v.getNormalized();
	}

	static void writeTrajectoryFile(const std::string& outFile, const SensorData& sd) {

		std::vector<mat4f> trajecotry;
		for (size_t i = 0; i < sd.m_frames.size(); i++) {
			trajecotry.push_back(sd.m_frames[i].getCameraToWorld());
		}

		BinaryDataStreamFile outStream(outFile, true);
		outStream << trajecotry;
		outStream.close();
	}

	static std::vector<mat4f> readTrajectoryFile(const std::string& inFile) {
		BinaryDataStreamFile in(inFile, false);
		std::vector<mat4f> trajectory;
		in >> trajectory;
		return trajectory;
	}

	static void removeInvalidIMUFrames(SensorData& sd) {
		const bool checkTimeStamp = true;	//at the moment only remove invalid time frames
		const bool checkGravity = false;
		const bool checkAccel = false;

		unsigned int removedInvalidFrames = 0;
		for (std::vector<SensorData::IMUFrame>::iterator iter = sd.m_IMUFrames.begin(); iter != sd.m_IMUFrames.end();) {
			bool rem = false;
			if (checkTimeStamp && iter->timeStamp == 0) rem = true;
			if (checkGravity && iter->gravity == vec3d::origin) rem = true;
			if (checkAccel && iter->gravity == vec3d::origin) rem = true;

			if (rem == true) {
				iter = sd.m_IMUFrames.erase(iter);
				removedInvalidFrames++;
			}
			else {
				iter++;
			}
		}

		if (removedInvalidFrames > 0) {
			std::cout << "removed " << removedInvalidFrames << " invalid IMUFrames" << std::endl;
		}
	}

	static void alignScan(const std::string& path, bool forceRealign = false) {

		const std::string processedFile(path + "/" + "processed.txt");
		if (!util::fileExists(processedFile)) {
			std::cout << "no reconstruction available for " << path << "\n\t -> skipping folder" << std::endl;
			return;
		}
		ParameterFile parameterFile(processedFile);
		ProcessedFile pf;	pf.readMembers(parameterFile);
		if (!pf.valid) {
			std::cout << "reconstruction was invalid for " << path << "\n\t -> skipping folder" << std::endl;
			return;
		}
		if (pf.aligned && !forceRealign) {
			std::cout << "reconstruction is already aligned " << path << "\n\t -> skipping folder" << std::endl;
			return;
		}

		Directory dir(path);
		const std::vector<std::string> tmp = ml::util::split(util::replace(path, "\\", "/"), "/");	//we assume forward slashes
		const std::string base = tmp.back();
		const std::string baseFile = path + "/" + base;

		const std::string sensFile = path + "/" + base + ".sens";
		const std::string plyFile = path + "/" + base + ".ply";
		const std::string trajFile = util::replace(sensFile, ".sens", ".traj");


		mat4f transform = mat4f::identity();

		SensorData sd(sensFile);
		removeInvalidIMUFrames(sd);

		if (sd.m_frames.size() == 0) throw MLIB_EXCEPTION("no frames found in the sensor file");

		if (sd.m_frames[0].getCameraToWorld() != mat4f::identity()) {
			std::cout << "already found a previous alignment -> reverting to original" << std::endl;

			if (sd.m_frames[0].getCameraToWorld()(0, 0) == -std::numeric_limits<float>::infinity()) {
				std::cout << "error can't revert due to an invalid transform in the first frame" << std::endl;
				std::cout << "\tskipping folder " << std::endl;
				return;
			}
			mat4f inverse = sd.m_frames[0].getCameraToWorld().getInverse();
			sd.applyTransform(inverse);

			//if we have a multiple ply files (e.g., if VH was already run):
			Directory dir(path);
			std::vector<std::string> plyFiles = dir.getFilesWithSuffix(".ply");
			for (const std::string& plyFile : plyFiles) {
				MeshDataf md = MeshIOf::loadFromFile(dir.getPath() + "/" + plyFile);
				md.applyTransform(inverse);
				MeshIOf::saveToFile(dir.getPath() + "/" + plyFile, md);
			}
		}

		MeshDataf md = MeshIOf::loadFromFile(plyFile);
		md.mergeCloseVertices(0.0005f, true);
		md.removeIsolatedPieces(5000);

		//compute approx up vector from camera views or gravity (if available)
		if (true) {
			vec3f upEstimateView = upVectorFromViews(sd);
			vec3f upEstimate = upEstimateView;
			if (hasGravity(sd)) { //try to use gravity if possible
				vec3f upEstimateGrav = upVectorFromGravity(sd);
				upEstimate = upEstimateGrav;
			}

			vec3f x = upEstimate ^ vec3f(upEstimate.y, -upEstimate.z, upEstimate.x);	x.normalize();
			vec3f y = upEstimate ^ x;	y.normalize();
			vec3f z = upEstimate;
			mat4f mat = mat4f(x, y, z);
			md.applyTransform(mat);
			transform = mat * transform;
		}

		md.computeVertexNormals();
		
		//attempting to find a ground plane (and align to it)
		if (true) {
			PlaneExtract pe(md);
			pe.cluster();
			pe.removeSmallClusters();
			pe.removeNonBoundingClusters(0.1f, 100);	//the 
			const auto& clusters = pe.getClusters();

			bool foundHorizontalPlane = false;
			const size_t maxClusters = clusters.size();
			size_t i = 0;
			for (const Cluster& c : clusters) {
				if ((c.m_rep.plane.getNormal() | vec3f(0.0f, 0.0f, 1.0f)) > 0.8f) {
					mat4f mat = c.reComputeNormalAlignment();
					md.applyTransform(mat);
					transform = mat * transform;
					foundHorizontalPlane = true;
					break;
				}
				i++;
				if (i >= maxClusters) break;
			}
			if (!foundHorizontalPlane) std::cout << "could not find a horizontal plane" << std::endl;

			//final alignment with xy plane (translation)
			if (true) {
				BoundingBox3f bb = md.computeBoundingBox();
				mat4f mat = mat4f::translation(-bb.getMinZ());
				transform = mat * transform;
				md.applyTransform(mat);

				bb = md.computeBoundingBox();
				mat = mat4f::translation(-vec3f(bb.getCenter().x, bb.getCenter().y, 0.0f));
				md.applyTransform(mat);
				transform = mat * transform;
			}

			//attempting to find a vertical plane to align with the x and y axis
			if (true) {
				if (true) {
					//first find the bounding box using cgal
					OrientedBoundingBox3f obb = CGALWrapperf::computeOrientedBoundingBox(md.m_Vertices, CGALWrapperf::CONSTRAIN_Z);
					mat4f mat = mat4f(obb.getAxisX().getNormalized(), obb.getAxisY().getNormalized(), obb.getAxisZ().getNormalized());
					md.applyTransform(mat);
					transform = mat * transform;
				}

				//make sure x and y are in the positive quadrant
				if (true) {
					BoundingBox3f bb = md.computeBoundingBox();
					mat4f mat = mat4f::translation(-vec3f(bb.getMin().x, bb.getMin().y, 0.0f));
					md.applyTransform(mat);
					transform = mat * transform;
				}
			}
		}

		md.m_Normals.clear();

		{
			//if we have a multiple ply files (e.g., if VH was already run):
			Directory dir(path);
			std::vector<std::string> plyFiles = dir.getFilesWithSuffix(".ply");
			for (const std::string& plyFile : plyFiles) {
				MeshDataf mesh = MeshIOf::loadFromFile(dir.getPath() + "/" + plyFile);
				mesh.applyTransform(transform);
				MeshIOf::saveToFile(dir.getPath() + "/" + plyFile, mesh);
			}


			sd.applyTransform(transform);
			sd.saveToFile(sensFile);
			pf.aligned = true;	//it's now aligned
			pf.saveToFile(processedFile);
		}
	}

	static mat4f readTransformFromAln(const std::string& filename)
	{
		std::ifstream s(filename);
		if (!s.is_open()) throw MLIB_EXCEPTION("failed to open file " + filename + " for read");
		std::string tmp;
		std::getline(s, tmp); std::getline(s, tmp); std::getline(s, tmp);//ignore header lines
		mat4f m;
		for (unsigned int i = 0; i < 16; i++) s >> m[i];
		s.close();
		return m;
	}

	static void alignScanFromAlnFile(const std::string& path) {

		const std::string processedFile(path + "/" + "processed.txt");
		if (!util::fileExists(processedFile)) {
			std::cout << "no reconstruction available for " << path << "\n\t -> skipping folder" << std::endl;
			return;
		}
		ParameterFile parameterFile(processedFile);
		ProcessedFile pf;	pf.readMembers(parameterFile);
		if (!pf.valid) {
			std::cout << "reconstruction was invalid for " << path << "\n\t -> skipping folder" << std::endl;
			return;
		}

		Directory dir(path);
		const std::vector<std::string> tmp = ml::util::split(util::replace(path, "\\", "/"), "/");	//we assume forward slashes
		const std::string base = tmp.back();
		const std::string baseFile = path + "/" + base;

		const std::string alnFile = path + "/alignment.aln";
		const std::string sensFile = path + "/" + base + ".sens";
		std::vector<std::string> plyFiles = dir.getFilesWithSuffix(".ply");

		const mat4f transform = readTransformFromAln(alnFile);

		SensorData sd(sensFile);
		if (sd.m_frames.size() == 0) throw MLIB_EXCEPTION("no frames found in the sensor file");
		sd.applyTransform(transform);

		{ //save out transformed
			for (const std::string& plyFile : plyFiles) {
				MeshDataf mesh = MeshIOf::loadFromFile(dir.getPath() + "/" + plyFile);
				mesh.applyTransform(transform);
				MeshIOf::saveToFile(dir.getPath() + "/" + plyFile, mesh);
			}
			sd.applyTransform(transform);
			sd.saveToFile(sensFile);
			pf.aligned = true;	//it's now aligned
			pf.saveToFile(processedFile);
		}
	}


private:
};