#include <stdio.h>
#include <math.h>
#include <stdlib.h>
#include <memory.h>
#include <assert.h>
#include "main.h"
#include <gsl/gsl_errno.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_multimin.h>

static int maxarg1, maxarg2;
#define FMAX(a,b) (maxarg1=(a), maxarg2=(b), (maxarg1)>(maxarg2) ? (maxarg1) : (maxarg2))

double **ppLikMat;
int **ppDataMat;
int *pDataInfo;
int numobs, numdata;


int MemAlloc() {

  int i;

  ppLikMat = (double**)malloc((numobs+1)*sizeof(double*));
  if (ppLikMat==NULL) {
    printf("Can't memalloc ppLikMat\n");
    return 0;
  }
  for (i=0; i<(numobs+1); i++) {
    ppLikMat[i]=(double*)malloc((2*numdata+1)*sizeof(double));
    if (ppLikMat[i]==NULL) {
      printf("Can't memalloc ppLikMat[%d]\n",i);
      return 0;
    }
  }

  ppDataMat = (int**)malloc((numobs+1)*sizeof(int*));
  if (ppDataMat==NULL) { 
    printf("Can't memalloc ppDataMat\n"); 
    return 0; 
  }
  for (i=0; i<(numobs+1); i++) {
    ppDataMat[i]=(int*)malloc(3*sizeof(int));
    if (ppDataMat[i]==NULL) { 
      printf("Can't memalloc ppDataMat[%d]\n",i); 
      return 0;
    }
  }

  pDataInfo = (int*)malloc(numdata*sizeof(int));
  if (pDataInfo==NULL) {
    printf("Can't memalloc pDataInfo\n");
    return 0;
  }

  return 1;

}


/*** Function to compute the log likelihood ***/

double my_f(const gsl_vector *v, void *params) {

  int i, j=0;
  int flag=0;
  int cumobs = pDataInfo[0];
  double loglik=0.0;
  double sum=0.0;
  double paramvec[5];

  for (i=0; i<5; i++) paramvec[i] = gsl_vector_get(v, i);

  for (i=0; i<5; i++) if (paramvec[i]<0.0 || paramvec[i]>1.0) flag=1;
  
  if (flag==1) {
    for (i=0;i<5;i++) {
      sum+=fabs(paramvec[i])*10000000000.0;
    }
    return sum;
    }


  for (i=0; i<pDataInfo[0]; i++) {

      ppLikMat[i][0] = ppDataMat[i][2]*log(pow(1-paramvec[2*j],ppDataMat[i][1])*pow(paramvec[2*j],ppDataMat[i][0]-ppDataMat[i][1])*paramvec[2*numdata]+pow(paramvec[2*j+1],ppDataMat[i][1])*pow(1-paramvec[2*j+1],ppDataMat[i][0]-ppDataMat[i][1])*(1-paramvec[2*numdata]));

  }
  
  for (j=1; j<numdata; j++) {

    cumobs += pDataInfo[j];
    
    for (i=pDataInfo[j-1]; i<cumobs; i++) {
      ppLikMat[i][0] = ppDataMat[i][2]*log(pow(1-paramvec[2*j],ppDataMat[i][1])*pow(paramvec[2*j],ppDataMat[i][0]-ppDataMat[i][1])*paramvec[2*numdata]+pow(paramvec[2*j+1],ppDataMat[i][1])*pow(1-paramvec[2*j+1],ppDataMat[i][0]-ppDataMat[i][1])*(1-paramvec[2*numdata]));
    }

  }

  for (i=0; i<numobs; i++) {
    /*printf("%f ",ppLikMat[i][0]);*/
    loglik += ppLikMat[i][0];
  }

  /* printf("The log likelihood is %f\n",loglik);
     fflush(0);*/

  for (i=0; i<numobs; i++) {
    for (j=0; j<2*numdata+1; j++) {
      ppLikMat[i][j] = 0.0;
    }
  }

  return -loglik;

}


int main() {

  FILE *datafile;
  int k, j;
  int temp, temp0, temp1, temp2, temp3;
  int numpar;
  double testlik;
  float startval;
  char inputfile[20];

  size_t np = 5;
  double par[2] = {1.0, 2.0};

  const gsl_multimin_fminimizer_type *T = gsl_multimin_fminimizer_nmsimplex;
  gsl_multimin_fminimizer *s = NULL;
  gsl_vector *ss, *x;
  gsl_multimin_function minex_func;

  size_t iter = 0, i;
  int rval = GSL_CONTINUE;
  int status = GSL_SUCCESS;
  double ssval;

  printf("\n\nPlease enter the name of the file containing your data: ");
  scanf("%s",inputfile);
  datafile = fopen(inputfile,"r");

  printf("\n\n\n        Program to find MLEs of the parameters nu, phi, and rho\n");
  printf("        using the Nelder-Mead simplex algorithm for two data sets \n\n");

  fscanf(datafile,"%d %d",&numdata, &numobs);

  if (!MemAlloc()) {
    printf("Can't allocate MemAlloc\n");
    exit(1);
  }

  for (k=0;k<numdata; k++) {
    fscanf(datafile,"%d",&temp);
    pDataInfo[k] = temp;
  }

  printf("Input data was read as follows:\n\n  There are %d datasets with \n",numdata);
  for (k=0; k<numdata; k++) {
   printf("    %d observations in dataset %d\n",pDataInfo[k],k+1);
  }
  printf("\n\n");
  fflush(0);

  for (k=0; k<numobs; k++) {
    fscanf(datafile,"%d %d %d %d",&temp0,&temp1,&temp2,&temp3);
 
    ppDataMat[k][0] = temp1;
    ppDataMat[k][1] = temp2;
    ppDataMat[k][2] = temp3;

  }

  printf("The data matrix is \n");
  for (k=0;k<numobs;k++) printf("%d %d %d\n",ppDataMat[k][0],ppDataMat[k][1],ppDataMat[k][2]);


  /* Initialize ppLikMat */

  for (k=0; k<numobs; k++) {
    for (j=0; j<2*numdata+1; j++) {
      ppLikMat[k][j] = 0.0;
     }
   }


 
  /******** GSL code to find MLE's **********/
	    
	  
/* Initial vertex size vector */
	    
  ss = gsl_vector_alloc (np);

  if (ss == NULL)
    {
      GSL_ERROR_VAL ("failed to allocate space for ss", GSL_ENOMEM, 0);
    }
  
  gsl_vector_set_all (ss, 1.0);

  /* Starting point */

  x = gsl_vector_alloc (np);

  if (x == NULL)
    {
      gsl_vector_free(ss);
      GSL_ERROR_VAL ("failed to allocate space for x", GSL_ENOMEM, 0);
    }
  
  printf("Enter a starting value for nu in dataset 1: ");
  scanf("%f",&startval);
  gsl_vector_set(x, 0, startval);

  printf("Enter a starting value for phi in dataset 1: ");
  scanf("%f",&startval);
  gsl_vector_set(x, 1, startval);

  printf("Enter a starting value for nu in dataset 2: ");
  scanf("%f",&startval);
  gsl_vector_set(x, 2, startval);

  printf("Enter a starting value for phi in dataset 2: ");
  scanf("%f",&startval);
  gsl_vector_set(x, 3, startval);

  printf("Enter a starting value for rho: ");
  scanf("%f",&startval);
  gsl_vector_set(x, 4, startval);


  /* Initialize method and iterate */

  minex_func.f = &my_f;
  minex_func.n = np;
  minex_func.params = (void *)&par;

  s = gsl_multimin_fminimizer_alloc (T, np);
  gsl_multimin_fminimizer_set (s, &minex_func, x, ss);

  printf("Iterations beginning .....\n\n");
  printf("Iter. #    Nu-1      Phi-1      Nu-2       Phi-2      Rho     -Log Likelihood     Simplex size\n");

  while (rval == GSL_CONTINUE)
    {

      iter++;
      status = gsl_multimin_fminimizer_iterate(s);

      if (status) printf("error: %s\n", gsl_strerror (status));
      fflush(0);

      if (status) 
        break;

      rval = gsl_multimin_test_size (gsl_multimin_fminimizer_size (s), 1e-6);
      ssval = gsl_multimin_fminimizer_size (s);

      if (rval == GSL_SUCCESS)
        printf ("converged to a local maximum at\n");

      printf("%5d ", iter);
      for (i = 0; i < np; i++)
        {
          printf ("%10.5f ", gsl_vector_get (s->x, i));
        }
      printf("f() = %-10.5f ssize = %.7f\n", s->fval, ssval);
    }
  
printf("\n\n Please note: Program should be run many times with varying starting points to detemine global maximum\n\n");

  gsl_vector_free(x);
  gsl_vector_free(ss);
  gsl_multimin_fminimizer_free (s);

  return status;
}