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

#include "global.h"
#include "tree.h"
#include "evolve.h"
#include "models.h"
#include "gamma.h"
#include "random.h"
#include "covarion.h"

int numTaxa, numBases, maxPartitions, numPartitions, fileFormat;
double covarionRate, ONfreq;
double CovRate[2];
int covarion;

char *onoff="-+";




/* prototypes */

char SetHiddenState();
void RandomHiddenStates(char *seq, int inNumSites);
void covarionSimulation(char *state, double *effectiveLength, double len);
void MutateHiddenStates(char *seq, double *effectiveLength, int inFromSite, int inNumSites, double len);
void EvolveHiddenStateNode(TNode *anc, TNode *des, int inFromSite, int inNumSites);

void WriteHiddenStatesNode(FILE *fv, TTree *tree, int *nodeNo, TNode *des);
void WriteAncestralEffLengthNode(FILE *fv, TTree *tree, int *nodeNo, TNode *des);


/* functions */


void CreateHiddenStates(TTree *tree, int inNumSites)
{
  /* Creates memory for the "hiddenStates" and the "effectiveLength0" pointers
     of all the nodes of the tree. 
  */

	TNode *P;
	
	P=tree->nodeList;
	while (P!=NULL) {
		if (P->hiddenStates != NULL)
			free(P->hiddenStates);
			
		P->hiddenStates=CAllocMem(inNumSites+1, "hiddenStates", "CreateHiddenStates", 0);

		if (P->effectiveLength0 != NULL)
		  free(P->effectiveLength0 );
		P->effectiveLength0=CAllocMem(sizeof(double)*(inNumSites+1),"hiddenLengths", "CreateHiddenStates", 0);
		P=P->next;
	}
}



char SetHiddenState()
{
/* creates a state on/off at random according to the probability ONfreq. */

	char j;
	double r;
	
	j='\1';   		/* ON */
	r=rndu();
	if (r > ONfreq) j='\0'; /* OFF */
	#ifdef DEBUG 
		print("Debugging: state created:",j);
	#endif
	return j;
}



void RandomHiddenStates(char *seq, int inNumSites)
{
/* creates a sequence of hidden states of length inNumSites,
   according to the probability ONfreq to be 'on'*/

	int i;
	char *P;
	
	P=seq;
	for (i=0; i<inNumSites; i++) {
		*P=SetHiddenState();
		P++;
	}
}

void covarionSimulation(char *state, double *effectiveLength, double len)
{
/* it simulates the process ON/OFF : times when the switches occur. It stores
   the amount of time spent on "ON". */

	char *P;
	double *L;
	double r;
	double T,Ton,Ttot;

	P=state;
	Ton=Ttot=0.0; 

	while ( Ttot <= len ) {
		r=rndu();
		T=-log(r)/CovRate[*P];
		Ttot +=T;
		if (*P) Ton += T;
		if (*P) *P='\0'; else (*P)++; /* changes the state *P from ON to OFF and vice-versa*/
	}
	
	if (*P) *P='\0'; else (*P)++;;
	if (*P) *effectiveLength=Ton-(Ttot-len);
	else *effectiveLength=Ton;

}


void MutateHiddenStates(char *states, double *effectiveLength, int inFromSite, int inNumSites, double len)
{
/* From site number inFromSite for inNumSites sites, it mutates the hidden states contained in
   seq (seq is changed). "len" is the length of evolution. The effective length = time spent on "ON"
   is stored in effectiveLength (each site has its own). */

	int i;
	double *L;
	short *S;
	char *P;
	
	P=states;
	L=effectiveLength;

	/* Here the hidden states ON /OFF are simulated for all the sites, including
	   those that will be invariable. It would speed up to skip those sites.
	   But: The variable "static short *invariable" that stores which site is
	   invariable and which is not, is a local variable of the "evolve" file.
	   Moreover, it is used for the first time (and initialized) with 
	   "EvolveSequence", that is called after the evolution of the hidden states.
	*/

	/*	if (invariableSites) {
		S = invariable + inFromSite;
		for (i=0; i<inNumSites; i++) {
			if (!(*S))
				covarionSimulation(P,L,len);
				P++;
				L++;
				S++;
		}
	} 
	else 
	*/
	{
		for (i=0; i<inNumSites; i++) {
			covarionSimulation(P,L,len);
			P++;
			L++;
		}
	}
}


void EvolveHiddenStateNode(TNode *anc, TNode *des, int inFromSite, int inNumSites)
{
/* recursive function. The hidden process evolves from node 'anc' to node 'des'. 
   Then the function calls itself twice to make the evolution from node 'des'
   to its 2 descendants (if it is not a tip).*/

	double len;
	
	len=des->length0;
		
	memcpy(des->hiddenStates, anc->hiddenStates, inNumSites);
	MutateHiddenStates(des->hiddenStates, des->effectiveLength0, inFromSite, inNumSites, len);
	
	if (des->tipNo==-1) {
		EvolveHiddenStateNode(des, des->branch1, inFromSite, inNumSites);
		EvolveHiddenStateNode(des, des->branch2, inFromSite, inNumSites);
	}
}


void EvolveHiddenStates(TTree *tree, int inFromSite, int inNumSites)
{
/* The evolution of the hidden process on/off is started from the root.*/


	RandomHiddenStates(tree->root->hiddenStates, inNumSites);

	CovRate[1]=covarionRate/(2*ONfreq);	/* rate of switch when the state is ON */
	CovRate[0]=covarionRate/(2*(1-ONfreq));		/* rate when the state id OFF. */
		
	EvolveHiddenStateNode(tree->root, tree->root->branch1, inFromSite, inNumSites);
	EvolveHiddenStateNode(tree->root, tree->root->branch2, inFromSite, inNumSites);
	if (!tree->rooted)
		EvolveHiddenStateNode(tree->root, tree->root->branch0, inFromSite, inNumSites);
}



void WriteHiddenStates(FILE *fv, TTree *tree)
{
	int i, j, n;
	char *P;

	fprintf(fv, "\n");
	n = numTaxa + 1;
	
	fprintf(fv, "%d\t", n);

       	if (covarion) {
       		P=tree->root->hiddenStates;
		for (j=0; j<numBases; j++) {
       			 fputc(onoff[*P], fv);
       			 P++;
       		}
       	}

	fputc('\n', fv);
	
	if (!tree->rooted)
		WriteHiddenStatesNode(fv, tree, &n, tree->root->branch0);
	WriteHiddenStatesNode(fv, tree, &n, tree->root->branch1);
	WriteHiddenStatesNode(fv, tree, &n, tree->root->branch2);
}


void WriteHiddenStatesNode(FILE *fv, TTree *tree, int *nodeNo, TNode *des)
{
	int j;
	char *P = des->hiddenStates;
	
	if (des->tipNo==-1) {
		(*nodeNo)++;
		fprintf(fv, "%d\t", *nodeNo);

		if (covarion) {
			for (j=0; j<numBases; j++) {
				 fputc(onoff[*P], fv);
				 P++;
			}
		}

		fputc('\n', fv);
		
		WriteHiddenStatesNode(fv, tree, nodeNo, des->branch1);
		WriteHiddenStatesNode(fv, tree, nodeNo, des->branch2);
	} else {
		fprintf(fv, "%s\t", tree->names[des->tipNo]);

		if (covarion) {
			for (j=0; j<numBases; j++) {
				 fputc(onoff[*P], fv);
				 P++;
			}
		}

		fputc('\n', fv);
	}
}





void WriteAncestralEffectiveLengths(FILE *fv, TTree *tree)
{
	int n;
	fprintf(fv, "\nThese effective lengths are re-scaled (by scale/(ONfreq*(1-Pinv))\n");
	fprintf(fv, "when they are used for the nucleotide evolution.\n");

	if (!tree->rooted)
		n = (2 * numTaxa) - 3;
	else
		n = (2 * numTaxa) - 2;
	fprintf(fv, " %d %d\n", n, numBases);

	n = numTaxa + 1;
	
	fprintf(fv, "%d\t", n);

	fputc('\n', fv);
	
	if (!tree->rooted)
		WriteAncestralEffLengthNode(fv, tree, &n, tree->root->branch0);
	WriteAncestralEffLengthNode(fv, tree, &n, tree->root->branch1);
	WriteAncestralEffLengthNode(fv, tree, &n, tree->root->branch2);
}

void WriteAncestralEffLengthNode(FILE *fv, TTree *tree, int *nodeNo, TNode *des)
{
	int j;
	double *P = des->effectiveLength0;
	
	if (des->tipNo==-1) {
		(*nodeNo)++;
		
		fprintf(fv, "%d\t", *nodeNo);
		for (j=0; j<numBases; j++) {
			fprintf(fv,"%G ",*P);
			P++;
		}

		fputc('\n', fv);
		
		WriteAncestralEffLengthNode(fv, tree, nodeNo, des->branch1);
		WriteAncestralEffLengthNode(fv, tree, nodeNo, des->branch2);
	} else {
		fprintf(fv, "%s\t", tree->names[des->tipNo]);
		for (j=0; j<numBases; j++) {
			fprintf(fv,"%G ",*P);
			P++;
		}

		fputc('\n', fv);
	}

}