/*
 * run_esprit_tree.cpp
 *
 *  Created on: Dec 19, 2014
 *      Author: qimao
 */

#include "esprit_tree.h"

#include <iostream>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/wait.h>
#include <unordered_set>
#include <string>
#include <vector>

#include <armadillo>
using namespace arma;


void create_fasta_files( std::vector<std::unordered_set<int> > clusters, seqan::StringSet<seqan::DnaString> &seqs,
		std::string &fasta_filename, std::vector<std::string> &fa_names, std::string &output_path)
{
	int nc = clusters.size();
	fa_names.clear();

	for(int i = 0;i<nc;i++)
	{
		std::string fname(output_path); // have to copy string
		fname.append("/epsrit_tree_"+std::to_string(i)+".fa");
		fa_names.push_back(fname);
	}

	for(int i = 0; i<nc; i++)
	{
		std::ofstream out(fa_names[i]);

		std::unordered_set<int>::iterator it = clusters[i].begin();
		for(;it != clusters[i].end();++it)
		{
			int idx = *it;
//			seqan::CharString id(std::to_string(idx));
//			seqan::writeRecord(out,id, seqs[idx],seqan::Fasta());
			out<<">"<<idx<<std::endl;
			out<<seqs[idx]<<std::endl;
		}
		out.close();
	}

}


void preproc(std::vector<std::string> &fa_names, int num_proc, uvec &sort_idx, std::string &exe_path)
{
	std::string exe_full_path(exe_path+"/preproc");

	int nc = fa_names.size();
	num_proc = std::min(num_proc, nc);

	int pos = 0;
	for(;pos<num_proc;++pos)
	{
		int cur_idx = sort_idx(pos);
		std::string fasta_file = fa_names[cur_idx];

		char * my_args[4] = { (char *)(exe_full_path.c_str()),
				(char *)"-a",
				(char *)(fasta_file.c_str() ),
				NULL};

		printf("\n++++++++call preproc %s +++++++\n",fasta_file.c_str());
		printf("execute command: %s %s %s\n", my_args[0],my_args[1],my_args[2]);

		pid_t pid = fork();
		if(pid<0){ perror("fork"); exit(1);}
		else if(pid==0){
			execv(my_args[0],my_args);
			exit(1);
		}
	}

	while(num_proc > 0)
	{
	  int status = -1;
		wait(&status);
		--num_proc;
		if(pos < nc)
		{
			int cur_idx = sort_idx(pos);
			std::string fasta_file = fa_names[cur_idx];

			char * my_args[4] = { (char *)(exe_full_path.c_str()),
					(char *)"-a",
					(char *)(fasta_file.c_str() ),
					NULL};


			printf("\n++++++++call preproc %s +++++++\n",fasta_file.c_str());
			printf("execute command: %s %s %s\n", my_args[0],my_args[1],my_args[2]);

			pid_t pid = fork();
			if(pid<0){ perror("fork"); exit(1);}
			else if(pid==0){
				execv(my_args[0],my_args);
				exit(1);
			}

			++pos;
			++num_proc;
		}
	}
}

void pbpcluster(std::vector<std::string> &fa_names, int num_proc, uvec &sort_idx, std::string &exe_path)
{

	std::string exe_full_path(exe_path+"/pbpcluster");

	int nc = fa_names.size();
	num_proc = std::min(num_proc, nc);

	int pos = 0;
	for(;pos<num_proc;++pos)
	{
		int cur_idx = sort_idx(pos);
		std::string fasta_file = fa_names[cur_idx];

		int idx = fasta_file.find(".fa");
		std::string prefix = fasta_file.substr(0,idx);
		std::string freq_file = prefix+"_Clean.frq";
		std::string clean_fasta_file = prefix+"_Clean.fa";

		char * my_args[11]={(char*)(exe_full_path.c_str()),
				(char*)"-l",(char*)"0.01",
				(char*)"-u",(char*)"0.1",
				(char*)"-m",(char*)"3",
				(char*)"-f", (char*)freq_file.c_str(),
				(char*)clean_fasta_file.c_str(),
				NULL};

		printf("\n++++++++call pbpcluster %s +++++++\n",fasta_file.c_str());
		printf("execute command: %s %s %s %s %s %s %s %s %s %s\n",
				my_args[0],my_args[1],my_args[2],my_args[3],my_args[4],my_args[5],
				my_args[6],my_args[7],my_args[8],my_args[9]);

		pid_t pid = fork();
		if(pid<0){ perror("fork"); exit(1);}
		else if(pid==0){
			execv(my_args[0],my_args);
			exit(1);
		}
	}

//	pid_t tmp;
	while(num_proc > 0)
	{
	  int status = -1;
		wait(&status);
		--num_proc;
		if(pos < nc)
		{
			int cur_idx = sort_idx(pos);
			std::string fasta_file = fa_names[cur_idx];

			int idx = fasta_file.find(".fa");
			std::string prefix = fasta_file.substr(0,idx);
			std::string freq_file = prefix+"_Clean.frq";
			std::string clean_fasta_file = prefix+"_Clean.fa";

			char * my_args[11]={(char*)(exe_full_path.c_str()),
					(char*)"-l",(char*)"0.01",
					(char*)"-u",(char*)"0.1",
					(char*)"-m",(char*)"2",
					(char*)"-f", (char*)freq_file.c_str(),
					(char*)clean_fasta_file.c_str(),
					NULL};

			printf("\n++++++++call pbpcluster %s +++++++\n",fasta_file.c_str());


			pid_t pid = fork();
			if(pid<0){ perror("fork"); exit(1);}
			else if(pid==0){
				execv(my_args[0],my_args);
				exit(1);
			}

			++pos;
			++num_proc;
		}
	}

}

void esprit_tree(std::vector<std::unordered_set<int> > &clusters, seqan::StringSet<seqan::DnaString> &seqs,
		std::string &fasta_filename, int num_proc, std::string &output_path, std::string &exe_path)
{
	std::vector<std::string> fa_names;
	create_fasta_files(clusters,seqs,fasta_filename,fa_names,output_path);

	int nc = clusters.size();
	ivec cluster_sizes(nc);
	for(int i = 0;i<nc;i++)
		cluster_sizes(i) = clusters[i].size();
	uvec sort_idx = sort_index(cluster_sizes,"descend");

	printf("\n++++++start preproc+++++++++++\n");
	preproc(fa_names, num_proc, sort_idx, exe_path);

	printf("\n++++++start pbpcluster++++++++\n");
	pbpcluster(fa_names, num_proc, sort_idx,exe_path);
}
//
//void test_esprit_tree(int argc, char **argv)
//{
//	std::string filename = "/Users/qimao/Documents/workspace/hc_seqs_cpp/data/HMP/Saliva_by_sample/Saliva.fa";
//
//	seqan::SequenceStream seqIO(filename.c_str());
//	seqan::StringSet<seqan::CharString> ids;
//	seqan::StringSet<seqan::DnaString> seqs;
//
//	int res = readAll(ids, seqs, seqIO);
//	if (res != 0){
//		std::cerr << "ERROR: Could not read records!\n";
//		exit(1);
//	}
//
//	std::vector<std::unordered_set<int> > clusters;
//	int nc = 4;
//	for(int i=0;i<nc;++i)
//	{
//		std::unordered_set<int> cluster;
//		for(int j=i; j<1000*nc; j+=nc)
//		{
//			cluster.insert(j);
//		}
//		clusters.push_back(cluster);
//	}
//
//	double t = omp_get_wtime();
//
//	std::string path = "Saliva.fa";
//	std::string result_path = "./test";
//
//	esprit_tree(clusters,seqs,path,4,result_path,"");
//
//	t = omp_get_wtime() - t;
//	printf ("It took me %f seconds.\n",t);
//}
