// Copyright 2008 Michiaki Hamada
// Adapted from public domain code by Yi-Kuo Yu, NCBI

/**
 * See lambda_calculator.h
 */
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "nrutil.h"
int Alphsize;
#include "nrutil.c"
#include "ludcmp.c"
#include "lubksb.c"
#include "lambda_calculator.h"
#define Epsilon 1.0e-36
#define E_bound 1.0e-12 
#define Infty   1000000.0
#define min(A,B) ((A) > (B) ? (B) : (A) )
#define max(A,B) ((A) > (B) ? (A) : (B) )
#define bool int
#define true 1
#define false 0
double Lambda_UB; //lambda upper bound
double r_max_m, c_max_m; //min of each row's (column's) max

void makematrix( const double** mat_b, double **a, double lambda);
typedef struct Lambda {
  double min;
  double max;
  int flag;    // 1 means there is a range, -1 means no solution possible.
} Lambda;
typedef struct Sum {
  double value;
  int flag;   // 1 means no negative bg_freq, -1 means there is negative bg_freq
} Sum;
Lambda Find_JP(const double** mat_b, double la_min, double la_max, double **JP, double *p_in, double *q_in);
Sum Check_root(const double** mat_b, double **a, double lambda, double *p, double *q);
double Check_det(const double** mat_b, double **a, double lambda);
Sum Nail_lambda(const double** mat_b, int flag_sign, double lambda_min, double lambda_max, double *p, double *q, double *la_add);
double Nail_det(const double** mat_b, int flag_sign, double lambda_min, double lambda_max);
bool Check_range( const double** mat_b );

double *Locate_det_zero( const double** mat_b, int * ); //pointer to root list are returned with how many of them by int 


double calculate_lambda ( const double** mat_b, int alpha_size,
                          double* p, double* q )
{
  double **JP/*, *q, *p*/;
  int k;
  double *root_location;
  int N_root; 
  Lambda Lambda_local;

  Alphsize = alpha_size;

  if( ! Check_range( mat_b ) ) return -1.0;

  root_location = Locate_det_zero(mat_b, &N_root);
  if( root_location == NULL && N_root > 0 ) return -1.0;

  //q=dvector(1,Alphsize);
  //p=dvector(1,Alphsize);
  JP = dmatrix(1,Alphsize,1,Alphsize);

  if (N_root == 0){
    Lambda_local = Find_JP(mat_b, 0, Lambda_UB, JP, p, q);
    if (1 == Lambda_local.flag) { // sensible solution found  
      // Remember to find the right place to free the vectors
      //free_dvector(p, 1,Alphsize);
      //free_dvector(q, 1,Alphsize);
      free( root_location );
      free_dmatrix( JP, 1, Alphsize, 1, Alphsize );
      return (Lambda_local.min + Lambda_local.max) / 2.0;
    }
    else if (-1 == Lambda_local.flag) {
      //printf("matrix pass first screening but no sensible solution found. :-( \n");
    }
  } 
  else if (N_root > 0) {
    //printf("N_root = %d for this matirx \n", N_root);
    //for (i=0;i<N_root;i++) printf("N_root[%d] = %lf\n",i,root_location[i]);
    for (k=0; k<= N_root; k++){
      if (k==0) {Lambda_local.min =0; Lambda_local.max = root_location[0];}
      else if (k== N_root) {Lambda_local.min = root_location[N_root-1]; Lambda_local.max = Lambda_UB+Epsilon;}
      else {Lambda_local.min = root_location[k-1]; Lambda_local.max = root_location[k]; }
      Lambda_local = Find_JP(mat_b, Lambda_local.min, Lambda_local.max, JP, p, q);
      if (1 == Lambda_local.flag) { // sensible solution found  
	//free_dvector(p, 1,Alphsize);
	//free_dvector(q, 1,Alphsize);
	free( root_location );
      free_dmatrix( JP, 1, Alphsize, 1, Alphsize );
	return (Lambda_local.min + Lambda_local.max) / 2.0;
      }
      else if (-1 == Lambda_local.flag) {
	//printf("matrix pass first screening but still no sensible solution found. :-( \n");
      }
    }
  } 
  // Remember to find the right place to free the vectors
  //free_dvector(p, 1,Alphsize);
  //free_dvector(q, 1,Alphsize);
  free( root_location );
  free_dmatrix( JP, 1, Alphsize, 1, Alphsize );
  return -1.0;
}

bool Check_range( const double** mat_b )
{
  int i,j;
  int pos_flag_r, neg_flag_r;
  int pos_flag_c, neg_flag_c;
  double r_max, c_max; //max of each row (or column)

  int L_r=0, L_c=0; // number of zero-score rows (columns)

  // First make sure each row and column have both pos and neg entries
  r_max_m = c_max_m = 100000000000.0;
  for (i=1;i<=Alphsize; i++) {
    r_max = 0; c_max = 0;
    pos_flag_r = -1; neg_flag_r = -1; pos_flag_c = -1; neg_flag_c = -1;
    for (j = 1; j<= Alphsize; j++){
      if (mat_b[i][j] > 0) {
	if (mat_b[i][j]> r_max) r_max = mat_b[i][j];
	pos_flag_r = 1;
      }
      else if (mat_b[i][j] < 0) neg_flag_r = 1;
      if (mat_b[j][i] > 0) {
	if (mat_b[j][i] > c_max) c_max = mat_b[j][i];
	pos_flag_c = 1;
      }
      else if (mat_b[j][i] < 0) neg_flag_c = 1;
    }
    if ((pos_flag_r == -1)||(neg_flag_r == -1)||(pos_flag_c == -1)||(neg_flag_c == -1)) {
      if ((pos_flag_r == -1)&&(neg_flag_r == -1)){
	printf ("only zero score at row  %d\n", i);
	L_r++;
      }
      else if ((pos_flag_c == -1)&& (neg_flag_c == -1)){
	printf ("only zero score at column %d\n", i);
	L_c++;
      }
      else{
	//printf("all positive or all negative at row or column %d\n", i);
	//printf("therefore invalid matrix. exit now. \n");
	return false;
	//exit(1);
      }
    }
    if ((r_max < r_max_m)&&(r_max > 0)) r_max_m = r_max;
    if ((c_max < c_max_m)&&(c_max > 0)) c_max_m = c_max;
  }


  // Find the upper bound for lambda
  if (r_max_m > c_max_m) {
    Lambda_UB = 1.1*log(1.0*Alphsize-L_r)/r_max_m ;
  }
  else {
    Lambda_UB = 1.1*log(1.0*Alphsize-L_c)/c_max_m;
  }
  //printf("the upper bound for lambda is %lf\n", Lambda_UB);
  return true;
}



double Check_det(const double** mat_b, double **a, double lambda)
{
  double d;
  int i, /*j,*/ *indx;

  indx = ivector(1,Alphsize);
  makematrix(mat_b, a,lambda);
  ludcmp(a,Alphsize,indx,&d);
  for(i=1;i<=Alphsize;i++) d *= a[i][i];
  free_ivector(indx,1,Alphsize);
  return d;  //returning the determinant
}


Sum Check_root(const double** mat_b, double **a, double lambda, double *p, double *q)
{
  double **y, /* *col,*/ d;
  //double sum = 0.0;
  int i,j;//, *indx;
  Sum Sum_here; 

  y=dmatrix(1,Alphsize,1,Alphsize);
  //indx = ivector(1,Alphsize);
  int indx[Alphsize+1];
  //col = dvector(1,Alphsize);
  double col[Alphsize+1];

  makematrix(mat_b, a,lambda);
  ludcmp(a,Alphsize,indx,&d);
  Sum_here.value = 0.0;
  for(i=1;i<=Alphsize;i++) q[i]=0.0;
  for(j=1;j<=Alphsize;j++){
    for (i=1;i<=Alphsize;i++) col[i]=0.0;
    col[j]=1.0;
    lubksb(a,Alphsize,indx,col);
    p[j]=0.0;
    for (i=1;i<=Alphsize;i++) {
      y[i][j]=col[i]; Sum_here.value += y[i][j]; 
      p[j]+=y[i][j]; q[i]+=y[i][j]; 
    }
  }
  
  Sum_here.flag = 1;
  for (i=1;i<Alphsize; i++) {
    if ((p[i] < 0) || (q[i] < 0)) {
      Sum_here.flag = -1;
      //printf("problematic freq. p[%d] = %.4f q[%d]=%.4f\n",i,p[i],i,q[i]);
    }
  }
  free_dmatrix(y,1,Alphsize,1,Alphsize);
  return Sum_here; 
}



double * Locate_det_zero(const double** mat_b, int *N_root_add)
{
  double **a/*,  *q, *p */; // a is the exponentiated matrix of socres, p and q are bg_freqs
  int i/*,j,k*/;
  int N;  // number of points for first round
  int flag_sign; 
  double lambda/*, l_tmp, sum,  sum_min, sum_max */;
  double lambda_root, dlambda /*, dsum=0.5 */; 
  //double *l_here, *s_here;
  double root[5000];
  double *root_temp; 
  //double error=0.000000000001;
  int zero_monitor = 0;  // record number of zeros found in the range 
  //int flag;  

  a=dmatrix(1,Alphsize,1,Alphsize);
  //Sum_local = (Sum *)malloc(sizeof(Sum));
  //Lambda_local = (Lambda *)malloc(sizeof(Lambda));
  
  N = 2 + max(400, ((int) (Lambda_UB-0)/0.005)); 
  //printf("N = %d in Locate_det_zero\n", N);
  dlambda = (Lambda_UB)/(N*1.0); 
  //l_here = (double *)malloc((N+1)*sizeof(double));
  //s_here = (double *)malloc((N+1)*sizeof(double));
  double l_here[N+1];
  double s_here[N+1];

  for (i=0; i< N ; i++){
    lambda = (i+1)*dlambda; 
    s_here[i] = Check_det(mat_b, a,lambda); 
    l_here[i] = lambda; 
  }

  if (s_here[0] < 0.0) flag_sign = -1;
  if (s_here[0] > 0.0 ) flag_sign = 1;
  if (fabs(s_here[0])/exp(l_here[0]*(r_max_m+c_max_m)/2.0) <= Epsilon) {
    root[zero_monitor++] = l_here[0];
    flag_sign = 0;
  }
  
  for (i=1;i<N; i++){
    if ((flag_sign != 0)&& (fabs(s_here[i])>Epsilon)) {
      if (s_here[i-1]*s_here[i] < 0){
	//printf("occurring at regular places\n");
	lambda_root = Nail_det(mat_b, flag_sign, l_here[i-1], l_here[i]);
	root[zero_monitor++] = lambda_root;
	flag_sign = -flag_sign;  // the flag switch sign after one sol found 
	//printf("a (regular) root of det found at %12.10f, i= %d\n", lambda_root,i);
      }
    }
    else {
      if (s_here[i] < 0.0) flag_sign = -1;
      if (s_here[i] > 0.0 ) flag_sign = 1;
      if (fabs(s_here[i])/exp(l_here[i]*(r_max_m+c_max_m)/2.0) <= Epsilon) {
	root[zero_monitor++] = l_here[i];
      }
    }
  }
  //printf("total number of solution found in range is %d\n", i_monitor);
  root_temp = (double *)malloc(zero_monitor*sizeof(double));
  *N_root_add = zero_monitor; 
  if (zero_monitor > 0) {
    if (zero_monitor >= N/4) {
      //printf("It is likely that uniform zero determinant is occurring.\n");
      //printf("number of small det points = %d out of %d, exit now....\n",zero_monitor, N);
      free( root_temp );
      return NULL;
      //exit(1);
    }
    for (i = 0; i < zero_monitor; i++){
      root_temp[i] = root[i];
      //printf("root_location[%d] = %lf\n",i,root_temp[i]);
    }
  }
  free_dmatrix( a, 1, Alphsize, 1, Alphsize );
  return root_temp;  

} 


Lambda Find_JP(const double** mat_b, double la_min, double la_max, double **JP, double *p_in, double *q_in)
{
  double **a, *q, *p; // a is the exponentiated matrix of socres, p and q are bg_freqs
  int i,j/*,k*/;
  int N;  // number of points for first round
  double lambda/*, l_tmp, sum, sum_min, sum_max*/;
  double lambda_max, lambda_min, lambda_final, dlambda/*, dsum=0.5*/; 
  //double *l_here, *s_here;
  //double error=0.000000000000001;
  //int validity_flag; // 1 means valid, -1 means not valid.
  int flag_sign;       // 1 means small lambda sum > 1, -1 means otherwise
  int flag_done = -1;       // 1 means find sensible solution, -1 means sensible not found 
  int i_monitor = 0;          // record number of solution found in the range, including nonsense ones 
  int j_monitor;

  Lambda Lambda_local;
  //Sum *Sum_local;   
  Sum Sum_local;   

  lambda_min = la_min; 
  lambda_max = la_max;
  q = q_in;
  p = p_in;
  a=dmatrix(1,Alphsize,1,Alphsize);
  //Sum_local = (Sum *)malloc(sizeof(Sum));
  //Lambda_local = (Lambda *)malloc(sizeof(Lambda));
  
  N = 2 + max(400, ((int) (lambda_max-lambda_min)/0.005)); 
  //printf("N = %d in Find_JP\n", N);
  dlambda = (lambda_max-lambda_min)/(N*1.0); 
  //l_here = (double *)malloc((N+1)*sizeof(double));
  //s_here = (double *)malloc((N+1)*sizeof(double));
  double l_here[N+1];
  double s_here[N+1];
  //printf("lambda_min enter = %12.10e, lambda_max = %12.10f\n", lambda_min, lambda_max);
  for (i=0; i< N-1 ; i++){
    lambda = lambda_min + (i+1)*dlambda; 
    makematrix(mat_b, a,lambda);
    Sum_local = Check_root(mat_b, a,lambda,p,q); 
    l_here[i] = lambda; 
    s_here[i] = Sum_local.value - 1.0;
    //printf("scan %d th time in Find_JP\n",i );
  }
  //printf("finish first time scanining in Find_JP\n");
  if (s_here[0] < 0.0) flag_sign = -1;
  else if (s_here[0] > 0.0 ) flag_sign = 1;
  else if (s_here[0] == 0.0) {  //needs refined definition on flag_sign
    printf("enter the exact hit, rarely occurs other than when lambda = 0 \n"); 
    j_monitor = 1;
    flag_sign = 0;
    while ((flag_sign == 0) &&(j_monitor < N)){
      Sum_local = Check_root(mat_b, a,l_here[0]+ j_monitor*dlambda/N,p,q);
      if (Sum_local.value > 1.0) {
	flag_sign  = 1;
      }
      else if (Sum_local.value <1.0){
	flag_sign = -1;
      }
      j_monitor++; 
    }
  }
  
  for (i=1;i<N; i++){  // should be N-1 ???
    if (flag_sign == 0) {printf("flag_sign = 0 \n"); exit(1);}
    if (s_here[i-1]*s_here[i] < 0){
      lambda_min = l_here[i-1];
      lambda_max = l_here[i];
      Sum_local = Nail_lambda(mat_b, flag_sign, lambda_min, lambda_max, p,q, &lambda_final);
      if (Sum_local.flag == 1) {
	i =N; flag_done = 1; Lambda_local.flag = 1;
	Lambda_local.min = lambda_final, Lambda_local.max = lambda_final;
      }
      flag_sign = -flag_sign;  // the flag switch sign after one sol found 
      i_monitor++ ;
    }
  }

  if (flag_done == 1){
    // Write correct JP to the matrix
    makematrix(mat_b, a,lambda_final);
    for (i=1;i<=Alphsize;i++){
      for (j=1;j<=Alphsize;j++){
	JP[i][j] =  a[i][j]*p[i]*q[j];
      }
    }     
    free_dmatrix(a,1,Alphsize,1,Alphsize);
    return Lambda_local;
  }
  else if (flag_done == -1){
    //printf("no sensible solution in the plausible x range: (%lf,%lf)\n", la_min, la_max);
    Lambda_local.flag = -1; Lambda_local.min = 0; Lambda_local.max = Infty;
    return Lambda_local; 
  }
  // never come here
  return Lambda_local;
} 


Sum Nail_lambda(const double** mat_b, int flag_sign, double lambda_min, double lambda_max, double *p, double *q, double *lam_add)
{
  double **a;
  double lambda;

  //Sum *Sum_local;
  Sum Sum_local;
  a = dmatrix(1,Alphsize,1,Alphsize);
  //Sum_local = (Sum *)malloc(sizeof(Sum));

  lambda = (lambda_min+lambda_max)/2.0;
  Sum_local = Check_root(mat_b, a, lambda, p, q);  
  while (fabs(Sum_local.value-1.0)>E_bound){
    if (flag_sign*(Sum_local.value-1.0)<0) lambda_max = lambda;
    else if (flag_sign*(Sum_local.value-1.0)> 0) lambda_min = lambda;

    // Added by MCF to avoid infinite loop:
    if (lambda == (lambda_min+lambda_max)/2.0){
      Sum_local.flag = -1;
      break;
    }

    lambda = (lambda_min+lambda_max)/2.0;
    Sum_local = Check_root(mat_b, a, lambda, p, q); 
  }
  free_dmatrix(a,1,Alphsize,1,Alphsize);
  *lam_add = lambda; 
  return Sum_local;
}


double Nail_det(const double** mat_b, int flag_sign, double lambda_min, double lambda_max)
{
  double **a;
  double lambda;
  double value;

  a = dmatrix(1,Alphsize,1,Alphsize);

  lambda = (lambda_min+lambda_max)/2.0;
  value = Check_det(mat_b, a, lambda);  
  while ((fabs(value)>E_bound)&&(lambda > 0)){
    if (flag_sign*(value)<0) lambda_max = lambda;
    else if (flag_sign*(value)> 0) lambda_min = lambda;
    lambda = (lambda_min+lambda_max)/2.0;
    value = Check_det(mat_b, a, lambda); 
  }
  free_dmatrix(a,1,Alphsize,1,Alphsize);
  return lambda; 
}

void makematrix(const double** mat_b, double **a, double lambda)
{

  int i,j;
 
  for (i=1;i<=Alphsize;i++)
    for(j=1;j<=Alphsize;j++){
      *(*(a+i)+j)=exp(lambda*mat_b[i][j]);       
    }
}



