/*  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/>.*/


using namespace std;

class PropOddsNewton {
	// a class for fitting a proportional odds model to genetic data
 public:
  double *x;
  int *y;
  double *beta;
  double *initTau;
  int *ipiv;
  double *p;
  double *U;
  double *I;
  double *c;
  int * missing;
  int Npred;
  int Nind;
  int Npheno;
  PropOddsNewton(int Npred, int Nind, int Npheno);
  ~PropOddsNewton();
  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 ();
  void zeroBetas();
};


void PropOddsNewton::setMissing(){
  // record which samples have either missing phenotype or missing genotype data
  int i, j;
  for (i = 0 ; i < Nind; i++){
    missing[i] = 0;
    if (y[i] < 0) missing[i] = 1;
    if (x[i] < 0) missing[i] = 1;
    for (j = 1; j < Npred; j++){
      if (custom_isnan(x[i + j*Nind])) missing[i] = 1;
    }
    // cout << i << " " << y[i] << " " << x[i] << " " << missing[i] << endl;
  }
      
}

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

void PropOddsNewton::printStuff(){
	//print some debugging information (not used under normal circumstances)
  
  int i, j;

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



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

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

}


PropOddsNewton::PropOddsNewton(int initNpred, int initNind, int initNpheno){
	//initialize the object (allocate space for varaibles and initailize them)
  Npred = initNpred;
  Nind = initNind;
  Npheno = initNpheno;
  int i;
  x = new double [Nind*Npred];
  missing = new int [Nind];
  for (i = 0 ; i < Nind; i++){
	  missing[i] = 0;
  }
  y = new int [Nind];
  beta = new double [Npheno + Npred];
  initTau = new double [Npheno];
  ipiv = new int [Npheno + Npred + 1];
  for (i = 0; i < (Npheno + Npred); i++){ beta[i] = 0;}
  p = new double [(Npheno + 1) * Nind];
  U = new double [Npheno + Npred];
  I = new double [(Npheno + Npred)*(Npheno + Npred)];
  c = new double [Npheno + Npred];
}


PropOddsNewton::~PropOddsNewton(){
	// clean up all the allocated data
  delete x;
  delete y;
  delete beta;
  delete p;
  delete U;
  delete I;
  delete c;
}

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

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

void PropOddsNewton::setCovars(covars *cov){
	// set the covariate data from a covar object
  for (int i = 0; i < Nind; i++) {
    for (int j = 0; j < cov->Ncovars ; j++)   x[i + (j+1)*Nind] = cov->covariates[j][i];
  }
}

void PropOddsNewton::setPheno(phenos *phen, int phenofocus){
	// set the phenotype data using the phenofocus-th phenotype in a phenotype object
  int i, counter=0;
  double cumul = 0;
  for (i = 0; i < Npheno; i++){
    initTau[i] = 0;
  }
  for (i = 0; i < Nind; i++) {
    y[i] = phen->phenotypes[phenofocus][i];
    if (y[i] >= 0){
	counter++;
	if (y[i] < Npheno) initTau[y[i]] += 1.0;
   }
  }
    
  for (i = 0; i < Npheno; i++){
	  cumul += initTau[i]/counter;
    initTau[i] = log(cumul/(1 - cumul));
  }
 
}

double PropOddsNewton::fitLogistic(int * constrain, int maxIter, double relTol, double absTol){
	//fit the model by Newton's method
	
  int i, j;
  double maxTol1, maxTol2;
  
  for (i = 0 ; i < maxIter; i++){
    NRiter(constrain);
    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]);
  		}
    }
    if (maxTol1 < relTol || maxTol2 < absTol){
      break;
    }
  }
  
  return(min(maxTol1,maxTol2));
}
  

void PropOddsNewton::calcU (){
	// calculate the gradiant
  int i, j, k, a;
  int tempj;
  double temp;
  for (a = 0; a < Npheno; a++){
    U[a] = 0;
    for (i = 0; i < Nind; i++){
      if (missing[i]) continue;
      if (y[i] == a) {temp = 1; tempj = a;} 
      else if (y[i] == (a + 1)){ temp = -1; tempj = a + 1;} 
      else continue;
      
      if (tempj == 0) temp *= p[i + 0*Nind];
      else if (tempj == Npheno) temp *= 1 - p[i + (Npheno - 1)*Nind];
      else temp *= p[i + tempj*Nind] - p[i + (tempj-1)*Nind];
      
      if (temp == 0) continue;
      U[a] += p[i + a*Nind]*(1 - p[i + a*Nind])/temp;
    }
  }
  for (j = Npheno; j < (Npheno + Npred); j++) U[j] = 0;
  for (i = 0; i < Nind; i++){
    if (missing[i]) continue;
    if (y[i] == 0) temp = (1 - p[i + 0*Nind]);
    else if (y[i] == Npheno) temp = -p[i + (Npheno - 1)*Nind];
    else temp = (1 - p[i + y[i]*Nind] - p[i + (y[i]-1)*Nind]);
    for (j = Npheno; j < (Npheno + Npred); j++) U[j] += temp*x[i + (j - Npheno)*Nind];
  }

}  

void PropOddsNewton::calcI(){
	//calculate the hessian

  int a, b, c, d, i, j;
  double temp1, temp2;


  for (a = 0; a < (Npheno + Npred); a++)
    for (b = 0; b < (Npheno + Npred); b++) 
      I[a + b*(Npheno + Npred)] = 0;

  // calc d2L/dTau_a dTau_b
  for (a = 0; a < Npheno; a++){
    for (b = 0; b < Npheno; b++){
      if (a > b){
	I[a + b*(Npheno + Npred)] = I[b + a*(Npheno + Npred)];
	continue;
      }
      if (a == b){
	for (i = 0; i < Nind; i++){
	  if (missing[i]) continue;
	  temp1 = -p[i + a*Nind]*(1 - p[i + a*Nind]);
	  if (y[i] == a){
	    if (a > 0 && a < Npheno){
	      temp2 = (p[i + a*Nind] - p[i + (a - 1)*Nind])*(p[i + a*Nind] - p[i + (a - 1)*Nind]); 
	      if (temp2 == 0) continue;
	      temp1 *= 1 + (p[i + (a - 1)*Nind]*(1 - p[i + (a - 1)*Nind]))/temp2;
	    }
	  } else if (y[i] == (a + 1)){
	    if (a < Npheno){
	      temp2 = (p[i + (a+1)*Nind] - p[i + a*Nind])*(p[i + (a+1)*Nind] - p[i + a*Nind]);
	      
		  if (temp2 == 0) continue;
	      
		  temp1 *= 1 + (p[i + (a + 1)*Nind]*(1 - p[i + (a + 1)*Nind]))/temp2;
	    }
	  } else {
	    continue;
	  }
	  I[a + b*(Npheno + Npred)] += temp1;
	}
      }
      if (a == b - 1){
	for (i = 0; i < Nind; i++){
	  if (missing[i]) continue;
	  if (y[i] == b){
	    temp2 = (p[i + b*Nind] - p[i + (b - 1)*Nind])*(p[i + b*Nind] - p[i + (b - 1)*Nind]);
	    if (temp2 == 0) continue;
	    I[a + b*(Npheno + Npred)] += p[i + b*Nind]*(1 - p[i + b*Nind])*p[i + (b-1)*Nind]*(1 - p[i + (b-1)*Nind])/temp2;
	  }
	}
      }
    }
  }

  // calc d2L/dBeta_a dBeta_b
  for (i = 0; i < Nind; i++){
    if (missing[i]) continue;
    if (y[i] == 0) temp1 = p[i + 0*Nind]*(1 - p[i + 0*Nind]);
    else if (y[i] == Npheno) temp1 = p[i + (Npheno - 1)*Nind]*(1 - p[i + (Npheno - 1)*Nind]);
    else temp1 = (1 - p[i + y[i]*Nind])*p[i + y[i]*Nind] + p[i + (y[i]-1)*Nind]*(1 - p[i + (y[i]-1)*Nind]);
    for (a = Npheno; a < (Npheno + Npred); a++)
      for (b = Npheno; b < (Npheno + Npred); b++) 
	I[a + b*(Npheno + Npred)] -= temp1*x[i + (a - Npheno)*Nind]*x[i + (b - Npheno)*Nind];
  }

  // calc d2L/ dTau_a dBeta_b

  for (a = 0; a < Npheno; a++){
    for (i = 0; i < Nind; i++){
      if (missing[i]) continue;
      if (y[i] == a || y[i] == (a + 1)){
	temp1 = p[i + a*Nind]*(1 - p[i + a*Nind]);
	for (b = Npheno; b < (Npheno + Npred); b++) I[a + b*(Npred + Npheno)] -= x[i + (b - Npheno)*Nind]*temp1;
      }
    }
  }

  for (b = 0; b < Npheno; b++){
    for (a = Npheno; a < (Npheno + Npred); a++){
      I[a + b*(Npred + Npheno)] = I[b + a*(Npred + Npheno)];
    }
  }
}

void PropOddsNewton::calcPs (){
	// calculate the category probabilities conditional on genotypes and covariates
  double temp1;
  int i, j, k;
  for (i = 0; i < Nind; i++) {
	  for (k = 0; k < Npheno + 1; k++){
		temp1 = beta[k];
      	for (j = 0; j < Npred; j++){
		  temp1 += beta[Npheno + j]*x[i + j*Nind];
		}
        p[i + k*Nind] = 1/(1 + exp(-temp1));
    }
  }
}


void PropOddsNewton::NRiter(int * constrain){
    //carry out an iteration of the algorithm
	
  int i, j;
  calcPs(); 
  calcU(); //calculate the gradiant
  calcI(); //calculate the hessian

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

  // solve the matrices to get the step update
  int N = (Npred + Npheno);
  int nrhs = 1;
  int info;
  dgesv_(&N, &nrhs, I, &N, ipiv, U, &N, &info);

  // update the parameter vector
  for (i = 0; i < (Npred + Npheno); i++){
    beta[i] -= U[i]; //The 0.75 makes it slower, but more stable
  }

}

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