/*  trinculo, Flexible tools for multinomial association
    Copyright (C) 2014  Luke Jostins

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.*/


#ifndef LINUX
#include <Accelerate/Accelerate.h>
#endif 
using namespace std;

#ifdef LINUX 
extern "C" {
	//lapack functions
void dgetrf_(int* M, int *N, double* A, int* lda, int* IPIV, int* INFO);
void dgetri_(int* N, double* A, int* lda, int* IPIV, double* WORK, int* lwork, int* INFO);
void dgesv_( int* n, int* nrhs, double* a, int* lda, int* ipiv, double* b, int* ldb, int* info );
}
#endif

double inverse(double* A, int N)
{
	//invert the matrix A
    int *IPIV = new int[N+1];
    int LWORK = N*N;
    double *WORK = new double[LWORK];
    double det = 1.0;
    int INFO;

    dgetrf_(&N,&N,A,&N,IPIV,&INFO);
    for (int i = 0; i < N; i++) det *= (A[i + i*N]);
    dgetri_(&N,A,&N,IPIV,WORK,&LWORK,&INFO);

    delete IPIV;
    delete WORK;
    return(det);
}

void multiplyVec(double *A, double *b, double *c, int N){
  // does c = A*b where A is square

  for (int i = 0; i < N; i++){
    c[i] = 0;
    for (int j = 0; j < N; j++){
      c[i] += A[i + j*N]*b[j];
    }
  }
}

void solve(double *A, double *b, int N, int *ipiv){
	// wrapper for the lapack solve function
  int info;
  int nrhs = 1;
  dgesv_(&N, &nrhs, A, &N, ipiv, b, &N, &info);
}

class MultinomNewton {
	// class for holding data required for, and fitting, a multinomial logistic model
 public:
  int Npred; //number of predictors (genotypes plus confounders)
  int Ncovar; //number of confounders
  int Nind; //number of individuals
  int Npheno; //number of phenotype categories
  int * missing; //vector indicating which individuals have missing genotype or phenotype data
  int bayesian; //a flag to indicate that a Bayesian analysis should be carried out (i.e. the prior used)
 
  double *x; //predictors (genotypes and confounders)  
  int *y; //phenotype data
  double *beta; //matrix of effect sizes
  double *sigma; //effect size covariance matrix prior 
  
  double *p; //probabilities for each individual for each category (conditional on genotypes and covariates)
  double *U; //gradiant
  double *I; //hessian
  double *c; //step change
  
  double *tempsigma; //temporary variables
  double *tempbeta; //temporary variables
  int *ipiv; //temporary variables
  double *invsigma; //temporary variables
  double *priortemp; //temporary variables
  double *priortemp2; //temporary variables
  double det; //temporary variables

  //double *initTau; //initial values of the baseline frequency parameter (speeds up model fitting, but can make it unstable. Not currently used.)

  MultinomNewton();
  MultinomNewton(int Npred, int Nind, int Npheno);
  ~MultinomNewton();
  void AllocateSpace(int initNpred, int initNind, int initNpheno);
  void setGenos(bedfile *bed, int SNPindex);
  void setDosage(dosagefile *dosage, int SNPindex);
  void setCovars(covars *cov);
  void setPheno(phenos *phen, int phenofocus);
  double fitLogistic(int * constrain, int maxIter = 20, double relTol = 1e-6, double absTol = 1e-8);
  void NRiter(int * constrain);
  void printStuff();
  void setMissing();
  void calcPs();
  void calcU();
  void calcI();
  double calcL ();
  //double calcP ();
  void zeroBetas();
  void setPrior(double *newSigma);
  double calcPrior(int *constrain);
  void setConditional(vector<int> conditional, dosagefile *dosage);
  void setConditional(vector<int> conditional, bedfile *bed);
};


void MultinomNewton::setMissing(){
	//figure out which individuals have missinge genotype, confounder or phenotype data
  int i, j;
  for (i = 0 ; i < Nind; i++){
    missing[i] = 0;
    if (y[i] < 0) missing[i] = 1;
    if (x[i + 1*Nind] < 0) missing[i] = 1;
    for (j = 2; j < (Npred + 1); j++){
      if (custom_isnan(x[i + j*Nind])) missing[i] = 1;
    }
  }
      
}

void MultinomNewton::setPrior(double *newSigma){
	// set the prior covariance matrix (and invert it)
  for (int i = 0; i < (Npheno * (Npred + 1)*Npheno * (Npred + 1)); i++) sigma[i] = newSigma[i];
  det = fabs(inverse(newSigma,Npheno * (Npred + 1)));
  for (int i = 0; i < (Npheno * (Npred + 1)*Npheno * (Npred + 1)); i++) invsigma[i] = newSigma[i];
  bayesian=1;
}

void MultinomNewton::zeroBetas(){
	// set the parameter vector to the initial values
 int i;
 for (i = 0; i < Npheno*(Npred + 1); i++){ beta[i] = 0;}
 for (i = 0; i < Npheno; i++){ beta[i*(Npred + 1)] = 0;}//initTau[i]; }
}

void MultinomNewton::printStuff(){
	// print some debugging information
  int i, j;

  for (i = 0 ; i < Nind; i++) {
    cout << "IND" << i;
    cout << " " << y[i];
    for ( j = 0 ; j < (Npred + 1); j++) cout << " " <<  x[i + j*Nind] ;
    cout << endl;
    }



  cout << "beta:";
  for ( i = 0 ; i < (Npred + 1)*(Npheno); i++) cout << beta[i] << " " ;
  cout << endl;
  
  cout << "U:";
  for ( i = 0 ; i < (Npred + 1)*(Npheno); i++) cout << U[i] << " " ;
  cout << endl;

  cout << "I:";
  for ( i = 0 ; i < (Npred + 1)*(Npheno)*(Npred + 1)*(Npheno); i++) cout << I[i] << " " ;
  cout << endl;

  cout << "invsigma:";
  for ( i = 0 ; i < (Npred + 1)*(Npheno)*(Npred + 1)*(Npheno); i++) cout << invsigma[i] << " " ;
  cout << endl;

}

void MultinomNewton::AllocateSpace(int initNpred, int initNind, int initNpheno){
	//initialize variables, allocate space, etc
  bayesian=0;
  Npred = initNpred;
  Nind = initNind;
  Npheno = initNpheno;
  x = new double [Nind*(Npred + 1)];
  missing = new int [Nind];
  int i;
  for (i = 0 ; i < Nind; i++){
    x[i] = 1;
    missing[i] = 0;
  }
  y = new int [Nind];
  //initTau = new double [Npheno];
  beta = new double [Npheno * (Npred + 1)];
  tempbeta = new double [Npheno * (Npred + 1)];
  for (i = 0; i < Npheno*(Npred + 1); i++){ beta[i] = 0;}
  p = new double [(Npheno + 1) * Nind];
  U = new double [Npheno * (Npred + 1)];
  I = new double [Npheno * (Npred + 1) * Npheno * (Npred + 1)];
  invsigma = new double [Npheno * (Npred + 1) * Npheno * (Npred + 1)];
  sigma = new double [Npheno * (Npred + 1) * Npheno * (Npred + 1)];
  tempsigma = new double [Npheno * (Npred + 1) * Npheno * (Npred + 1)];
  priortemp = new double [Npheno * (Npred + 1)];
  priortemp2 = new double [Npheno * Npheno];
  c = new double [(Npheno + 1) * Nind];
  ipiv = new int [1 + Npheno * (Npred + 1)];
}

MultinomNewton::MultinomNewton(int initNpred, int initNind, int initNpheno){
	AllocateSpace(initNpred, initNind, initNpheno);
}


MultinomNewton::MultinomNewton(){
	return;
}


MultinomNewton::~MultinomNewton(){
	// free up all the allocated variables
  delete x;
  delete missing;
  delete y;
  delete beta;
  delete tempbeta;
  delete p;
  delete U;
  delete I;
  delete c;
  delete invsigma;
  delete sigma;
  delete tempsigma;
  delete priortemp;
  delete priortemp2;
  delete ipiv;
}

void MultinomNewton::setGenos (bedfile *bed, int SNPindex){
	// set genotypes from a bed object
  for (int i = 0; i < bed->Ninds; i++) x[i + bed->Ninds] = bed->getDosage(SNPindex,i);
}

void MultinomNewton::setDosage (dosagefile *dosage, int SNPindex){
	// set genotypes from a dosage object
  for (int i = 0; i < dosage->Ninds; i++) x[i + dosage->Ninds] = dosage->getDosage(SNPindex,i);
}


void MultinomNewton::setCovars(covars *cov){
	// set covariates
  for (int i = 0; i < Nind; i++) {
    for (int j = 0; j < cov->Ncovars ; j++)   x[i + (j+2)*Nind] = cov->covariates[j][i];
  }
  Ncovar = cov->Ncovars;
}

void MultinomNewton::setConditional(vector<int> conditional, dosagefile *dosage){
	//set additional SNPs to condition on from a dosage file
  for (int i = 0; i < Nind; i++) {
    for (int j = 0; j < conditional.size() ; j++){
      x[i + (j+2+Ncovar)*Nind] = dosage->getDosage(conditional[j],i);
      if (x[i + (j+2+Ncovar)*Nind] < 0) x[i + (j+2+Ncovar)*Nind] = 0.0/0.0;
    }
  }
}

void MultinomNewton::setConditional(vector<int> conditional, bedfile *bed){
	//set additional SNPs to condition on from a bed file
  for (int i = 0; i < Nind; i++) {
    for (int j = 0; j < conditional.size() ; j++){
      x[i + (j+2+Ncovar)*Nind] = bed->getDosage(conditional[j],i);
      if (x[i + (j+2+Ncovar)*Nind] < 0) x[i + (j+2+Ncovar)*Nind] = 0.0/0.0;
    }
  }
}

void MultinomNewton::setPheno(phenos *phen, int phenofocus){
	// set the phenotypes and calculate an initial guess at the baseline frequency parameters
  int i, counter=0;
  for (i = 0; i < Nind; i++) {
    y[i] = phen->phenotypes[phenofocus][i];
    if (y[i] >= 0){
      counter++;
      //if (y[i] > 0) initTau[y[i] - 1]++;
   }
  }

 //for (i = 0; i < Npheno; i++){
//	 initTau[i] = log(initTau[i]/counter);
 //}

}

double MultinomNewton::fitLogistic(int * constrain, int maxIter, double relTol, double absTol){
	// carry out the NR algorithm
	
  int i, j;
  double maxTol1, maxTol2;

  for (i = 0 ; i < maxIter; i++){
    NRiter(constrain); // carry out an iteration
    maxTol1 = 0;
	maxTol2 = 0;
    for (j = 0; j < (Npred + 1)*Npheno; j++){
      if (fabs(U[j]/beta[j]) > maxTol1 && !constrain[j]){
		  maxTol1 = fabs(U[j]/beta[j]);
		}
        if (fabs(U[j]) > maxTol2 && !constrain[j]){
  		  maxTol2 = fabs(U[j]);
  		}
    }
	//cout << maxTol1 << " " << maxTol2 << endl;
	//cout << max_j << " " << U[max_j] << " " << beta[max_j] << " " << maxTol << endl;
	//stop if the algorithm converged
    if (maxTol1 < relTol || maxTol2 < absTol){
      break;
    }
  }
  
  return(min(maxTol1,maxTol2));
}
  

void MultinomNewton::calcU (){
	// calculate the gradient
  int i, j, k;
  int yi;
  double temp;
  for (i = 0; i < Npheno*(Npred + 1); i++) U[i] = 0;

  for (k = 0; k < Npheno; k++){
    for (i = 0; i < Nind; i++){
       if (missing[i]) continue;
       if (y[i] == (k+1)) yi = 1;
       else yi = 0;
       temp = (yi - p[i + (k+1)*Nind]);
       for (j = 0; j < (Npred + 1); j++) U[j + k*(Npred+1)] += x[i + j*Nind]*temp;
   }
  }

  if (bayesian){
    multiplyVec(invsigma, beta, priortemp, Npheno*(Npred + 1));
    for (i = 0; i < Npheno*(Npred + 1); i++) U[i] -= priortemp[i];
  }

 /* for (j = 0; j < (Npred + 1); j++){
    for (k = 0; k < Npheno; k++){
      U[j + k*(Npred+1)] = 0;
      for (i = 0; i < Nind; i++) {
	if (missing[i]) continue;
	if (y[i] == (k + 1)) yi = 1;
	else yi = 0;
	U[j + k*(Npred+1)] += x[i + j*Nind]*(yi - p[i + (k+1)*Nind]);
      }
      
    }
  }*/

}  

void MultinomNewton::calcI(){
	//calculate the hessian
  int a, b, c, d, i;
  double temp, temp2;
  int temp3, temp4;
  for (i = 0; i < (Npred + 1)*(Npred + 1)*Npheno*Npheno; i++) I[i] = 0;      
      
  //simple version:
  /*for (a = 0; a < (Npred + 1); a++){
    for (b = 0; b < Npheno; b++){
      for (c = 0; c < (Npred + 1); c++){
	for (d = 0; d < Npheno; d++){
	  cout << "Map: " << a << " " << b << " " << c << " " << d << " = " << a + b*(Npred + 1) + c*(Npred + 1)*(Npheno) + d*(Npred + 1)*(Npred + 1)*Npheno << endl;
	  //I[a + b*(Npred + 1) + c*(Npred + 1)*(Npheno) + d*(Npred + 1)*(Npred + 1)*Npheno] = 0;
	  for (i = 0; i < Nind; i++){
	    if (missing[i]) continue;
	    if (b == d) temp = p[i + (b + 1)*Nind];
	    else temp = 0;
	    I[a + b*(Npred + 1) + c*(Npred + 1)*(Npheno) + d*(Npred + 1)*(Npred + 1)*Npheno] += (temp -  p[i + (b + 1)*Nind]* p[i + (d + 1)*Nind])*x[i + a*Nind]*x[i + c*Nind];
	  }
	}
      }
    }
    }*/

  // below is an optimized for of the above 
  int temp5 = (Npred + 1);
  int temp6 = Npheno*(Npred + 1);
  int temp7, temp8;
  
  for (b = 0; b < Npheno; b++){
    for (d = b; d < Npheno; d++){
      temp3 = b*temp5 + d*temp5*temp6;
      if (b == d){
	temp7 = (b + 1)*Nind;
	for (i = 0; i < Nind; i++){
          if (missing[i]) continue;
          temp = p[i + temp7];
          temp = temp*(1 - temp);
          for (c = 0; c < (Npred + 1); c++){ 
	    temp2 = temp*x[i + c*Nind];
	    temp4 = temp3 + c*temp6;
	    for (a = c; a < temp5; a++){
			I[temp4 + a] += temp2*x[i + a*Nind];
	    }
          }
	}
      } else {
	temp7 = (b + 1)*Nind;
	temp8 = (d + 1)*Nind;
	for (i = 0; i < Nind; i++){
	  if (missing[i]) continue;
	  temp =  - p[i + temp7]*p[i + temp8];
	  for (c = 0; c < temp5; c++){
	    temp2 = temp*x[i + c*Nind];
	    temp4 = temp3 + c*temp6;
	    for (a = c; a < temp5; a++){
			I[temp4 + a] += temp2*x[i + a*Nind];
	    }
	  }
	}
      }
    }
  }
 

  for (b = 0; b < Npheno; b++){
    for (d = b; d < Npheno; d++){
      temp3 = b*temp5 + d*temp5*temp6;
      for (c = 0; c < (Npred + 1); c++)
	for (a = 0; a < c; a++){
		I[temp3 + a + c*(Npred + 1)*(Npheno)] = I[temp3 + c + a*(Npred + 1)*(Npheno)];
	}
    }
  }
 
  for (b = 0; b < Npheno; b++)
    for (d = 0; d < b; d++)
      for (a = 0; a < (Npred + 1); a++)
	for (c = 0; c < (Npred + 1); c++){
          I[a + b*(Npred + 1) + c*(Npred + 1)*(Npheno) + d*(Npred + 1)*(Npred + 1)*Npheno] = I[a + d*(Npred + 1) + c*(Npred + 1)*(Npheno) + b*(Npred + 1)*(Npred + 1)*Npheno];
	}
  
  if (bayesian) for (i = 0; i < Npheno*Npheno*(Npred + 1)*(Npred + 1); i++) {
      I[i] += invsigma[i];
    }

}

void MultinomNewton::calcPs (){
	//calculate the probabilities of each phenotype category for each individual conditional on genotypes and confounders

  double temp1, temp2;
  int i, j, k;
  for (i = 0; i < Nind; i++) {
    p[i] = 1;
    temp1 = 1;
    for (k = 1; k < (Npheno + 1); k++){
      temp2 = 0;
      for (j = 0; j < (Npred + 1); j++) { temp2 += x[i + j*Nind]*beta[j + (k-1)*(Npred+1)]; /* cout << "i:" << i << " j:" << j << " k:" << k << " " << x[i + j*Nind] << " * " << beta[j + (k-1)*Npheno] << endl;*/}
      p[i + k*Nind] = exp(temp2);
      temp1 += p[i + k*Nind];
    }
    for (k = 0; k < (Npheno + 1); k++) p[i + k*Nind] *= 1/temp1;
  }

}


void MultinomNewton::NRiter(int * constrain){
	// carry out an interation of the NR algorithm
  
  calcPs(); //calculate probabilities
  calcU(); //calculate gradiant
  calcI(); //calculate hessian
  int i;

  // apply constraints
  for (int i = 0; i < (Npred+1)*Npheno; i++){
    if (constrain[i]) {
     U[i] = 0;
     for (int j = 0; j < (Npred + 1)*Npheno; j++) { I[i + j*(Npred + 1)*Npheno] = 0; I[j + i*(Npred + 1)*Npheno] = 0;}
     I[i + i*(Npred + 1)*Npheno] = 1;
    }
  }

  //solve the matrices to get the update
  int N = (Npred + 1)*Npheno;
  int nrhs = 1;
  int info;
  dgesv_(&N, &nrhs, I, &N, ipiv, U, &N, &info);
  
  //apply the update
  for (int i = 0; i < (Npred + 1)*Npheno; i++){
	  beta[i] += U[i];
  }

}

double MultinomNewton::calcL (){
	//calculate the likelihood
  calcPs();
  double L = 0;
  for (int i = 0; i < Nind; i++) {
    if (missing[i]) continue;
    L += log(p[i + y[i]*Nind]);
  }
  return(L);
}

double MultinomNewton::calcPrior (int *constrain){
	//calculate the value of the prior
  int i, j, k;
  double detcon = 1;
  int consum = 0, consi = 0, consj;

  for (i = 0; i < (Npred + 1)*Npheno; i++) if (constrain[i] == 0) consum++;
  
  for (i = 0; i < (Npred + 1)*Npheno; i++){
    if (constrain[i] == 0){
      consj = 0;
      tempbeta[consi] = beta[i];
      for (j = 0; j < (Npred + 1)*Npheno; j++){
	if (constrain[j] == 0) {
	  tempsigma[consi + consum*consj] = sigma[i + (Npred + 1)*Npheno*j];
	  consj++;
	}
      }
      consi++;
    }
  }
  detcon = inverse(tempsigma,consum);

  double P = -log(2*PI)*(consum)*0.5 - 0.5*log(detcon);
  multiplyVec(tempsigma, tempbeta, priortemp,consum);
  for (i = 0; i < consum; i++) P -= 0.5*priortemp[i]*tempbeta[i];

  return(P);
}

