#include <math.h>
#include "pseud.h"

/*global variables*/
int   nR=4;
double  Cijk[64], Root[4];

static int
RootTN93 (int model, double *kappa1, double *kappa2, double *pi, double *f, double Root[]);
static int
EngenTN93(int model, double *kappa, double *pi, int *nR, double Root[], double Cijk[]);
static int
PMatCijk(double *PMat, double t);
static void
DisplayPMat(double *PMat, FILE *fout);

char  *models[]={"JC69","K80","F81","F84","HKY85","TN93","REV","UNREST"};
enum {JC69,K80,F81,F84,HKY85,TN93,REV,UNREST} MODELS;

static int
RootTN93 (int model, double *kappa1, double *kappa2, double *pi, double *f, double Root[])
{
   double T=pi[0],C=pi[1],A=pi[2],G=pi[3],Y=T+C,R=A+G;

   if (model==F84) { *kappa2=1+*kappa1/R; *kappa1=1+*kappa1/Y; }
   else if (model==HKY85) *kappa2=*kappa1;

   *f=1/(2*T*C*(*kappa1)+2*A*G*(*kappa2) + 2*Y*R);

   Root[0] = 0;
   Root[1] = - (*f);
   Root[2] = -(Y+R*(*kappa2)) * (*f);
   Root[3] = -(Y*(*kappa1)+R) * (*f);
   return (0);
}

static int
EngenTN93(int model, double *kappa, double *pi, int *nR, double Root[], double Cijk[])
{
   int i,j,k, nr;
   double kappa1,kappa2;
   double f, U[16],V[16], t;
   double T=pi[0],C=pi[1],A=pi[2],G=pi[3],Y=T+C,R=A+G;

   if (model==JC69 || model==F81) kappa1=kappa2=*kappa=1;
   else if (model<=HKY85)     kappa2=kappa1=*kappa;
   else if (model==TN93)     kappa2=kappa1=*kappa;
   RootTN93 (model, &kappa1, &kappa2, pi, &f, Root);

   *nR=nr = 2+(model==K80||model>=F84)+(model>=HKY85);
   U[0*4+0]=U[1*4+0]=U[2*4+0]=U[3*4+0]=1;
   U[0*4+1]=U[1*4+1]=1/Y;   U[2*4+1]=U[3*4+1]=-1/R;
   U[0*4+2]=U[1*4+2]=0;  U[2*4+2]=G/R;  U[3*4+2]=-A/R;
   U[2*4+3]=U[3*4+3]=0;  U[0*4+3]=C/Y;  U[1*4+3]=-T/Y;

   xtoy (pi, V, 4);
   V[1*4+0]=R*T;   V[1*4+1]=R*C;
   V[1*4+2]=-Y*A;  V[1*4+3]=-Y*G;
   V[2*4+0]=V[2*4+1]=0;  V[2*4+2]=1;   V[2*4+3]=-1;
   V[3*4+0]=1;  V[3*4+1]=-1;   V[3*4+2]=V[3*4+3]=0;

   for(i=0;i<4;i++) for(j=0;j<4;j++) {
      Cijk[i*4*nr+j*nr+0]=U[i*4+0]*V[0*4+j];
      switch (model) {
      case JC69:
      case F81:
         for (k=1,t=0; k<4; k++) t += U[i*4+k]*V[k*4+j];
         Cijk[i*4*nr+j*nr+1] = t;
         break;
      case K80:
      case F84:
         Cijk[i*4*nr+j*nr+1]=U[i*4+1]*V[1*4+j];
         for (k=2,t=0; k<4; k++) t += U[i*4+k]*V[k*4+j];
         Cijk[i*4*nr+j*nr+2]=t;
         break;
      case HKY85:   case TN93:
         for (k=1; k<4; k++)  Cijk[i*4*nr+j*nr+k] = U[i*4+k]*V[k*4+j];
         break;
      default:
         printf("\nerror: err:model\n"); exit(-1);
      }
   }
   return(0);


}



/*P(t)ij=SUM Cijk*exp{Root*t}*/
static int PMatCijk(double *PMat, double t)
{
	int  i,j,k, n=4, nr=nR;
   double expt[4];

   if(t<-1e-3)  printf("\nt=%.5f in PMatCijk",t);
   if(t<1e-7){
   	for(i=0;i<n;i++)  for(j=0;j<n;j++) { PMat[i*n+j]=0; PMat[i*n+i]=1; }
		return (0);
   }

   for(k=0;k<nr;k++)expt[k]=exp(t*Root[k]);

   for(i=0;i<n;i++){
   	for(j=0;j<n;j++){
      	for(k=1,t=Cijk[i*n*nr+j*nr+0];k<nr;k++){
         	t+= Cijk[i*n*nr+j*nr+k]*expt[k];
         }
         PMat[i*n+j]=t;
      }
   }
   return (0);
}


static void DisplayPMat(double *PMat, FILE *fout)
{
   int  i,j;

   fprintf(fout,"PMat=(");
	for(i=0;i<4;i++){
   	for(j=0;j<4;j++){fprintf(fout,"%1.3e, ",PMat[i*4+j]);}
      fprintf(fout,"\n");
   }
   fprintf(fout,")\n");


   return;
}


void
CalculatePMat(double *PMat, OptionsBlkPtr options, FILE *fout)
{
   /*int    i; */
   double t;
   /*double *pi; */

   t = options->branch_length;
   fprintf(fout,"t=%1.3e,\n",t);
	EngenTN93(options->model, &(options->kappa), options->pi, &nR, Root,Cijk);
	PMatCijk(PMat, t);

   /*fprintf(fout,"kappa=%1.2e, model=%d, t=%1.2e\n",options->kappa, options->model,t);
   pi = options->pi;
   for(i=0;i<4;i++) fprintf(fout,"pi[%d]=%1.3f,",i,pi[i]);
   fprintf(fout,"\n");
	DisplayPMat(PMat, fout);  */

	return;
}