/****************************************************************************

For a cost function of
C(x) = 0.5*x'*Lxx*x + 0.5*u'*Luu*u
The value function is
  V(x) = V0 + Vx*x + 0.5*x'*Vxx*x

/****************************************************************************/
/*
TODO:

Change A to Fx, B to Fu

Better to make Zxx, Zuu symmetric?
make Fxx, Fuu symmetric
make Fux = Fxu'

rationalize arguments

Implement adaptive shifting to fix Zuu not full rank.

Levenberg Marquardt (what does that mean in this context?

use space-time opt techniques in create_trajectory
DIRCOL
MUSMOL
*/
/****************************************************************************/
/*
CHANGES:

handle time-varying costs and non-time varying costs
This is taken care of by supplying a cost function argument
Have Lxx and Luu indexed by time: double **Lxx -> double ***Lxx, ...
if time_varying_criteria = 1
in initialize_lqrd();

add discount from opt-2nd-order.c

switch from inverting Zuu to solving set of linear equations.

Rework to include time index  -> lqrd2
0 -> ...; -1 means steady state.

Remove all non essential stuff from init_lqrd:
 cost_time_varying
 cost_linear_terms
 cost_cross_terms

bidirectional numerical derivatives.

Have all global variables in structure.

Get rid of X0, and just have first point in trajectory be initial
point.

initialize_lqrd():
method, eps removed as arguments
added time_varying_criteria to initialize_lqrd() arguments;

added set_xxx() functions to set variables xxx.

reset_lqrd() is now set_n_points()

n_steps -> n_points

added argument to one_step_dynamics and compute_linear_model
*/
/****************************************************************************/
/* INCLUDES */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include "../useful/useful.h"
#include "../cholesky/ldlt.h"
#include "../dm/dm.h"
#include "lqrd.h"

/****************************************************************************/
/* DEFINES */

/****************************************************************************/
/****************************************************************************/

LQRD *initialize_lqrd( int nstates, /* Dimensionality of state */
		       int ncontrols, /* Dimensionality of controls */
		       int max_npoints, /* Maximum number of points in trajectory */
		       int (*one_step_dynamics)(),
		       double (*cost_function)(),
		       void *argument /* argument to dynamics and cost functions */
		       )
{
  int i;
  LQRD *r;

  r = (LQRD *) malloc( sizeof( LQRD ) );
  if ( r == NULL )
    {
      fprintf( stderr, "allocation failure 1 in LQRD()");
      exit( -1 );
    }

  r->n_states = nstates;
  r->n_controls = ncontrols;
  r->max_n_points = max_npoints;
  r->n_points = max_npoints; 

  r->use_U_limits = FALSE;
  r->U_min = dv(ncontrols);
  r->U_max = dv(ncontrols);

  r->one_step_dynamics_function = one_step_dynamics;
  r->cost_function = cost_function;
  r->argument = argument;

  printf( "Initializing LQRD: %d %d %d 0x%x 0x%x 0x%x\n",
	  r->n_states, r->n_controls, r->max_n_points,
	  r->one_step_dynamics_function, 
	  r->cost_function,
	  r->argument ); 

  r->terminal_cost_function = NULL;
  r->compute_linear_model_function = NULL;

  r->X = dm(max_npoints, nstates);
  dm_zero(r->X, max_npoints, nstates);
  r->X_d = dm(max_npoints, nstates);
  dm_zero(r->X_d, max_npoints, nstates);
  r->X_best = dm(max_npoints, nstates);
  dm_zero(r->X_best, max_npoints, nstates);
  /*
  r->X0 = dv(nstates);
  dv_zero(r->X0, nstates);
  */
  r->X_new = dv(nstates);
  dv_zero(r->X_new, nstates);
  r->X_d_final = dv(nstates);
  dv_zero(r->X_d_final, nstates);

  r->A = dm(nstates, nstates);
  dm_zero(r->A, nstates, nstates);
  r->B = dm(nstates, ncontrols);
  dm_zero(r->B, nstates, ncontrols);
  r->K = dm(ncontrols, nstates);
  dm_zero(r->K, ncontrols, nstates);
  r->K_ss = dm(ncontrols, nstates);
  dm_zero(r->K_ss, ncontrols, nstates);

  r->P = dm(nstates, nstates);
  dm_zero(r->P, nstates, nstates);
  r->P_new = dm(nstates, nstates);
  dm_zero(r->P_new, nstates, nstates);
  r->P_ss = dm(nstates, nstates);
  dm_zero(r->P_ss, nstates, nstates);

  r->Lx = dv( nstates );
  dv_zero( r->Lx, nstates );
  r->Lu = dv( ncontrols );
  dv_zero( r->Lu, ncontrols );

  r->Lxx = dm( nstates, nstates );
  dm_zero( r->Lxx, nstates, nstates );
  r->Lxu = dm( nstates, ncontrols );
  dm_zero( r->Lxu, nstates, ncontrols );
  r->Lux = dm( ncontrols, nstates );
  dm_zero( r->Lux, ncontrols, nstates );
  r->Luu = dm( ncontrols, ncontrols );
  dm_zero( r->Luu, ncontrols, ncontrols );
 
  r->Vx = dv(nstates);
  dv_zero(r->Vx, nstates);
  r->Vxx = dm(nstates, nstates);
  dm_zero(r->Vxx, nstates, nstates);
  r->Zx = dv(nstates);
  dv_zero(r->Zx, nstates);
  r->Zu = dv(ncontrols);
  dv_zero(r->Zu, ncontrols);
  r->Zxu = dm(nstates, ncontrols);
  dm_zero(r->Zxu, nstates, ncontrols);
  r->Zux = dm(ncontrols, nstates);
  dm_zero(r->Zux, ncontrols, nstates);
  r->Zuu = dm(ncontrols, ncontrols);
  dm_zero(r->Zuu, ncontrols, ncontrols);
  r->Zuu_inv = dm(ncontrols, ncontrols);
  dm_zero(r->Zuu_inv, ncontrols, ncontrols);
  r->Zuu_L = dm(ncontrols, ncontrols);
  dm_zero(r->Zuu_L, ncontrols, ncontrols);
  r->Zuu_D = dv(ncontrols);
  dv_zero(r->Zuu_D, ncontrols);
  r->Zxx = dm(nstates, nstates);
  dm_zero(r->Zxx, nstates, nstates);

  r->U = dm(max_npoints, ncontrols);
  dm_zero(r->U, max_npoints, ncontrols);
  r->U_d = dm(max_npoints, ncontrols);
  dm_zero (r->U_d, max_npoints, ncontrols);
  r->U_best = dm(max_npoints, ncontrols);
  dm_zero(r->U_best, max_npoints, ncontrols);
  r->U_new = dv(ncontrols);
  dv_zero(r->U_new, ncontrols);
  r->K_array = d3( max_npoints, ncontrols, nstates );
  d3_zero(r->K_array, max_npoints, ncontrols, nstates );
  r->V_array = dv( max_npoints );
  dv_zero(r->V_array, max_npoints );
  r->Vx_array = dm( max_npoints, nstates );
  dm_zero(r->Vx_array, max_npoints, nstates );
  r->Vxx_array = d3( max_npoints, nstates, nstates );
  d3_zero(r->Vxx_array, max_npoints, nstates, nstates );

  r->delta_U_array = dm(max_npoints, ncontrols);
  dm_zero(r->delta_U_array, max_npoints, ncontrols);

  r->temps1 = dv(nstates);
  dv_zero(r->temps1, nstates);
  r->temps2 = dv(nstates);
  dv_zero(r->temps2, nstates);
  r->temps3 = dv(nstates);
  dv_zero(r->temps3, nstates);
  r->tempc1 = dv(ncontrols);
  dv_zero(r->tempc1, ncontrols);
  r->tempc2 = dv(ncontrols);
  dv_zero(r->tempc2, ncontrols);
  r->tempc3 = dv(ncontrols);
  dv_zero(r->tempc3, ncontrols);
  r->tempss1 = dm(nstates, nstates);
  dm_zero(r->tempss1, nstates, nstates);
  r->tempss2 = dm(nstates, nstates);
  dm_zero(r->tempss2, nstates, nstates);
  r->tempss3 = dm(nstates, nstates);
  dm_zero(r->tempss3, nstates, nstates);
  r->tempsc1 = dm(nstates, ncontrols);
  dm_zero(r->tempsc1, nstates, ncontrols);
  r->tempsc2 = dm(nstates, ncontrols);
  dm_zero(r->tempsc2, nstates, ncontrols);
  r->tempcc1 = dm(ncontrols, ncontrols);
  dm_zero(r->tempcc1, ncontrols, ncontrols);
  r->tempcs1 = dm(ncontrols, nstates);
  dm_zero(r->tempcs1, ncontrols, nstates);
  r->tempcs2 = dm(ncontrols, nstates);
  dm_zero(r->tempcs2, ncontrols, nstates);

  r->Fxx = d3( nstates, nstates, nstates );
  d3_zero(r->Fxx, nstates, nstates, nstates );
  r->Fux = d3( nstates, ncontrols, nstates );
  d3_zero(r->Fux, nstates, ncontrols, nstates );
  r->Fxu = d3( nstates, nstates, ncontrols );
  d3_zero(r->Fxu, nstates, nstates, ncontrols );
  r->Fuu = d3( nstates, ncontrols, ncontrols );
  d3_zero(r->Fuu, nstates, ncontrols, ncontrols );
  
  r->epsilon_gradient = 1e-3;
  r->epsilon_second_order = 1.0;
  r->delta = 1e-4;
  r->ss_printout_interval = 10000;
  r->ss_min_iterations = 100;
  r->ss_max_iterations = 1000000;
  r->ss_convergence_threshold = 1e-20;
  r->discount = 1.0;

  // printf( "Initialization done.\n" );

  return r;
}

/****************************************************************************/

int compute_linear_model( LQRD *lqrd,
			  double *x, 
			  double *u,
			  int index,
			  /* return values */
			  double **a, 
			  double **b )
{
  int i, j;
  double save;

  if ( lqrd->compute_linear_model_function != NULL )
    {
      return lqrd->compute_linear_model_function( lqrd->argument,
						  x, u, index, a, b );
    }

  /* Do it numerically */

  for( i = 0; i < lqrd->n_states; i++ )
    {
      save = x[i];
      x[i] += lqrd->delta;
      lqrd->one_step_dynamics_function( lqrd->argument, x, u, index, lqrd->temps2 );
      x[i] = save - lqrd->delta;
      lqrd->one_step_dynamics_function( lqrd->argument, x, u, index, lqrd->temps1 );
      x[i] = save;
      for( j = 0; j < lqrd->n_states; j++ )
	a[j][i] = (lqrd->temps2[j] - lqrd->temps1[j])/(2*lqrd->delta);
    }

  for( i = 0; i < lqrd->n_controls; i++ )
    {
      save = u[i];
      u[i] += lqrd->delta;
      lqrd->one_step_dynamics_function( lqrd->argument, x, u, index, lqrd->temps2 );
      u[i] = save - lqrd->delta;
      lqrd->one_step_dynamics_function( lqrd->argument, x, u, index, lqrd->temps1 );
      u[i] = save;
      for( j = 0; j < lqrd->n_states; j++ )
	b[j][i] = (lqrd->temps2[j] - lqrd->temps1[j])/(2*lqrd->delta);
    }
  return 0;
}

/****************************************************************************/

int compute_cost_first_derivatives( LQRD *lqrd,
				    double *x, 
				    double *u,
				    int index,
				    /* return values */
				    double *Lx, 
				    double *Lu )
{
  int i, j;
  double c1, c2;
  double save;

  /*
  if ( lqrd->compute_linear_model_function != NULL )
    {
      return lqrd->compute_linear_model_function( lqrd->argument,
						  x, u, index, a, b );
    }
  */

  /* Do it numerically */

  for( i = 0; i < lqrd->n_states; i++ )
    {
      save = x[i];
      x[i] += lqrd->delta;
      c1 = lqrd->cost_function( lqrd->argument, x, u, index );
      // printf( "1: %g %g %g -> %g\n", x[0], x[1], u[0], c1 );
      x[i] = save - lqrd->delta;
      c2 = lqrd->cost_function( lqrd->argument, x, u, index );
      // printf( "1: %g %g %g -> %g\n", x[0], x[1], u[0], c2 );
      x[i] = save;
      Lx[i] = (c1 - c2)/(2*lqrd->delta);
    }

  for( i = 0; i < lqrd->n_controls; i++ )
    {
      save = u[i];
      u[i] += lqrd->delta;
      c1 = lqrd->cost_function( lqrd->argument, x, u, index );
      u[i] = save - lqrd->delta;
      c2 = lqrd->cost_function( lqrd->argument, x, u, index );
      u[i] = save;
      Lu[i] = (c1 - c2)/(2*lqrd->delta);
    }
  return 0;
}

/****************************************************************************/

int compute_cost_second_derivatives( LQRD *z,
				     double *x, 
				     double *u,
				     int index,
				     /* return values */
				     double **Lxx, 
				     double **Lxu, 
				     double **Lux, 
				     double **Luu )
{
  int i, j;
  double save;

  /*
  if ( z->compute_linear_model_function != NULL )
    {
      return z->compute_linear_model_function( z->argument,
						  x, u, index, a, b );
    }
  */

  /* Do it numerically */

  for( i = 0; i < z->n_states; i++ )
    {
      save = x[i];
      x[i] += z->delta;
      compute_cost_first_derivatives( z, x, u, index, z->temps1, z->tempc1 );
      // printf( "2: %g %g %g -> %g %g %g\n", x[0], x[1], u[0],
      //      z->temps1[0], z->temps1[1], z->tempc1[0] );
      x[i] = save - z->delta;
      compute_cost_first_derivatives( z, x, u, index, z->temps2, z->tempc2 );
      // printf( "2: %g %g %g -> %g %g %g\n", x[0], x[1], u[0],
      //      z->temps2[0], z->temps2[1], z->tempc2[0] );
      x[i] = save;
      for( j = 0; j < z->n_states; j++ )
	Lxx[j][i] = (z->temps1[j] - z->temps2[j])/(2*z->delta);
      for( j = 0; j < z->n_controls; j++ )
	Lux[j][i] = (z->tempc1[j] - z->tempc2[j])/(2*z->delta);
    }

  /* Let's make sure Lxx is symmetric. */
  dm_transpose( Lxx, z->tempss2, z->n_states, z->n_states );
  dm_acc( Lxx, z->tempss2, z->n_states, z->n_states );
  dm_scale( z->tempss2, 0.5, Lxx, z->n_states, z->n_states );

  for( i = 0; i < z->n_controls; i++ )
    {
      save = u[i];
      u[i] += z->delta;
      compute_cost_first_derivatives( z, x, u, index,
				     z->temps1, z->tempc1 );
      u[i] = save - z->delta;
      compute_cost_first_derivatives( z, x, u, index,
				     z->temps2, z->tempc2 );
      u[i] = save;
      for( j = 0; j < z->n_states; j++ )
	Lxu[j][i] = (z->temps1[j] - z->temps2[j])/(2*z->delta);
      for( j = 0; j < z->n_controls; j++ )
	Luu[j][i] = (z->tempc1[j] - z->tempc2[j])/(2*z->delta);
    }

  /* Let's make sure Luu is symmetric. */
  dm_transpose( Luu, z->tempcc1, z->n_controls, z->n_controls );
  dm_acc( Luu, z->tempcc1, z->n_controls, z->n_controls );
  dm_scale( z->tempcc1, 0.5, Luu, z->n_controls, z->n_controls );

  /* Let's make sure Lxu = Lux'. */
  for( i = 0; i < z->n_states; i++ )
    {
      for( j = 0; j < z->n_controls; j++ )
	{
	  save = (Lxu[i][j] + Lux[j][i])/2;
	  Lxu[i][j] = Lux[j][i] = save;
	}
    }

  return 0;
}

/************************************************************************/

int update_P( LQRD *z,
	      double **a, 
	      double **b, 
	      double **Lxx, 
	      double **Lxu, 
	      double **Lux, 
	      double **Luu, 
	      double **p, 
	      double **p_new, 
	      double **k )
{
  double value;
  int i, j;

  /* This is from page 46 of Dyer and McReynolds */
  /* Zxx = A'PA + Lxx */
  dm_mult( p, a, z->tempss1, z->n_states, z->n_states, z->n_states );
  dm_transpose( a, z->tempss2, z->n_states, z->n_states );
  dm_mult( z->tempss2, z->tempss1, z->Zxx, 
	   z->n_states, z->n_states, z->n_states );
  dm_acc( Lxx, z->Zxx, z->n_states, z->n_states );

  /* Let's make sure Zxx is symmetric. */
  dm_transpose( z->Zxx, z->tempss2, z->n_states, z->n_states );
  dm_acc( z->Zxx, z->tempss2, z->n_states, z->n_states );
  dm_scale( z->tempss2, 0.5, z->Zxx, z->n_states, z->n_states );

  /* Zux = B'PA + Lux */
  /* tempss1 is still PA */
  dm_transpose( b, z->tempcs2, z->n_states, z->n_controls );
  dm_mult( z->tempcs2, z->tempss1, z->Zux,
	   z->n_controls, z->n_states, z->n_states );
  dm_acc( Lux, z->Zux, z->n_controls, z->n_states );
  dm_transpose( z->Zux, z->Zxu, z->n_controls, z->n_states );

  /* Zuu = B'PB + Luu */
  dm_mult( p, b, z->tempsc1, z->n_states, z->n_states, z->n_controls );
  /* tempcs2 is still B' */
  dm_mult( z->tempcs2, z->tempsc1, z->Zuu,
	   z->n_controls, z->n_states, z->n_controls );
  dm_acc( Luu, z->Zuu, z->n_controls, z->n_controls );

  /* Let's make sure Zuu is symmetric. */
  dm_transpose( z->Zuu, z->tempcc1, z->n_controls, z->n_controls );
  dm_acc( z->Zuu, z->tempcc1, z->n_controls, z->n_controls );
  dm_scale( z->tempcc1, 0.5, z->Zuu, z->n_controls, z->n_controls );

  /* Zuu_inv */
  dm_invert( z->Zuu, z->Zuu_inv, z->n_controls );
  /* Check quality of inverse */
  dm_mult( z->Zuu_inv, z->Zuu, z->tempcc1, z->n_controls,
	   z->n_controls, z->n_controls );
  for( i = 0; i < z->n_controls; i++ )
    {
      for( j = 0; j < z->n_controls; j++ )
	{
	  if ( i == j )
	    value = 1;
	  else
	    value = 0;
	  if ( fabs( z->tempcc1[i][j] - value ) > 1e-8 )
	    printf( "Bad inverse: %g %g\n", z->tempcc1[i][j],
		    z->tempcc1[i][j] - value );
	}
    }
      
  /* P_new = Zxx - Zxu*Zuu_inv*Zux */
  dm_mult( z->Zuu_inv, z->Zux, k, z->n_controls, z->n_controls, z->n_states );
  dm_mult( z->Zxu, k, z->tempss1, z->n_states, z->n_controls, z->n_states );
  dm_subtract( z->Zxx, z->tempss1, p_new, z->n_states, z->n_states );

  /* Let's make sure P is symmetric. */
  dm_transpose( p_new, z->tempss1, z->n_states, z->n_states );
  dm_acc( p_new, z->tempss1, z->n_states, z->n_states );
  dm_scale( z->tempss1, 0.5, p_new, z->n_states, z->n_states );

  return 0;
}

/****************************************************************************/
/*
Initialize lqrd->P_ss as you wish. For example:
dm_zero( lqrd->P_ss, lqrd->n_states, lqrd->n_states );
*/

int design_steady_state_lqr_controller( LQRD *lqrd, double *x, double *u )
{
  int i, j, k;
  double P_norm, Pd_norm;

  dv_copy( x, lqrd->X_d_final, lqrd->n_states );
  compute_linear_model( lqrd, x, u, STEADY_STATE_INDEX, lqrd->A, lqrd->B );
  compute_cost_first_derivatives( lqrd, x, u, STEADY_STATE_INDEX, lqrd->Lx, lqrd->Lu );
  compute_cost_second_derivatives( lqrd, x, u, STEADY_STATE_INDEX, lqrd->Lxx, lqrd->Lxu,
				   lqrd->Lux, lqrd->Luu );
  if ( lqrd->ss_printout_interval > 0 )
  {
    printf( "A\n" );
    for( i = 0; i < lqrd->n_states; i++ )
    {
      for( j = 0; j < lqrd->n_states; j++ )
	printf( " %10.5g ", lqrd->A[i][j] );
      printf( "\n" );
    }
    printf( "B\n" );
    for( i = 0; i < lqrd->n_states; i++ )
    {
      for( j = 0; j < lqrd->n_controls; j++ )
	printf( " %10.5g ", lqrd->B[i][j] );
      printf( "\n" );
    }
    printf( "Lx\n" );
    for( i = 0; i < lqrd->n_states; i++ )
      printf( " %10.5g ", lqrd->Lx[i] );
    printf( "\n" );
    printf( "Lu\n" );
    for( i = 0; i < lqrd->n_controls; i++ )
      printf( " %10.5g ", lqrd->Lu[i] );
    printf( "\n" );
    printf( "Lxx\n" );
    for( i = 0; i < lqrd->n_states; i++ )
    {
      for( j = 0; j < lqrd->n_states; j++ )
	printf( " %10.5e ", lqrd->Lxx[i][j] );
      printf( "\n" );
    }
    printf( "Lxu\n" );
    for( i = 0; i < lqrd->n_states; i++ )
    {
      for( j = 0; j < lqrd->n_controls; j++ )
	printf( " %10.5e ", lqrd->Lxu[i][j] );
      printf( "\n" );
    }
    printf( "Luu\n" );
    for( i = 0; i < lqrd->n_controls; i++ )
    {
      for( j = 0; j < lqrd->n_controls; j++ )
	printf( " %10.5e ", lqrd->Luu[i][j] );
      printf( "\n" );
    }
  } 
  /* Initialize P from P_ss */
  dm_copy( lqrd->P_ss, lqrd->P, lqrd->n_states, lqrd->n_states );
  for( i = 0; 1; i++ )
  {
    update_P( lqrd, lqrd->A, lqrd->B,
	      lqrd->Lxx, lqrd->Lxu, lqrd->Lux, lqrd->Luu,
	      lqrd->P, lqrd->P_new, lqrd->K );
    dm_subtract( lqrd->P, lqrd->P_new, lqrd->tempss1,
		 lqrd->n_states, lqrd->n_states );
    dm_copy( lqrd->P_new, lqrd->P, lqrd->n_states, lqrd->n_states );
    dm_norm( lqrd->P, &P_norm, lqrd->n_states, lqrd->n_states );
    dm_norm( lqrd->tempss1, &Pd_norm, lqrd->n_states, lqrd->n_states );
    if ( i%lqrd->ss_printout_interval == 0
	 && i > 0 && (lqrd->ss_printout_interval > 0) )
    {
      printf (" %d iterations\n", i);
      for (j=0; j<lqrd->n_states; j++)
      {
	for (k=0; k<lqrd->n_states; k++)
	  printf (" %10.6g ", lqrd->P[j][k]);
	printf ("\n");
      }
      for (j=0; j<lqrd->n_states; j++)
      {
	for (k=0; k<lqrd->n_states; k++)
	  printf (" %10.6g ", lqrd->tempss1[j][k]);
	printf ("\n");
      }
      printf( "P_norm: %g; Pd_norm: %g; ratio: %g\n",
	      P_norm, Pd_norm,
	      Pd_norm/P_norm );
    }
    if ( ( Pd_norm/P_norm < lqrd->ss_convergence_threshold ) &&
	( i > lqrd->ss_min_iterations ) )
      break;
    if ( i > lqrd->ss_max_iterations )
      break;
  }
  dm_copy( lqrd->P, lqrd->P_ss, lqrd->n_states, lqrd->n_states );
  dm_copy( lqrd->K, lqrd->K_ss, lqrd->n_controls, lqrd->n_states );
  if ( lqrd->ss_printout_interval > 0 )
  {
    printf (" computed P_ss in %d iterations\n", i);
    for (j=0; j<lqrd->n_states; j++)
    {
      for (k=0; k<lqrd->n_states; k++)
	printf (" %10.6g ", lqrd->P_ss[j][k]);
      printf ("\n");
    }
    printf( "K_ss\n" );
    for (j=0; j<lqrd->n_controls; j++)
    {
      for (k=0; k<lqrd->n_states; k++)
	printf (" %10.6g ", lqrd->K_ss[j][k]);
      printf ("\n");
    }
  }
  return 0;
}

/****************************************************************************/
/* A test routine to see if the steady state gains work.
Need to have set:
X_d
K_ss
 */

int integrate_ss( LQRD *z, double *x0 )
{
int j;

/* X_new = X0 */

   dv_copy( x0, z->X_new, z->n_states );
   for ( j = 0; j < z->n_points; j++ )
     {
       /* X[j] = X_new */
       dv_copy( z->X_new, z->X[j], z->n_states );

       /* U_new = - K_ss * ( X_new - X_d[j] ) */
       dv_subtract ( z->X_new, z->X_d[j], z->temps1, z->n_states );
       dmv_mult ( z->K_ss, z->temps1, z->tempc1, z->n_controls, z->n_states );
       dv_scale( z->tempc1, -1.0, z->U_new, z->n_controls );

       /* U[j] = U_new */
       dv_copy( z->U_new, z->U[j], z->n_controls);

       z->one_step_dynamics_function( z->argument, z->X[j], z->U[j], STEADY_STATE_INDEX,
				      z->X_new );
     }
   return 0;
}

/****************************************************************************/

double trajectory_cost( LQRD *z, 
			double **x_array, double **u_array, int n_points )
{
  int i;
  double cost = 0;

  if ( z->terminal_cost_function )
    cost += z->terminal_cost_function( z->argument, x_array[n_points-1] );
  else
    cost += z->cost_function( z->argument, x_array[n_points-1], 
				u_array[n_points-1], STEADY_STATE_INDEX );

  z->V_array[n_points-1] = cost;

  for ( i = n_points - 2; i >= 0; i-- )
    {
      cost = z->cost_function( z->argument, x_array[i], u_array[i], i )
	+ z->discount*cost;
      z->V_array[i] = cost;
    }

  if ( !finite( cost ) )
    cost = LQRD_BAD_VALUE;

  z->V_array[0] = cost;

  return( cost );
}

/****************************************************************************/

double total_cost( LQRD *z )
{
  return trajectory_cost( z, z->X, z->U, z->n_points );
}

/****************************************************************************/

int matlab_traj( LQRD *z, char *file_name )
{
  int i, j;
  FILE *stream;

  stream = fopen( file_name, "w" );
  if ( stream == NULL )
    {
      fprintf( stderr, "Can't open %s for output.\n", file_name );
      exit( -1 );
    }
  for( i = 0; i < z->n_points; i++ )
    {
      fprintf( stream, "%d ", i );
      for( j = 0; j < z->n_states; j++ )
	fprintf( stream, "%g ", z->X[i][j] );
      for( j = 0; j < z->n_controls; j++ )
	fprintf( stream, "%g ", z->U[i][j] );
      fprintf( stream, "\n" );
    }
  fclose( stream );
  printf( "Wrote %d points to %s\n", z->n_points, file_name );

  return 0;
}

/****************************************************************************/

int gradient( LQRD *z )
{
  int i, j;

  /* This is from page 58 of Dyer and McReynolds */
  /* Vx = P_final * (X[n_points-1] - X_d_final) */
  dv_subtract ( z->X[z->n_points-1], z->X_d_final, z->temps1, z->n_states ); 
  dmv_mult ( z->P_ss, z->temps1, z->Vx, z->n_states, z->n_states ); 
  /*
  printf( "%g %g; %g %g\n", z->temps1[0], z->temps1[1], 
	  z->Vx[0], z->Vx[1] );
  */

  for ( i = z->n_points - 1; i >= 0; i-- )
    {
      compute_linear_model( z, z->X[i], z->U[i], i, z->A, z->B );
      compute_cost_first_derivatives( z, z->X[i], z->U[i], i, z->Lx, z->Lu );

#ifdef COMMENT
      /* Zx = Vx * A + (Lxx * (X[i] - X_d[i])) */
      dv_subtract( z->X[i], z->X_d[i], z->temps1, z->n_states );
      dmv_mult( z->Lxx, z->temps1, z->temps2, z->n_states, z->n_states );
      dvm_mult( z->Vx, z->A, z->temps1, z->n_states, z->n_states );
      dv_add( z->temps1, z->temps2, z->Zx, z->n_states );
#endif
      /* Zx = Vx * A + Lx */
      dvm_mult( z->Vx, z->A, z->temps1, z->n_states, z->n_states );
      dv_scale( z->temps1, z->discount, z->temps1, z->n_states );
      dv_add( z->temps1, z->Lx, z->Zx, z->n_states );

#ifdef COMMENT
      /* Zu = Vx * B + Luu * (U[i] - U_d[i]) */
      dvm_mult( z->Vx, z->B, z->tempc2, z->n_states, z->n_controls );
      dv_subtract( z->U[i], z->U_d[i], z->tempc1, z->n_controls );
      dmv_mult( z->Luu, z->tempc1, z->tempc3,
		z->n_controls, z->n_controls );
      dv_add( z->tempc3, z->tempc2, z->Zu, z->n_controls );
#endif
      /* Zu = Vx * B + Lu */
      dvm_mult( z->Vx, z->B, z->tempc1, z->n_states, z->n_controls );
      dv_scale( z->tempc1, z->discount, z->tempc1, z->n_controls );
      dv_add( z->tempc1, z->Lu, z->Zu, z->n_controls );

      /* delta_U_array(i,:) = Zu */
      dv_copy( z->Zu, z->delta_U_array[i], z->n_controls );

      /* Vx = Zx */
      dv_copy( z->Zx, z->Vx, z->n_states );

      /*
	printf( "%d: %g %g %g\n", i, z->delta_U_array[i][0], 
	z->Vx[0], z->Vx[1] );
      */
    }
  return 0;
}

/****************************************************************************/

int integrate_gradient( LQRD *z )
{
  int j;

  /* X_new = X0 */
  dv_copy ( z->X[0], z->X_new, z->n_states );
  for ( j = 0; j < z->n_points; j++ )
    {
      /* U_new = U[j] + epsilon * delta_U_array[j] */
      dv_scale( z->delta_U_array[j], z->epsilon_gradient, 
		z->tempc1, z->n_controls );
      dv_subtract( z->U[j], z->tempc1, z->U_new, z->n_controls );

      /* X[j] = X_new */
      dv_copy( z->X_new, z->X[j], z->n_states );

      /* U[j] = U_new */
      dv_copy( z->U_new, z->U[j], z->n_controls );

      z->one_step_dynamics_function( z->argument, z->X[j], z->U[j], j, z->X_new );
    }
  return 0;
}

/****************************************************************************/

int compute_dynamics_second_derivatives( LQRD *z, double *x, double *u, int index )
{
  int i, j, k;
  double save;

  for( i = 0; i < z->n_states; i++ )
    {
      save = x[i];
      x[i] += z->delta;
      compute_linear_model( z, x, u, index, z->tempss1, z->tempsc1 );
      x[i] = save - z->delta;
      compute_linear_model( z, x, u, index, z->tempss2, z->tempsc2 );
      x[i] = save;
      for( j = 0; j < z->n_states; j++ )
	{
	  for( k = 0; k < z->n_states; k++ )
	    z->Fxx[j][k][i] = (z->tempss1[j][k] - z->tempss2[j][k])
	      /(2*z->delta);
	  for( k = 0; k < z->n_controls; k++ )
	    z->Fux[j][k][i] = (z->tempsc1[j][k] - z->tempsc2[j][k])
	      /(2*z->delta);
	}
    }

  for( i = 0; i < z->n_controls; i++ )
    {
      save = u[i];
      u[i] += z->delta;
      compute_linear_model( z, x, u, index, z->tempss1, z->tempsc1 );
      u[i] = save - z->delta;
      compute_linear_model( z, x, u, index, z->tempss2, z->tempsc2 );
      u[i] = save;
      for( j = 0; j < z->n_states; j++ )
	{
	  for( k = 0; k < z->n_states; k++ )
	    z->Fxu[j][k][i] = (z->tempss1[j][k] - z->tempss2[j][k])/
	      (2*z->delta);
	  for( k = 0; k < z->n_controls; k++ )
	    z->Fuu[j][k][i] = (z->tempsc1[j][k] - z->tempsc2[j][k])/
	      (2*z->delta);
	}
    }
  return 0;
}

/****************************************************************************/

int second_order_default_initialization( LQRD *z )
{
  /* This is from page 71 of Dyer and McReynolds */
  /* Vx = P_final * (X[n_points-1] - X_final) */
  dv_subtract( z->X[z->n_points-1], z->X_d_final, z->temps1, z->n_states ); 
  dmv_mult( z->P_ss, z->temps1, z->Vx, z->n_states, z->n_states ); 

  /* Vxx = P_final */
  dm_copy( z->P_ss, z->Vxx, z->n_states, z->n_states );
  return 0;
}

/****************************************************************************/

int second_order_sweep( LQRD *z )
{
  int i, j, k, state, ok;

  /* This is from page 71 of Dyer and McReynolds */
  for ( i = z->n_points - 1; i >= 0; i-- )
    {
      compute_linear_model( z, z->X[i], z->U[i], i, z->A, z->B );
      compute_dynamics_second_derivatives( z, z->X[i], z->U[i], i );
      compute_cost_first_derivatives( z, z->X[i], z->U[i], i, z->Lx, z->Lu );
      compute_cost_second_derivatives( z, z->X[i], z->U[i], i, z->Lxx, z->Lxu,
				       z->Lux, z->Luu );

#ifdef COMMENT
      /* Zx = Vx * A + (Lxx * (X[i] - X_d[i])) */
      dv_subtract( z->X[i], z->X_d[i], z->temps1, z->n_states );
      dmv_mult( z->Lxx, z->temps1, z->temps2, z->n_states, z->n_states );
      dvm_mult( z->Vx, z->A, z->temps1, z->n_states, z->n_states );
      dv_add( z->temps1, z->temps2, z->Zx, z->n_states );
#endif
      /* Zx = Vx * A + Lx */
      dvm_mult( z->Vx, z->A, z->temps1, z->n_states, z->n_states );
      dv_scale( z->temps1, z->discount, z->temps1, z->n_states );
      dv_add( z->temps1, z->Lx, z->Zx, z->n_states );

#ifdef COMMENT
      /* Zu = Vx * B + Luu * (U[i] - U_d[i]) */
      dvm_mult( z->Vx, z->B, z->tempc2, z->n_states, z->n_controls );
      dv_subtract( z->U[i], z->U_d[i], z->tempc1, z->n_controls );
      dmv_mult( z->Luu, z->tempc1, z->tempc3,
		z->n_controls, z->n_controls );
      dv_add( z->tempc3, z->tempc2, z->Zu, z->n_controls );
#endif
      /* Zu = Vx * B + Lu */
      dvm_mult( z->Vx, z->B, z->tempc1, z->n_states, z->n_controls );
      dv_scale( z->tempc1, z->discount, z->tempc1, z->n_controls );
      dv_add( z->tempc1, z->Lu, z->Zu, z->n_controls );

      /* This is from page 69 of Dyer and McReynolds */
      /* Zxx = A' * Vxx * A + Vx * Fxx + Lxx */
      /* A' * Vxx * A term */
      dm_mult( z->Vxx, z->A, z->tempss1,
	       z->n_states, z->n_states, z->n_states );
      dm_transpose( z->A, z->tempss2, z->n_states, z->n_states );
      dm_mult( z->tempss2, z->tempss1, z->Zxx,
	       z->n_states, z->n_states, z->n_states );
      /* Vx*Fxx term */
      for( state = 0; state < z->n_states; state++ )
	{
	  dm_scale( z->Fxx[state], z->Vx[state], z->tempss2,
		    z->n_states, z->n_states );
	  dm_acc( z->tempss2, z->Zxx, z->n_states, z->n_states );
	}
      dm_scale( z->Zxx, z->discount, z->Zxx, z->n_states, z->n_states );
      dm_acc( z->Lxx, z->Zxx, z->n_states, z->n_states );

      /* Let's make sure Zxx is symmetric. */
      dm_transpose( z->Zxx, z->tempss2, z->n_states, z->n_states );
      dm_acc( z->Zxx, z->tempss2, z->n_states, z->n_states );
      dm_scale( z->tempss2, 0.5, z->Zxx, z->n_states, z->n_states );

      /* Zux = B' * Vxx * A + Vx * Fux + Lux */
      /* B' * Vxx * A term: tempss1 is still Vxx*A */
      dm_transpose( z->B, z->tempcs2, z->n_states, z->n_controls );
      dm_mult( z->tempcs2, z->tempss1, z->Zux, 
	       z->n_controls, z->n_states, z->n_states );
      /* Vx*Fux term */
      for( state = 0; state < z->n_states; state++ )
	{
	  dm_scale( z->Fux[state], z->Vx[state], z->tempcs1,
		    z->n_controls, z->n_states );
	  dm_acc( z->tempcs1, z->Zux, z->n_controls, z->n_states );
	}
      dm_scale( z->Zux, z->discount, z->Zux, z->n_controls, z->n_states );
      dm_acc( z->Lux, z->Zux, z->n_controls, z->n_states );
      dm_transpose( z->Zux, z->Zxu, z->n_controls, z->n_states );

#ifdef COMMENT
  /*
    Is Zux really the transpose of Zxu?
    compute it and check
  */
  /* Zxu = Fx' * Vxx * Fu + Vx * Fxu + Lxu */
  dm_mult( traj->next->Vxx, Fu, temp_xu1,
	   traj->next->n_x, traj->next->n_x, traj->n_u );
  dm_transpose( Fx, temp_xx1, traj->next->n_x, traj->n_x );
  dm_mult( temp_xx1, temp_xu1, Zxu, traj->n_x, traj->next->n_x, traj->n_u );
  /* Vx * Fxu term */
  for( state = 0; state < traj->next->n_x; state++ )
    {
      dm_scale( Fxu[state], traj->next->Vx[state], temp_xu1, traj->n_x, traj->n_u );
      dm_acc( temp_xu1, Zxu, traj->n_x, traj->n_u );
    }
  dm_scale( z->Zxu, z->discount, z->Zxu, z->n_states, z->n_controls );
  dm_acc( Lxu, Zxu, traj->n_x, traj->n_u );
#endif      

      /* Zuu = B' * Vxx * B + Vx * Fuu + Luu */
      dm_mult( z->Vxx, z->B, z->tempsc1,
	       z->n_states, z->n_states, z->n_controls );
      /* tempcs2 is still B' */
      dm_mult( z->tempcs2, z->tempsc1, z->Zuu,
	       z->n_controls, z->n_states, z->n_controls );
      /* Vx*Fuu term */
      for( state = 0; state < z->n_states; state++ )
	{
	  dm_scale( z->Fuu[state], z->Vx[state], z->tempcc1,
		    z->n_controls, z->n_controls );
	  dm_acc( z->tempcc1, z->Zuu, z->n_controls, z->n_controls );
	}
      dm_scale( z->Zuu, z->discount, z->Zuu, z->n_controls, z->n_controls );
      dm_acc( z->Luu, z->Zuu, z->n_controls, z->n_controls );

      /* Let's make sure Zuu is symmetric. */
      dm_transpose( z->Zuu, z->tempcc1, z->n_controls, z->n_controls );
      dm_acc( z->Zuu, z->tempcc1, z->n_controls, z->n_controls );
      dm_scale( z->tempcc1, 0.5, z->Zuu, z->n_controls, z->n_controls );

#ifdef COMMENT
      /* Zuu_inv */
      dm_invert( z->Zuu, z->Zuu_inv, z->n_controls );

      /* Check out if inverse is correct */
      dm_mult( z->Zuu_inv, z->Zuu, z->tempcc1,
	       z->n_controls, z->n_controls, z->n_controls );
      // dm_print( stdout, "%g ", "\n", z->tempcc1, z->n_controls, z->n_controls );
      // printf( "\n" );
      for ( j = 0; j < z->n_controls; j++ )
	for ( k = 0; k < z->n_controls; k++ )
	  {
	    if ( j == k )
	      {
		if ( fabs( z->tempcc1[j][k] - 1.0 ) > 1e-10 )
		  printf( "diagonal error: %g %g\n", z->tempcc1[j][k], z->tempcc1[j][k] - 1.0 );
	      }
	    else
	      {
		if ( fabs( z->tempcc1[j][k] ) > 1e-4 )
		  printf( "off diagonal error: %g\n", z->tempcc1[j][k] );
	      }
	  }
#endif

#ifdef COMMENT
      /* delta_U_array(i,:) = Zuu_inv*Zu */
      dmv_mult( z->Zuu_inv, z->Zu, z->tempc1, z->n_controls, z->n_controls );

      /* K(t) = Zuu_inv*Zux */
      dm_mult( z->Zuu_inv, z->Zux, z->K,
	       z->n_controls, z->n_controls, z->n_states );
#endif

      /* Let's do an LDL' decomposition of Zuu */
      ok = ldlt_decompose( z->Zuu, z->Zuu_L, z->Zuu_D,
			   z->tempc1, z->n_controls );
      if ( !ok )
	{
	  printf( "index %d\n", i );
	  printf( "Luu:\n" );
	  dm_print( stdout, "%g ", "\n",
		    z->Luu, z->n_controls, z->n_controls );
	  printf( "Zuu:\n" );
	  dm_print( stdout, "%g ", "\n",
		    z->Zuu, z->n_controls, z->n_controls );
	  exit( -1 );
	}

      /* delta_U_array(i,:) = Zuu_inv*Zu */
      ldlt_solve_vector( z->Zuu_L, z->Zuu_D, z->Zu, z->tempc1,
			 z->tempc2, z->n_controls );
      dv_copy( z->tempc1, z->delta_U_array[i], z->n_controls );

      /* K(t) = Zuu_inv*Zux */
      ldlt_solve_matrix( z->Zuu_L, z->Zuu_D, z->Zux, z->K,
			 z->tempc2, z->n_controls, z->n_states );
      dm_copy ( z->K, z->K_array[i], z->n_controls, z->n_states );

      /* Vx = Zx - Zu*K */
      dvm_mult( z->Zu, z->K, z->temps1, z->n_controls, z->n_states );
      dv_subtract( z->Zx, z->temps1, z->Vx, z->n_states );

      /* Vxx = Zxx - Zxu*K */
      dm_mult( z->Zxu, z->K, z->tempss1,
	       z->n_states, z->n_controls, z->n_states );
      dm_subtract( z->Zxx, z->tempss1, z->Vxx, z->n_states, z->n_states );

      /* Let's make sure Vxx is symmetric. */
      dm_transpose( z->Vxx, z->tempss1, z->n_states, z->n_states );
      dm_acc( z->Vxx, z->tempss1, z->n_states, z->n_states );
      dm_scale( z->tempss1, 0.5, z->Vxx, z->n_states, z->n_states );

      dv_copy ( z->Vx, z->Vx_array[i], z->n_states );
      dm_copy ( z->Vxx, z->Vxx_array[i], z->n_states, z->n_states );
    }
  return 0;
}

/****************************************************************************/

int second_order( LQRD *z )
{
  second_order_default_initialization( z );
  second_order_sweep( z );
  return 0;
}

/****************************************************************************/

int integrate_second_order( LQRD *z )
{
  int i, j, k;

  // printf( "z->epsilon_second_order: %g\n", z->epsilon_second_order );

  /* X_new = X0 */
  dv_copy ( z->X[0], z->X_new, z->n_states );
  /* dv_copy (X0, X[0], n_states); This is unnecessary,
     since X[j] = X_new later. */
  for ( j = 0; j < z->n_points; j++ )
    {
      /* U_new = U[j] - epsilon * delta_U_array[j] - Ux[j] * (X_new - X[j])  */
      dv_subtract( z->X_new, z->X[j], z->temps1, z->n_states );
      /*
      if ( j < 660 )
	{
	  printf( "%d delta_x: ", j );
	  for( i = 0; i < z->n_states; i++ )
	    printf( "%g ", z->temps1[i] );
	  printf( "\n" );
	}
      */
      /*
      if ( j < 4 )
	{
	  printf( "K\n" );
	  dm_print( stdout, "%g ", "\n", z->K_array[j], z->n_controls, z->n_states );
	}
      */
      dv_zero( z->tempc1, z->n_controls );
      dmv_mult( z->K_array[j], z->temps1, z->tempc1,
		z->n_controls, z->n_states );
      dv_scale( z->delta_U_array[j], z->epsilon_second_order, z->tempc2,
		z->n_controls);
      dv_subtract( z->U[j], z->tempc2, z->tempc2, z->n_controls );
      dv_subtract( z->tempc2, z->tempc1, z->U_new, z->n_controls );

      if ( z->use_U_limits )
	for( i = 0; i < z->n_controls; i++ )
	  {
	    if ( z->U_new[i] > z->U_max[i] )
	      z->U_new[i] = z->U_max[i];
	    else if ( z->U_new[i] < z->U_min[i] )
	      z->U_new[i] = z->U_min[i];
	  }

      /*
      if ( j < 4 )
	{
	  printf( "U: " );
	  for( i = 0; i < z->n_controls; i++ )
	    printf( "%g ", z->U_new[i] );
	  printf( "\n" );
	}
      */

      /* X[j] = X_new */
      dv_copy( z->X_new, z->X[j], z->n_states );

      /* U[j] = U_new */
      dv_copy( z->U_new, z->U[j], z->n_controls );

      z->one_step_dynamics_function( z->argument, z->X[j], z->U[j], j, z->X_new );
    }
  return 0;
}

/****************************************************************************/

int save_best_trajectory( LQRD *z )
{
  int j;

  for ( j = 0; j < z->n_points; j++ )
    {
      dv_copy( z->X[j], z->X_best[j], z->n_states );
      dv_copy( z->U[j], z->U_best[j], z->n_controls );
    }
  return 0;
}

/****************************************************************************/

int restore_best_trajectory( LQRD *z )
{
  int j;

  for ( j = 0; j < z->n_points; j++ )
    {
      dv_copy( z->X_best[j], z->X[j], z->n_states );
      dv_copy( z->U_best[j], z->U[j], z->n_controls );
    }
  return 0;
}

/****************************************************************************/

set_n_points( LQRD *z, int npoints )
{

  if (( npoints <= z->max_n_points ) && ( npoints > 0 ))
    z->n_points = npoints;
  else
    {
      printf( "LQRD: npoints invalid %d; %d\n", npoints, z->max_n_points );
      exit( -1 );
    }
  /* printf( "N_points set to: %d\n", n_points ); */
}

/****************************************************************************/

set_U_limits( LQRD *z, double *U_min, double *U_max )
{
  dv_copy ( U_min, z->U_min, z->n_controls );
  dv_copy ( U_max, z->U_max, z->n_controls );
  z->use_U_limits = TRUE;
}

/****************************************************************************/
#ifdef COMMENT
/****************************************************************************/
/****************************************************************************/
/****************************************************************************/

set_lqrd_dynamics( one_step_dynamics )
     int (*one_step_dynamics)();
{
  if ( !initialized() )
    return;

  one_step_dynamics_function = one_step_dynamics;
}

/****************************************************************************/

set_lqrd_derivatives( compute_linear_model )
     int (*compute_linear_model)();
{
  int compute_linear_model_numerically();

  if ( !initialized() )
    return;

  if ( compute_linear_model )
    compute_linear_model_function = compute_linear_model;
  else
    compute_linear_model_function = compute_linear_model_numerically;
}

/****************************************************************************/

set_delta( arg )
     double arg;
{
  if ( !initialized() )
    return;

  delta = arg;
}

/****************************************************************************/

set_epsilon_gradient( epsilon )
     double epsilon;
{
  if ( !initialized() )
    return;

  epsilon_gradient = epsilon;
}

/****************************************************************************/

set_epsilon_second_order( epsilon )
     double epsilon;
{
  if ( !initialized() )
    return;

  epsilon_second_order = epsilon;
}

/****************************************************************************/

set_steady_state_parameters( convergence_threshold, min_iterations, 
			    printout_interval )
     double convergence_threshold;
     int min_iterations;
     int printout_interval;
{
  if ( !initialized() )
    return;
  
  if ( convergence_threshold > 0 )
    ss_convergence_threshold = convergence_threshold;
  if ( min_iterations > 0 )
    ss_min_iterations = min_iterations;
  if ( printout_interval >= 0 )
    ss_printout_interval = printout_interval;
}

/****************************************************************************/

copy_trajectory( x1_array, u1_array, x2_array, u2_array )
     double **x1_array, **u1_array;
     double **x2_array, **u2_array;
{
  dm_copy( x1_array, x2_array, n_points, n_states );
  dm_copy( u1_array, u2_array, n_points, n_controls );
}

/****************************************************************************/

copy_K( k1_array, k2_array )
     double ***k1_array, ***k2_array;
{
  d3_copy( k1_array, k2_array, n_points, n_controls, n_states );
}

/****************************************************************************/

save_trajectory( x_array, u_array )
     double **x_array, **u_array;
{
  int j;

  for ( j = 0; j < n_points; j++ )
    {
      dv_copy( X[j], x_array[j], n_states );
      dv_copy( U[j], u_array[j], n_controls );
    }
}

/****************************************************************************/

restore_trajectory( x_array, u_array )
     double **x_array, **u_array;
{
  int j;

  for ( j = 0; j < n_points; j++ )
    {
      dv_copy( x_array[j], X[j], n_states );
      dv_copy( u_array[j], U[j], n_controls );
    }
}

/****************************************************************************/

second_order_sweep_old()
{
  int i;
  double determinant;

  /* This version did not include Fxx, Fux, Fxu, or Fuu terms */
  /* This is from page 71 of Dyer and McReynolds */

  for ( i = n_points - 1; i >= 0; i-- )
    {
      compute_linear_model_function( the_argument, X[i], U[i], A, B );

      /* Zx = Vx * A + (Q * (X[i] - X_d[i])) */
      dv_subtract( X[i], X_d[i], temps1, n_states );
      dmv_mult( Q, temps1, temps2, n_states, n_states );
      dvm_mult( Vx, A, temps1, n_states, n_states );
      dv_add( temps1, temps2, Zx, n_states );

      /* Zu = Vx * B + R * (U[i] - U_d[i]) */
      dvm_mult( Vx, B, tempc2, n_states, n_controls );
      dv_subtract( U[i], U_d[i], tempc1, n_controls );
      dmv_mult( R, tempc1, tempc3, n_controls, n_controls );
      dv_add( tempc3, tempc2, Zu, n_controls );

      /* This is from page 69 of Dyer and McReynolds */
      /* Zxx = A' * Vxx * A + Q */
      dm_mult( Vxx, A, tempss1, n_states, n_states, n_states );
      dm_transpose( A, tempss2, n_states, n_states );
      dm_mult( tempss2, tempss1, Zxx, n_states, n_states, n_states );
      dm_acc( Q, Zxx, n_states, n_states );

      /* Zux = B' * Vxx * A */
      /* tempss1 is still Vxx*A */
      dm_transpose( B, tempcs2, n_states, n_controls );
      dm_mult( tempcs2, tempss1, Zux, n_controls, n_states, n_states );
      dm_transpose( Zux, Zxu, n_controls, n_states );

      /* Zuu = B' * Vxx * B + R */
      dm_mult( Vxx, B, tempsc1, n_states, n_states, n_controls );
      /* tempcs2 is still B' */
      dm_mult( tempcs2, tempsc1, Zuu, n_controls, n_states, n_controls );
      dm_acc( R, Zuu, n_controls, n_controls );

      /* Zuu_inv */
      if ( n_controls == 1 ) 
	{
	  /* 1D case: */ 
	  Zuu_inv[0][0] = 1.0/Zuu[0][0];
	}
      else if ( n_controls == 2 ) 
	{
	  /* 2D case: */ 
	  determinant = Zuu[0][0]*Zuu[1][1] - Zuu[1][0]*Zuu[0][1];
	  Zuu_inv[0][0] = Zuu[1][1]/determinant;
	  Zuu_inv[0][1] = -Zuu[0][1]/determinant;
	  Zuu_inv[1][0] = -Zuu[1][0]/determinant;
	  Zuu_inv[1][1] = Zuu[0][0]/determinant;
	}
      else
	{
	  printf( "Need to handle inverting Zuu! of dimension %d\n",
		 n_controls );
	  exit(-1);
	}

      /* delta_U_array(i,:) = Zuu_inv*Zu */
      dmv_mult( Zuu_inv, Zu, tempc1, n_controls, n_controls );
      dv_copy( tempc1, delta_U_array[i], n_controls );

      /* K(t) = Zuu_inv*Zux */
      dm_mult( Zuu_inv, Zux, K, n_controls, n_controls, n_states );
      dm_copy ( K, K_array[i], n_controls, n_states );

      /* Vx = Zx - Zu*K */
      dvm_mult( Zu, K, temps1, n_controls, n_states );
      dv_subtract( Zx, temps1, Vx, n_states );

      /* Vxx = Zxx - Zxu*K */
      dm_mult( Zxu, K, tempss1, n_states, n_controls, n_states );
      dm_subtract( Zxx, tempss1, Vxx, n_states, n_states );

      /* Let's make sure Vxx is symmetric. */
      dm_transpose( Vxx, tempss1, n_states, n_states );
      dm_acc( Vxx, tempss1, n_states, n_states );
      dm_scale( tempss1, 0.5, Vxx, n_states, n_states );
    }
}

/****************************************************************************/
#endif
