/////////////////////////////////////////////////////////////////////////////////
//
//  Levenberg - Marquardt non-linear minimization algorithm
//  Copyright (C) 2004-06  Manolis Lourakis (lourakis at ics forth gr)
//  Institute of Computer Science, Foundation for Research & Technology - Hellas
//  Heraklion, Crete, Greece.
//
//  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 2 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.
//
/////////////////////////////////////////////////////////////////////////////////

/*******************************************************************************
 * This file implements combined box and linear equation constraints.
 *
 * Note that the algorithm implementing linearly constrained minimization does
 * so by a change in parameters that transforms the original program into an
 * unconstrained one. To employ the same idea for implementing box & linear
 * constraints would require the transformation of box constraints on the
 * original parameters to box constraints for the new parameter set. This
 * being impossible, a different approach is used here for finding the minimum.
 * The trick is to remove the box constraints by augmenting the function to
 * be fitted with penalty terms and then solve the resulting problem (which
 * involves linear constrains only) with the functions in lmlec.c
 *
 * More specifically, for the constraint a<=x[i]<=b to hold, the term C[i]=
 * (2*x[i]-(a+b))/(b-a) should be within [-1, 1]. This is enforced by adding
 * the penalty term w[i]*max((C[i])^2-1, 0) to the objective function, where
 * w[i] is a large weight. In the case of constraints of the form a<=x[i],
 * the term C[i]=a-x[i] has to be non positive, thus the penalty term is
 * w[i]*max(C[i], 0). If x[i]<=b, C[i]=x[i]-b has to be non negative and
 * the penalty is w[i]*max(C[i], 0). The derivatives needed for the Jacobian
 * are as follows:
 * For the constraint a<=x[i]<=b: 4*(2*x[i]-(a+b))/(b-a)^2 if x[i] not in [a, b],
 *                                0 otherwise
 * For the constraint a<=x[i]: -1 if x[i]<=a, 0 otherwise
 * For the constraint x[i]<=b: 1 if b<=x[i], 0 otherwise
 *
 * Note that for the above to work, the weights w[i] should be large enough;
 * depending on your minimization problem, the default values might need some
 * tweaking (see arg "wghts" below).
 *******************************************************************************/

#ifndef LM_REAL // not included by lmblec.c
#error This file should not be compiled directly!
#endif


#define __MAX__(x, y)   (((x)>=(y))? (x) : (y))
#define __BC_WEIGHT__   LM_CNST(1E+04)

#define __BC_INTERVAL__ 0
#define __BC_LOW__      1
#define __BC_HIGH__     2

/* precision-specific definitions */
#define LEVMAR_BOX_CHECK LM_ADD_PREFIX(levmar_box_check)
#define LMBLEC_DATA LM_ADD_PREFIX(lmblec_data)
#define LMBLEC_FUNC LM_ADD_PREFIX(lmblec_func)
#define LMBLEC_JACF LM_ADD_PREFIX(lmblec_jacf)
#define LEVMAR_LEC_DER LM_ADD_PREFIX(levmar_lec_der)
#define LEVMAR_LEC_DIF LM_ADD_PREFIX(levmar_lec_dif)
#define LEVMAR_BLEC_DER LM_ADD_PREFIX(levmar_blec_der)
#define LEVMAR_BLEC_DIF LM_ADD_PREFIX(levmar_blec_dif)
#define LEVMAR_COVAR LM_ADD_PREFIX(levmar_covar)

struct LMBLEC_DATA{
  LM_REAL *x, *lb, *ub, *w;
  int *bctype;
  void (*func)(LM_REAL *p, LM_REAL *hx, int m, int n, void *adata);
  void (*jacf)(LM_REAL *p, LM_REAL *jac, int m, int n, void *adata);
  void *adata;
};

/* augmented measurements */
static void LMBLEC_FUNC(LM_REAL *p, LM_REAL *hx, int m, int n, void *adata)
{
struct LMBLEC_DATA *data=(struct LMBLEC_DATA *)adata;
int nn;
register int i, j, *typ;
register LM_REAL *lb, *ub, *w, tmp;

  nn=n-m;
  lb=data->lb;
  ub=data->ub;
  w=data->w;
  typ=data->bctype;
  (*(data->func))(p, hx, m, nn, data->adata);

  for(i=nn, j=0; i<n; ++i, ++j){
    switch(typ[j]){
      case __BC_INTERVAL__:
        tmp=(LM_CNST(2.0)*p[j]-(lb[j]+ub[j]))/(ub[j]-lb[j]);
        hx[i]=w[j]*__MAX__(tmp*tmp-LM_CNST(1.0), LM_CNST(0.0));
      break;

      case __BC_LOW__:
        hx[i]=w[j]*__MAX__(lb[j]-p[j], LM_CNST(0.0));
      break;

      case __BC_HIGH__:
        hx[i]=w[j]*__MAX__(p[j]-ub[j], LM_CNST(0.0));
      break;
    }
  }
}

/* augmented Jacobian */
static void LMBLEC_JACF(LM_REAL *p, LM_REAL *jac, int m, int n, void *adata)
{
struct LMBLEC_DATA *data=(struct LMBLEC_DATA *)adata;
int nn, *typ;
register int i, j;
register LM_REAL *lb, *ub, *w, tmp;

  nn=n-m;
  lb=data->lb;
  ub=data->ub;
  w=data->w;
  typ=data->bctype;
  (*(data->jacf))(p, jac, m, nn, data->adata);

  /* clear all extra rows */
  for(i=nn*m; i<n*m; ++i)
    jac[i]=0.0;

  for(i=nn, j=0; i<n; ++i, ++j){
    switch(typ[j]){
      case __BC_INTERVAL__:
        if(lb[j]<=p[j] && p[j]<=ub[j])
          continue; // corresp. jac element already 0

        /* out of interval */
        tmp=ub[j]-lb[j];
        tmp=LM_CNST(4.0)*(LM_CNST(2.0)*p[j]-(lb[j]+ub[j]))/(tmp*tmp);
        jac[i*m+j]=w[j]*tmp;
      break;

      case __BC_LOW__: // (lb[j]<=p[j])? 0.0 : -1.0;
        if(lb[j]<=p[j])
          continue; // corresp. jac element already 0

        /* smaller than lower bound */
        jac[i*m+j]=-w[j];
      break;

      case __BC_HIGH__: // (p[j]<=ub[j])? 0.0 : 1.0;
        if(p[j]<=ub[j])
          continue; // corresp. jac element already 0

        /* greater than upper bound */
        jac[i*m+j]=w[j];
      break;
    }
  }
}

/*
 * This function seeks the parameter vector p that best describes the measurements
 * vector x under box & linear constraints.
 * More precisely, given a vector function  func : R^m --> R^n with n>=m,
 * it finds p s.t. func(p) ~= x, i.e. the squared second order (i.e. L2) norm of
 * e=x-func(p) is minimized under the constraints lb[i]<=p[i]<=ub[i] and A p=b;
 * A is kxm, b kx1. Note that this function DOES NOT check the satisfiability of
 * the specified box and linear equation constraints.
 * If no lower bound constraint applies for p[i], use -DBL_MAX/-FLT_MAX for lb[i];
 * If no upper bound constraint applies for p[i], use DBL_MAX/FLT_MAX for ub[i].
 *
 * This function requires an analytic Jacobian. In case the latter is unavailable,
 * use LEVMAR_BLEC_DIF() bellow
 *
 * Returns the number of iterations (>=0) if successful, or an error code (<0) on failure.
 *
 * For more details on the algorithm implemented by this function, please refer to
 * the comments in the top of this file.
 *
 */
int LEVMAR_BLEC_DER(
  void (*func)(LM_REAL *p, LM_REAL *hx, int m, int n, void *adata), /* functional relation describing measurements. A p \in R^m yields a \hat{x} \in  R^n */
  void (*jacf)(LM_REAL *p, LM_REAL *j, int m, int n, void *adata),  /* function to evaluate the Jacobian \part x / \part p */
  LM_REAL *p,         /* I/O: initial parameter estimates. On output has the estimated solution */
  LM_REAL *x,         /* I: measurement vector. NULL implies a zero vector */
  int m,              /* I: parameter vector dimension (i.e. #unknowns) */
  int n,              /* I: measurement vector dimension */
  LM_REAL *lb,        /* I: vector of lower bounds. If NULL, no lower bounds apply */
  LM_REAL *ub,        /* I: vector of upper bounds. If NULL, no upper bounds apply */
  LM_REAL *A,         /* I: constraints matrix, kxm */
  LM_REAL *b,         /* I: right hand constraints vector, kx1 */
  int k,              /* I: number of constraints (i.e. A's #rows) */
  LM_REAL *wghts,     /* mx1 weights for penalty terms, defaults used if NULL */
  int itmax,          /* I: maximum number of iterations */
  LM_REAL opts[4],    /* I: minim. options [\mu, \epsilon1, \epsilon2, \epsilon3]. Respectively the scale factor for initial \mu,
                       * stopping thresholds for ||J^T e||_inf, ||Dp||_2 and ||e||_2. Set to NULL for defaults to be used
                       */
  LM_REAL info[LM_INFO_SZ],
					           /* O: information regarding the minimization. Set to NULL if don't care
                      * info[0]= ||e||_2 at initial p.
                      * info[1-4]=[ ||e||_2, ||J^T e||_inf,  ||Dp||_2, mu/max[J^T J]_ii ], all computed at estimated p.
                      * info[5]= # iterations,
                      * info[6]=reason for terminating: 1 - stopped by small gradient J^T e
                      *                                 2 - stopped by small Dp
                      *                                 3 - stopped by itmax
                      *                                 4 - singular matrix. Restart from current p with increased mu
                      *                                 5 - no further error reduction is possible. Restart with increased mu
                      *                                 6 - stopped by small ||e||_2
                      *                                 7 - stopped by invalid (i.e. NaN or Inf) "func" values. This is a user error
                      * info[7]= # function evaluations
                      * info[8]= # Jacobian evaluations
                      * info[9]= # linear systems solved, i.e. # attempts for reducing error
                      */
  LM_REAL *work,     /* working memory at least LM_BLEC_DER_WORKSZ() reals large, allocated if NULL */
  LM_REAL *covar,    /* O: Covariance matrix corresponding to LS solution; mxm. Set to NULL if not needed. */
  void *adata)       /* pointer to possibly additional data, passed uninterpreted to func & jacf.
                      * Set to NULL if not needed
                      */
{
  struct LMBLEC_DATA data;
  int ret;
  LM_REAL locinfo[LM_INFO_SZ];
  register int i;

  if(!jacf){
    PRINT_ERROR(RCAT("No function specified for computing the Jacobian in ", LEVMAR_BLEC_DER)
      RCAT("().\nIf no such function is available, use ", LEVMAR_BLEC_DIF) RCAT("() rather than ", LEVMAR_BLEC_DER) "()\n");
    return LM_ERROR_NO_JACOBIAN;
  }

  if(!lb && !ub){
    PRINT_ERROR(RCAT(LCAT(LEVMAR_BLEC_DER, "(): lower and upper bounds for box constraints cannot be both NULL, use "),
          LEVMAR_LEC_DER) "() in this case!\n");
    return LM_ERROR_NO_BOX_CONSTRAINTS;
  }

  if(!LEVMAR_BOX_CHECK(lb, ub, m)){
    PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): at least one lower bound exceeds the upper one\n"));
    return LM_ERROR_FAILED_BOX_CHECK;
  }

  /* measurement vector needs to be extended by m */
  if(x){ /* nonzero x */
    data.x=(LM_REAL *)malloc((n+m)*sizeof(LM_REAL));
    if(!data.x){
      PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): memory allocation request #1 failed\n"));
      return LM_ERROR_MEMORY_ALLOCATION_FAILURE;
    }

    for(i=0; i<n; ++i)
      data.x[i]=x[i];
    for(i=n; i<n+m; ++i)
      data.x[i]=0.0;
  }
  else
    data.x=NULL;

  data.w=(LM_REAL *)malloc(m*sizeof(LM_REAL) + m*sizeof(int)); /* should be arranged in that order for proper doubles alignment */
  if(!data.w){
    PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): memory allocation request #2 failed\n"));
    if(data.x) free(data.x);
    return LM_ERROR_MEMORY_ALLOCATION_FAILURE;
  }
  data.bctype=(int *)(data.w+m);

  /* note: at this point, one of lb, ub are not NULL */
  for(i=0; i<m; ++i){
    data.w[i]=(!wghts)? __BC_WEIGHT__ : wghts[i];
    if(!lb) data.bctype[i]=__BC_HIGH__;
    else if(!ub) data.bctype[i]=__BC_LOW__;
    else if(ub[i]!=LM_REAL_MAX && lb[i]!=LM_REAL_MIN) data.bctype[i]=__BC_INTERVAL__;
    else if(lb[i]!=LM_REAL_MIN) data.bctype[i]=__BC_LOW__;
    else data.bctype[i]=__BC_HIGH__;
  }

  data.lb=lb;
  data.ub=ub;
  data.func=func;
  data.jacf=jacf;
  data.adata=adata;

  if(!info) info=locinfo; /* make sure that LEVMAR_LEC_DER() is called with non-null info */
  ret=LEVMAR_LEC_DER(LMBLEC_FUNC, LMBLEC_JACF, p, data.x, m, n+m, A, b, k, itmax, opts, info, work, covar, (void *)&data);

  if(data.x) free(data.x);
  free(data.w);

  return ret;
}

/* Similar to the LEVMAR_BLEC_DER() function above, except that the Jacobian is approximated
 * with the aid of finite differences (forward or central, see the comment for the opts argument)
 */
int LEVMAR_BLEC_DIF(
  void (*func)(LM_REAL *p, LM_REAL *hx, int m, int n, void *adata), /* functional relation describing measurements. A p \in R^m yields a \hat{x} \in  R^n */
  LM_REAL *p,         /* I/O: initial parameter estimates. On output has the estimated solution */
  LM_REAL *x,         /* I: measurement vector. NULL implies a zero vector */
  int m,              /* I: parameter vector dimension (i.e. #unknowns) */
  int n,              /* I: measurement vector dimension */
  LM_REAL *lb,        /* I: vector of lower bounds. If NULL, no lower bounds apply */
  LM_REAL *ub,        /* I: vector of upper bounds. If NULL, no upper bounds apply */
  LM_REAL *A,         /* I: constraints matrix, kxm */
  LM_REAL *b,         /* I: right hand constraints vector, kx1 */
  int k,              /* I: number of constraints (i.e. A's #rows) */
  LM_REAL *wghts,     /* mx1 weights for penalty terms, defaults used if NULL */
  int itmax,          /* I: maximum number of iterations */
  LM_REAL opts[5],    /* I: opts[0-3] = minim. options [\mu, \epsilon1, \epsilon2, \epsilon3, \delta]. Respectively the
                       * scale factor for initial \mu, stopping thresholds for ||J^T e||_inf, ||Dp||_2 and ||e||_2 and
                       * the step used in difference approximation to the Jacobian. Set to NULL for defaults to be used.
                       * If \delta<0, the Jacobian is approximated with central differences which are more accurate
                       * (but slower!) compared to the forward differences employed by default.
                       */
  LM_REAL info[LM_INFO_SZ],
					           /* O: information regarding the minimization. Set to NULL if don't care
                      * info[0]= ||e||_2 at initial p.
                      * info[1-4]=[ ||e||_2, ||J^T e||_inf,  ||Dp||_2, mu/max[J^T J]_ii ], all computed at estimated p.
                      * info[5]= # iterations,
                      * info[6]=reason for terminating: 1 - stopped by small gradient J^T e
                      *                                 2 - stopped by small Dp
                      *                                 3 - stopped by itmax
                      *                                 4 - singular matrix. Restart from current p with increased mu
                      *                                 5 - no further error reduction is possible. Restart with increased mu
                      *                                 6 - stopped by small ||e||_2
                      *                                 7 - stopped by invalid (i.e. NaN or Inf) "func" values. This is a user error
                      * info[7]= # function evaluations
                      * info[8]= # Jacobian evaluations
                      * info[9]= # linear systems solved, i.e. # attempts for reducing error
                      */
  LM_REAL *work,     /* working memory at least LM_BLEC_DIF_WORKSZ() reals large, allocated if NULL */
  LM_REAL *covar,    /* O: Covariance matrix corresponding to LS solution; mxm. Set to NULL if not needed. */
  void *adata)       /* pointer to possibly additional data, passed uninterpreted to func.
                      * Set to NULL if not needed
                      */
{
  struct LMBLEC_DATA data;
  int ret;
  register int i;
  LM_REAL locinfo[LM_INFO_SZ];

  if(!lb && !ub){
    PRINT_ERROR(RCAT(LCAT(LEVMAR_BLEC_DIF, "(): lower and upper bounds for box constraints cannot be both NULL, use "),
          LEVMAR_LEC_DIF) "() in this case!\n");
    return LM_ERROR_NO_BOX_CONSTRAINTS;
  }

  if(!LEVMAR_BOX_CHECK(lb, ub, m)){
    PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): at least one lower bound exceeds the upper one\n"));
    return LM_ERROR_FAILED_BOX_CHECK;
  }

  /* measurement vector needs to be extended by m */
  if(x){ /* nonzero x */
    data.x=(LM_REAL *)malloc((n+m)*sizeof(LM_REAL));
    if(!data.x){
      PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): memory allocation request #1 failed\n"));
      return LM_ERROR_MEMORY_ALLOCATION_FAILURE;
    }

    for(i=0; i<n; ++i)
      data.x[i]=x[i];
    for(i=n; i<n+m; ++i)
      data.x[i]=0.0;
  }
  else
    data.x=NULL;

  data.w=(LM_REAL *)malloc(m*sizeof(LM_REAL) + m*sizeof(int)); /* should be arranged in that order for proper doubles alignment */
  if(!data.w){
    PRINT_ERROR(LCAT(LEVMAR_BLEC_DER, "(): memory allocation request #2 failed\n"));
    if(data.x) free(data.x);
    return LM_ERROR_MEMORY_ALLOCATION_FAILURE;
  }
  data.bctype=(int *)(data.w+m);

  /* note: at this point, one of lb, ub are not NULL */
  for(i=0; i<m; ++i){
    data.w[i]=(!wghts)? __BC_WEIGHT__ : wghts[i];
    if(!lb) data.bctype[i]=__BC_HIGH__;
    else if(!ub) data.bctype[i]=__BC_LOW__;
    else if(ub[i]!=LM_REAL_MAX && lb[i]!=LM_REAL_MIN) data.bctype[i]=__BC_INTERVAL__;
    else if(lb[i]!=LM_REAL_MIN) data.bctype[i]=__BC_LOW__;
    else data.bctype[i]=__BC_HIGH__;
  }

  data.lb=lb;
  data.ub=ub;
  data.func=func;
  data.jacf=NULL;
  data.adata=adata;

  if(!info) info=locinfo; /* make sure that LEVMAR_LEC_DIF() is called with non-null info */
  ret=LEVMAR_LEC_DIF(LMBLEC_FUNC, p, data.x, m, n+m, A, b, k, itmax, opts, info, work, covar, (void *)&data);

  if(data.x) free(data.x);
  free(data.w);

  return ret;
}

/* undefine all. THIS MUST REMAIN AT THE END OF THE FILE */
#undef LEVMAR_BOX_CHECK
#undef LMBLEC_DATA
#undef LMBLEC_FUNC
#undef LMBLEC_JACF
#undef LEVMAR_COVAR
#undef LEVMAR_LEC_DER
#undef LEVMAR_LEC_DIF
#undef LEVMAR_BLEC_DER
#undef LEVMAR_BLEC_DIF