#include "iostream"
#include "fstream"
#include "sstream"
#include "string.h"
#include "cmath"
#include "vector"
#include "iomanip"
#include "algorithm"
#include "ctime"

#include "gctl/core.h"
#include "gctl/geometry.h"
#include "gctl/io.h"

using namespace std;

struct vertex;
struct ele;

struct vertex : public gctl::vertex2dc
{
	int tag = 0; //0 = far away, 1 = close, 2 = active
	double time = 1e+30;
	double syn_time = 1e+30; //synthetic direct arrive time (only for error test)
	vector<vertex*> v_neigh; //neighbor vertex unknown amount
	vector<ele*> e_neigh; //neighbor element unknown amount
	vector<int> ve_order; // 此顶点在相接三角形中的顶点序号(0、1或者2)，与e_neigh中的三角形一一对应，应该能提高计算效率
};

struct ele
{
	int id, tag = 0; //0 = not calculated, 1 = already calculated
	vertex* vec_ptr[3];
	double slow;
	double area;
	double edge_l[3];
};

double local_update_triangle(double area, double length_ab, double length_ac, 
	double length_bc, double a_time, double b_time, double slow)
{
	///< 沿边ac或bc传播的走时值
	double direct_time = GCTL_MIN(a_time + slow*length_ac, 
		b_time + slow*length_bc);

	double u = b_time - a_time;
	double w = slow*slow*length_ab*length_ab - u*u;
	if (w < 0) return direct_time;

	w = sqrt(w);
	double rho0 = 2.0*area/length_ab;
	double ksi0 = sqrt(length_ac*length_ac - rho0*rho0)/length_ab;
	double ksi  = ksi0 - u*rho0/(length_ab*w);

	if (ksi >= 1 || ksi <= 0) return direct_time;
	else
	{
		double trial_time = a_time + u*ksi0 + w*rho0/length_ab;
		if (trial_time > a_time && trial_time > b_time)
			return GCTL_MIN(trial_time, direct_time);
		else return direct_time;
	}
}

void local_update(vertex* vert_ptr)
{
	int j;
	double time1, time2;
	// loop over all neighboring triangles. If any one have two known vertice, update time
	for (int i = 0; i < vert_ptr->e_neigh.size(); i++)
	{
		if (vert_ptr->e_neigh[i]->tag == 0)
		{
			j = vert_ptr->ve_order[i];
			if (vert_ptr->e_neigh[i]->vec_ptr[(j+1)%3]->tag == 2 &&
				vert_ptr->e_neigh[i]->vec_ptr[(j+2)%3]->tag == 2)
			{
				time1 = vert_ptr->e_neigh[i]->vec_ptr[(j+1)%3]->time;
				time2 = vert_ptr->e_neigh[i]->vec_ptr[(j+2)%3]->time;

				if (time1 <= time2)
					vert_ptr->time = GCTL_MIN(vert_ptr->time, local_update_triangle(
						vert_ptr->e_neigh[i]->area, vert_ptr->e_neigh[i]->edge_l[j],
						vert_ptr->e_neigh[i]->edge_l[(j+2)%3], vert_ptr->e_neigh[i]->edge_l[(j+1)%3],
						time1, time2, vert_ptr->e_neigh[i]->slow));
				else
					vert_ptr->time = GCTL_MIN(vert_ptr->time, local_update_triangle(
						vert_ptr->e_neigh[i]->area, vert_ptr->e_neigh[i]->edge_l[j],
						vert_ptr->e_neigh[i]->edge_l[(j+1)%3], vert_ptr->e_neigh[i]->edge_l[(j+2)%3],
						time2, time1, vert_ptr->e_neigh[i]->slow));
				vert_ptr->e_neigh[i]->tag = 1;
			}
		}
	}
	return;
}

class FMM_2D_TRIANGLE
{
public:
	FMM_2D_TRIANGLE(){}
	~FMM_2D_TRIANGLE(){}
	// read mesh files generated by triangle
	int ReadFiles(char* filename);
	// output a Gmsh (.msh) file
	int OutMsh(char* filename, bool tag);
	// add abnormal slowness
	void SetRectangeSlowness(double ab_slow,double xmin,double xmax,double ymin,double ymax);
	// initialize triangle's slowness
	void InitEleSlowness(double in_slow);
	// calculate synthetic direct arrive time and fmm time
	void CalculateSolution(double bkg_slow);
	// set source parameters
	void set_init_source(const gctl::point2dc &loc);
private:
	size_t node_num_, ele_num_, src_id_;
	vector<vertex> nodes_;
	vector<ele> elements_;

	vector<vertex*> nodes_ptr_;
};

int FMM_2D_TRIANGLE::ReadFiles(char* filename)
{
	ifstream infile;
	string temp_str;
	stringstream temp_ss;

	char full_name[1024];
	//read node file first
	strcpy(full_name,filename);
	strcat(full_name,".node");
	infile.open(full_name);
	if (!infile)
	{
		cerr << "file open error: " << full_name << endl;
		return -1;
	}
	else
	{
		clog << "reading node file:\t" << full_name;

		getline(infile,temp_str);
		gctl::str2ss(temp_str, temp_ss);
		temp_ss >> node_num_;

		nodes_.resize(node_num_);
		for (int i = 0; i < node_num_; i++)
		{
			getline(infile,temp_str);
			gctl::str2ss(temp_str, temp_ss);
			temp_ss >> nodes_[i].id >> nodes_[i].x >> nodes_[i].y;
		}
		infile.close();

		clog << " ... done!" << endl;
	}

	//read element file
	int tmp_ids[3];
	strcpy(full_name,filename);
	strcat(full_name,".ele");
	infile.open(full_name);
	if (!infile)
	{
		cerr << "file open error: " << full_name << endl;
		return -1;
	}
	else
	{
		clog << "reading element file:\t" << full_name;

		getline(infile,temp_str);
		gctl::str2ss(temp_str, temp_ss);
		temp_ss >> ele_num_;

		elements_.resize(ele_num_);
		for (int i = 0; i < ele_num_; i++)
		{
			getline(infile,temp_str);
			gctl::str2ss(temp_str, temp_ss);
			temp_ss >> elements_[i].id >> tmp_ids[0] >> tmp_ids[1] >> tmp_ids[2];
			elements_[i].vec_ptr[0] = &nodes_[tmp_ids[0]];
			elements_[i].vec_ptr[1] = &nodes_[tmp_ids[1]];
			elements_[i].vec_ptr[2] = &nodes_[tmp_ids[2]];
		}
		infile.close();

		clog << " ... done!" << endl;
	}

	// reorder the vertice sequence to anti-clockwise
	gctl::point2dc v01, v02;
	vertex *tmp_ptr;
	for (int i = 0; i < ele_num_; i++)
	{
		v01.x = elements_[i].vec_ptr[1]->x - elements_[i].vec_ptr[0]->x;
		v01.y = elements_[i].vec_ptr[1]->y - elements_[i].vec_ptr[0]->y;

		v02.x = elements_[i].vec_ptr[2]->x - elements_[i].vec_ptr[0]->x;
		v02.y = elements_[i].vec_ptr[2]->y - elements_[i].vec_ptr[0]->y;

		if (v01.x*v02.y - v01.y*v02.x < 0)
		{
			tmp_ptr = elements_[i].vec_ptr[1];
			elements_[i].vec_ptr[1] = elements_[i].vec_ptr[2];
			elements_[i].vec_ptr[2] = tmp_ptr;
		}
	}

	//sort neighboring vertice and triangles for a vertex
	double p;
	for (int i = 0; i < ele_num_; i++)
	{
		for (int j = 0; j < 3; j++)
		{
			// opposite edge length of vertex j
			elements_[i].edge_l[j] = gctl::distance(*elements_[i].vec_ptr[(j+1)%3], *elements_[i].vec_ptr[(j+2)%3]);
		}

		// calculate triangle's area
		p = 0.5*(elements_[i].edge_l[0] + elements_[i].edge_l[1] + elements_[i].edge_l[2]);
		elements_[i].area = sqrt(p*(p-elements_[i].edge_l[0])*(p-elements_[i].edge_l[1])*(p-elements_[i].edge_l[2]));

		for (int j = 0; j < 3; j++)
		{
			// add current element to each vertex element's neighboring list
			elements_[i].vec_ptr[j]->e_neigh.push_back(&elements_[i]);
			elements_[i].vec_ptr[j]->ve_order.push_back(j);
			// add the other two vertice to current vertex's node neighboring list
			// this operation will create duplicate enters for nodes's neighbors
			elements_[i].vec_ptr[j]->v_neigh.push_back(elements_[i].vec_ptr[(j+1)%3]);
			elements_[i].vec_ptr[j]->v_neigh.push_back(elements_[i].vec_ptr[(j+2)%3]);
		}
	}

	// remove duplicate enters for neighboring vertice
	vector<vertex*>::iterator pos;
	for (int i = 0; i < node_num_; i++)
	{
		sort(nodes_[i].v_neigh.begin(),nodes_[i].v_neigh.end()); //对顶点序列由小到大排序
		pos = unique(nodes_[i].v_neigh.begin(),nodes_[i].v_neigh.end()); //获取重复序列开始的位置
		nodes_[i].v_neigh.erase(pos,nodes_[i].v_neigh.end()); //删除重复点
	}

	//output statistics to screen
	clog << "node number:\t" << node_num_ << endl;
	clog << "element number:\t" << ele_num_ << endl;

	return 0;
}

int FMM_2D_TRIANGLE::OutMsh(char* filename, bool tag)
{
	ofstream outfile;
	if (tag)
		strcat(filename,"_linear.msh");
	else
		strcat(filename,".msh");

	outfile.open(filename);
	if (!outfile)
	{
		cerr << "file open error: " << filename << endl;
		return -1;
	}
	else clog << "writing .msh file:\t" << filename << endl;

	outfile<<"$MeshFormat"<<endl<<"2.2 0 8"<<endl<<"$EndMeshFormat"<<endl<<"$Nodes"<<endl<<node_num_<<endl;
	// set the first vertex index to 1 to avoid a display bug in Gmsh
	for (int i = 0; i < node_num_; i++)
	{
		outfile << nodes_[i].id + 1 << " " << setprecision(16) << nodes_[i].x << " " << nodes_[i].y << " 0" << endl;
	}
	outfile<<"$EndNodes"<<endl;
	outfile<<"$Elements"<<endl<<ele_num_<<endl;
	for (int i = 0; i < ele_num_; i++)
	{
		outfile << elements_[i].id + 1 <<" 2 1 0";
		for (int j = 0; j < 3; j++)
			outfile << " " << elements_[i].vec_ptr[j]->id + 1;
		outfile << endl;
	}
	outfile << "$EndElements"<< endl;
	if (tag)
	{
		int tmp_node_num = 0;
		for (int i = 0; i < node_num_; i++)
		{
			if (nodes_[i].time < 1e+10)
				tmp_node_num++;
		}
		outfile <<"$NodeData"<< endl;
		outfile<<1<<endl<<"\"fmm arrive time (ms)\""<<endl<<1<<endl<<0.0<<endl
		<<3<<endl<<0<<endl<<1<<endl<<tmp_node_num<<endl;
		for (int i = 0; i < node_num_; i++)
		{
			if (nodes_[i].time < 1e+10)
			{
				outfile << nodes_[i].id + 1 << " " << setprecision(16) << nodes_[i].time << endl;
			}
		}
		outfile<<"$EndNodeData"<< endl;

		outfile <<"$NodeData"<< endl;
		outfile<<1<<endl<<"\"synthetic time (ms)\""<<endl<<1<<endl<<0.0<<endl
		<<3<<endl<<0<<endl<<1<<endl<<tmp_node_num<<endl;
		for (int i = 0; i < node_num_; i++)
		{
			if (nodes_[i].time < 1e+10) // this condition seems unnecessary, we will keep it here until we know better
			{
				outfile << nodes_[i].id + 1 << " " << setprecision(16) << nodes_[i].syn_time << endl;
			}
		}
		outfile<<"$EndNodeData"<< endl;

		// 这里我们顺带统计一下误差的平均值和标准差
		vector<double> tmp_err(tmp_node_num, 0.0);
		double mean_err = 0.0;
		outfile <<"$NodeData"<< endl;
		outfile<<1<<endl<<"\"fmm - synthetic time (ms)\""<<endl<<1<<endl<<0.0<<endl
		<<3<<endl<<0<<endl<<1<<endl<<tmp_node_num<<endl;
		for (int i = 0; i < node_num_; i++)
		{
			if (nodes_[i].time < 1e+10)
			{
				tmp_err[i] = nodes_[i].time - nodes_[i].syn_time;
				mean_err += tmp_err[i];

				outfile << nodes_[i].id + 1 << " " << setprecision(16) << tmp_err[i] << endl;
			}
		}
		outfile<<"$EndNodeData"<< endl;

		mean_err /= 1.0*tmp_node_num;
		double rms_err = 0.0;
		for (int i = 0; i < tmp_node_num; i++)
		{
			rms_err += pow(tmp_err[i] - mean_err,2);
		}
		clog << "mean error  = " << mean_err << endl;
		clog << "rms of error= " << sqrt(rms_err/tmp_node_num) << endl;

		outfile <<"$ElementData"<< endl;
		outfile<<1<<endl<<"\"model slowness (s/m)\""<<endl<<1<<endl<<0.0<<endl
		<<3<<endl<<0<<endl<<1<<endl<<ele_num_<<endl;
		for (int i = 0; i < ele_num_; i++)
		{
			outfile << elements_[i].id + 1 << " " << setprecision(16) << elements_[i].slow << endl;
		}
		outfile<<"$EndElementData"<< endl;
	}
	outfile.close();
	return 0;
}

// set slowness within a rectangular area.
void FMM_2D_TRIANGLE::SetRectangeSlowness(double ab_slow,double xmin,double xmax,double ymin,double ymax)
{
	gctl::point2dc cen;
	for (int i = 0; i < ele_num_; i++)
	{
		cen = 1.0/3.0 * (*elements_[i].vec_ptr[0] + *elements_[i].vec_ptr[1] + *elements_[i].vec_ptr[2]);
		if (cen.x >= xmin && cen.x <= xmax && cen.y >= ymin && cen.y <= ymax)
		{
			elements_[i].slow = ab_slow;
		}
	}
	return;
}

// set element's slowness as the mean value of three vertice's slowness.
void FMM_2D_TRIANGLE::InitEleSlowness(double in_slow)
{
	for (int i = 0; i < ele_num_; i++)
	{
		elements_[i].slow = in_slow;
	}
	return;
}

void FMM_2D_TRIANGLE::CalculateSolution(double bkg_slow)
{
	time_t start_time, stop_time;
	start_time = time(NULL);
	clog << "calculating first arrive time ... ";

	//calculate synthetic direct arrival time for all nodes
	//this calculation is only valid for testing when there is no abnormal slowness or obstacles.
	for (int i = 0; i < node_num_; i++)
	{
		nodes_[i].syn_time = gctl::distance(nodes_[i], nodes_[src_id_]) * bkg_slow;
	}

	//set all close vertice's tag to 'close' and add them to wave front list
	for (int i = 0; i < node_num_; i++)
	{
		if (nodes_[i].tag == 2)
		{
			for (int j = 0; j < nodes_[i].v_neigh.size(); j++)
			{
				if (nodes_[i].v_neigh[j]->tag == 0)
				{
					nodes_[i].v_neigh[j]->tag = 1;
					nodes_ptr_.push_back(nodes_[i].v_neigh[j]);
				}
			}
		}
	}

	//calculate trial time for all close nodes
	for (int i = 0; i < nodes_ptr_.size(); i++)
	{
		local_update(nodes_ptr_[i]);
	}

	//marching forward and updating the close nodes set
	while (!nodes_ptr_.empty())
	{
		// heap sort close nodes pointers to put a node at the first place if it has the smallest time
		std::sort(nodes_ptr_.begin(), nodes_ptr_.end(), [](vertex *a, vertex *b)->bool{
			if (a->time < b->time) return true; return false;
		});

		//change the first node's tag to 2 and update it's neighbor's time if it is not active (tag != 2)
		//this step is needed as a vertex may possess smaller time once it's neighbor time is updated.
		nodes_ptr_[0]->tag = 2;
		for (int i = 0; i < nodes_ptr_[0]->v_neigh.size(); i++)
		{
			// if a neighboring node is not added to the wave front list.
			// Add it to the list and update it's time
			if (nodes_ptr_[0]->v_neigh[i]->tag == 0)
			{
				nodes_ptr_[0]->v_neigh[i]->tag = 1;
				local_update(nodes_ptr_[0]->v_neigh[i]);
				nodes_ptr_.push_back(nodes_ptr_[0]->v_neigh[i]);
			}
			// already in the wave front list, update it's time
			else if (nodes_ptr_[0]->v_neigh[i]->tag == 1)
			{
				local_update(nodes_ptr_[0]->v_neigh[i]);
			}
		}
		// remove the first node from the wave front list
		nodes_ptr_.erase(nodes_ptr_.begin());
	}

	stop_time = time(NULL);
	clog << stop_time - start_time << " s" << endl;
	return;
}

void FMM_2D_TRIANGLE::set_init_source(const gctl::point2dc &loc)
{
	// Find the closest node and make it the source
    double dist, mini_dist = 1e+30;
    for (size_t i = 0; i < node_num_; i++)
    {
        dist = gctl::distance(loc, nodes_[i]);
        if (dist < mini_dist)
        {
            mini_dist = dist;
            src_id_ = i;
        }
    }

    // initiate the source time
    nodes_[src_id_].time = 0.0;
    nodes_[src_id_].tag = 2;
    for (size_t i = 0; i < nodes_[src_id_].e_neigh.size(); i++)
    {
        nodes_[src_id_].e_neigh[i]->tag = 1;

        for (size_t j = 0; j < 3; j++)
        {
            if (nodes_[src_id_].e_neigh[i]->vec_ptr[j]->id != src_id_)
            {
				dist = gctl::distance(nodes_[src_id_], *nodes_[src_id_].e_neigh[i]->vec_ptr[j]);
                nodes_[src_id_].e_neigh[i]->vec_ptr[j]->time = nodes_[src_id_].e_neigh[i]->slow * dist;
                nodes_[src_id_].e_neigh[i]->vec_ptr[j]->tag = 2;
            }
        }
    }
    return;
}

/*******************************************main function here*****************************/
// argv[1] -> running mode
// argv[2] -> name of the input triangle file
// argv[3] -> source parameter
// argv[4] -> abnormal slownesses separated by commas

int main(int argc, char* argv[])
{
	FMM_2D_TRIANGLE instance;
	if (!strcmp(argv[1],"convert"))
	{
		if(!instance.ReadFiles(argv[2]))
			instance.OutMsh(argv[2],false);
	}
	else if (!strcmp(argv[1],"calculate"))
	{
		if(instance.ReadFiles(argv[2])) return 0;

		double s_x, s_y, s_slow;
		if (3 != sscanf(argv[3], "%lf/%lf/%lf", &s_x, &s_y, &s_slow))
		{
			cerr << "wrong source parameter." << endl;
			return 0;
		}
		//set slowness here
		instance.InitEleSlowness(s_slow);
		instance.set_init_source(gctl::point2dc(s_x, s_y));

		//add abnormal slowness
		double ab_slow, ab_xmin, ab_xmax, ab_ymin, ab_ymax;
		string all_ab_para, old_ab_para;
		string tmp_str;
		stringstream tmp_ss;
		if (argc >= 5)
		{
			old_ab_para = argv[4];
			gctl::replace_all(all_ab_para, old_ab_para, ",", " ");
			gctl::str2ss(all_ab_para, tmp_ss);
			while (tmp_ss >> tmp_str)
			{
				if (5 != sscanf(tmp_str.c_str(), "%lf/%lf/%lf/%lf/%lf", 
					&ab_slow, &ab_xmin, &ab_xmax, &ab_ymin, &ab_ymax))
				{
					cerr << "wrong source parameter." << endl;
					return 0;
				}
				instance.SetRectangeSlowness(ab_slow, ab_xmin, ab_xmax, ab_ymin, ab_ymax);
			}
		}

		instance.CalculateSolution(s_slow);
		if(instance.OutMsh(argv[2],true)) return 0;
	}
	return 0;
}