/*
 * Copyright (C) 1996-2011 Daniel Waggoner and Tao Zha
 *
 * This 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.
 *
 * It 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.
 *
 * If you did not received a copy of the GNU General Public License
 * with this software, see <http://www.gnu.org/licenses/>.
 */

#include "dw_switch_opt.h"
#include "dw_std.h"

#include <stdlib.h>
#include <string.h>
#include <math.h>

//#define DW_NPSOL_PRESENT
//#define DW_CSMINWEL_PRESENT

//====== Static Global Variables ======
static TStateModel *Model=(TStateModel*)NULL;
static PRECISION *buffer=(PRECISION*)NULL;
static PRECISION *ModifiedFreeParameters=(PRECISION*)NULL;
static PRECISION *FreeParameters_Q=(PRECISION*)NULL;
static int NumberFreeParameters_Q=0;
static PRECISION *FreeParameters_Theta=(PRECISION*)NULL;
static int NumberFreeParameters_Theta=0;

/* 
   Sets up global variables for optimization.  The buffer used to store the free 
   parameters is allocated and deallocated outside these routines.  The buffers 
   FreeQ and FreeTheta do not have to be contiguous, but FreeQ must be at least
   NumberFreeParametersQ(model)*sizeof(PRECISION) bytes and FreeTheta must be at
   least NumberFreeParametersTheta(model)*sizeof(PRECISION) bytes.  The region
   from Modified + 0 through Modified + n - 1 must be completely contained in 
   FreeQ and FreeTheta.  Here n is the parameter passed to MLEObjectiveFunction() 
   or PosteriorObjectiveFunction().
*/
void SetupObjectiveFunction(TStateModel *model, PRECISION *Modified, PRECISION *FreeQ, PRECISION *FreeTheta)
{
  if (buffer) dw_free(buffer);
  buffer=(PRECISION*)NULL;
  Model=model;
  FreeParameters_Q=FreeQ;
  NumberFreeParameters_Q=NumberFreeParametersQ(model);
  FreeParameters_Theta=FreeTheta;
  NumberFreeParameters_Theta=NumberFreeParametersTheta(model);
  ModifiedFreeParameters=Modified;
}

/* 
   Sets up global variables for optimization.  The buffer used to store the free 
   parameters is allocated and deallocated in this routine.  One of FreeTheta_Idx 
   or FreeQ_Idx must be zero and the other must be NumberFreeParametersQ(model) 
   or NumberFreeParameters(Theta).  Modified_Idx and Modified_Idx + n -1 must be 
   between 0 and NumberFreeParametersQ(model) + NumberFreeParameters(Theta) - 1, 
   inclusive.  Here n is the parameter passed to PosteriorObjectiveFunction() or
   MLEObjectiveFunction().
   
*/
void SetupObjectiveFunction_AllocateMemory(TStateModel *model, int FreeTheta_Idx, int FreeQ_Idx, int Modified_Idx)
{
  if (buffer) dw_free(buffer);
  Model=model;
  NumberFreeParameters_Q=NumberFreeParametersQ(model);
  NumberFreeParameters_Theta=NumberFreeParametersTheta(model);
  buffer=(PRECISION*)dw_malloc((NumberFreeParameters_Q + NumberFreeParameters_Theta)*sizeof(PRECISION));

  FreeParameters_Q=buffer+FreeQ_Idx;
  FreeParameters_Theta=buffer+FreeTheta_Idx;
  ModifiedFreeParameters=buffer+Modified_Idx;
}

PRECISION PosteriorObjectiveFunction(PRECISION *x, int n)
{
  if (x != ModifiedFreeParameters) memmove(ModifiedFreeParameters,x,n*sizeof(PRECISION));
  ConvertFreeParametersToQ(Model,FreeParameters_Q);
  ConvertFreeParametersToTheta(Model,FreeParameters_Theta);
  return -LogPosterior_StatesIntegratedOut(Model);
}

PRECISION PosteriorObjectiveFunction_csminwel(double *x, int n, double **args, int *dims)
{
  return PosteriorObjectiveFunction(x,n);
}

void PosteriorObjectiveFunction_npsol(int *mode, int *n, double *x, double *f, double *g, int *nstate)
{
  *f=PosteriorObjectiveFunction(x,*n);
}

PRECISION MLEObjectiveFunction(PRECISION *x, int n)
{
  if (x != ModifiedFreeParameters) memmove(ModifiedFreeParameters,x,n*sizeof(PRECISION));
  ConvertFreeParametersToQ(Model,FreeParameters_Q);
  ConvertFreeParametersToTheta(Model,FreeParameters_Theta);
  return -LogLikelihood_StatesIntegratedOut(Model);
}

PRECISION MLEObjectiveFunction_csminwel(double *x, int n, double **args, int *dims)
{
  return MLEObjectiveFunction(x,n);
}

void MLEObjectiveFunction_npsol(int *mode, int *n, double *x, double *f, double *g, int *nstate)
{
  *f=MLEObjectiveFunction(x,*n);
}

#ifdef DW_NPSOL_PRESENT
//=== NPSOL declarations ===
#define npsol npsol_
#define npoptn npoptn_
#define npsol_options npsol_options_
void npsol(int *n, int *nclin, int *ncnln, int *ldA, int *ldJ, int *ldR, double *A, double *bl, double *bu,
	   void (*funcon)(int *mode, int *ncnln, int *n, int *ldJ, int *needc, double *x, double *c, double *cJac, int *nstate),
	   void (*funobj)(int *mode, int *n, double *x, double *f, double *g, int *nstate),
	   int *inform, int *iter, int *istate, double *c, double *cJac, double *clamda, double *f, double *g, double *R, double *x,
	   int *iw, int *leniw, double *w, int *lenw);
void npoptn(char *,int);
static void confun(int *mode, int *ncnln, int *n, int *ldJ, int *needc, double *x, double *c, double *cJac, int *nstate)
{
  *mode=-1;
}
int PosteriorOptimizationQ_npsol(TStateModel *model, int max_iteration)
{
  int j, size_theta, pos_theta, size_Q, pos_Q, opt;
  double objective, objective_last, likelihood, prior;
  char  option_string[256];

  // npsol arguments
  int n, nclin=0, ncnln=0, ldA=1, ldJ=1, ldR, nctotl, *istate, *iw, leniw, lenw, inform, iter;
  double A[1], cJac[1], *bl, *bu, *clamda, *R, *x, *w, f, *g, c[1], tolerance;

  //==== Allocate memory  ===
  size_theta=NumberFreeParametersTheta(model);
  size_Q=NumberFreeParametersQ(model);
  pos_theta=0;
  pos_Q=size_theta;

  n=size_Q;
  nctotl=n+nclin+ncnln;
  leniw=3*n + nclin + 2*ncnln;
  if (ncnln == 0)
    if (nclin == 0)
      lenw=20*n;
    else
      lenw=2*n*n + 11*nclin;
  else
    lenw=2*n*n + n*nclin + 2*n*ncnln + 20*n + 11*nclin + 21*ncnln;
  istate=(int*)dw_malloc(nctotl*sizeof(int));
  iw=(int*)dw_malloc(leniw*sizeof(int));
  bl=(double*)dw_malloc(nctotl*sizeof(double));
  bu=(double*)dw_malloc(nctotl*sizeof(double));
  clamda=(double*)dw_malloc(nctotl*sizeof(double));
  R=(double*)dw_malloc(n*n*sizeof(double));
  x=(double*)dw_malloc((size_Q + size_theta)*sizeof(double));
  g=(double*)dw_malloc(n*sizeof(double));
  w=(double*)dw_malloc(lenw*sizeof(double));
  //====================================================================

  //=== Set starting value ===
  ConvertQToFreeParameters(model,x+pos_Q);
  ConvertThetaToFreeParameters(model,x+pos_theta);

  SetupObjectiveFunction(model,x+pos_Q,x+pos_Q,x+pos_theta);
  ldR=n;
  for (j=0; j < n; j++)
    {
      bl[j]=0.0;
      bu[j]=1.0;
    }

  //=== Set NPSOL options
  npoptn("Derivative level = 0",20);
  sprintf(option_string,"Optimality tolerance = %le",tolerance);
  npoptn(option_string,strlen(option_string));
  sprintf(option_string,"Major iterations limit = %d",max_iteration);
  npoptn(option_string,strlen(option_string));

  npsol(&n,&nclin,&ncnln,&ldA,&ldJ,&ldR,A,bl,bu,&confun,&PosteriorObjectiveFunction_npsol,
	&inform,&iter,istate,c,cJac,clamda,&f,g,R,x+pos_Q,iw,&leniw,w,&lenw);

  ConvertFreeParametersToQ(model,x+pos_Q);
  ConvertFreeParametersToTheta(model,x+pos_theta);

  //=== Free memory ===
  dw_free(istate);
  dw_free(iw);
  dw_free(bl);
  dw_free(bu);
  dw_free(clamda);
  dw_free(R);
  dw_free(x);
  dw_free(g);
  dw_free(w);
}
#endif

/* #ifdef DW_CSMINWEL_PRESENT */
/* /\* */
/*    Maximizes the posterior kernel with respect to a subblock of theta or q given  */
/*    by offset and length.  It must be the case that offset plus length be less than */
/*    or equal to the number of free parameters in theta or q.  If flag is 1, then */
/*    subblocks of theta are used; if flag is 2, then subblocks of q are used, if */
/*    flag is 3, then maximization over both theta and q is performed. */

/*    Returns 1 upon completion and 0 if the arguments are invalid. */
/* *\/ */
/* int PosteriorOptimization_csminwel(TStateModel *model, int offset, int length, int flag, int max_iteration, PRECISION tolerance) */
/* { */
/*   int size_theta, pos_theta, size_Q, pos_Q; */

/*   // csminwel arguments  */
/*   int itct, fcount, retcodeh; */
/*   double *x, fh; */
/*   TMatrix H; */
/*   TVector g; */

/*   //==== Sizes */
/*   size_theta=NumberFreeParametersTheta(model); */
/*   size_Q=NumberFreeParametersQ(model); */
/*   pos_theta=0; */
/*   pos_Q=size_theta; */
/*   switch (flag) */
/*     { */
/*     case 1: */
/*       offset+=pos_theta; */
/*       if (offset+length > size_theta) return 0; */
/*       break; */
/*     case 2: */
/*       offset+=pos_Q; */
/*       if (offset+length > size_Q) return 0; */
/*       break; */
/*     case 3: */
/*       offset=0; */
/*       length=size_Q+size_theta; */
/*     } */

/*   //=== Allocate memory */
/*   x=(PRECISION*)dw_malloc((size_Q + size_theta)*sizeof(double)); */
/*   g=CreateVector(length); */
/*   H=IdentityMatrix((TMatrix)NULL,length); */
/*   ProductMS(H,H,1e-5); */

/*   //=== Setup objective function */
/*   SetupObjectiveFunction(model,x+offset,x+pos_Q,x+pos_theta); */

/*   //=== Set starting value */
/*   ConvertQToFreeParameters(model,x+pos_Q); */
/*   ConvertThetaToFreeParameters(model,x+pos_theta); */

/*   //=== Optimize */
/*   csminwel(PosteriorObjectiveFunction_csminwel,x+offset,length,pElementM(H),pElementV(g),NULL, */
/* 	   &fh,tolerance,&itct,max_iteration,&fcount,&retcodeh,NULL,NULL); */

/*   //=== Get terminal values */
/*   ConvertFreeParametersToQ(model,x+pos_Q); */
/*   ConvertFreeParametersToTheta(model,x+pos_theta); */

/*   //=== Free memory === */
/*   FreeVector(g); */
/*   FreeMatrix(H); */
/*   dw_free(x); */

/*   return 1; */
/* } */
/* #endif */

#ifdef DW_CSMINWEL_PRESENT

#include "csminwel.h"

/*
   Maximizes the posterior kernel with respect to a subblock of theta or q given 
   by offset and length.  It must be the case that offset plus length be less than
   or equal to the number of free parameters in theta or q.  If flag is 1, then
   subblocks of theta are used; if flag is 2, then subblocks of q are used, if
   flag is 3, then maximization over both theta and q is performed.

   Returns 1 upon completion and 0 if the arguments are invalid.
*/
int PosteriorOptimization_csminwel(TStateModel *model, int offset, int length, int flag, int max_iteration, 
					    PRECISION tolerance, PRECISION stepsize, PRECISION hessian_scale)
{
  int size_theta, pos_theta, size_Q, pos_Q;
  PRECISION *parameters;

  // csminwel arguments 
  int itct, fcount, retcodeh;
  double *x, fh, original_stepsize=GRADSTPS_CSMINWEL;
  TMatrix H;
  TVector g;

  //==== Sizes
  size_theta=NumberFreeParametersTheta(model);
  size_Q=NumberFreeParametersQ(model);
  pos_theta=0;
  pos_Q=size_theta;
  if (flag == 3)
    if (length > 0)
      {
	if (offset+length > size_theta + size_Q) return 0;
      }
    else
      {
	offset=0;
	length=size_Q + size_theta;
      }
  else if (flag == 2)
    if (length > 0)
      {
	offset+=pos_Q;
	if (offset+length > size_Q) return 0;
      }
    else
      {
	offset=pos_Q;
	length=size_Q;
      }
  else if (flag == 1)
    if (length > 0)
      {
	offset+=pos_theta;
	if (offset+length > size_theta) return 0;
      }
    else
      {
	offset=pos_theta;
	length=size_theta;
      }
  else
    {
      printf("PosteriorOptimization_csminwel(): Unknown flag\n");
      dw_exit(0);
    }

  //=== Allocate memory and set Hessian scale and step size
  parameters=(PRECISION*)dw_malloc((size_Q + size_theta)*sizeof(double));
  x=(PRECISION*)dw_malloc(length*sizeof(PRECISION));
  g=CreateVector(length);
  H=IdentityMatrix((TMatrix)NULL,length);
  ProductMS(H,H,hessian_scale);
  GRADSTPS_CSMINWEL=stepsize;

  //=== Setup objective function
  SetupObjectiveFunction(model,parameters+offset,parameters+pos_Q,parameters+pos_theta);

  //=== Set starting value
  ConvertQToFreeParameters(model,parameters+pos_Q);
  ConvertThetaToFreeParameters(model,parameters+pos_theta);
  memcpy(x,parameters+offset,length*sizeof(PRECISION));

  //=== Optimize
  csminwel(PosteriorObjectiveFunction_csminwel,x,length,pElementM(H),pElementV(g),NULL,
	   &fh,tolerance,&itct,max_iteration,&fcount,&retcodeh,NULL,NULL);

  //=== Get terminal values
  ConvertFreeParametersToQ(model,parameters+pos_Q);
  ConvertFreeParametersToTheta(model,parameters+pos_theta);

  //=== Free memory
  FreeMatrix(H);
  FreeVector(g);
  dw_free(x);
  dw_free(parameters);

  //=== Reset global step size
  GRADSTPS_CSMINWEL=original_stepsize;

  return 1;
}

int Estimate_csminwel(TStateModel *model, int iter_block, int iter_start, PRECISION iter_inc, 
			       PRECISION crit_start, PRECISION crit_end, PRECISION crit_inc,
			       PRECISION stepsize_theta, PRECISION hessian_scale_theta,
			       PRECISION stepsize_q, PRECISION hessian_scale_q, 
			       PRECISION stepsize_full, PRECISION hessian_scale_full,FILE *f_out)
{
  int nq, ntheta, i, total_iteration=1, iteration, iter=iter_start;
  PRECISION crit=crit_start, posterior, posterior_last, likelihood, likelihood_last, original_stepsize=GRADSTPS_CSMINWEL;
  TVector theta=(TVector)NULL, q=(TVector)NULL;
  time_t *currenttime;
   
  nq=NumberFreeParametersQ(model);
  ntheta=NumberFreeParametersTheta(model);
  if (nq > 0) q=CreateVector(nq);
  if (ntheta > 0) theta=CreateVector(ntheta);

  posterior=LogPosterior_StatesIntegratedOut(model);

  if (f_out)
    {
      fprintf(f_out,"\n\n//=== Initial Value ===//\n");
      fprintf(f_out,"likelihood value:  %22.14le\n",LogLikelihood_StatesIntegratedOut(model));
      fprintf(f_out,"posterior value:  %22.14le\n",posterior);
      if (ntheta > 0) 
	{
	  fprintf(f_out,"theta:\n");
	  ConvertThetaToFreeParameters(model,pElementV(theta));
	  dw_PrintVector(f_out,theta,"%.13le ");
	}
      if (nq > 0)
	{
	  fprintf(f_out,"q:\n");
	  ConvertQToFreeParameters(model,pElementV(q));
	  dw_PrintVector(f_out,q,"%.13le ");
	}
      fflush(f_out);
    }

  currenttime=(time_t*)dw_malloc(4*sizeof(time_t));

  for ( ; crit >= crit_end; crit*=crit_inc, iter=(int)((PRECISION)iter*iter_inc))
    {
      for (iteration=1; iteration <= iter_block; iteration++)
	{
	  posterior_last=posterior;
	  time(currenttime);
        
	  // optimize over theta for model i
	  PosteriorOptimization_csminwel(model,0,0,1,iter,crit,stepsize_theta,hessian_scale_theta);
	  time(currenttime+1);

	  // optimize over q
	  PosteriorOptimization_csminwel(model,0,0,2,iter,crit,stepsize_q,hessian_scale_q);
	  time(currenttime+2);

	  // full optimization
	  PosteriorOptimization_csminwel(model,0,0,3,iter,crit,stepsize_full,hessian_scale_full); 
	  time(currenttime+3);

	  posterior=LogPosterior_StatesIntegratedOut(model);

	  if (f_out)
	    {
	      fprintf(f_out,"\n//=== Iteration %d ===//\n",total_iteration);
	      fprintf(f_out,"Criterion/Max Iteration:  %le  %d\n",crit,iter);
	      fprintf(f_out,"Elapsed time for theta: %0.4f\n",difftime(currenttime[1],currenttime[0]));
	      fprintf(f_out,"Elapsed time for q: %0.4f\n",difftime(currenttime[2],currenttime[1]));
	      fprintf(f_out,"Elapsed time for full: %0.4f\n",difftime(currenttime[3],currenttime[2]));
	      fprintf(f_out,"likelihood value:  %22.14le\n",LogLikelihood_StatesIntegratedOut(model));
	      fprintf(f_out,"posterior value:  %22.14le\n",posterior);
	      fprintf(f_out,"%s",ctime(&currenttime[3]));
	      if (ntheta > 0) 
		{
		  fprintf(f_out,"theta:\n");
		  ConvertThetaToFreeParameters(model,pElementV(theta));
		  dw_PrintVector(f_out,theta,"%.13le ");
		}
	      if (nq > 0)
		{
		  fprintf(f_out,"q:\n");
		  ConvertQToFreeParameters(model,pElementV(q));
		  dw_PrintVector(f_out,q,"%.13le ");
		}
	      fflush(f_out);
	    }

	  total_iteration++;
	  if (fabs(posterior - posterior_last) <= crit) break;
	}
    }

  FreeVector(q);
  FreeVector(theta);
  dw_free(currenttime);

  return 1;
}
#endif



