#include "file_io.h"

#include "../geometry/point_set.h"
#include "../geometry/vertex_group.h"
#include "../basic/stop_watch.h"
#include "../basic/file_system.h"
#include "color.h"
#include "../widget/logger.h"
#include "../widget/message_box.h"
#include <fstream>



PointSet* FileIO::read_point_cloud(const std::string& file_name)
{
	StopWatch w;
	Logger::output(title()) << "reading file..." << Logger::endl();

	const std::string& ext = FileSystem::extension(file_name);

	PointSet* point_set = new PointSet;

	if (ext == "bxyz") 
		load_bxyz(point_set, file_name);

	else if (ext == "xyz") 
		load_xyz(point_set, file_name);

	else if (ext == "bpn") 
		load_bpn(point_set, file_name);

	else if (ext == "pn") 
		load_pn(point_set, file_name);

	else if (ext == "blab") 
		load_blab(point_set, file_name);

	else if (ext == "lab") 
		load_lab(point_set, file_name);

	else if (ext == "pcd") 
		load_pcd(point_set, file_name);

	else {
		Logger::output(title()) << "reading file failed (unknown file format)" << Logger::endl();
		delete point_set;
		return nil;
	}

	if (point_set->nb_vertices() < 1) {
		Logger::output(title()) << "reading file failed (no data exist)" << Logger::endl();
		delete point_set;
		return nil;
	}

	Logger::output(title()) << "reading file done. Time: "
		<< w.elapsed() << " seconds" << Logger::endl();

	return point_set;
}



bool FileIO::save_point_cloud(PointSet* point_set, const std::string& file_name) {
	StopWatch w;
	const std::string& ext = FileSystem::extension(file_name);

	if (ext == "bxyz")
		save_bxyz(point_set, file_name);

	else if (ext == "xyz")
		save_xyz(point_set, file_name);

	else if (ext == "pcd")
		save_pcd(point_set, file_name);

	else if (ext == "bpn") 
		save_bpn(point_set, file_name);

	else if (ext == "pn") 
		save_pn(point_set, file_name);

	else if (ext == "blab")
		save_blab(point_set, file_name);

	else if (ext == "lab")
		save_lab(point_set, file_name);

	else {
		Logger::output(title()) << "saving file failed (unknown file format)" << Logger::endl();
		return false;
	}

	Logger::output(title()) << "saving file done. Time: "
		<< w.elapsed() << " seconds" << Logger::endl();

	return true;
}


void FileIO::load_bxyz(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str(), std::fstream::binary) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	// check size of types
	if (sizeof(float) != 4) {
		Logger::output(title()) << "wrong type size, sizeof(float) != 4" << Logger::endl();
		return ;
	}

	int begin_pos = input.tellg();
	input.seekg(0, std::ios::end);
	int end_pos = input.tellg();
	// num of points in the file
	int num = (end_pos - begin_pos) / 12;  // 12 bytes per point

	input.seekg(0, std::ios::beg); 

	float* coords = new float[num * 3];
	// read large blocks
	input.read((char*)coords, num * 12);

	Vertex** vertices = new Vertex*[num];

	for (int i=0; i<num; ++i) {
		vertices[i] = new Vertex(coords[i*3], coords[i*3+1], coords[i*3+2]);
	}

	delete [] coords;
	
	pointSet->set_vertices(vertices, num);
	pointSet->assign_color_by_height();
}


void FileIO::load_bpn(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str(), std::fstream::binary) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	// check size of types
	unsigned int size_int   = sizeof(int);
	unsigned int size_float = sizeof(float);

	if (size_float != 4) {
		Logger::output(title()) << "wrong type size, sizeof(float) != 4"  << Logger::endl();
		return ;
	}

	// num of points in the file
	int num;
	input.read((char*)(&num), size_int);

	float* data = new float[num * 6]; // 6 coords per points
	// read large blocks
	input.read((char*)data, num * 6 * size_float); 

	Vertex** vertices = new Vertex*[num];
	for (int i=0; i<num; ++i) {
		vertices[i] = new Vertex(data[i*6], data[i*6 + 1], data[i*6 + 2], data[i*6 + 3], data[i*6 + 4], data[i*6 + 5]);
	}
	delete [] data;

	pointSet->set_vertices(vertices, num);
	pointSet->set_normals_exist(true);
	pointSet->assign_color_by_normal();
}

void FileIO::load_xyz(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str()) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	std::list<Vertex*> tmp;

	// first point
	float x0, y0, z0;
	if(!input.eof()) {
		input >> x0 >> y0 >> z0;

		if (input.good())
			tmp.push_back(new Vertex(0, 0, 0));
	}

	while(!input.eof()) {
		float x, y, z;
		input >> x >> y >> z;

		if (input.good()) {
			Vertex* v = new Vertex(x - x0, y - y0, z - z0) ;
			tmp.push_back(v);
		}
	}

	//////////////////////////////////////////////////////////////////////////

	Vertex** vertices = new Vertex*[tmp.size()];
	std::list<Vertex*>::iterator itr = tmp.begin();
	for (int i = 0; itr != tmp.end(); ++itr) {
		vertices[i] = *itr;
		vertices[i]->set_index(i);
		++i;
	}

	pointSet->set_vertices(vertices, tmp.size());
	pointSet->assign_color_by_height();
}


void FileIO::save_bxyz(PointSet* pointSet, const std::string& file_name) {
	// open file
	std::ofstream output(file_name.c_str(), std::fstream::binary) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	Vertex** vertices = pointSet->vertices();
	unsigned int num = pointSet->nb_vertices();

	for (unsigned int i=0; i<num; ++i) {
		const Point3f& p = vertices[i]->point();
		float x = p.x();    output.write((char*)&x, 4);
		float y = p.y();    output.write((char*)&y, 4);
		float z = p.z();    output.write((char*)&z, 4);
	}
}


void FileIO::save_xyz(PointSet* pointSet, const std::string& file_name) {
	// open file
	std::ofstream output(file_name.c_str()) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}
	output.precision(16);

	Vertex** vertices = pointSet->vertices();
	unsigned int num = pointSet->nb_vertices();

	for (unsigned int i=0; i<num; ++i) {
		const Point3f& p = vertices[i]->point();
		output << p << std::endl;
	}
}


void FileIO::save_pn(PointSet* pointSet, const std::string& file_name) {
	// open file
	std::ofstream output(file_name.c_str()) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}
	output.precision(16);

	Vertex** vertices = pointSet->vertices();
	unsigned int num = pointSet->nb_vertices();

	for (unsigned int i=0; i<num; ++i) {
		const Point3f& p = vertices[i]->point();
		const Vector3f& n = vertices[i]->normal();

		output << p << " " << n << std::endl;
	}
}


void FileIO::save_bpn(PointSet* pointSet, const std::string& file_name) {
	unsigned int size_int   = sizeof(int);
	unsigned int size_float = sizeof(float);

	// open file
	std::ofstream output(file_name.c_str(), std::fstream::binary) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	Vertex** vertices = pointSet->vertices();
	unsigned int num = pointSet->nb_vertices();

	unsigned int num_write = 0;
	// Leon: we don't know how many to be saved now. So write point number 
	// after all the points are saved. But now we need write something, 
	// otherwise the first float point number will be replaced. 
	output.write((char*)&num, size_int); 

	for (unsigned int i=0; i<num; ++i) {
		const Point3f& p = vertices[i]->point();
		float x = p.x();    output.write((char*)&x, 4);
		float y = p.y();    output.write((char*)&y, 4);
		float z = p.z();    output.write((char*)&z, 4);

		const Vector3f& n = vertices[i]->normal();
		x = n.x();    output.write((char*)&x, 4);
		y = n.y();    output.write((char*)&y, 4);
		z = n.z();    output.write((char*)&z, 4);

		++num_write;
	}

	output.seekp(std::ios::beg);
	output.write((char*)&num_write, 4);
}


void FileIO::load_pn(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str()) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	std::list<Vertex*> tmp;

	// first point
	float x0, y0, z0, nx, ny, nz;
	if(!input.eof()) {
		input >> x0 >> y0 >> z0 >> nx >> ny >> nz;

		if (input.good())
			tmp.push_back(new Vertex(0, 0, 0, nx, ny, nz));
	}

	while(!input.eof()) {
		float x, y, z;
		input >> x >> y >> z >> nx >> ny >> nz;

		if (input.good()) {
			Vertex* v = new Vertex(x - x0, y - y0, z - z0, nx, ny, nz) ;
			tmp.push_back(v);
		}
	}

	//////////////////////////////////////////////////////////////////////////
	
	Vertex** vertices = new Vertex*[tmp.size()];
	std::list<Vertex*>::iterator itr = tmp.begin();
	for (int i = 0; itr != tmp.end(); ++itr) {
		vertices[i] = *itr;
		vertices[i]->set_index(i);
		++i;
	}

	pointSet->set_vertices(vertices, tmp.size());
	pointSet->set_normals_exist(true);
	pointSet->assign_color_by_normal();
}


VertexGroup* FileIO::lab_read_group(std::ifstream& input, Vertex** vts) {
	std::string dumy;
	std::string label;
	input >> dumy >> label;

	float r, g, b;
	input >> dumy >> r >> g >> b;
	Colorf color(r, g, b);

	int num_points;
	input >> dumy >> num_points;

	VertexGroup* grp = new VertexGroup;
	for (int i=0; i<num_points; ++i) {
		int idx;
		input >> idx;
		Vertex* v = vts[idx];
		grp->push_back(v);
	}

	String::to_lowercase(label);
	grp->set_label(label);
	grp->set_color(color);

	return grp;
}


// string are stored as array of chars in binary file
std::string FileIO::read_binary_string(std::ifstream& input) {
	int size_int = sizeof(int);
	int size_char = sizeof(char);

	int num_char = 0;    
	input.read((char*)&num_char, size_int);  

	std::string str;
	for (int i=0; i<num_char; ++i) {
		char c;
		input.read(&c, size_char);
		if (c != ' ')
			str.push_back(c);
		else
			str.push_back('-');
	}

	return str;
}

// string are stored as array of chars in binary file
void FileIO::write_binary_string(std::ofstream& output, const std::string& str) {
	int size_int = sizeof(int);
	int size_char = sizeof(char);

	int num_char = str.size();
	output.write((char*)&num_char, size_int);

	for (int i=0; i<num_char; ++i) {
		char c = str[i];
		if (c == ' ')
			c = '-';
		output.write(&c, size_char);
	}
}


// for binary file format, no string stuff except labels. we add size info before each label
VertexGroup* FileIO::blab_read_group(std::ifstream& input, Vertex** vts) {
	int float_size = sizeof(float);
	int int_size = sizeof(int);

	std::string label = read_binary_string(input);
	
	////////////////////////////////////////////////////////////////////////// 
	// Leon: check why this doesn't work in release mode
	//int label_size = 0;   
	//std::string label;
	//input.read((char*)&label_size, int_size);  
	//input.read((char*)&label, label_size); 
	//////////////////////////////////////////////////////////////////////////
	
	float arr[3];
	input.read((char*)arr, 3*float_size);
	Colorf color(arr[0], arr[1], arr[2]);

	int num_points = 0;
	input.read((char*)&num_points, int_size);  

	int* indices = new int[num_points];
	if (indices == nil) {
		Logger::output(title()) << "bad allocation of memory." << Logger::endl();
		return nil;
	}
	
	input.read((char*)indices, num_points * int_size); 

	VertexGroup* grp = new VertexGroup;
	for (int i=0; i<num_points; ++i) {
		int idx = indices[i];
		Vertex* v = vts[idx];
		grp->push_back(v);
	}
	delete [] indices;

	String::to_lowercase(label);
	grp->set_label(label);
	grp->set_color(color);

	return grp;
}

/*
group_label: label  // the first group info
group_color: color
group_num_points: num       
idx ...
*/

void FileIO::lab_write_group(std::ofstream& output, VertexGroup* g) {
	std::string label = g->label();
	output << "group_label: " << label << std::endl;

	Colorf c = g->color();
	output << "group_color: " << c.r() << " " << c.g() << " " << c.b() << std::endl;

	unsigned int num_point = g->size();
	output << "group_num_point: " << num_point << std::endl;

	for (unsigned int j=0; j<g->size(); ++j) {
		output << g->at(j)->index() << " ";
	}
	output << std::endl;
}


// for binary file format, no string stuff except labels. we add size info before each label
void FileIO::blab_write_group(std::ofstream& output, VertexGroup* group) {
	int float_size = sizeof(float);
	int int_size = sizeof(int);
	
	std::string label = group->label();
	write_binary_string(output, label);

	//////////////////////////////////////////////////////////////////////////
	// Leon: check why not work
	//std::string label = group->label();
	//int label_size = sizeof(label);
	//output.write((char*)&label_size, int_size);
	//output.write((char*)&label, label_size);
	//////////////////////////////////////////////////////////////////////////

	Colorf c = group->color();
	float r = c.r();
	float g = c.g();
	float b = c.b();
	output.write((char*)&r, float_size);
	output.write((char*)&g, float_size);
	output.write((char*)&b, float_size);

	int num_point = group->size();
	output.write((char*)&num_point, int_size);

	for (unsigned int j=0; j<group->size(); ++j) {
		int idx = group->at(j)->index();
		output.write((char*)&idx, int_size);
	}
}

/*  
// file format definition
num_points: num
x  y  z  nx  ny  nz r g b
...

num_groups: num

group_label: label  // the first group info
group_color: color
group_num_points: num
idx ...

num_children: num 

group_label: label  // 0th child of group 0    
group_color: color
group_num_points: num
idx ...          

group_label: label  // 1th child of group 0 
group_color: color
group_num_points: num
idx ...


*/

void FileIO::save_lab(PointSet* pointSet, const std::string& file_name) {
	Vertex** vts = pointSet->vertices();
	int num = pointSet->nb_vertices();
	for (int i=0; i<num; ++i) {
		vts[i]->set_index(i);
	}
	
	// open file
	std::ofstream output(file_name.c_str()) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}
	output.precision(16);

	//////////////////////////////////////////////////////////////////////////

	output << "num_points: " << num << std::endl;
	for (int i=0; i<num; ++i) {
		Vertex* v = vts[i];
		Point3f p = v->point();
		Vector3f n = v->normal();
		Colorf c = v->color();
		output << p << " " << n << " " << c.r() << " " << c.g() << " " << c.b() << std::endl;
	}

	const std::vector<VertexGroup*>& vertex_groups = pointSet->vertex_groups();
	unsigned int num_groups = vertex_groups.size();
	output << "num_groups: " << num_groups << std::endl;

	for (unsigned int i=0; i<num_groups; ++i) {
		VertexGroup* g = vertex_groups[i];
		lab_write_group(output, g);

		// children
		std::vector<VertexGroup*> children = g->children();
		output << "num_children: " << children.size() << std::endl;
		for (unsigned int j=0; j<children.size(); ++j) {
			VertexGroup* chld = children[j];
			lab_write_group(output, chld);
		}
	}
}


void FileIO::save_blab(PointSet* pointSet, const std::string& file_name) {
	Vertex** vts = pointSet->vertices();
	int num = pointSet->nb_vertices();
	for (int i=0; i<num; ++i) {
		vts[i]->set_index(i);
	}

	// open file
	std::ofstream output(file_name.c_str(), std::fstream::binary) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	//////////////////////////////////////////////////////////////////////////
	
	int float_size = sizeof(float);
	int int_size = sizeof(int);

	output.write((char*)&num, int_size);
	for (int i=0; i<num; ++i) {
		Vertex* v = vts[i];
		Point3f p = v->point();
		Vector3f n = v->normal();
		Colorf c = v->color();

		float values[9] = {p.x(), p.y(), p.z(), n.x(), n.y(), n.z(), c.r(), c.g(), c.b()};
		output.write((char*)values, 9 * float_size);
	}

	const std::vector<VertexGroup*>& vertex_groups = pointSet->vertex_groups();
	int num_groups = vertex_groups.size();
	output.write((char*)&num_groups, int_size); 

	for (int i=0; i<num_groups; ++i) {
		VertexGroup* g = vertex_groups[i];
		blab_write_group(output, g);

		// children
		std::vector<VertexGroup*> children = g->children();
		int chld_num = children.size();
		output.write((char*)&chld_num, int_size);
		for (int j=0; j<chld_num; ++j) {
			VertexGroup* chld = children[j];
			blab_write_group(output, chld);
		}
	}
}


void FileIO::load_lab(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str()) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	std::string dumy;
	int num_points;
	input >> dumy >> num_points;
	Vertex** vertices = new Vertex*[num_points];
	for (int i=0; i<num_points; ++i) {
		float x, y, z, nx, ny, nz, r, g, b;
		input >> x >> y >> z >> nx >> ny >> nz >> r >> g >> b;
		vertices[i] = new Vertex(x, y, z, nx, ny, nz);
		vertices[i]->set_color(Colorf(r, g, b));
		vertices[i]->set_index(i);
	}

	pointSet->set_vertices(vertices, num_points);
	pointSet->set_normals_exist(true);
	
	//////////////////////////////////////////////////////////////////////////

	int num_groups = 0;
	input >> dumy >> num_groups;
	for (int i=0; i<num_groups; ++i) {
		VertexGroup* g = lab_read_group(input, vertices);
		
		if (g->nb_vertice() >0) {
			g->set_point_set(pointSet);
			pointSet->add_vertex_group(g);
		}
		else
			delete g;

		int num_children = 0;
		input >> dumy >> num_children;
		for (int j=0; j<num_children; ++j) {
			VertexGroup* chld = lab_read_group(input, vertices);
			if (chld->nb_vertice() > 0) {
				chld->set_point_set(pointSet);
				g->add_child(chld);
			}
			else
				delete chld;
		}
	}
}


void FileIO::load_blab(PointSet* pointSet, const std::string& file_name) {
	std::ifstream input(file_name.c_str(), std::fstream::binary) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	int float_size = sizeof(float);
	int int_size = sizeof(int);

	int num_points = 0;
	input.read((char*)&num_points, int_size); 

	float* data = new float[num_points * 9]; // 10 coords per point.
	input.read((char*)data, num_points * 9 * float_size); 

	Vertex** vertices = new Vertex*[num_points];
	if (vertices == nil) {
		Logger::output(title()) << "bad allocation of memory." << Logger::endl();
		return ;
	}

	for (int i=0; i<num_points; ++i) {
		vertices[i] = new Vertex(data[i*9], data[i*9 + 1], data[i*9 + 2], data[i*9 + 3], data[i*9 + 4], data[i*9 + 5]);
		vertices[i]->set_color(Colorf(data[i*9 + 6], data[i*9 + 7], data[i*9 + 8]));
		vertices[i]->set_index(i);
	}
	delete [] data;

	pointSet->set_vertices(vertices, num_points);
	pointSet->set_normals_exist(true);

	//////////////////////////////////////////////////////////////////////////

	int num_groups = 0;
	input.read((char*)&num_groups, int_size); 
	for (int i=0; i<num_groups; ++i) {
		VertexGroup* g = blab_read_group(input, vertices);

		if (g->nb_vertice() > 0) {
			g->set_point_set(pointSet);
			pointSet->add_vertex_group(g);
		}
		else
			delete g;

		int num_children = 0;
		input.read((char*)&num_children, int_size); 
		for (int j=0; j<num_children; ++j) {
			VertexGroup* chld = blab_read_group(input, vertices);
			if (chld->nb_vertice() > 0) {
				chld->set_point_set(pointSet);
				g->add_child(chld);
			}
			else
				delete chld;
		}
	}
}



void FileIO::save_pcd(PointSet* pointSet, const std::string& file_name)
{
	std::ofstream output(file_name.c_str()) ;
	if(output.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}
	output.precision(16);
	unsigned int num = pointSet->nb_vertices();
	output<<"# .PCD v.7 - Point Cloud Data file format\n"
		<<"VERSION .7\n"
		<<"FIELDS x y z rgb normal_x normal_y normal_z curvature\n"
		<<"SIZE 4 4 4 4 4 4 4 4\n"
		<<"TYPE F F F F F F F F\n"
		<<"COUNT 1 1 1 1 1 1 1 1\n"
		<<"WIDTH " << num <<"\n"
		<<"HEIGHT 1\n"
		<<"VIEWPOINT 0 0 0 1 0 0 0\n"
		<<"POINTS "	<< num <<"\n"
		<<"DATA ascii\n";
	Vertex** vertices = pointSet->vertices();

	for (unsigned int i=0; i<num; ++i) {
		const Vertex*  v = vertices[i];
		const Point3f& p = v->point();

		unsigned int  r = v->color().r() * 255;
		unsigned int  g = v->color().g() * 255;
		unsigned int  b = v->color().b() * 255;

		unsigned int rgb = (r << 16 | g << 8 | b /*<< 8 | 255*/);

		output 
			<< p.x() << " " << p.y() << " " <<  p.z() << " " 
			<< *reinterpret_cast<float*>(&rgb) << " " 
			<< v->normal().x() << " " << v->normal().y() << " " << v->normal().z() << " 0" << std::endl;
	}
	output.close();
}


void FileIO::load_pcd(PointSet* pointSet, const std::string& file_name)
{
	std::ifstream input(file_name.c_str()) ;
	if(input.fail()) {
		Logger::output(title()) << "could not open file\'" << file_name << "\'" << Logger::endl();
		return ;
	}

	unsigned int num_header = 11;
	unsigned int num_points = 0;
	// this line contains point number
	unsigned int line_num = 9;

	for (unsigned int i=0; i<num_header; ++i) {
		std::string line;
		getline(input, line);

		if (i==line_num) {
			std::istringstream line_input(line);
			std::string dummy;
			line_input >> dummy >> num_points;
		}
	}

	Vertex** vertices = new Vertex*[num_points];
	for (unsigned int i=0; i<num_points; ++i) {
		float x, y, z, rgb, nx, ny, nz, cur;
		input >> x >> y >> z >> rgb >> nx >> ny >> nz >> cur;
		vertices[i] = new Vertex(x, y, z, nx, ny, nz);

		// unpack rgb into r/g/b
		unsigned int rrggbb = *reinterpret_cast<int*>(&rgb);
		unsigned int rr = (rrggbb >> 16) & 0x0000ff;
		unsigned int gg = (rrggbb >> 8)  & 0x0000ff;
		unsigned int bb = (rrggbb)       & 0x0000ff;
		float r = rr / 255.0f;
		float g = gg / 255.0f;
		float b = bb / 255.0f;
		vertices[i]->set_color(Colorf(r, g, b));
	}

	pointSet->set_vertices(vertices, num_points);
	pointSet->set_normals_exist(true);
	pointSet->assign_color_by_normal();
}