/***************************************************************************************
ESPRIT: Estimating Species Richness Using Large Collections of 16S rRNA Shotgun Sequences  
Yijun Sun, Yunpeng Cai, Li Liu, Fahong Yu, Michael L. Farrell, William McKendree, William Farmerie
Nucleic Acids Research 2009

Programed by Yunpeng Cai & Yubo Cheng
Copyright Hold by University of Florida, Gainesville, FL32610

****************************************************************************************/
/********************************Filtering Sequences************************/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <time.h>
#include <math.h>

#include "FASTA.h"
#include "util.h"
#include "needle.h"
int SeqCount=10000000;

int **SeqTabs;
char **SeqStrs;//**OriginalSeqStrs;
char **LabelStrs;
int *SeqLens;
int *SeqFreq;
int *SeqIDs;
char *primer;
int mis_allowed=1;
double var_allowed=1.0;
int do_trim=0;
bool exact_match=false;

bool var_cut=true;
double len_min=0;
double len_max=65535;


#define PRIMER_LEN 1000
#define MAXPATH 2000


struct Sequence
{
	char *SeqStr;
	char *LabelStr;
	int ID;
};

int CheckPrimer(char *seq)
{
	char dup[2* PRIMER_LEN];
	char str1[2* PRIMER_LEN],str2[2* PRIMER_LEN];
	int i,tail,cntmiss;
	Needle needle(10,0.5);

	if (primer==NULL) return 1;
		
	strncpy(dup,seq,strlen(primer)+mis_allowed);
	needle.Align(primer,dup,str1, str2,1.0);
	tail=strlen(str1)-1;
	while (tail > 0 && str1[tail]=='-') tail--;
  cntmiss=0;
  for (i=0;i<=tail;i++)
  	if (str2[i]!=str1[i]) cntmiss++;
  if (cntmiss >mis_allowed) 
  {
  	return 0;
  } 
  else 
  {
  	return 1;
  }
}

void CheckLen(int &seqnum, char **SeqStrs, char **LabelStrs, int *SeqLens, int *SeqIDs)
{
	int i,j;
	double ave_len,std_len;
	int minlen,maxlen;
	char *swap;
	int nswap;
	
	ave_len=0;
	std_len=0;
	minlen=65535;
	maxlen=0;
	for (i=0;i<seqnum;i++)
	{
		ave_len+=SeqLens[i];
		minlen=Min(minlen,SeqLens[i]);
		maxlen=Max(maxlen,SeqLens[i]);
	}
	ave_len/=(seqnum+0.0);
	for (i=0;i<seqnum;i++)
		std_len+=Sqr(SeqLens[i]-ave_len);
	std_len/=(seqnum+0.0);
	std_len=sqrt(std_len);
	
	printf("Sequence Average Len %.1f Std %.1f Min %d Max %d\n",ave_len,std_len,minlen,maxlen);

	if (var_cut)
	{
		len_min=Max(len_min,(ave_len-var_allowed*std_len));
		len_max=Min(len_max,(ave_len+var_allowed*std_len));
	}
	#ifdef DEBUG
	 fprintf(stderr,"Av len %f Std %f\n",ave_len,std_len);
	#endif
	for (i=0;i<seqnum;i++)
	{
		if (SeqLens[i]< len_min || SeqLens[i] > len_max)
		{
			#ifdef DEBUG
				fprintf(stderr,"Len Discard %d %s\n",SeqLens[i],LabelStrs[i]);
			#endif
			SeqLens[i]=0;
		}
	}	
	
	i=0;
	while (i<seqnum && SeqLens[i]>0) i++;

	j=i+1;
	while (j<seqnum)
	{
		if (SeqLens[j]>0)
		{
					swap=SeqStrs[j];
					SeqStrs[j]=SeqStrs[i];
					SeqStrs[i]=swap;
					swap=LabelStrs[j];
					LabelStrs[j]=LabelStrs[i];
					LabelStrs[i]=swap;
					SeqLens[i]=SeqLens[j];
					SeqLens[j]=0;
					nswap=SeqIDs[j];
					SeqIDs[j]=SeqIDs[i];
					SeqIDs[i]=nswap;
					i++;
		}
		j++;
	}
	for (j=i;j<seqnum;j++)
	{
		free(LabelStrs[j]);
		free(SeqStrs[j]);
	}
	seqnum=i;
}

//int CompSeq(const void * a, const void * b)
//{
//	char *seq1,*seq2;
//	seq1=(char *)(*(char **)a);
//	seq2=(char *)(*(char **)b);
//	int i,len1,len2;
//	len1=strlen(seq1);
//	len2=strlen(seq2);
//	for (i=0;i<Min(len1,len2);i++)
//	{
//		if (seq1[i]>seq2[i]) return 1;
//		if (seq1[i]<seq2[i]) return -1;
//	}
//	return len1-len2;
//}

int CompSeq(const void * a, const void * b)
{
	struct Sequence *Seq1=(struct Sequence *)a;
	struct Sequence *Seq2=(struct Sequence *)b;
	char *seq1,*seq2;
	seq1=Seq1->SeqStr;
	seq2=Seq2->SeqStr;
	int i,len1,len2;
	len1=strlen(seq1);
	len2=strlen(seq2);
	for (i=0;i<Min(len1,len2);i++)
	{
		if (seq1[i]>seq2[i]) return 1;
		if (seq1[i]<seq2[i]) return -1;
	}
	return len1-len2;
}

int SubSeq(char *seq, char*subseq)
{
	int i,len1,len2;
	len1=strlen(seq);
	len2=strlen(subseq);
	if (len2 > len1) return 0;
	for (i=0;i<len2;i++)
	{
		if (seq[i]!=subseq[i]) return 0;
	}
	return 1;
}

void DoPurge(char *ifile,char *ofile,char *ffile,char *rfile, char *mfile)
{
	FASTA fasta;
	FASTA fastaf;
	char *seq;
	char *label;
	int seqlen;
	int idx=0;
	int i,j,numdel;
	FILE *fp;
	FILE *fmap;
		
	SeqStrs=(char **)Malloc(SeqCount*sizeof(char *));
//	OriginalSeqStrs=(char **)Malloc(SeqCount*sizeof(char *));
	LabelStrs=(char **)Malloc(SeqCount*sizeof(char *));
	SeqFreq=(int *)Malloc(SeqCount*sizeof(int));
	SeqLens=(int *)Malloc(SeqCount*sizeof(int));
	SeqIDs=(int *)Malloc(SeqCount*sizeof(int));
	

	int idx_org=0;	
	fasta.OpenRead(ifile);
	while ((seqlen=fasta.GetFastaSeq(&label,&seq,do_trim,1))>0)
	{
//		fprintf(stderr,"seq %d len %d %s\n",idx,seqlen,seq);	
		if (seq[0] !=0)
		{
			if (CheckPrimer(seq))
				{
					LabelStrs[idx]=label;
					//strcpy(LabelStrs[idx],label);
					SeqLens[idx]=seqlen;
					SeqIDs[idx]=idx_org;
					SeqStrs[idx++]=seq;
					//strcpy(SeqStrs[idx++],seq);
				}
			else
				{
					#ifdef DEBUG
						fprintf(stderr,"Primer Discard  %s\n",label);
					#endif
					free(seq);
					free(label);
				}
		}
		else{
			seq[0]='1';
			free(seq);
			free(label);		
		}		
		idx_org++;
	}
	fasta.Close();
//	#ifdef CONSOLE
	printf("%d Seqs Match Primer\n",idx);
	//#endif
	
	CheckLen(idx,SeqStrs,LabelStrs,SeqLens,SeqIDs);
	//#ifdef CONSOLE
	printf("%d Seqs Valid Len\n",idx);
	//#endif
//	printf("passed!%d;%d\n",idx,sizeof(struct Sequence));

	struct Sequence* Seq=new struct Sequence[idx];

	for (i=0;i<idx;i++)
	{
		Seq[i].SeqStr=SeqStrs[i];
		Seq[i].LabelStr=LabelStrs[i];
		Seq[i].ID=SeqIDs[i];
	}

	//OriginalSeqStrs=SeqStrs;
	// Merge Duplicated Reads
	qsort(Seq,idx,sizeof(struct Sequence),CompSeq);
	
	for (i=0;i<idx;i++) SeqFreq[i]=1;
	
	i=idx-1;j=i-1;
	while (i>0)
	{
		while (j >=0 && (((!exact_match) && SubSeq(Seq[i].SeqStr,Seq[j].SeqStr)) || strcmp(Seq[i].SeqStr,Seq[j].SeqStr)==0))
		{
			SeqFreq[i]++;
			SeqFreq[j--]=0;
		}
		i=j;j--;
	}
	
	numdel=0;
	fasta.OpenWrite(ofile);
	fastaf.OpenWrite(rfile);
	
	if ((fp=fopen(ffile,"w"))==NULL)
		{
			fprintf(stderr,"Create Frequency File Failed\n");
			exit(1);
		}
	if ((fmap=fopen(mfile,"w"))==NULL)
		{
			fprintf(stderr,"Create Mapping File Failed\n");
			exit(1);
		}
	
	for (i=0;i<idx;i++)
	{
		fastaf.WriteFastaSeq(Seq[i].LabelStr,Seq[i].SeqStr);
		fprintf(fmap,"%d",Seq[i].ID);
		if (SeqFreq[i]>0)
		{
			fasta.WriteFastaSeq(Seq[i].LabelStr,Seq[i].SeqStr);
			fprintf(fp,"%s %d\n",Seq[i].LabelStr,SeqFreq[i]);
			fprintf(fmap,"\n");
		}
		else
		{
			fprintf(fmap," ");
			numdel++;
		}
	}
	fasta.Close();
	fastaf.Close();
	fclose(fp);
	fclose(fmap);
	
	SeqCount=idx-numdel;
	
	for (i=0;i<idx;i++)
	{
		free(SeqStrs[i]);
		free(LabelStrs[i]);	
	}
	
	free(SeqIDs);
	free(SeqStrs);
	free(LabelStrs);
	free(SeqFreq);	
  free(SeqLens);
  //#ifdef CONSOLE
	printf("\n%d Seqs After Process\n",SeqCount);
	//#endif
}


void ReadPrimer(char *ifile)
{
	FASTA fasta;
	char *label;
	int len;
	
	fasta.OpenRead(ifile);
	if ((len=fasta.GetFastaSeq(&label,&primer,0))<=0)
		{
			fprintf(stderr,"Reading Primer File Failed\n");
			primer=NULL;
		}
	else
		free(label);
	fasta.Close();
}


void Usage(){
	printf("Preprocess : Removes low quality reads and merge duplicated sequences\n");
	printf("preproc [-p primer_file] [-e] [-m mis_allowed] [-a] [-v var_allowed] [-w]  [-l minlen] [-u maxlen] <input[.fas]> [<output>] [<freq_output>]\n");
	printf("\t primer_file:\t a FASTA file contains sequencing primers\n");
	printf("\t -w:\t\t remove seqeuences with ambiguous char\n");
	printf("\t -e:\t\t only merges exactly identical sequences.\n");
	printf("\t mis_allowed:\t maximum allowed mismatch with the primer sequence, default 1\n");
	printf("\t -a:\t\t turn off default length filtering.\n");
	printf("\t var_allowed:\t maximum allowed senquence length variance, default 1.0\n");
	printf("\t minlen:\t minimal allowed sequence length.\n");
	printf("\t maxlen:\t maximial allowed sequence length.\n");
	printf("\t output:\t processed sequence file, default <input>_Clean.fas\n");
	printf("\t freq_output:\t sequence frequency output, default <output>.frq\n");
}

int main(int argc, char **argv)
{
	char *ifile, *pfile,*ofile,*ffile,*mfile,*rfile;
	char c;
	clock_t time_1,time_2;
	
	ifile=(char *)Malloc(sizeof(char)*MAXPATH);
	ofile=(char *)Malloc(sizeof(char)*MAXPATH);
	pfile=(char *)Malloc(sizeof(char)*MAXPATH);
	ffile=(char *)Malloc(sizeof(char)*MAXPATH);
	rfile=(char *)Malloc(sizeof(char)*MAXPATH);
	mfile=(char *)Malloc(sizeof(char)*MAXPATH);
	

	primer=NULL;
	
  while ((c = getopt(argc, argv, "hweal:u:p:m:v:")) != -1)
    {
        switch (c)
        {
        case 'h':
        		Usage();
        		exit(0);
        		break;
        case 'a':
                // turn of length filtering
                var_cut=false;
        case 'w':
        		do_trim=1;
        		break;
        case 'p':
				strcpy(pfile,optarg);
				ReadPrimer(pfile);
        	break;
       	case 'm':
       		mis_allowed = atoi(optarg);
       		break;
       	case 'v':
       		var_allowed = atof(optarg);
			var_cut=true;
       		break;	
		case 'l':
			len_min=atoi(optarg);
			break;
		case 'u':
			len_max=atoi(optarg);
			break;
		case 'e':
				exact_match=true;
				break;
        default:
        	fprintf(stderr,"Unknown option -%c\n",c);
        	break;
        }
    }
    argc -= optind;
    argv += optind;

	if (argc==0)
	{
		Usage();
		exit(0);	
	}
		
	sprintf(ifile,"%s",argv[0]);
	if (argc<2)
		AddSuff(ofile,ifile,"_Clean");
	else
		strcpy(ofile,argv[1]);
	
	if (argc<3)
	{
		strcpy(ffile,ofile);
		RepExt(ffile,"frq");
	}
	else
		strcpy(ffile,argv[2]);
	AddSuff(rfile,ofile,"_All");
	strcpy(mfile,ffile);
	RepExt(mfile,"map");
	
	time_1=clock();
	DoPurge(ifile,ofile,ffile,rfile,mfile);
	time_2=clock();
	#ifdef CONSOLE
	printf("\r%10.3g secs in Purging Strings. \n",(double)(time_2-time_1)/CLOCKS_PER_SEC);
	#endif
	free(ifile);
	free(pfile);
	free(ofile);
	free(ffile);
	free(rfile);
	free(mfile);
	
	if (primer!=NULL)
		free(primer);
	return 0;		
}

