#include <string.h>
#include <stdio.h>
#include <ctype.h>
#include <math.h>
#include "kmer.h"
#include "util.h"
#include <algorithm>
#include <iostream>
#include <fstream>
using namespace std;
#define MINPROB 1e-10

Kmer::Kmer(int klen, int clen)
	{
		CodeLen=clen;
		KmerLen=klen;
		TableLen=(int ) pow(CodeLen,KmerLen);
		debug=0;
	}

int Kmer::GetKmerLen(){return TableLen;}

Kmer::~Kmer()
{
}
//need to override if the sequence is not nuclocide
int Kmer::KmerCode(char c) 
{
	switch (toupper(c))
	{
		case 'A':
				return 0;
		case 'G':
				return 1;
		case 'C':
			  return 2;
		case 'T':
			  return 3;
		default:
				return -1;
	}
}

int Kmer::KmerIndex(char *s)
{
	int idx=0;
	int code;
	if (strlen(s)<KmerLen)
		return -(KmerLen+1);
	for (int i=0;i<KmerLen;i++)
	{
	  code=KmerCode(s[i]);
	  if (code ==-1)
	  	 return -(i+1);
	  else
	  	{	 
	  		idx=idx*CodeLen+code; 
	  	}
	 } 	
	return idx;	
}

int Kmer::KmerIndexNext(int lastidx, char *s)
{
	int code;
	if (strlen(s)<KmerLen)
		return -(KmerLen+1);
	code = KmerCode(*(s+KmerLen-1));
	if (code ==-1)
	 	  return -KmerLen;
	 else
	 		return (lastidx*CodeLen % TableLen) + code;	
}

int* Kmer::AllocCodeTable()
{
	int *ptr;
	ptr= (int *)Malloc(TableLen*sizeof(int));
	memset(ptr,0,TableLen*sizeof(int));
	return ptr;
}
	
REAL* Kmer::AllocCodeTableP()
{
	REAL *ptr;
	ptr= (REAL *)Malloc(TableLen*sizeof(REAL));
	for (int i=0;i<TableLen;i++) ptr[i]=0.0;
	return ptr;
}


void Kmer::KmerCount(char *seq, int *tab, int *kseq)
{
	  int idx;
	  int slen=strlen(seq)-KmerLen+1;
		int *oldkseq=kseq;
	  idx=KmerIndex(seq);
	  char *seq_end=seq+slen;
	  
	  if (idx <0)
	  	{
	  		seq-=idx;
	  	}
	  else
	  	{	
	  		tab[idx]++;
	  		if (kseq!=NULL) *(kseq++)=idx;
	  		seq++;
	  	}
	  while (seq<seq_end)
	  {
	  	if (idx <0)
	  	{
	  		idx=KmerIndex(seq);
	  	}
	  	else
	  	{	
		  	idx=KmerIndexNext(idx,seq);
		  }
			if (idx <0)
		  {
	  		seq-=idx;
	  	}
	  	else
	  	{
		  	tab[idx]++;
	  		if (kseq!=NULL) *(kseq++)=idx;
				seq++;
			}
	  }
	  if (kseq!=NULL)
	  	{
	  	while (kseq <oldkseq+slen)
	  		*(kseq++)=-1;
	  	}
}

long int Kcomp=0;	
REAL Kmer::KmerComp(int *tab1, int *tab2, int *kseq1,int seqlen)
/* assume that seq1 is shorter than seq2 */
{
	int ucomm=0;
	int i;
	int idx;
	
	Kcomp++;
	if (kseq1==NULL)
	{
		for (i=0;i<TableLen;i++)
			ucomm+=min(tab1[i],tab2[i]);
	}	
	else
	{
		for (i=0;i<seqlen-KmerLen+1;i++)
		{
			idx=kseq1[i];
			if (idx >=0 && idx <TableLen)
				if (tab1[idx] >0 && tab1[idx] <=seqlen)
				{
					ucomm+=min(tab1[idx],tab2[idx]);
					tab1[idx]+=seqlen;
				}		
		}
		for (i=0;i<seqlen-KmerLen+1;i++)
		{
			idx=kseq1[i];
			if (idx >=0 && idx <TableLen)
				if (tab1[idx] > seqlen)
				{
					tab1[idx]-=seqlen;
				}		
		}
	}
	return (REAL) ucomm/(seqlen-KmerLen+1.0);
}

void	Kmer::CodeNorm(ProbChar &pc,REAL *vect)
{
	int i;
	REAL sum=0;
	
	for (i=1;i<=Num_RNAType;i++)
	{
		vect[i-1]=pc[i];
		sum+=pc[i];
	} 
	if (sum >MINPROB)
	{
		for (i=0;i<Num_RNAType;i++)
			vect[i]/=sum;
	}
}

void Kmer::CodeFirst(ProbString &seq,REAL *code,int *nonz)
{
	int i,j;
	REAL **codev=(REAL **)Malloc(KmerLen*sizeof(REAL *));
	int pos[KmerLen];
	int ptr;
	int nzptr=0;
	int idxnz;
	REAL prob;

	for (i=0;i<KmerLen;i++)
	{
		codev[i]=(REAL *)Malloc(CodeLen*sizeof(REAL));
		CodeNorm(seq[i],codev[i]);
		pos[i]=0;
	}	
 	ptr=0;
	
	while (ptr>=0)
	{
		while (pos[ptr]< CodeLen && codev[ptr][pos[ptr]] <= MINPROB)
			pos[ptr]++;
		if (pos[ptr]>=CodeLen)
		{
			pos[ptr]=0;
			ptr--;
			if (ptr>=0) pos[ptr]++;
		}
		else
		{
			ptr++;
			if (ptr >=KmerLen)
			{
				prob=1;
				idxnz=0;
				for (j=0;j<KmerLen;j++)
				{
					idxnz=idxnz*CodeLen+pos[j];
					prob*=codev[j][pos[j]];
				}
				if (idxnz>=TableLen)
						cerr <<"Overflow in Kmer Index"<<endl;
				code[idxnz]=prob;
				nonz[nzptr++]=idxnz;
				ptr--;
				pos[ptr]++;
			}
		}
	}
	nonz[nzptr]=-1;
	for (i=0;i<KmerLen;i++) free(codev[i]);
	free(codev);
}

void Kmer::ShiftCode(REAL *code1,int *nonz1,REAL *probv,	REAL *code2,int *nonz2)
{
	int i,j;
	int ptr=0;
	
	int relen=(int) (TableLen/CodeLen);
	REAL *tmp=(REAL *)Malloc(relen*sizeof(REAL));
	for (i=0;i<relen;i++)
		tmp[i]=0.0;
	while (nonz1[ptr]>=0)
	{
		tmp[nonz1[ptr] % relen]+=code1[nonz1[ptr]];
		ptr++;
	}
	ptr=0;
	for (i=0;i<relen;i++)
	{
		if (tmp[i] >MINPROB)
		{
			for (j=0;j<CodeLen;j++)
			{
				if (probv[j]>MINPROB)
				{
					nonz2[ptr]=i*CodeLen+j;
					if (nonz2[ptr]>=TableLen)
						cerr <<"Overflow Kmer Index"<<endl;
					code2[nonz2[ptr]]=tmp[i]*probv[j];
					ptr++;
				}
			} 
		}
	}
	nonz2[ptr]=-1;
	free(tmp);
}
void Kmer::AddCode(REAL *res,REAL *code,int *nonz)
{
	int ptr=0;
	while (nonz[ptr]>=0)
	{
		res[nonz[ptr]]+=code[nonz[ptr]];
		ptr++;
	}
}

REAL *Kmer::KmerCount(ProbString & seq)
{

	REAL *codeswp;
	int *nzswp;
	REAL *res=AllocCodeTableP();
	REAL *code1=AllocCodeTableP();
	REAL *code2=AllocCodeTableP();
	int *nonz1=(int *)Malloc((TableLen+1)*sizeof(int));
	int *nonz2=(int *)Malloc((TableLen+1)*sizeof(int));

	CodeFirst(seq,code1,nonz1);
	int idx=KmerLen;
	REAL *probv= new REAL[CodeLen];
	
	AddCode(res,code1,nonz1);
	while (idx < seq.Len())
	{
		CodeNorm(seq[idx],probv);
		ShiftCode(code1,nonz1,probv,code2,nonz2);
		AddCode(res,code2,nonz2);
		codeswp=code1;
		code1=code2;
		code2=codeswp;
		nzswp=nonz1;
		nonz1=nonz2;
		nonz2=nzswp;
		idx++;
	}

	free(nonz1);
	free(nonz2);
	delete [] probv;

	free(code1);
	free(code2);
	return res;
}  

long int Kcomp2=0;
REAL Kmer::KmerComp(REAL *tab1, REAL *tab2)
{
	REAL ucomm=0;
	REAL tot1=0;
	REAL tot2=0;
	Kcomp2++;
		for (int i=0;i<TableLen;i++)
		{
			tot1+=tab1[i];
			tot2+=tab2[i];
			ucomm+=min(tab1[i],tab2[i]);
		}	
	if (tot1 <=MINPROB || tot2 <=MINPROB)
		return 0;
	return ucomm/min(tot1,tot2);
}


REAL Kmer::KmerComp(int *tab1, REAL *tab2, int *kseq1,int seqlen,int seqlen2)
{
	REAL ucomm=0;
	int i;
	int idx;
	
	Kcomp++;
	if (kseq1==NULL)
	{
		for (i=0;i<TableLen;i++)
			ucomm+=min((REAL)tab1[i],tab2[i]);
	}	
	else
	{
		for (i=0;i<seqlen-KmerLen+1;i++)
		{
			idx=kseq1[i];
			if (idx >=0 && idx <TableLen)
				if (tab1[idx] >0 && tab1[idx] <=seqlen)
				{
					ucomm+=min((REAL)tab1[idx],tab2[idx]);
					tab1[idx]+=seqlen;
				}		
		}
		for (i=0;i<seqlen-KmerLen+1;i++)
		{
			idx=kseq1[i];
			if (idx >=0 && idx <TableLen)
				if (tab1[idx] > seqlen)
				{
					tab1[idx]-=seqlen;
				}		
		}
	}
	return (REAL) ucomm/(min(seqlen,seqlen2)-KmerLen+1.0);
}


REAL Kmer::NormalizedKmerComp(REAL *tab1, REAL *tab2)
{
	REAL ucomm=0;
		for (int i=0;i<TableLen;i++)
		{
			ucomm+=min(tab1[i],tab2[i]);
		}	
	return ucomm;
}

REAL *Kmer::KmerAdd(REAL *tab1,REAL *tab2,int frq1,int frq2)
{
	REAL *tab=AllocCodeTableP();
	for (int i=0;i<TableLen;i++)
	{
		tab[i]=(tab1[i]*frq1+tab2[i]*frq2)/(frq1+frq2);
	}
	return tab;
}

REAL *Kmer::KmerAdd(int *tab1,REAL *tab2,int frq1,int frq2)
{
	REAL *tab=AllocCodeTableP();
	for (int i=0;i<TableLen;i++)
	{
		tab[i]=(tab1[i]*frq1+tab2[i]*frq2)/(frq1+frq2);
	}
	return tab;
}

REAL *Kmer::KmerAdd(int *tab1,int *tab2,int frq1,int frq2)
{
	REAL *tab=AllocCodeTableP();
	for (int i=0;i<TableLen;i++)
	{
		tab[i]=(tab1[i]*frq1+tab2[i]*frq2+0.0)/(frq1+frq2+0.0);
	}
	return tab;
}

REAL *Kmer::DupCode(REAL *code)
{
	REAL *res=AllocCodeTableP();
	memcpy(res,code,TableLen*sizeof(REAL));
	return res;
}



