#include "src/complementarity/impls/ssls/ssls.h"

// Context for ISXLS
//   -- interior-point  - smoothed fisched function
//   -- active-set      - reduced matrices formed
//			  - inherit properties of original system
//   -- semismooth (S)  - function not differentiable
//                      - merit function continuously differentiable
//                      - Fischer-Burmeister reformulation of complementarity
//                        - Billups composition for two finite bounds
//   -- infeasible (I)  - iterates not guaranteed to remain within bounds
//   -- feasible (F)    - iterates guaranteed to remain within bounds
//   -- linesearch (LS) - Armijo rule on direction
//
// Many other reformulations are possible and combinations of 
// feasible/infeasible and linesearch/trust region are possible.
//
// Basic theory 
//   Fischer-Burmeister reformulation is semismooth with a continuously
//   differentiable merit function and strongly semismooth if the F has
//   lipschitz continuous derivatives.
//
//   Every accumulation point generated by the algorithm is a stationary
//   point for the merit function.  Stationary points of the merit function
//   are solutions of the complementarity problem if
//     a.  the stationary point has a BD-regular subdifferential, or
//     b.  the Schur complement F'/F'_ff is a P_0-matrix where ff is the
//         index set corresponding to the free variables.
//
//   If one of the accumulation points has a BD-regular subdifferential then
//     a.  the entire sequence converges to this accumulation point at
//         a local q-superlinear rate
//     b.  if in addition the reformulation is strongly semismooth near
//         this accumulation point, then the algorithm converges at a
//         local q-quadratic rate.
//
// The theory for the feasible version follows from the feasible descent
// algorithm framework.
//
// References:
//   Billups, "Algorithms for Complementarity Problems and Generalized
//     Equations," Ph.D thesis, University of Wisconsin - Madison, 1995.
//   De Luca, Facchinei, Kanzow, "A Semismooth Equation Approach to the
//     Solution of Nonlinear Complementarity Problems," Mathematical
//     Programming, 75, pages 407-439, 1996.
//   Ferris, Kanzow, Munson, "Feasible Descent Algorithms for Mixed
//     Complementarity Problems," Mathematical Programming, 86,
//     pages 475-497, 1999.
//   Fischer, "A Special Newton-type Optimization Method," Optimization,
//     24, pages 269-284, 1992
//   Munson, Facchinei, Ferris, Fischer, Kanzow, "The Semismooth Algorithm
//     for Large Scale Complementarity Problems," Technical Report 99-06,
//     University of Wisconsin - Madison, 1999.

/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "Tao_ISLS_SFunction"
int Tao_ISLS_SFunction(TAO_SOLVER tao, TaoVec *X, double *fcn, void *solver) 
{
  TAO_SSLS *isls = (TAO_SSLS *)solver;
  TaoVec *l, *u;
  int info;

  info = TaoGetVariableBounds(tao, &l, &u); CHKERRQ(info);

  info = TaoComputeConstraints(tao, X, isls->f); CHKERRQ(info);
  info = isls->ff->SFischer(X, isls->f, l, u, isls->mu); CHKERRQ(info);
  info = isls->ff->Norm2squared(&isls->merit_eqn); CHKERRQ(info);

  // Add in the constraint for mu
  isls->mucon = exp(isls->mu) - 1.0;
  isls->merit_mu = isls->mucon*isls->mucon;

  isls->merit = isls->merit_eqn + isls->merit_mu;

  // Calculate the objective and merit function
  *fcn = 0.5*isls->merit;
  isls->merit = sqrt(isls->merit);
  return 0;
}

/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "Tao_ISLS_SFunctionGradient"
int Tao_ISLS_SFunctionGradient(TAO_SOLVER tao, TaoVec *X, double *fcn, 
                               TaoVec *G, void *solver)
{
  TAO_SSLS *isls = (TAO_SSLS *)solver;
  TaoVec *l, *u;
  TaoMat *J;
  int info;

  info = TaoGetVariableBounds(tao, &l, &u); CHKERRQ(info);
  info = TaoGetJacobian(tao, &J); CHKERRQ(info);

  info = TaoComputeConstraints(tao, X, isls->f); CHKERRQ(info);
  info = isls->ff->SFischer(X, isls->f, l, u, isls->mu); CHKERRQ(info);
  info = isls->ff->Norm2squared(&isls->merit_eqn); CHKERRQ(info);

  // Add in the constraint for mu
  isls->mucon = exp(isls->mu) - 1.0;
  isls->merit_mu = isls->mucon*isls->mucon;

  isls->merit = isls->merit_eqn + isls->merit_mu;

  // Calculate the objective and merit function
  *fcn = 0.5*isls->merit;
  isls->merit = sqrt(isls->merit);

  info = TaoComputeJacobian(tao, X, J); CHKERRQ(info);
  info = J->D_SFischer(X, isls->f, l, u, isls->mu, 
		       isls->t1, isls->t2, 
		       isls->da, isls->db, isls->dm); CHKERRQ(info);
  isls->d_mucon = exp(isls->mu);

  info = isls->t1->PointwiseMultiply(isls->ff, isls->db);
  info = J->MultiplyTranspose(isls->t1, G); CHKERRQ(info);
  info = isls->t1->PointwiseMultiply(isls->ff, isls->da); CHKERRQ(info);
  info = G->Axpy(1.0, isls->t1); CHKERRQ(info);
  info = isls->dm->Dot(isls->ff, &isls->g_mucon); CHKERRQ(info);
  isls->g_mucon += exp(isls->mu)*isls->mucon;
  return 0;
}

/*------------------------------------------------------------*/
#undef __FUNCT__  
#define __FUNCT__ "TaoSolve_ISILS"
static int TaoSolve_ISILS(TAO_SOLVER tao, void *solver)
{
  TAO_SSLS *isls = (TAO_SSLS *)solver;

  TaoVec *x;			// Solution vector
  TaoVec *l, *u;		// Problem bounds
  TaoMat *J;			// Jacobian matrix

  TaoIndexSet *FreeVariableSet;	// Free variables
  TaoIndexSet *FixedVariableSet;// Fixed variables

  TaoMat *Jsub;			// Reduced Jacobian
  TaoVec *r1, *r2, *r3;		// Temporary vectors

//  TaoLinearSolver *lsolver;

  double psi, ndpsi, normd, innerd, t=0;
  double two_norm;
  double sigma = 1e-4, beta = 0.5, minstep = TAO_EPSILON;
  double muval, psival;
  int iter=0, step_type, info, nn, nf;
  TaoTerminateReason reason;
  TaoTruth flag;

  TaoFunctionBegin;

  // Assume that Setup has been called!

  // Start by obtaining the solution vector, Jacobian,
  // and problem size
  info = TaoGetSolution(tao, &x); CHKERRQ(info);
  info = TaoGetJacobian(tao, &J); CHKERRQ(info);
  info = x->GetDimension(&nn); CHKERRQ(info);

  // Evaluate the bounds
  info = TaoGetVariableBounds(tao, &l, &u); CHKERRQ(info);
  info = TaoEvaluateVariableBounds(tao, l, u); CHKERRQ(info);

  // Create vector for derivative wrt smoothing parameter
  info = x->Clone(&(isls->dm)); CHKERRQ(info);

  // Create temporary vectors and index sets (for active set)
  info = x->CreateIndexSet(&FreeVariableSet); CHKERRQ(info);
  info = x->CreateIndexSet(&FixedVariableSet); CHKERRQ(info);
  info = x->Clone(&r1); CHKERRQ(info);
  info = x->Clone(&r2); CHKERRQ(info);
  info = x->Clone(&r3); CHKERRQ(info);

  // Create submatrix
  info = J->CreateReducedMatrix(FreeVariableSet, FreeVariableSet, &Jsub); CHKERRQ(info);

  // Initialize the merit function
  //   Used in the linesearch
  //   Termination based on different function
  info = TaoSetMeritFunction(tao, Tao_ISLS_SFunction, Tao_ISLS_SFunctionGradient,
                             TAO_NULL, TAO_NULL, isls); CHKERRQ(info);

  // Initialize the value for the smoothing parameter
  isls->mu = isls->mu_init;

#if 0
  // Choice of mu:
  //   Simple search

  info = TaoComputeMeritFunctionGradient(tao, x, &psi, isls->dpsi); CHKERRQ(info);

  psival = fabs(isls->merit_eqn - nn*isls->merit_mu);
  printf("mu: %5.4e merit_eqn: %5.4e merit_mu: %5.4e\n", 
	 isls->mu, isls->merit_eqn, nn*isls->merit_mu);

  double fact;
  if (isls->merit_eqn < nn*isls->merit_mu) {
    fact = 0.5;
  } else
    {
      fact = 2.0;
    }

  isls->mu *= fact;
  info = TaoComputeMeritFunctionGradient(tao, x, &psi, isls->dpsi); CHKERRQ(info);
  printf("mu: %5.4e merit_eqn: %5.4e merit_mu: %5.4e\n", 
	 isls->mu, isls->merit_eqn, nn*isls->merit_mu);
  
  while (fabs(isls->merit_eqn - nn*isls->merit_mu) < psival) {
    psival = fabs(isls->merit_eqn - nn*isls->merit_mu);

    isls->mu *= fact;
    info = TaoComputeMeritFunctionGradient(tao, x, &psi, isls->dpsi); CHKERRQ(info);
    printf("mu: %5.4e merit_eqn: %5.4e merit_mu: %5.4e\n", 
	   isls->mu, isls->merit_eqn, nn*isls->merit_mu);
  }

  isls->mu /= fact;
  printf("mu_choice : %5.4e\n", isls->mu);
#endif

  // DONE:
  //   1.  Eliminated smoothing parameter vector.  Do not know how to deal
  //       with mixture of differentiable and nondifferentiable functions.
  //       Maybe let the smoothing parameter be a vector later.
  //   2.  Rewrote smoothed functions to provide the following information:
  //       - derivative of the function with respect to smoothing parameter
  //       - gives nonsmooth function when smoothing parameter == 0
  //   3.  Rewrote the function and function/gradient routines to add
  //       the smoothing constraint (e^{mu}-1 = 0) and the appropriate
  //       derivatives.
  //   4.  Modified the right-hand side of the direction calculation to
  //       account for the smoothing constraint.  I.e. the direction in 
  //       the smoothing parameter is substituted out of the problem.
  //   5.  Modified descent direction test to deal with the augmented
  //       problem with the smoothing parameter.
  //   6.  Hacked in a monotone linesearch to deal with the augmented system.
  //   7.  When smoothing parameter == 0, make sure we get the same 
  //       behavior as asils.
  // TODO:
  //   8.  Add method to optimally choose initial value for mu.

  
  // Calculate the fischer function at the initial iterate
  Tao_ASLS_FunctionGradient(tao, x, &psi, isls->dpsi, isls);
  info = isls->dpsi->Norm2(&ndpsi); CHKERRQ(info);

  while (1) {

    // Assert: 
    //   Fischer function for terminatation
    //   Evaluated at current iterate
    //   Norm gradient of Fischer function calculated

    printf("nd_merit : %5.4e\n", isls->merit);
    printf("nd_ndpsi : %5.4e\n", ndpsi);

    // Check the termination criteria

    info = TaoMonitor(tao, iter++, isls->merit, ndpsi, 0.0, t, &reason); CHKERRQ(info);
    if (TAO_CONTINUE_ITERATING != reason) break;

    // Compute the smoothed Fischer function.
    // This forms the basis for the interior-point method.

    info = TaoComputeMeritFunctionGradient(tao, x, &psi, isls->dpsi); CHKERRQ(info);

    printf("merit    : %5.4e\n", isls->merit);
    printf("merit_eqn: %5.4e\n", isls->merit_eqn);
    printf("merit_mu : %5.4e\n", isls->merit_mu);
    printf("mu       : %5.4e\n", isls->mu);
    printf("mucon    : %5.4e\n", isls->mucon);
    printf("d_mucon  : %5.4e\n", isls->d_mucon);
    printf("g_mucon  : %5.4e\n", isls->g_mucon);

#if 0
    // Compute tolerances for linear system solver

    // We want to set tolerances so that we maintain a superlinear
    // asymptotic rate of convergence.  Note: the tolerances set are 
    // for the reduced system.  We really need to make sure that the 
    // full system satisfies the termination conditions.

    // Superlinear asymptotic convergence:
    // isls->atol = TaoMin(0.5, isls->merit*sqrt(isls->merit));
    // isls->rtol = 0.0;

    // Quadratic asymptotic convergence:
    isls->atol = TaoMin(0.5, isls->merit*isls->merit);
    isls->rtol = 0.0;

    // Set the convergence tolerances.
    info = TaoGetLinearSolver(tao, &lsolver); CHKERRQ(info);
    // info = lsolver->SetTolerances(isls->rtol, isls->atol, 1e30, 10*nn); CHKERRQ(info);
    // info = lsolver->SetTolerances(0.0, 0.0, 1e30, 10*nn); CHKERRQ(info);
    // info = lsolver->SetTolerances(0.0, 0.0, 1e30, 10*nn); CHKERRQ(info);
#endif

    // Determine whether we are taking a predictor or corrector step
    // Calculate the direction in the smoothing parameter
    // Substitute out of the problem

    if ((isls->merit_eqn <= 1e-10) ||
        (isls->merit_eqn <= nn*isls->merit_mu)) {
      // Only take a predictor step when ||eqn|| <= const*||mu||
      printf("Predictor step.\n");
      step_type = 1;
      isls->dmu = isls->mucon / isls->d_mucon;
      printf("dmu    : %5.4e\n", isls->dmu);
    }
    else {
      // Otherwise take a correct step
      printf("Corrector step.\n");
      step_type = 0;
      isls->dmu = 0;
    }
    info = isls->w->Waxpby(1.0, isls->ff, -isls->dmu, isls->dm);

    // Assert: w vector now contains right-hand side
    // We solve the system:
    //    [DPhi DMu ][dx ] = Phi
    //    [   0 DRho][dmu] = Rho
    // where Rho is the forcing function.

    // Now calculate a free and fixed set of variables.  The fixed set of
    // variables are those for which d_b is approximately equal to zero.
    // The definition of approximately changes as we approact the 
    // solution.

    // No one rule is guaranteed to work in all cases.  The following
    // definition is based on the norm of the Jacobian matrix.  If the
    // norm is large, the tolerance becomes smaller.

    info = J->Norm1(&(isls->identifier)); CHKERRQ(info);
    isls->identifier = TaoMin(isls->merit, 1e-2) / (1 + isls->identifier);

    info = isls->t1->SetToConstant(-isls->identifier); CHKERRQ(info);
    info = isls->t2->SetToConstant( isls->identifier); CHKERRQ(info);

    info = isls->t1->SetToConstant(0.0); CHKERRQ(info);
    info = isls->t2->SetToConstant(0.0); CHKERRQ(info);

    info = FixedVariableSet->WhichBetweenOrEqual(isls->t1, isls->db, isls->t2); CHKERRQ(info);
    info = FreeVariableSet->ComplementOf(FixedVariableSet); CHKERRQ(info);

    info = FixedVariableSet->GetSize(&nf);
    if (nf) printf("Fixed size: %d\n", nf);

    // Assert:
    //   Partition created.
    // Now calculate the direction in the fixed variable space.
    info = r1->SetReducedVec(isls->w, FixedVariableSet); CHKERRQ(info);
    info = r2->SetReducedVec(isls->da, FixedVariableSet); CHKERRQ(info);
    info = r1->PointwiseDivide(r1, r2); CHKERRQ(info);

    // Throw the fixed variable direction into a global direction
    info = isls->d->SetToZero(); CHKERRQ(info);
    info = isls->d->ReducedXPY(r1, FixedVariableSet); CHKERRQ(info);

    // Our direction in the FixedVariableSet is fixed.  Calculate the
    // information needed for the step in the FreeVariableSet.  To
    // do this, we need to know the diagonal perturbation and the
    // right hand side.
    info = r1->SetReducedVec(isls->da, FreeVariableSet); CHKERRQ(info);
    info = r2->SetReducedVec(isls->w,  FreeVariableSet); CHKERRQ(info);
    info = r3->SetReducedVec(isls->db, FreeVariableSet); CHKERRQ(info);

    info = r1->PointwiseDivide(r1, r3); CHKERRQ(info);
    info = r2->PointwiseDivide(r2, r3); CHKERRQ(info);

    // Assert:
    //   r1 is the diagonal perturbation
    //   r2 is the right hand side
    //   r3 is no longer needed
    // Now need to modify r2 for our direction choice in the fixed
    // variable set:  calculate t1 = J*d, take the reduced vector
    // of t1 and modify r2.
    info = J->Multiply(isls->d, isls->t1); CHKERRQ(info);
    info = r3->SetReducedVec(isls->t1, FreeVariableSet); CHKERRQ(info);
    info = r2->Axpy(-1.0, r3);

    // Calculate the reduced problem matrix
    info = Jsub->SetReducedMatrix(J, FreeVariableSet, FreeVariableSet); CHKERRQ(info);
    info = Jsub->AddDiagonal(r1); CHKERRQ(info);

    // Assert:
    //   r1 is no longer needed
    // Calculate the reduced direction.  (Really negative of Newton
    // direction.  Therefore, rest of the code uses -d.)
    info = TaoPreLinearSolve(tao, Jsub); CHKERRQ(info);
    info = TaoLinearSolve(tao, Jsub, r2, r1, &flag); CHKERRQ(info);

    // Add the direction in the free variables back into the real direction.
    info = isls->d->ReducedXPY(r1, FreeVariableSet); CHKERRQ(info);

    // Calculate the norm of the relative residual in the real space
    info = isls->t1->PointwiseMultiply(isls->da, isls->d); CHKERRQ(info);
    info = J->Multiply(isls->d, isls->t2); CHKERRQ(info);
    info = isls->t2->PointwiseMultiply(isls->db, isls->t2); CHKERRQ(info);
    info = isls->t1->Axpy( 1.0, isls->t2); CHKERRQ(info);
    info = isls->t1->Axpy(-1.0, isls->w); CHKERRQ(info);
    info = isls->t1->Norm2(&two_norm); CHKERRQ(info);

    printf("eqn_resid: %5.4e\n", two_norm);

    // Check direction for descent.
    info = isls->d->Norm2squared(&normd); CHKERRQ(info);
    normd += isls->dmu*isls->dmu;
    normd = sqrt(normd);

    info = isls->d->Dot(isls->dpsi, &innerd); CHKERRQ(info);
    innerd += isls->g_mucon*isls->dmu;

    printf("normd  : %5.4e\n", normd);
    printf("innerd : %5.4e\n", innerd);

    if (innerd <= isls->delta*pow(normd, isls->rho)) {
      // Descent direction test failed, use gradient direction
      info = PetscInfo1(tao, "TaoSolve_ISILS: %d: newton direction not descent\n", iter); CHKERRQ(info);
      printf("Gradient direction\n");

      info = isls->d->CopyFrom(isls->dpsi); CHKERRQ(info);
      if (1 == step_type) {
	// Predictor step gradient
	isls->dmu = isls->g_mucon;
      }
      else {
	// Corrector step gradient
	isls->dmu = 0;
      }

      info = isls->d->Dot(isls->dpsi, &innerd); CHKERRQ(info);
      innerd += isls->g_mucon*isls->dmu;

      printf("g_innerd : %5.4e\n", innerd);
    }

    // We have a descent direction.  We need to take the negative to
    // correct the sign of the direction.
    info = isls->d->Negate(); CHKERRQ(info);
    isls->dmu = -isls->dmu;
    innerd = -innerd;

    // Apply an armijo linesearch to find the new iterate.
    innerd *= sigma;
    muval = isls->mu;
    t = 1;

    while (t >= minstep) {
      // Calculate iterate
      info = isls->w->Waxpby(1.0, x, t, isls->d); CHKERRQ(info);
      isls->mu = muval + t*isls->dmu;
 
      // Calculate function at new iterate
      info = TaoComputeMeritFunction(tao, isls->w, &psival); CHKERRQ(info);

      printf("psival: %5.4e psi+t*innerd: %5.4e\n", psival, psi + t*innerd);

      // Check descent condition
      if (psival <= psi + t*innerd) {
        break;
      }
      t *= beta;
    }

    x->CopyFrom(isls->w);

    printf("t      : %5.4e\n", t);
    printf("psi    : %5.4e\n", psi);
    printf("psival : %5.4e\n", psival);
    printf("diff   : %5.4e\n", psival - psi);

    // Have an improvement for the merit function.  Calculate the function
    // values for the global merit function.

    printf("ip_merit: %5.4e\n", isls->merit);

    Tao_ASLS_FunctionGradient(tao, x, &psi, isls->dpsi, isls);
    info = isls->dpsi->Norm2(&ndpsi); CHKERRQ(info);
  }

  info = TaoVecDestroy(isls->dm); CHKERRQ(info);
  isls->dm = TAO_NULL;

  info = TaoMatDestroy(Jsub); CHKERRQ(info);
  info = TaoVecDestroy(r1); CHKERRQ(info);
  info = TaoVecDestroy(r2); CHKERRQ(info);
  info = TaoVecDestroy(r3); CHKERRQ(info);

  info = TaoIndexSetDestroy(FreeVariableSet); CHKERRQ(info);
  info = TaoIndexSetDestroy(FixedVariableSet); CHKERRQ(info);
  TaoFunctionReturn(0);
}

/* ---------------------------------------------------------- */
EXTERN_C_BEGIN
#undef __FUNCT__  
#define __FUNCT__ "TaoCreate_ISILS"
int TaoCreate_ISILS(TAO_SOLVER tao)
{
  TAO_SSLS *isls;
  int        info;

  TaoFunctionBegin;

  info = TaoNew(TAO_SSLS,&isls); CHKERRQ(info);
  info = PetscLogObjectMemory(tao, sizeof(TAO_SSLS)); CHKERRQ(info);

  isls->delta = 1e-10;
  isls->rho = 2.1;
  isls->mu_init = 0.1;

  isls->identifier = 1e-5;

  info=TaoSetTaoSolveRoutine(tao,TaoSolve_ISILS,(void*)isls); CHKERRQ(info);
  info=TaoSetTaoSetUpDownRoutines(tao,TaoSetUp_SSLS,TaoSetDown_SSLS); CHKERRQ(info);
  info=TaoSetTaoOptionsRoutine(tao,TaoSetOptions_SSLS); CHKERRQ(info);
  info=TaoSetTaoViewRoutine(tao,TaoView_SSLS); CHKERRQ(info);

  info = TaoSetMaximumIterates(tao,2000); CHKERRQ(info);
  info = TaoSetMaximumFunctionEvaluations(tao,4000); CHKERRQ(info);

  info = TaoSetTolerances(tao,0,0,0,0); CHKERRQ(info);
  info = TaoSetGradientTolerances(tao,1.0e-16,0.0,0.0); CHKERRQ(info);
  info = TaoSetFunctionLowerBound(tao,1.0e-8); CHKERRQ(info);

  TaoFunctionReturn(0);
}
EXTERN_C_END