#include "snopt.h"
#include <stdlib.h>
#include <stdio.h>

/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */
/*  springConstraint    springObjective    springInit                          */
/*  main                                                                       */
/* +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ */

static cint T;


void springConstraint
(
 cint   *mode,
 cint   *nncon,
 cint   *nnjac,
 cint   *nejac,
 double *x,
 double *fCon,
 double *gCon,
 cint   *nState,
 char   *cu,
 cint   *lencu,
 double *iu,
 cint   *leniu,
 double *ru,
 cint   *lenru
)
{

  cint    i, jg, jy;
  double yt, ytp1;

  jy = -1;   /* Counts y(t) variables. */
  jg = -2;   /* Counts nonlinear Jacobian elements */

  for ( i = 1; i <= T; i++) {
    jy = jy+1;
    jg = jg+2;
    yt = x[jy];
    ytp1 = x[jy+1];
    fCon[i-1] = 0.01*yt*yt + (ytp1 - yt);
    gCon[jg] = 0.02*yt - 1.0;
    /* gCon[jg+1] = 2.0;    Constant term set by SpringInit. */
  }

}

/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/

void springObjective
(
 cint*   mode,
 cint*   nnobj,
 double* x,
 double* fObj,
 double* gObj,
 cint*   nState,
 char*   cu,
 cint*   lencu,
 double* iu,
 cint*   leniu,
 double* ru,
 cint*   lenru
)
{
  double u;
  cint   i, jy, jx;


  *fObj = 0.0;
  jy = -1;
  jx = T;

  for ( i = 0; i <= T; i++ ) {
    jy = jy+1;
    jx = jx+1;
    u = x[jx];
    *fObj = *fObj + u*u;
    gObj[jy] = 0.0;
    gObj[jx] = u;
  }

  *fObj = *fObj/2.0;

}

/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/

void springInit( snProblem* spring )
{
  cint    i, j, k, neJ;
  double  objective_add =  0.0;
  double  infty         = 1e20;
  double  dummy         = 0.11;
  double *Jcol;
  cint   *indJ;
  cint   *locJ;

  snSetProblemName     ( spring, "     Spring" );
  snSetOptionString    ( spring, "Cold start          "       );

  /* For the variables y(t) and x(t), respectively. */
  for (i = 0; i <= T ; i++ ) {
    snSetVariableBounds  ( spring, i    ,  -1.0 , infty );
    snSetVariable( spring, i, SN_BASIS_COLD_ELIGIBLE_0, -1.0 );

    snSetVariableBounds  ( spring, T+1+i, -infty, infty );
    snSetVariable( spring, T+1+i , SN_BASIS_COLD_ELIGIBLE_3 , 0.0 );
  }

  /* For the variables u(t). */
  for (i = 0; i <= T-1 ; i++ ) {
    snSetVariableBounds  ( spring, 2*T+2+i, -0.2, 0.2 );
    snSetVariable( spring, 2*T+2+i, SN_BASIS_COLD_ELIGIBLE_3, 0.0 );
    }

  /* Set the bounds on the slacks. */
  for (i = 0 ; i <= 2*T-1 ; i++ ) {
    snSetConstraintBounds( spring, i,   0.0 ,  0.0 );
  }

  for ( i = 0 ; i <= (T-1); i++ ) {
    snSetMultiplier( spring, i, 0.0 );
    snSetMultiplier( spring, T+i , 0.0 );
  }

  Jcol = snGetNonzeroValues        ( spring );
  indJ = snGetNonzeroRowIndices    ( spring );
  locJ = snGetNonzeroColumnPointers( spring );

  /* Populate the Jacobian data structure */

  j =   -1;
  neJ =  0;

  /* Generate columns for y(t). */
  for ( k = 0 ; k<= T ; k++ ) {
    j = j+1;
    locJ[j] = neJ + 1;
    if ( k > 0 ) {
      neJ = neJ+1;
      indJ[neJ-1] = k;
      Jcol[neJ-1] = 1.0;
    }
    if ( k < T ) {
      neJ = neJ+1;
      indJ[neJ-1] = k+1;
      Jcol[neJ-1] = dummy;
    }
    if ( k < T ) {
      neJ = neJ+1;
      indJ[neJ-1] = T+k+1;
      Jcol[neJ-1] = -0.2;
    }
  }
  /* Generate columns for x(t). */
  for ( k = 0 ; k<= T ; k++ ) {
    j = j+1;
    locJ[j] = neJ + 1;
    if ( k < T ) {
      neJ = neJ+1;
      indJ[neJ-1] = k+1;
      Jcol[neJ-1] = 0.004;
    }
    if ( k > 0 ) {
      neJ = neJ+1;
      indJ[neJ-1] = T+k;
      Jcol[neJ-1] = 1.0;
    }
    if ( k < T ) {
      neJ = neJ+1;
      indJ[neJ-1] = T+k+1;
      Jcol[neJ-1] = -1.0;
    }
  }

  /* Generate columns for u(t). */
  for ( k = 0 ; k <= (T-1) ; k++ ) {
    j = j+1;
    locJ[j] = neJ + 1;
    neJ = neJ+1;
    indJ[neJ-1] = k+1;
    Jcol[neJ-1] = -0.2;
  }

  /* locJ (*) has one extra element.*/
  locJ[spring->dimension.n] = neJ+1;

  snSetObjectiveFunction ( spring, springObjective  );
  snSetConstraintFunction( spring, springConstraint );
  snSetObjectiveAdd      ( spring, objective_add );

}

/*+++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++*/

int main( int argc , char* argv[] )
{

  snProblem spring;

  if ( argc != 2 ) {
    printf( "usage: %s <num_periods>\n" , argv[0] );
    exit(EXIT_FAILURE);
  }

  T = atoi( argv[1] );

  cint    linear_objective_row          =  0;
  cint    nonzeros                      =  7*T;
  cint    jacobian_nonzeros             =  2*T;
  cint    variables                     =  3*T+2;
  cint    constraints                   =  2*T;
  cint    nonlinear_constraints         =  T;
  cint    nonlinear_jacobian_variables  =  T+1;
  cint    nonlinear_objective_variables =  2*T+2;
  double  objective_add                 =  0.0;

  cint    sn_error;
  char   *error_msg;

  /* Allocate snopt workspace and various arrays that define problem spring. */

  sn_error =
    snInit( &spring,
	    variables,
	    constraints,
	    nonzeros,
	    jacobian_nonzeros,
	    nonlinear_constraints,
	    nonlinear_objective_variables,
	    nonlinear_jacobian_variables,
	    linear_objective_row,
            "springb.out",
            "stdout" );

  if ( sn_error != SN_OK ) {
    snGetError( &spring , &error_msg );
    printf( "%s: error occurred (%d).\n%s", argv[0] , sn_error, error_msg );

    snDelete( &spring );
    free(error_msg);

    exit(EXIT_FAILURE);
  }

  /* Populate the arrays that define problem spring. */
  /* Set some optional parameters.                */

  snSetOptionInteger( &spring, "Verify Level", 3 ); 

  springInit( &spring );

  /* Solve problem spring.                           */

  sn_error = snSolve ( &spring );

  if ( sn_error != SN_OK ) {

    snGetError( &spring , &error_msg );
    printf( "%s: error occurred (%d).\n%s", argv[0] , sn_error, error_msg );

    free(error_msg);
    snDelete( &spring );

    exit(EXIT_FAILURE);

  }

  snDelete( &spring );
  exit(EXIT_SUCCESS);

  return 0;

}
