#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <unistd.h>
#include "FASTA.h"
#include "util.h"
#include "global.h"
#include "ProbCluster.h"
#include <math.h>
#include <iostream>
#include <string>
using namespace std;

int SeqCount=8000000;

char **SeqStrs;
char **LabelStrs;
char **AlignedSeqs;
int *SeqLens;
int *Freq;
int SeqNum;
int **KmerTabs;
int **KmerSeqs;
REAL **KmerCodes;
int *SeqID;
double AveSeqLen;
int SumFreq;

double MemLim=1.0;


typedef struct{
	char *seq;
	int len;
	int frq;
	int id;
}SeqRec;

extern int NumAlign1;
extern int NumAlign2;
extern int NumAlign3;
extern long int Kcomp;
extern long int Kcomp2;
extern int MaxHeaps;

void LoadSeqs(char *seqfile)
{
	FASTA fasta;
	char *seq;
	char *label;
	int idx=0;
	int seqlen;
	
	SeqStrs=(char **)Malloc(SeqCount*sizeof(char *));
	SeqLens=(int *)Malloc(SeqCount*sizeof(int));

	if (Global::ShowAlign)
	{
		LabelStrs=(char **)Malloc(SeqCount*sizeof(char *));
		AlignedSeqs=(char **)Malloc(SeqCount*sizeof(char *));
	}
	
	Freq=(int *) Malloc(SeqCount*sizeof(int));
	fasta.OpenRead(seqfile);
	while ((seqlen=fasta.GetFastaSeq(&label,&seq,0))>0)
	{
			SeqLens[idx]=seqlen;
			if (Global::ShowAlign)
			{
				LabelStrs[idx]=label;
			}
			else
			{	
				free(label);		
			}	
			SeqStrs[idx++]=seq;
			if (idx >=SeqCount)
			{
				printf("The free software supports up to 4M unique reads. Please contact the authors for larger data.\n");
				exit(0);
			}
	}
	fasta.Close();
	SeqNum=idx;
	for (idx=0;idx<SeqNum;idx++)
	{
		Freq[idx]=1;
	}

	if (Global::ShowAlign)
	{	
		for (int i=0;i<SeqNum;i++)
		{
			AlignedSeqs[i]=(char *)Malloc((SeqLens[i]+1)*sizeof(char));
			strcpy(AlignedSeqs[i],SeqStrs[i]);
		}
	}
}


void LoadFreq(char *freqfile)
{
	int i,fq;
	FILE *fp;
	char buf[65535];
	
	if (freqfile==NULL)
		return;
	
	if ((fp=fopen(freqfile,"r"))==NULL)
		{
			fprintf(stderr,"Cannot find frequency file\n");
			return;
		}
		i=0;
	while (fscanf(fp,"%s %d",buf,&fq)!=EOF)
	{
		Freq[i++]=fq;
	}			
	fclose(fp);
}


int CompSeq(const void * a, const void * b)
{
	SeqRec *Seq1=(SeqRec *)a;
	SeqRec *Seq2=(SeqRec *)b;
	if (Seq1->len > AveSeqLen && Seq1->frq > Seq2->frq) return -1;
	if (Seq2->len > AveSeqLen && Seq2->frq > Seq1->frq) return 1;
	
	if (Seq1->len >Seq2->len) return -1;
	if (Seq1->len <Seq2->len) return 1;
	return 0;	
}

void SortSeqs()
{
	SeqRec *AllRec=(SeqRec *)Malloc(SeqNum*sizeof(SeqRec));
	
	AveSeqLen=0.0;
	for (int i=0;i<SeqNum;i++)
		SumFreq+=Freq[i];
		
	for (int i=0;i<SeqNum;i++)
	{
		AveSeqLen+=SeqLens[i]*Freq[i]+0.0;
		
		AllRec[i].seq=SeqStrs[i];
		AllRec[i].len=SeqLens[i];
		AllRec[i].frq=Freq[i];
		AllRec[i].id=i;
	}
	AveSeqLen/=SumFreq;
	qsort(AllRec,SeqNum,sizeof(SeqRec),CompSeq);
	for (int i=0;i<SeqNum;i++)
	{
		SeqStrs[i]=AllRec[i].seq;
		
		SeqLens[i]=AllRec[i].len;
		Freq[i]=AllRec[i].frq;
		SeqID[i]=AllRec[i].id;
	}
	free(AllRec);
}


void BuildKmer()
{
	KmerTabs=(int**)Malloc(SeqNum*sizeof(int *));
	KmerSeqs=(int**)Malloc(SeqNum*sizeof(int *));
	KmerCodes=(REAL **) Malloc(SeqNum*sizeof(REAL *));
	for (int i=0;i<SeqNum;i++)
	{
		KmerSeqs[i]=(int *)Malloc((SeqLens[i]-Global::Kmer_Len+1)*sizeof(int));
		KmerTabs[i]=Global::kmer->AllocCodeTable();
		KmerCodes[i]=NULL;
		Global::kmer->KmerCount(SeqStrs[i],KmerTabs[i],KmerSeqs[i]);
	}

}


void Clear()
{
	int i;
	for (int i=0;i<SeqNum;i++)
	{
		free(SeqStrs[i]);
		if (Global::ShowAlign)
		{
			free(AlignedSeqs[i]);
			free(LabelStrs[i]);
		}
		if (KmerTabs[i]) free(KmerTabs[i]);
		if (KmerSeqs[i]) free(KmerSeqs[i]);
		if (KmerCodes[i]) free(KmerCodes[i]);
	}
	free(SeqStrs);
	if (Global::ShowAlign)
	{
		free(AlignedSeqs);
		free(LabelStrs);
	}
	free(SeqLens);
	free(KmerTabs);
	free(KmerSeqs);
	free(KmerCodes);
	free(Freq);
}

void DoCluster(fstream &flist,fstream &fstat,fstream &fgroup,fstream &ftree)
{
	clock_t time_1,time_2,time_3,time_4;

	ProbCluster *pcluster=new ProbCluster(Global::level_min,Global::level_max);
	pcluster->InitCluster(SeqNum);
	
	time_1=clock();
	for (int i=0;i<SeqNum;i++)
	{
//		printf("\rAdding Seq %d  of %d",i+1,SeqNum);
//		fflush(stdout);
		pcluster->AddSeq(i,Freq[i]);
	}
	//pcluster->Brief(flist);
	time_2=clock();
	printf("\nNum Aligns %d Num Kmer-Comp %ld\n", NumAlign1+NumAlign2+NumAlign3,Kcomp+Kcomp2);

	printf("\nTree Building Time %10.3g secs\n",(REAL)(time_2-time_1)/CLOCKS_PER_SEC);

//	pcluster->Brief(flist);
	
	if (Global::UseApprox)
		pcluster->CondenseBottom();
	
	pcluster->BuildNNlist();
	
	time_3=clock();
	printf("\nNum Aligns %d Num Kmer-Comp %ld\n", NumAlign1+NumAlign2+NumAlign3,Kcomp+Kcomp2);
	
	printf("\nNN Heap Building Time %10.3g secs\n",(REAL)(time_3-time_2)/CLOCKS_PER_SEC);

 	printf("\n Generating Clusters\n");
	pcluster->HCluster(flist,fstat, fgroup);
	time_4=clock();

	printf("\nNum Aligns %d Num Kmer-Comp %ld\n", NumAlign1+NumAlign2+NumAlign3,Kcomp+Kcomp2);
	printf("\nClustering Time %10.3g secs\n",(REAL)(time_4-time_3)/CLOCKS_PER_SEC);
	printf("Heap Used Size %d\n",MaxHeaps);
	pcluster->Finish();
	flist.close();
	fstat.close();
	fgroup.close();

	delete pcluster;
}


void DetermMem(int num, double lim)
{
	double memsize=lim*1024*1024*1024;
	memsize-=pow(4,Global::Kmer_Len)*num*sizeof(REAL);
	memsize-=(4*sizeof(REAL)+sizeof(char)+sizeof(int))*AveSeqLen*num;
	memsize*=0.9;
	if (memsize <1e6)
		Global::CacheSize=1000;
	else
		Global::CacheSize=(int) floor(sqrt(memsize/sizeof(DistRec) ));
	Global::CacheSize=min(Global::CacheSize,num);
	printf("Set cache dimension to %d\n",Global::CacheSize);
}

void WriteAlign(char *alfile)
{
	FASTA fasta;
	fasta.OpenWrite(alfile);
	for (int i=0;i<SeqNum;i++)
	{
		fasta.WriteFastaSeq(LabelStrs[i],AlignedSeqs[i]);
	}
	fasta.Close();
}

void Usage()
{
	printf("pbpcluster [-p] [-l lower] [-s intv] [-u upper] [-g gap_open] [-e gap_ext] [-m meml] [-r errorrate] [-a <alfile>] [-o output] [-f <freqfile>] <seqfile>  [<kfile>]\n");
	printf("\t output:\t path and filename prefix for generating output files.\n");
	printf("\t frqfile:\t file recording the duplication of sequences.\n");
	printf("\t kfile:\t\t choose the k-mer configure file. see manual.\n");
	printf("\t lower:\t\t lowest distance level for clustering. Default 0.01 . \n");
	printf("\t intv:\t\t interval of distance levels to show. Default 0.01 .\n");
	printf("\t upper:\t\t top distance level for clustering. Default 0.15 .\n");
	printf("\t gap_open/ext:\t parameters for NW alignment, see manual.\n");
	printf("\t meml:\t\t maximum allowed memory in gigabytes. Default 1.0 . \n");
	printf("\t\t\t Using larger cache can speed up the program. \n");
	printf("\t errorrate:\t base-wise sequencing error rate for single-base\n");
	printf("\t -p\t\t Use pre-clustering\n");
	printf("\t\t\t correcting. Default 0(No correcting).\n\t\t\t Error correcting may improve clustering quality.\n");
}

int main (int argc, char **argv)
{
	clock_t time_1,time_2;
	fstream flist,fstat,fgroup,ftree;
	char buf[2000];
	char c;
	char *frfile = NULL;
	char *prefix = NULL;
	char mypath[2000];
	char *alfile = NULL;
	
	strcpy(mypath,argv[0]);
	printf("PBPCluster execution path %s\n",mypath);
	
	while ((c = getopt(argc, argv, "phm:l:s:u:g:e:k:f:o:r:a:q:")) != -1)
    {
        switch (c)
        {
        case 'h':
        	Usage();
        	exit(0);
        	break;
        case 'm':
        	MemLim=atof(optarg);
        	break;
        case 'l':
			Global::level_min = atof(optarg);
			break;
		case 's':
			Global::level_step = atof(optarg);
			break;	
        case 'u':
			Global::level_max = atof(optarg);
			break;
        case 'r':
			Global::ErrorRate = atof(optarg)/100.0;
			break;
		case 'g':
			Global::gap_o = atof(optarg);
			break;	
		case 'e':
			Global::gap_e = atof(optarg);
			break;	
        case 'k':
			Global::Kmer_Len = atoi(optarg);
			break;
		case 'f':
			frfile = optarg;	
			break;
		case 'a':
			Global::ShowAlign=true;
			alfile = optarg;
			break;
		case 'q':
			Global::ProbSeqOut=optarg;
			break;
		case 'o':
			prefix = optarg;
			break;
		case 'p':
			Global::UseApprox=true;
			break;
        default:
        	fprintf(stderr,"Unknown option -%c\n",c);
					break;
        }
    }
	argc -= optind;
	argv += optind;

	if (argc==0)
	{
    	Usage();
     	exit(0);
	}
	
	char *ifile=argv[0];
	char *kfile;
		
		
	if (argc>=2)
		kfile = argv[1];
	else
		kfile = NULL;

			
	printf("Input %s Freq %s from %.3f to %.3f\n",ifile,frfile,Global::level_min,Global::level_max);

	time_1=clock();

	LoadSeqs(ifile);
	LoadFreq(frfile);
	SeqID=(int *)Malloc(SeqNum*sizeof(int));
	SortSeqs();
	
	if (prefix==NULL)
	{
		strcpy(buf,ifile);
	}
	else
	{	
		strcpy(buf,prefix);
	}
	
	RepExt(buf,"OTU");
	flist.open(buf,ios_base::out | ios_base::trunc);	
	if (flist.fail())
	{
		cerr << "Cannot Create Output" <<endl;
		exit(1);
	}
	RepExt(buf,"Cluster_List");
	fstat.open(buf,ios_base::out | ios_base::trunc);	
	if (fstat.fail())
	{
		cerr << "Cannot Create Output" <<endl;
		exit(1);
	}
	RepExt(buf,"Clusters");
	fgroup.open(buf,ios_base::out | ios_base::trunc);	
	if (fgroup.fail())
	{
		cerr << "Cannot Create Output" <<endl;
		exit(1);
	}
	//RepExt(buf,"tree");
	//ftree.open(buf,ios_base::out | ios_base::trunc);	
	printf("Total %d Reads Unique %d Average Len %.1f\n",SumFreq,SeqNum,AveSeqLen);
//	Global::SetParam(10,0.5,6,50);
	DetermMem(SeqNum,MemLim);
	
	if (kfile==NULL)
	{
		kfile=mypath;
		char *ptr=kfile+strlen(kfile)-1;
		while (ptr >kfile && *ptr!='/' && *ptr!='\\') ptr--;
		if (*ptr=='/' || *ptr=='\\') ptr++;

		switch (Global::Kmer_Len)
		{
		case 5:
			if (AveSeqLen < 100)
			{
				sprintf(ptr,"l50k5.krate");
			}
			else if (AveSeqLen <200)
			{
				sprintf(ptr,"l100k5.krate");
			}
			else if (AveSeqLen <400)
			{
				sprintf(ptr,"l200k5.krate");
			}
			else if (AveSeqLen <800)
			{
				sprintf(ptr,"l400k5.krate");
			}
			else
			{
				sprintf(ptr,"fulk5.krate");
			}
			printf("Used Kmer configure file: %s\n",kfile);
			break;
		case 6:
			if (AveSeqLen < 100)
			{
				sprintf(ptr,"l50k6.krate");
			}
			else if (AveSeqLen <200)
			{
				sprintf(ptr,"l100k6.krate");
			}
			else if (AveSeqLen <400)
			{
				sprintf(ptr,"l200k6.krate");
			}
			else if (AveSeqLen <800)
			{
				sprintf(ptr,"l400k6.krate");
			}
			else
			{
				sprintf(ptr,"fulk6.krate");
			}
			printf("Used Kmer configure file: %s\n",kfile);
			break;
		default:
			printf("Cannot find a K-mer configure file. Please specify it.\n");
			exit(0);
			break;
		}
	}
	else
	{
		printf("Used Kmer configure file: %s\n",kfile);
	}
	Global::Init(SeqNum);	
	Global::LoadKList(kfile);	
	BuildKmer();
	DoCluster(flist,fstat,fgroup,ftree);
	if (Global::ShowAlign)
		WriteAlign(alfile);
		
	Global::Finish();
	Clear();
	time_2=clock();
	printf("\r%10.3g secs in Statistical Analysis. \n",(REAL)(time_2-time_1)/CLOCKS_PER_SEC);
	free(SeqID);
	//free(mypath);
 	return 0;
}

