#include "matrix.h"
#include "aaTools.h"


// Initializes Matrix
Matrix::Matrix(int len) {
  if(DB) printf("entering Matrix constructor \n");
  N = len;
  cell = new int*[len];
  for(int i = 0; i < len; i++)
    cell[i] = new int[len];
  label = new char[len+1];
  if(len==20)
    strcpy(label, "ARNDCQEGHILKMFPSTWYV");
  if(len==7)
    strcpy(label, "BEGHSTC");
  if(DB) printf("leaving Matrix constructor \n");
}


// Default constructor
Matrix::Matrix() {
  if(DB) printf("entering default Matrix constructor \n");
  int len=20; // set by default
  N = len;
  cell = new int*[len];
  for(int i = 0; i < len; i++)
    cell[i] = new int[len];
  label = new char[len+1];
  if(len==20)
    strcpy(label, "ARNDCQEGHILKMFPSTWYV");
  if(len==7)
    strcpy(label, "BEGHSTC");
  if(DB) printf("leaving default Matrix constructor \n");
}


// Destructor
Matrix::~Matrix() {
  if(DB) printf("entering Matrix destructor \n");
  if(cell!=NULL)
    for(int i=0; i<N; i++)
      delete[] cell[i];
  if(label!=NULL) delete[] label;
  if(DB) printf("leaving Matrix destructor \n");
}


// copy constructor
Matrix::Matrix(const Matrix& m) {
  if(DB) printf("entering Matrix copy constructor \n");
  expectedValue = m.expectedValue;
  relativeEntropy = m.relativeEntropy;
  N = m.N;
  if(m.cell==NULL) cell=NULL;
  else {
    cell = new int*[N];
    for(int i=0; i<N; i++)
      cell[i] = new int[N];
    for(int i=0; i<N; i++)
      for(int j=0; j<N; j++)
        cell[i][j] = m.cell[i][j];
  }
  minimumCellValue = m.minimumCellValue;
  if(m.label==NULL) label=NULL;
  else {
    label = new char[N+1];
    strcpy(label, m.label);
  }
  if(DB) printf("leaving Matrix copy constructor \n");
}

// Assignment
Matrix& Matrix::operator= (const Matrix& m) {
  if(DB) printf("entering Matrix::operator= \n");
  if(this != &m) {
    if(m.cell==NULL) cell=NULL;
    else {
      if(cell!=NULL) 
        for(int i=0; i<N; i++)
          delete[] cell[i];
      cell = new int*[N];
      for(int i=0; i<N; i++)
        cell[i] = new int[N];
      for(int i=0; i<N; i++)
        for(int j=0; j<N; j++)
          cell[i][j] = m.cell[i][j];
    }
    if(m.label==NULL) label=NULL;
    else {
      if(label!=NULL) delete[] label;
      label = new char[N+1];
      strcpy(label, m.label);
    }
    expectedValue = m.expectedValue;
    relativeEntropy = m.relativeEntropy;
    N = m.N;
    minimumCellValue = m.minimumCellValue;
  }
  if(DB) printf("leaving Matrix::operator= \n");
  return *this;
}


// Matrix addition, returns a summation Matrix
const Matrix Matrix::operator+ (const Matrix& m) {
  if(DB) printf("entering Matrix::operator+ \n");

  if(m.N!=N) {
    printf("ERROR: attempted to add matrices of differnet dimensions\n");
    return *this;
  }

  Matrix n(m.N);

  for(int i=0; i<m.N; i++)
    for(int j=0; j<m.N; j++)
      n.cell[i][j] = cell[i][j] + m.cell[i][j];

  strcpy(n.label, m.label);

  n.expectedValue = expectedValue + m.expectedValue;
  n.relativeEntropy = relativeEntropy + m.relativeEntropy;
  n.N = m.N;
  n.N = m.N;
  n.minimumCellValue = 0;
  if(DB) printf("leaving Matrix::operator+ \n");
  return n;
}


// Matrix subtraction, returns a difference Matrix
const Matrix Matrix::operator- (const Matrix& m) {
  if(DB) printf("entering Matrix::operator- \n");

  if(m.N!=N) {
    printf("ERROR: attempted to subtract matrices of differnet dimensions\n");
    return *this;
  }

  Matrix n(m.N);

  for(int i=0; i<m.N; i++)
    for(int j=0; j<m.N; j++)
      n.cell[i][j] = cell[i][j] - m.cell[i][j];

  strcpy(n.label, m.label);

  n.expectedValue = expectedValue - m.expectedValue;
  n.relativeEntropy = relativeEntropy - m.relativeEntropy;
  n.N = m.N;
  n.N = m.N;
  n.minimumCellValue = 0;
  if(DB) printf("leaving Matrix::operator- \n");
  return n;
}


// Matrix multiplication, returns a entry-wise product Matrix
const Matrix Matrix::operator* (const Matrix& m) {
  if(DB) printf("entering Matrix::operator* \n");

  if(m.N!=N) {
    printf("ERROR: attempted to multiply matrices of differnet dimensions\n");
    return *this;
  }

  Matrix n(m.N);

  for(int i=0; i<m.N; i++)
    for(int j=0; j<m.N; j++)
      n.cell[i][j] = cell[i][j] * m.cell[i][j];

  strcpy(n.label, m.label);

  n.expectedValue = expectedValue * m.expectedValue;
  n.relativeEntropy = relativeEntropy * m.relativeEntropy;
  n.N = m.N;
  n.N = m.N;
  n.minimumCellValue = 0;
  if(DB) printf("leaving Matrix::operator* \n");
  return n;
}


// Matrix multiplication, returns a weighted Matrix
const Matrix Matrix::operator* (float w) {
  if(DB) printf("entering Matrix::operator* \n");

  Matrix n(N);

  for(int i=0; i<N; i++)
    for(int j=0; j<N; j++)
      if(w*cell[i][j]+.1-floor(w*cell[i][j]+.1)<.5)
        n.cell[i][j] = (int)floor(w*cell[i][j]+.1);
      else
        n.cell[i][j] = (int)ceil(w*cell[i][j]+.1);

  strcpy(n.label, label);

  n.expectedValue = w*expectedValue;
  n.relativeEntropy = w*relativeEntropy;
  n.N = N;
  n.minimumCellValue = 0;
  if(DB) printf("leaving Matrix::operator* \n");
  return n;
}


// Computes the substitution Matrix
void Matrix::computeFromAlignment(Alignment aln, float maxPGaps)
{
  if(DB) printf("entering Matrix::computeFromAlignment(Alignment)\n");
  int i, j;
  // int k;

  float **f = new float*[N]; for(i = 0; i < N; i++) f[i] = new float[N];
  float *p = new float[N];
  //float *aaFreq = new float[N];
  float **q = new float*[N]; for(i = 0; i < N; i++) q[i] = new float[N];
  float **e = new float*[N]; for(i = 0; i < N; i++) e[i] = new float[N];
  float **s = new float*[N]; for(i = 0; i < N; i++) s[i] = new float[N];
  float **q2 = new float*[N]; for(i = 0; i < N; i++) q2[i] = new float[N];
  for(i=0;i<N;i++) p[i]=0;
  for(i=0;i<N;i++) for(j=0;j<N;j++) {
    f[i][j]=0; q[i][j]=0; e[i][j]=0; s[i][j]=0; q2[i][j]=0;
  }

  computeFrequencies(aln, f, maxPGaps); //fills f

  float sum=0;
  for(i=0; i<N; i++) for(j=0; j<=i; j++) sum+=f[i][j];
  for(i=0; i<N; i++) for(j=0; j<N; j++) q[i][j]=f[i][j]/sum;


  for(i=0; i<N; i++) {
    sum=0;
    for(j=0; j<N; j++) if(i!=j) sum+=q[i][j];
    p[i] = q[i][i]+sum/2;
  }

  // try extended probability P(A|B) ~ P(A|k)*P(k|B)
//  for(i=0; i<N; i++) for(j=0; j<N; j++) for(k=0; k<N; k++)
//    q2[i][j]+=q[i][k]*q[k][j]/p[k];
//  for(i=0; i<N; i++) for(j=0; j<N; j++) printf("%d %d %f %f\n",i,j,q[i][j],q2[i][j]);
//  for(i=0; i<N; i++) for(j=0; j<N; j++) q[i][j]=q2[i][j];

  for(i=0; i<N; i++) for(j=0; j<N; j++)
      if(i==j)  e[i][j] =   p[i]*p[j];
      else      e[i][j] = 2*p[i]*p[j];
  for(i=0; i<N; i++) for(j=0; j<N; j++) {
    if(q[i][j]!=0) s[i][j] = log(q[i][j]/e[i][j])/log((float)2);
    else s[i][j]=-33;
  }

/*  for(i=0; i<N; i++) for(j=0; j<N; j++) {
    printf("f[%d][%d] = %f\n",i,j,f[i][j]);
    printf("p[%d] = %f p[%d] = %f\n",i,p[i],j,p[j]);
    printf("q[%d][%d] = %f\n",i,j,q[i][j]);
    printf("e[%d][%d] = %f\n",i,j,e[i][j]);
    printf("s[%d][%d] = %f\n",i,j,s[i][j]);
  }
*/
  relativeEntropy=0; expectedValue=0;
  for(i=0; i<N; i++) for(j=0; j<=i; j++) {
      relativeEntropy+=q[i][j]*s[i][j];
      expectedValue+=p[i]*p[j]*s[i][j];
    }

  // compute actual Matrix
  for(i=0; i<N; i++) for(j=0; j<N; j++)
    if(SCALE*s[i][j]-floor(SCALE*s[i][j])<.5)
      cell[i][j] = (int)floor(SCALE*s[i][j]);
    else
      cell[i][j] = (int)ceil(SCALE*s[i][j]);

  minimumCellValue=0;
  for(i=0;i<N;i++) for(j=0;j<N;j++)
    if(minimumCellValue>cell[i][j]) minimumCellValue=cell[i][j];

  delete[] p;
  for(int i=0; i<N; i++) delete[] f[i];
  for(int i=0; i<N; i++) delete[] q[i];
  for(int i=0; i<N; i++) delete[] e[i];
  for(int i=0; i<N; i++) delete[] s[i];

  if(DB) printf("leaving Matrix::computeFromAlignment(Alignment)\n");
}

// By default allows all gaps
void Matrix::computeFromAlignment(Alignment aln) {
  computeFromAlignment(aln,1);
}

// computes the f values for constructing a substitution
// Matrix
void Matrix::computeFrequencies(Alignment aln, float **f, float maxPGaps) {
  if(DB) printf("entering Matrix computeFrequencies \n");
  int i,j,k,A,B;
  float *gapP = new float[aln.maximumSequenceLength];
  for(i=0; i<aln.maximumSequenceLength; i++) gapP[i]=0;

  if(aln.sequences[1].length==aln.sequences[2].length) {
    for(i=0; i<aln.maximumSequenceLength; i++)
      for(k=0; k<aln.nSequences; k++)
        if(aln.sequences[k].residues[i]=='-') gapP[i]++;
    for(i=0; i<aln.maximumSequenceLength; i++) gapP[i]=gapP[i]/aln.nSequences;
    for(k=0; k<aln.nSequences; k++)
      for(j=0; j<k; j++)
        for(i=0; i<aln.maximumSequenceLength; i++) {
          A = residueCharToInt(aln.sequences[k].residues[i]);
          B = residueCharToInt(aln.sequences[j].residues[i]);
          if(A>-1 && B>-1 && gapP[i]<=maxPGaps) {
            f[A][B]++; f[B][A]++;
          }
        }
  }
  else
    for(k=0; k<aln.nSequences/2; k++) {
      for(i=0; i<aln.sequences[2*k].length; i++) {
        A = residueCharToInt(aln.sequences[2*k].residues[i]);
        B = residueCharToInt(aln.sequences[2*k+1].residues[i]);
        if(A>-1 && B>-1) {
          f[A][B]++; f[B][A]++;
        }
      }
    }

  if(DB) printf("leaving Matrix computeFrequencies \n");
}




// Loads the Blosum50 Matrix into the Matrix::cell fields
void Matrix::computeFromBlosumFile(char *s)
{
  if(DB) printf("entering Matrix computeFromBlosumFile \n");
  int i, j;
  FILE *blosFile;
  char* temp = new char[100];
  char* temp2 = new char[100];

  strcpy(temp, BLOSUM_PATH);
  strcat(temp, s);
  blosFile = fopen(temp, "r"); // Open input file

  for(i=0; i<6; i++)
    fgets(temp, 80+2, blosFile);
  strcpy(temp2, temp+13);
  temp2[9]=0;
  relativeEntropy = charToFloat(temp2);
  strcpy(temp2, temp+34);
  temp2[9]=0;
  expectedValue = charToFloat(temp2);

  fgets(temp, 80+2, blosFile);
  for(i=0; i<N; i++) {
    fgets(temp, 90, blosFile);
    for(j=0; j<N; j++) {
      strcpy(temp2, temp+1+3*j);
      temp2[4] = '\0';
      cell[i][j] = charToInt(temp2);
    }
  }

  delete[] temp;
  delete[] temp2;

  if(DB) printf("leaving Matrix computeFromBlosumFile \n");
}


// Prints Matrix info
void Matrix::print() {
  if(DB) printf("entering Matrix::print\n");
  int i, j;
  int min=0,max=0;
  for(i=0; i<N; i++) for(j=0; j<N; j++) if(cell[i][j]<min) min=cell[i][j];
  for(i=0; i<N; i++) for(j=0; j<N; j++) if(cell[i][j]>max) max=cell[i][j];
  printf("#  Matrix made by aaTools\n");
  printf("#  * column uses minimum score\n");
  printf("#  Scoring Matrix in 1/%d Bit Units\n",SCALE);
  printf("#  \n");
  printf("#  \n");
  printf("#  Entropy =   %1.4f, Expected =   %1.4f\n", 
      relativeEntropy,expectedValue);
  printf("   ");
  for(i=0; i<N; i++)
    printf("%c  ", label[i]);
  printf("\n");
  for(i=0; i<N; i++) {
    printf("%c", label[i]);
    for(j=0; j<N; j++) {
      printf("%3d", cell[j][i]);
    }
    printf("\n");
  }
  if(DB) printf("leaving Matrix::print\n");
}

// Prints Matrix info
void Matrix::printColor() {
  if(DB) printf("entering Matrix::print\n");
  int i, j;
#define RED printf("\033[22;31m");
#define YEL printf("\033[22;33m");
#define GRE printf("\033[22;32m");
#define IND printf("\033[22;36m");
#define BLU printf("\033[22;34m");
#define VIO printf("\033[22;35m");
  int min=0,max=0;
  for(i=0; i<N; i++) for(j=0; j<N; j++) if(cell[i][j]<min) min=cell[i][j];
  for(i=0; i<N; i++) for(j=0; j<N; j++) if(cell[i][j]>max) max=cell[i][j];
  printf("Relative Entropy (H) = %f\n", relativeEntropy);
  printf("Expected Score (E) = %f\n", expectedValue);
  printf("   ");
  for(i=0; i<N; i++)
    printf("%c  ", label[i]);
  printf("\n");
  for(i=0; i<N; i++) {
    printf("%c", label[i]);
    for(j=0; j<N; j++) {
/*      if(cell[j][i]>max*2/3) RED
        else if(cell[j][i]>max*1/3) YEL
          else if(cell[j][i]>0) GRE
            else if(cell[j][i]>min*1/3) IND
              else if(cell[j][i]>min*2/3) BLU
                else VIO
*/
      if(cell[j][i]>10) RED
        else if(cell[j][i]>5) YEL
          else if(cell[j][i]>0) GRE
            else if(cell[j][i]>-5) IND
              else if(cell[j][i]>-10) BLU
                else VIO

      printf("%3d", cell[j][i]);
    BLK
    }
    printf("\n");
  }
  BLK
  if(DB) printf("leaving Matrix::print\n");
}


// Maps each character in a set to consecutive integers
int Matrix::residueCharToInt(char residue) {
  for(int i=0; i<N; i++)
    if(residue==label[i]) return i;
  if(residue=='-') 
    return -1;
  else
    return -2;
}

// Returns the contents of cell [i,j]
int Matrix::getCell(int i, int j) {
  if(i<N && j<N && cell!=0) 
    return cell[i][j];
  else {
    if(cell==0)
      fprintf(stderr, "ERROR:  In function  int Matrix::getCell(int,int), array cell is not defined\n");
    else {
      if(i>=N)
        fprintf(stderr,"ERROR:  In function  int Matrix::getCell(int,int), index i=%d is greater than matrix dimension N=%d\n",i,N); 
      if(j>=N)
        fprintf(stderr,"ERROR:  In function  int Matrix::getCell(int,int), index j=%d is greater than matrix dimension N=%d\n",j,N); 
    }
    return -999;
  }
}