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

#include "FASTA.h"
#include "util.h"
#include "ProbModel.h"

int clnum[20];

#define MAX_SEQS 1000000

int SeqCount=10*MAX_SEQS;

char **SeqStrs;
char **LabelStrs;
char buf[100];
int groupid[10*MAX_SEQS];
int Freq[10*MAX_SEQS];

FASTA fasta;
char decoration[65535];
char frqout[65535];
char fcname[200];
char cmd[200];
double level;
char prefix[2000];
double leveln=0.0;
double bound=0.95;
bool do_trim=false;
int MaxLen=0;
FILE *fpout;

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");
			exit(1);
		}
		i=0;
	while (fscanf(fp,"%s %d",buf,&fq)!=EOF)
	{
		Freq[i++]=fq;
	}			
	fclose(fp);
}

static char RNAList[]="-AGCTN_";

void CorrectGroup(int groupidx,int gpptr)
{
	int i,j,k;
	int sumfrq=Freq[groupid[0]];
	int slen=strlen(SeqStrs[groupid[0]]);

	for (i=1;i<gpptr;i++)
	{
		sumfrq+=Freq[groupid[i]];
	}
	float codecount[7];
	char *pch;
	
	for (i=0;i<slen;i++)
	{	
		for (k=0;k<7;k++) codecount[k]=0;
		for (j=0;j<gpptr;j++)
		{
			pch=strchr(RNAList,SeqStrs[groupid[j]][i]);
			if (pch!=NULL)
				codecount[pch-RNAList] +=Freq[groupid[j]];
		}
		for (k=0;k<7;k++) codecount[k]/=sumfrq;
		int maxid=0;
		for (k=1;k<7;k++)
		{		
			if (codecount[k] > codecount[maxid])
				maxid=k;
		}
		if (codecount[maxid] >bound)
		{
			for (j=0;j<gpptr;j++)
			{
				pch=strchr(RNAList,SeqStrs[groupid[j]][i]);
				if (pch-RNAList !=maxid)
					SeqStrs[groupid[j]][i]=RNAList[maxid];
			}
			
		}
	}	
}

void Usage(){
printf("correction [-d consrate] <fastafile> <frqfile> <clusterfile> <level>\n");
printf("Making base corrections for aligned sequences.\n\n");
printf("\tconsrate:\t expected base rate to identify a conservative site. Default 0.95.\n"); 
printf("\tfastafile:\t trimmed FASTA sequences.\n");
printf("\tfrqfile:\t frequency file for trimmed FASTA sequences.\n");
printf("\tclusterfile:\t corresponding clustering result (.Cluster).\n");
printf("\tlevel:\t Reference level.\n");
}

int main(int argc, char **argv)
{
	FILE *fp,*fbuf;
	int i,numsq;
	char *label,*seq;
	int rlines=0;
	int bufptr,gpptr,groupidx;
	char gc;
	char c;
	
	while ((c = getopt(argc, argv, "thr:")) != -1)
    {
        switch (c)
        {
        case 'h':
        	Usage();
        	exit(0);
        	break;
        case 't':
        	do_trim=true;
        	break;
		case 'r':
			bound=atof(optarg);
			break;
        default:
        	fprintf(stderr,"Unknown option -%c\n",c);
			Usage();
			exit(0);
			break;
        }
    }
	argc -= optind;
	argv += optind;

	
	if (argc <4) 
	{
		Usage();
		exit(0);
	}
	char *seqfile=argv[0];
	char *frqfile=argv[1];
	char *clfile=argv[2];
	leveln=atof(argv[3]);

	SeqStrs=(char **)Malloc(SeqCount*sizeof(char *));
	LabelStrs=(char **)Malloc(SeqCount*sizeof(char *));

	fasta.OpenRead(seqfile);
	i=0;
	
	while (fasta.GetFastaSeq(&label,&seq,0,0))
	{
		LabelStrs[i]=label;
		SeqStrs[i++]=seq;
		if (strlen(seq) >MaxLen)
			MaxLen=strlen(seq);
	}
	fasta.Close();
	numsq=i;

	
	fp=fopen(clfile,"r");
	if (fp==NULL)
	{
		fprintf(stderr,"Cannot Open Cluster File\n");
		exit(0);
	}
	LoadFreq(frqfile);

	while(1)	
	{
		
		level=-1;
		fscanf(fp,"%lf",&level);
		if (level < 0) break;
		if (level >1)
		{
			fprintf(stderr,"Should not get here.\n");
		}
		
		if (fabs(level-leveln) <= 1e-4)
		{
			printf("Level=%.2f\n",level);
		}
		else
		{
			do {(gc=fgetc(fp));} while (gc!='\n');
			continue;
		}
		bufptr=0;
		gpptr=0;
		groupidx=0;
		
		while (1)
		{
			gc=fgetc(fp);
			if (gc=='\n' || gc=='\r')
			{
				if (bufptr>0)
				{
					buf[bufptr]='\0';
					groupid[gpptr++]=atoi(buf);
					bufptr=0;
				}

				if (gpptr >0)
				{
					CorrectGroup(groupidx,gpptr);
					groupidx++;
				}
				printf("Level %.2f Groups %d\n",level, groupidx);
				break;
			}					
			if (gc=='|')
			{
				if (bufptr>0)
				{
					buf[bufptr]='\0';
					groupid[gpptr++]=atoi(buf);
					bufptr=0;
				}

				if (gpptr >0)
				{
					CorrectGroup(groupidx,gpptr);
					groupidx++;
				}
				gpptr=0;			
			}
			else if (gc==' ')
			{
				if (bufptr>0)
				{
					buf[bufptr]='\0';
					groupid[gpptr++]=atoi(buf);
					bufptr=0;
				}
			}
			else
			{
				buf[bufptr++]=gc;
			}
		}			
	}
	fclose(fp);

	AddSuff(fcname,seqfile, ".cr");
	printf("Writing Solutions to %s\n",fcname);
	fasta.OpenWrite(fcname);
	for (i=0;i<numsq;i++)
	{
		fasta.WriteFastaSeq(LabelStrs[i],SeqStrs[i]);
	}
	fasta.Close();
	
	for (i=0;i<numsq;i++)
	{
		free(SeqStrs[i]);
		free(LabelStrs[i]);
	}
	free(SeqStrs);
	free(LabelStrs);
	return 0;
}
