#include "kmer.h"
#include "ProbTree.h"
#include "needle.h"
#include <iostream>
#include <fstream>
#include "global.h"
using namespace std;
bool logon=false;
REAL cutoff=0.1;

extern int *Freq;


Tree::Tree()
{
	Num_Seqs=0;
	Num_Children=0;
	thres =0.0;
	Child = NULL;
	Parent = NULL;
	TailChild=NULL;
	Brother =this;
}

Tree::Tree(ProbString & seq, int frq)
{
	Num_Seqs=frq;
	Num_Children=0;
	thres =0.0;
	Child = NULL;
	TailChild = NULL;
	Parent = NULL;
	Brother =this;
	TCenter=seq;
	SingleSeq=0;
}

Tree::Tree(unsigned int uid, int frq)
{
	UID=uid;
	Num_Seqs=frq;
	Num_Children=0;
	thres =0.0;
	Child = NULL;
	TailChild = NULL;
	Parent = NULL;
	Brother =this;
	SingleSeq=1;
}


Tree::~Tree()
{
	Tree *ptr,*ptr2;
	if (this->Child !=NULL)
	{
		ptr=this->Child;
		while (ptr->Brother !=ptr)
		{
			ptr2=ptr->Brother;
			ptr->Brother=ptr2->Brother;
			delete ptr2;
		}
		delete ptr;
	}
}

void Tree::AppendProbVect(ProbString &my_al,ProbString &seq_al,int numseq)
{
	REAL wt1=(Num_Seqs+0.0)/(Num_Seqs+numseq);
	REAL wt2=(numseq+0.0)/(Num_Seqs+numseq);
	my_al *=wt1;
	seq_al *=wt2;
	
/*	if (my_al[0].isGap()) 
	{
		int ptr=0;
		while (ptr < my_al.Len() && my_al[ptr].isGap())
		{
		  REAL norm=my_al[ptr].Norm();
		  my_al[ptr]=seq_al[ptr];
		  my_al[ptr] *=norm/seq_al[ptr].Norm();
		  ptr++;
		}
	}
	if (my_al[my_al.Len()-1].isGap()) 
	{
		int ptr=my_al.Len()-1;
		while (ptr>=0 && my_al[ptr].isGap())
		{
		  REAL norm=my_al[ptr].Norm();
		  my_al[ptr]=seq_al[ptr];
		  my_al[ptr] *=norm/seq_al[ptr].Norm();
		  ptr--;
		}
	}
	if (seq_al[0].isGap()) 
	{
		int ptr=0;
		while (ptr < seq_al.Len() && seq_al[ptr].isGap())
		{
		  REAL norm=seq_al[ptr].Norm();
		  seq_al[ptr]=my_al[ptr];
		  seq_al[ptr] *=norm/my_al[ptr].Norm();
		  ptr++;
		}
	}
	if (seq_al[seq_al.Len()-1].isGap()) 
	{
		int ptr=seq_al.Len()-1;
		while (ptr>=0 && seq_al[ptr].isGap())
		{
		  REAL norm=seq_al[ptr].Norm();
		  seq_al[ptr]=my_al[ptr];
		  seq_al[ptr] *=norm/my_al[ptr].Norm();
		  ptr--;
		}
	}*/
	
	my_al +=seq_al;
	TCenter=my_al;
}

void Tree::SubstractProbVect(ProbString &my_al,ProbString &seq_al,int numseq)
{
	if (Num_Seqs <=numseq) return;
	REAL wt1=(Num_Seqs+0.0)/(Num_Seqs-numseq);
	REAL wt2=-(numseq+0.0)/(Num_Seqs-numseq);
	my_al *=wt1;
	seq_al *=wt2;
	my_al +=seq_al;
	TCenter=my_al;
}

void Tree::AddChild(Tree *subtree,bool recount)
{
	if (subtree->Brother !=subtree)
	{
		cerr <<" Cannot insert multiple branch at the same time." <<endl;
	}
	if (Child==NULL)
	{
		Child=subtree;
		TailChild=subtree;
	}
	else
	{
		TailChild->Brother= subtree;
		subtree->Brother=this->Child;			
		TailChild=subtree;
		if (Child->Brother ==Child) cerr <<"FT here " <<endl;
	}
	
	subtree->Parent=this;
	if (recount)
		this->Num_Seqs+=subtree->NumSeqs();
	this->Num_Children++;
}

void Tree::DeleteChild(Tree *subtree)
{
		if (this->Child ==NULL)
		{
			cerr <<"Deleting an Empty Tree" <<endl;
			return;
		}
			
		Tree *ptr=this->Child;

		if (ptr->Brother==ptr)
		{
			if (ptr !=subtree)
			{
				cerr <<"Deleting Non-exist Child" <<endl;
			}
			else
			{
				this->Child=NULL;
				this->Num_Children=0;
				delete ptr;
			}
			return;
		}
		
		while (ptr->Brother != subtree)
		{
				ptr=ptr->Brother;
				if (ptr == this->Child)
				{
					cerr <<"Deleting Non-exist Child" <<endl;
					return;
				}	
		}
		ptr->Brother=subtree->Brother;
		this->Num_Children--;
		if (TailChild ==subtree)
				TailChild = ptr;		
		if (this->Child == subtree)
				this->Child =subtree->Brother;
		delete subtree;
}


REAL Tree::KmerDist(REAL *kcode,int slen)
{
	if (SingleSeq)
		return 1 - Global::kmer->KmerComp(KmerTabs[UID],kcode,KmerSeqs[UID],SeqLens[UID],slen);
	else
	{
		if (KmerCodes[UID-SeqNum]==NULL) BuildKmerCode();
		return 1- Global::kmer->KmerComp(kcode,KmerCodes[UID-SeqNum]);
	}
}
REAL Tree::KmerDist(int *kcode,int *kseqs,int slen)
{
	if (SingleSeq)
	{
		if (slen >= SeqLens[UID])
				return 1 - Global::kmer->KmerComp(KmerTabs[UID],kcode,KmerSeqs[UID],SeqLens[UID]);
		else
				return 1 - Global::kmer->KmerComp(kcode,KmerTabs[UID],kseqs,slen);
	}
	else
	{
	  if (KmerCodes[UID-SeqNum]==NULL) BuildKmerCode();
	  return 1- Global::kmer->KmerComp(kcode,KmerCodes[UID-SeqNum],kseqs,slen,TCenter.Len());
	}
}	

REAL Tree::KmerDist(Tree *node)
{
	if (node->SingleSeq)
		return this->KmerDist(KmerTabs[node->UID],KmerSeqs[node->UID],SeqLens[node->UID]);
	else
	{
		if (KmerCodes[node->UID-SeqNum]==NULL) node->BuildKmerCode();
		return this->KmerDist(KmerCodes[node->UID-SeqNum],node->Center().Len());
	}
}


REAL Tree::NeedleDist(Tree *node)
{
	REAL dist;
//	if (Global::cache->Get(this->UID,node->UID,dist))
//		return dist;
	//if (logon)
	//	cerr << "Ndist " << this->UID << " w " << node->UID <<endl;
	if (this->SingleSeq)
	{
		if (node->SingleSeq)
		{
			Global::needle->Align(SeqStrs[this->UID],SeqStrs[node->UID],Global::buf1,Global::buf2,Global::DiagRate);
			dist=calcDistance(Global::buf1,Global::buf2);
		}
		else
		{
			ProbString al;
			Global::needle->Align(SeqStrs[this->UID],node->Center(),Global::buf1,al,Global::DiagRate);
			if (Global::Dist_Style==DS_Ave)
				dist=al.AveDist(Global::buf1);
			else
				dist=al.ProbMaxDist(Global::buf1,cutoff/(this->NumSeqs()*node->NumSeqs()));
		}
	}
	else
	{
		if (node->SingleSeq)
		{
			ProbString al;
			Global::needle->Align(SeqStrs[node->UID],TCenter,Global::buf1,al,Global::DiagRate);
			if (Global::Dist_Style==DS_Ave)
				dist=al.AveDist(Global::buf1);
			else
				dist=al.ProbMaxDist(Global::buf1,cutoff/(this->NumSeqs()*node->NumSeqs()));
		}
		else
		{
			ProbString al1,al2;
			Global::needle->Align(TCenter,node->Center(),al1,al2,Global::DiagRate);
			if (Global::Dist_Style==DS_Ave)
				dist=al1.AveDist(al2);
			else
				dist=al1.ProbMaxDist(al2,cutoff/(this->NumSeqs()*node->NumSeqs()));
		}
	}
	Global::cache->Store(UID,node->UID,dist);
	return dist;	
}

void Tree::BuildKmerCode()
{
	if (UID < SeqNum) return;
	if (KmerCodes[UID-SeqNum] ==NULL)
		KmerCodes[UID-SeqNum]=Global::kmer->KmerCount(TCenter);  
}

REAL Tree::FindSpanChild(Tree *node,Tree *&subtree,int id1,int id2,int sq1,int sq2)
{
	
	if (this->Child ==NULL)
	{
		subtree=NULL;
		return 1.0;
	}

	Tree *ptrend=this->Child;
	Tree *ptr=this->Child;
	int clptr=0;
	do{
		if (clptr >=this->Num_Children)
			cerr << "Wrong Num Children" <<endl;
		Global::childlist[clptr++]=ptr;	
		ptr=ptr->Brother;
	}while (ptr !=ptrend);	
	
	REAL dist;
	ProbString al1,al2;
		
	for (int i=0;i<clptr;i++)
	{
		if (QuickDist(id1,id2,Global::childlist[i]->UID,node->UID,sq1,sq2, dist))
		{
			if (dist <this->FirstChild()->GetThres())
			{
				subtree=Global::childlist[i];
				return dist;
			}
		}
		else
		{
			dist=Global::childlist[i]->KmerDist(node);
			if (dist < KdistBound(Global::childlist[i]->GetThres()))
			{
				dist=Global::childlist[i]->NeedleDist(node);
				if (dist <this->FirstChild()->GetThres())
				{
					subtree=Global::childlist[i];
					return dist;
				}
			}
		}
	}
	subtree=NULL;
	return 1.0;
}

REAL Tree::FindSpanChild(unsigned int uid,Tree * &subtree)
{
	if (this->Child ==NULL)
	{
		subtree=NULL;
		return 1.0;
	}

	Tree *ptrend=this->Child;
	Tree *ptr=this->Child;
	int clptr=0;
	do{
		if (clptr >=this->Num_Children)
			cerr << "Wrong Num Children" <<endl;
		Global::childlist[clptr++]=ptr;	
		ptr=ptr->Brother;
	}while (ptr !=ptrend);	
	
	REAL dist;
	ProbString al2;
	 
	for (int i=0;i<clptr;i++)
		{
			if (Global::cache->Get(Global::childlist[i]->UID,uid,dist))
			{
				if (dist <this->FirstChild()->GetThres())
				{
					subtree=Global::childlist[i];
					return dist;
				}
			}
			else
			{
				dist=Global::childlist[i]->KmerDist(KmerTabs[uid],KmerSeqs[uid],SeqLens[uid]);
				if (dist < KdistBound(Global::childlist[i]->GetThres()) )
				{
					if (Global::childlist[i]->SingleSeq)
					{	
						Global::needle->Align(SeqStrs[Global::childlist[i]->UID],SeqStrs[uid],Global::buf1,Global::buf2,this->FirstChild()->GetThres());
						dist=calcDistance(Global::buf1,Global::buf2);
					}
					else
					{
						Global::needle->Align(SeqStrs[uid],Global::childlist[i]->Center(),Global::buf1,al2,this->FirstChild()->GetThres());
						if (Global::Dist_Style==DS_Ave)
							dist=al2.AveDist(Global::buf1);
						else
				  		dist=al2.ProbMaxDist(Global::buf1,cutoff/(Freq[uid]*Global::childlist[i]->NumSeqs()));		
					}
					Global::cache->Store(Global::childlist[i]->UID,uid,dist);

					if (dist <this->FirstChild()->GetThres())
					{
						subtree=Global::childlist[i];
						return dist;
					}
				}
			}
		}
	subtree=NULL;
	return 1.0;
}

void Tree::Print( fstream & out)
{
	
	out << "Tree Node PTR " << hex <<(long )this <<endl;
	out << "UID : ";
	out << dec << UID;
	out << " Thres :" <<  thres << " Num Seqs :" <<dec << Num_Seqs <<" Num Children" << Num_Children << endl;
	out << " Parent :" << hex << (long) Parent <<endl;
	out << " First Child " << hex << (long)Child <<endl;
	out << " Next Brother " << hex << (long)Brother <<endl;
	
	if (this->Child !=NULL)
	{
		Tree *ptrend=this->Child;
		Tree *ptr=this->Child;
	
		do{
			ptr->Print(out);
			ptr=ptr->Brother;
		}while (ptr !=ptrend);
	}
}


int Tree::Brief( fstream & out,int myid, int childid)
{
	if (this->Child !=NULL)
	{
		Tree *ptrend=this->Child;
		Tree *ptr=this->Child;
		int nextchild=childid+Num_Children;
		do{
			out << "0\t" << myid << '\t' << UID << "(" << thres << ",+" << Num_Seqs << "," << Num_Children << ")\t"<< childid << endl;
			nextchild=ptr->Brief(out,childid,nextchild);
			ptr=ptr->Brother;
			childid++;
		}while (ptr !=ptrend);
		return nextchild;
	}
	else{
		out << "1\t" << myid << '\t' << UID <<",+" <<Num_Seqs <<endl;
		return childid;
	}
}

void Tree::ListChildren(vector<Tree *> & tvec)
{
	if (this->Child ==NULL)
		return;
	Tree *ptr=this->Child;
	do{
		tvec.push_back(ptr);
		ptr=ptr->Brother;
	}while (ptr !=this->Child);
}

void Tree::ListChildrenAt(REAL thres, vector<Tree *> & tvec)
{
	if (this->Child ==NULL)
	{
		tvec.push_back(this);
		return;
	}
	Tree *ptrend=this->Child;
	Tree *ptr=this->Child;
	do{
			if (ptr->GetThres() <=thres+MIN_THRES/3) 
			{
				tvec.push_back(ptr);
			}
			else
			{
				ptr->ListChildrenAt(thres,tvec);
			}
			ptr=ptr->Brother;
	}while (ptr !=ptrend);
}

void Tree::ListLeaf(vector<Tree *> & tvec)
{
	this->ListChildrenAt(0.0,tvec);
}

void Tree::PrintAllSeqs( fstream & out)
{
	if (this->Child==NULL)
	{
		char *buf;
		out << '>' << UID <<endl;
		buf=TCenter.ToString();
		out << buf << endl;
		free(buf);
		return;
	}
	Tree *ptrend=this->Child;
	Tree *ptr=this->Child;
	do{
			ptr->PrintAllSeqs(out);
			ptr=ptr->Brother;
	}while (ptr !=ptrend);
}

void Tree::CountUID(int *counter)
{
	if (UID >=0) counter[UID]=1;
	if (this->Child !=NULL)
	{
		Tree *ptrend=this->Child;
		Tree *ptr=this->Child;
		do{
			ptr->CountUID(counter);
			ptr=ptr->Brother;
		}while (ptr !=ptrend);
	}
}


void Tree::Collapse(unsigned int uid)
// Assume that all Children are single sequences
{
	if (Num_Children<=1) return;
	Tree *ptr=this->Child;
	
	TCenter.FromString(SeqStrs[ptr->UID]);
	int totalseq=ptr->NumSeqs();
	REAL wt1,wt2;


	
	REAL *ktab1, *ktab2;
	ktab1=Global::kmer->KmerCount(TCenter);
	ktab2=NULL;
	
	while (ptr != TailChild){
		ptr=ptr->Brother;
		ktab2=Global::kmer->KmerAdd(KmerTabs[ptr->UID],ktab1,ptr->NumSeqs(),totalseq);
		Global::kmer->FreeCodeTable(ktab1);
		ktab1=ktab2;
		
		wt1=(totalseq+0.0)/(totalseq+ptr->NumSeqs());
		wt2=(ptr->NumSeqs()+0.0)/(totalseq+ptr->NumSeqs());
		
		ProbString al1,al2;
		Global::needle->Align(SeqStrs[ptr->UID],TCenter,Global::buf1,al1,Global::DiagRate);
		al1*=wt1;
		al2.FromString(Global::buf1);
		al2*=wt2;
	    al1+=al2;
		TCenter=al1;
		totalseq+=ptr->NumSeqs();	
	}
	//cerr <<"merged " <<endl;

	if (totalseq !=Num_Seqs) cerr <<"Warning Num of Seqs Not Match" <<endl;
	KmerCodes[uid-SeqNum]=ktab1;
	
	SingleSeq=0;
	this->UID=uid;
	
	Tree *ptr2;
	
	ptr=this->Child;
	while (ptr->Brother !=ptr)
	{
		ptr2=ptr->Brother;
		ptr->Brother=ptr2->Brother;
		delete ptr2;
	}
	delete ptr;
	Child=NULL;
	Num_Children=0;
	
	Tree *newnode=new Tree(TCenter,Num_Seqs);
	newnode->UID=uid;
	this->AddChild(newnode,false);
}

/*
void Tree::Expand(REAL minlevel)
{
	if (Num_Children <=1) return;
	
	Tree *lhead=Child;
	Tree *ltail=TailChild;
	Tree *ptr;
	Tree *bestchild;
	REAL nrdist;
	
	Tree *newnode=new Tree;
	newnode->AddChild(Child);
	newnode->SetThres(thres/Global::level_inc);
	
	if (Child->SingleSeq)
	{
		newnode->SingleSeq=1;
		newnode->UID=Child->UID;
	}
	else
	{
		newnode->Center()=Child->Center();
		newnode->SingleSeq=0;
		newnode->AssignKmerCode(Child->GetKmerCode());
	}
	lhead=Child->Brother;
	ltail->Brother=lhead;
	Child->Brother=Child;
	Child=newnode;
	TailChild=newnode;
	Num_Children=1;

	ptr=lhead;
	Tree *last=ltail;
	//First Round, Find Major Branches
	do{
		if (!ptr->IsLeaf()) cerr << "Warning Expanding None Leaf" <<endl;
		
		if (Num_Children==1)
		{
			if (Child->KmerDist(ptr) < KdistBound(Child->GetThres()))
			{
				REAL dist;
				if (!Global::cache->Get(Child->UID,ptr->UID,dist))
					 dist=Child->NeedleDist(ptr);
				if (dist < Child->GetThres())
				{
					last=ptr;
					ptr=ptr->Brother;
					continue;
				}
			}	
			if (ptr==lhead)	lhead=lhead->Brother;
			last->Brother=ptr->Brother;
			ptr->Brother=ptr;
			
			newnode=new Tree;
			newnode->AddChild(ptr);
			newnode->SetThres(thres/Global::level_inc);
			if (ptr->SingleSeq)
			{
				newnode->SingleSeq=1;
				newnode->UID=ptr->UID;
			}
			else
			{
				newnode->Center()=ptr->Center();
				newnode->SingleSeq=0;
				newnode->AssignKmerCode(ptr->GetKmerCode());
			}
			AddChild(newnode,false);
		}
		else
		{ 
			if (bestchild !=NULL)
			{
				last=ptr;
				ptr=ptr->Brother;
				continue;
			}
			last->Brother=ptr->Brother;
			ptr->Brother=ptr;
			
			newnode=new Tree;
			newnode->AddChild(ptr);
			newnode->SetThres(thres/Global::level_inc);
			if (ptr->SingleSeq)
			{
				newnode->SingleSeq=1;
				newnode->UID=ptr->UID;
			}
			else
			{
				newnode->Center()=ptr->Center();
				newnode->SingleSeq=0;
				newnode->AssignKmerCode(ptr->GetKmerCode());
			}
			AddChild(newnode,false);
		}
	}while (ptr !=lhead);
	
	//Second Round, Assign Other Leafs to Branches
	if (Num_Children==1)
	{
		last=lhead;
		do{
			ptr=last;
			last=last->Brother;
			ptr->Brother=ptr;
			Child->AddChild(ptr);
		}while(ptr !=ltail);
	}
	else
	{
		last=lhead;
		do{
			ptr=last;
			last=last->Brother;
			ptr->Brother=ptr;
			if (bestchild == NULL || nrdist > Child->GetThres()*1.001) cerr <<"FT No Correct Branch Found" <<endl;
			bestchild->AddChild(ptr);	
		}while (ptr !=ltail);
	}
}
*/
