/*****************************************************************************/
/*
  controller.c: control strategy.
*/
/*****************************************************************************/

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

#include "../useful/useful/useful.h"
#include "../useful/trajectory/trajectory.h"
#include "main.h"
#include "main2.h"
#include "sdfast/sra.h"

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

/* Controller states */
#define UNKNOWN_STATE 0
#define WAITING 1
#define DS1 2  // double support, left leg in front
#define LIFTOFF1 3  // double support, lifting right leg
#define SSL1 4 // first phase of single support, left leg in stance
#define SSL2 5 // second phase of single support, left leg in stance
#define DS2 6  // double support, right leg in front
#define LIFTOFF2 7  // double support, lifting left leg
#define SSR1 8 // first phase of single support, right leg in stance
#define SSR2 9 // second phase of single support, right leg in stance
#define LAUNCH 10 // launch state. Like DS1. largest value so plot is pretty
#define N_CSTATES 11

/* Joint servo modes */
#define PD_MODE 0
#define TORQUE_MODE 1

/*****************************************************************************/
/* Storage. */

// Knot point queues for trajectories.
DOFS desired_trajectory[N_BIPED_DOFS];
DOFS current_desireds[N_BIPED_DOFS];

// Previous states. Used to penalize asymmetry.
int last_state_valid[N_CSTATES];
double last_states[N_CSTATES][MAX_N_SDFAST_STATE];

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

void print_sdfast_state( SIM *s )
{
  int i;

  printf( "%d ", s->sdfast_model );
  if ( s->sdfast_model == SINGLE_SUPPORT )
    for( i = 0; i < SS_N_STATES; i++ )
      printf( "%8.4f ", s->sdfast_state[i] );
  if ( s->sdfast_model == DOUBLE_SUPPORT )
    for( i = 0; i < DS_N_STATES; i++ )
      printf( "%8.4f ", s->sdfast_state[i] );
  printf( "\n" );
}

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

void print_standard_state( SIM *s )
{
  int i;

  printf( "%d ", s->standard_state_type );
  if ( s->standard_state_type == SINGLE_SUPPORT )
    for( i = 0; i < SS_N_STATES; i++ )
      printf( "%8.4f ", s->standard_state[i] );
  if ( s->standard_state_type == DOUBLE_SUPPORT )
    for( i = 0; i < DS_N_STATES; i++ )
      printf( "%8.4f ", s->standard_state[i] );
  printf( "\n" );
}

/*****************************************************************************/
/* This should be the only routine that knows about the sdfast
representations in dynamics_xxx().
Right now this assumes dynamics_impact.c is used.
*/

void compute_standard_state( SIM *s )
{
  int i;

  if ( s->simulation_type != SDFAST_IMPACTS )
    {
      fprintf( stderr,
	       "compute_standard_state assumes dynamics_impact.c is used.\n" );
      exit( -1 );
    }

  switch( s->sdfast_model )
    {
    case IN_AIR:
      fprintf( stderr, "compute_standard_state can't handle IN_AIR\n" );
      exit( -1 );
      break;
    case SINGLE_SUPPORT:
      s->standard_state_type = SINGLE_SUPPORT;
      for ( i = 0; i < SS_N_STATES; i++ )
	s->standard_state[i] = s->sdfast_state[i];
      if ( s->ss_foot_down == LEFT )
	{
	}
      else // right foot down, need to flip.
	{
	  s->standard_state[SS_ANKLE1] *= -1;
	  s->standard_state[SS_HIP1] *= -1;
	  s->standard_state[SS_HIP2] *= -1;
	  s->standard_state[SS_ANKLE1D] *= -1;
	  s->standard_state[SS_HIP1D] *= -1;
	  s->standard_state[SS_HIP2D] *= -1;
	}
      break;
    case DOUBLE_SUPPORT:
      s->standard_state_type = DOUBLE_SUPPORT;
      for ( i = 0; i < DS_N_STATES; i++ )
	s->standard_state[i] = s->sdfast_state[i];
      if ( s->controller_state == SSL2 
	   || s->controller_state == DS2 
	   || s->controller_state == LIFTOFF2 )
	{
	  s->standard_state[DS_ANKLE1] 
	    = -(s->sdfast_state[DS_ANKLE1] + s->sdfast_state[DS_HIP1]
		+ s->sdfast_state[DS_HIP2]);
	  s->standard_state[DS_KNEE1] = s->sdfast_state[DS_KNEE2];
	  s->standard_state[DS_HIP1] = s->sdfast_state[DS_HIP2];
	  s->standard_state[DS_HIP2] = s->sdfast_state[DS_HIP1];
	  s->standard_state[DS_KNEE2] = s->sdfast_state[DS_KNEE1];
	  s->standard_state[DS_ANKLE1D] 
	    = -(s->sdfast_state[DS_ANKLE1D] + s->sdfast_state[DS_HIP1D]
		+ s->sdfast_state[DS_HIP2D]);
	  s->standard_state[DS_KNEE1D] = s->sdfast_state[DS_KNEE2D];
	  s->standard_state[DS_HIP1D] = s->sdfast_state[DS_HIP2D];
	  s->standard_state[DS_HIP2D] = s->sdfast_state[DS_HIP1D];
	  s->standard_state[DS_KNEE2D] = s->sdfast_state[DS_KNEE1D];
	}
      break;
    default:
      fprintf( stderr, "compute_standard_state unknown sdfast_model: %d\n",
	       s->sdfast_model );
      exit( -1 );
    }
}

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

void reset_score( SIM *s )
{
  s->summed_score = 0.0; // sum of penalties over time
  s->discounted_score = 0.0; // sum of discounted penalties over time
  s->discount_to_power = 1.0; // current discount factor
  s->crashed_score = 0; // penalty on crashing
  s->torque_score = 0; // sum of torque penalties
  s->knee_score = 0; // penalty on bending knee wrong way.
  s->speed_score = 0; // penalty on not achieving desired speed
  s->clearance_score = 0; // penalty on having foot too low
  s->f_x_score = 0; // penalty on horizontal foot forces
  s->f_z_score = 0; // penalty on vertical foot forces
  s->time_score = 0; // penalty on transition time errors
  s->impulse_score = 0; // penalty on impulses during touchdown.
  s->asymmetry_score = 0; // penalty on an asymmetric gait.
}

/*****************************************************************************/
/*
Could normalize torque and other penalties w.r.t. distance
Could penalize deviations from symmetry (if necessary).
*/

double one_step_cost( SIM *s )
{
  int i;
  double force;
  double total_z_force;

  s->one_step_total_score = 0;

  s->one_step_torque_score = 0;
  for ( i = LEFT; i <= RIGHT; i++ )
    {
      s->one_step_torque_score += s->hip_torque[i]*s->hip_torque[i];
      s->one_step_torque_score += s->ankle_torque[i]*s->ankle_torque[i];
    }
  if ( s->dynamics_type == IN_AIR )
    {
      for ( i = LEFT; i <= RIGHT; i++ )
	s->one_step_torque_score 
	  += s->leg_force_scale_factor*s->knee_force[i]*s->knee_force[i];
    }
  else if ( s->dynamics_type == SINGLE_SUPPORT )
    {
      if ( s->ss_foot_down == LEFT )
	{
	  force = s->knee_force[LEFT] - s->total_weight_in_newtons/2;
	  s->one_step_torque_score 
	    += s->leg_force_scale_factor*force*force;
	  s->one_step_torque_score 
	    += s->leg_force_scale_factor
	    *s->knee_force[RIGHT]*s->knee_force[RIGHT];
	}
      else
	{
	  s->one_step_torque_score 
	    += s->leg_force_scale_factor
	    *s->knee_force[LEFT]*s->knee_force[LEFT];
	  force = s->knee_force[RIGHT] - s->total_weight_in_newtons/2;
	  s->one_step_torque_score 
	    += s->leg_force_scale_factor*force*force;
	}
    }
  else if ( s->dynamics_type == DOUBLE_SUPPORT )
    {
      for ( i = LEFT; i <= RIGHT; i++ )
	{
	  force = s->knee_force[i] - s->total_weight_in_newtons/2;
	  s->one_step_torque_score += s->leg_force_scale_factor*force*force;
	}
    }
  s->one_step_torque_score *= s->torque_penalty_weight;

  s->one_step_knee_score = 0;

  s->one_step_clearance_score = 0;
  for ( i = LEFT; i <= RIGHT; i++ )
    {
      if ( s->ground_force[i][ZZ] <= 0 && s->foot[i][ZZ] < s->clearance_d )
	s->one_step_clearance_score += (s->foot[i][ZZ] - s->clearance_d)
	  *(s->foot[i][ZZ] - s->clearance_d);
    }
  s->one_step_clearance_score *= s->clearance_penalty_weight;

  s->one_step_f_x_score = 0;
  for ( i = LEFT; i <= RIGHT; i++ )
    s->one_step_f_x_score += s->ground_force[i][XX]*s->ground_force[i][XX];
  s->one_step_f_x_score *= s->f_x_penalty_weight;

  total_z_force = s->ground_force[LEFT][ZZ] + s->ground_force[RIGHT][ZZ];
  s->one_step_f_z_score = (total_z_force - s->total_weight_in_newtons)*
    (total_z_force - s->total_weight_in_newtons);
  s->one_step_f_z_score *= s->f_z_penalty_weight;

  s->one_step_crashed_score = 0;
  // This actually only gets called once since simulation stops on crash.
  // Let's make this quadratic in time rather than linear.
  if ( s->status == CRASHED )
    s->one_step_crashed_score = (s->duration - s->time)*(s->duration - s->time)
      *s->crashed_penalty_weight;

  s->one_step_total_score = 
    s->one_step_torque_score 
    + s->one_step_clearance_score
    + s->one_step_f_x_score
    + s->one_step_f_z_score
    + s->one_step_crashed_score;
  
  return s->one_step_total_score;
}

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

double transition_costs( SIM *s )
{
  int i;
  double duration;
  double diff;

  s->transition_time_score = 0;
  s->transition_impulse_score = 0;
  s->transition_asymmetry_score = 0; // penalty on an asymmetric gait.

  // Compute transition time penalty
  switch( s->controller_state )
    {
    case WAITING:
    case LAUNCH:
    case LIFTOFF1:
    case LIFTOFF2:
    case SSL2:
    case SSR2:
      break;
    case DS1:
    case DS2:
      if ( s->ss_start_time > 0 )
	{
	  duration = s->time - s->ss_start_time;
	  diff = duration - s->ss_time_d;
	  s->transition_time_score += diff*diff*s->timing_weight;
	  /*
	  printf( "%g %d %g %g %g \n", s->time, s->controller_state, duration,
		  s->ss_time_d, diff );
	  */
	}
      s->ss_start_time = -1;
      break;
    case SSL1:
    case SSR1:
      if ( s->ds_start_time > 0 )
	{
	  duration = s->time - s->ds_start_time;
	  diff = duration - s->ds_time_d;
	  s->transition_time_score += diff*diff*s->timing_weight;
	  /*
	  printf( "%g %d %g %g %g \n", s->time, s->controller_state, duration,
		  s->ds_time_d, diff );
	  */
	}
      s->ds_start_time = -1;
      break;
    default:
      fprintf( stderr, "Unknown state %d in transition_costs 1.\n",
	       s->controller_state );
      exit( -1 );
    }

  // Compute impulse penalty
  if ( s->last_impulse_time > 0 )
    {
      s->transition_impulse_score += ( s->last_impulse[XX]*s->last_impulse[XX]
	       + s->last_impulse[ZZ]*s->last_impulse[ZZ] )
	*s->impulse_weight;
    }
  s->last_impulse_time = -1;

  // Compute asymmetry penalty
  switch( s->controller_state )
    {
    case WAITING:
    case LAUNCH:
    case LIFTOFF1:
    case LIFTOFF2:
      break;
    case DS1:
      if ( !last_state_valid[DS2] )
	break;
      for ( i = 0; i < DS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[DS2][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[DS2] = FALSE;
      break;
    case DS2:
      if ( !last_state_valid[DS1] )
	break;
      for ( i = 0; i < DS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[DS1][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[DS1] = FALSE;
      break;
    case SSL1:
      if ( !last_state_valid[SSR1] )
	break;
      for ( i = 0; i < SS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[SSR1][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[SSR1] = FALSE;
      break;
    case SSR1:
      if ( !last_state_valid[SSL1] )
	break;
      for ( i = 0; i < SS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[SSL1][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[SSL1] = FALSE;
      break;
    case SSL2:
      if ( !last_state_valid[SSR2] )
	break;
      for ( i = 0; i < SS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[SSR2][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[SSR2] = FALSE;
      break;
    case SSR2:
      if ( !last_state_valid[SSL2] )
	break;
      for ( i = 0; i < SS_N_STATES; i++ )
	{
	  diff = s->standard_state[i] - last_states[SSL2][i];
	  s->transition_asymmetry_score += diff*diff;
	}
      last_state_valid[SSL2] = FALSE;
      break;
    default:
      fprintf( stderr, "Unknown state %d in transition_costs 2:\n",
	       s->controller_state );
      exit( -1 );
    }
  s->transition_asymmetry_score *= s->asymmetry_weight;

  return s->transition_time_score + s->transition_impulse_score
    + s->transition_asymmetry_score;
}

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

void update_score( SIM *s )
{
  double cost = 0;

  cost += one_step_cost( s );
  // printf( "%g\n", cost );
  cost += transition_costs( s );
  // printf( "%g\n", cost );

  s->summed_score += cost;
  // printf( "%g %g %g\n", s->time, s->summed_score, cost );
  s->discounted_score += s->discount_to_power * cost;
  s->discount_to_power = s->discount * s->discount_to_power;
  /* This is useful for debugging */
  s->torque_score += s->one_step_torque_score;
  s->clearance_score += s->one_step_clearance_score;
  s->f_x_score += s->one_step_f_x_score;
  s->f_z_score += s->one_step_f_z_score;
  s->crashed_score += s->one_step_crashed_score;
  s->time_score += s->transition_time_score;
  s->impulse_score += s->transition_impulse_score;
  s->asymmetry_score += s->transition_asymmetry_score;
}

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

double get_score( SIM *s )
{
  printf( "cost: %g, %g; crashed %g (%d); torque %g; clearance %g; f_x %g; f_z %g; time %g; impulse %g; asymmetry %g\n",
	  s->summed_score, s->discounted_score, s->crashed_score, s->status,
	  s->torque_score, s->clearance_score,
	  s->f_x_score, s->f_z_score,
	  s->time_score, s->impulse_score, s->asymmetry_score );
  return s->summed_score;
}

/*****************************************************************************/
/* call this many times to restart a simulation */

int reinit_controller( SIM *s )
{
  int i;

  s->controller_state = WAITING;
  s->cstate_start_time = s->time;
  s->cstate_elapsed_time = 0.0;

  s->hip_angle_d[LEFT] = s->hip_angle[LEFT];
  s->hip_angle_d[RIGHT] = s->hip_angle[RIGHT];
  s->knee_length_d[LEFT] = s->knee_length[LEFT];
  s->knee_length_d[RIGHT] = s->knee_length[RIGHT];
  s->ankle_angle_d[LEFT] = s->ankle_angle[LEFT];
  s->ankle_angle_d[RIGHT] = s->ankle_angle[RIGHT];
  s->hip_angled_d[LEFT] = s->hip_angled[LEFT];
  s->hip_angled_d[RIGHT] = s->hip_angled[RIGHT];
  s->knee_lengthd_d[LEFT] = s->knee_lengthd[LEFT];
  s->knee_lengthd_d[RIGHT] = s->knee_lengthd[RIGHT];
  /* sidestep angle perturbation for now
  s->ankle_angled_d[LEFT] = s->ankle_angled[LEFT];
  s->ankle_angled_d[RIGHT] = s->ankle_angled[RIGHT];
  */
  s->ankle_angled_d[LEFT] = 0;
  s->ankle_angled_d[RIGHT] = 0;

  s->hip_servo_mode[LEFT] = PD_MODE;
  s->hip_servo_mode[RIGHT] = PD_MODE;
  s->knee_servo_mode[LEFT] = PD_MODE;
  s->knee_servo_mode[RIGHT] = PD_MODE;
  s->ankle_servo_mode[LEFT] = PD_MODE;
  s->ankle_servo_mode[RIGHT] = PD_MODE;

  s->hip_command_ff[LEFT] = 0;
  s->hip_command_ff[RIGHT] = 0;
  s->knee_command_ff[LEFT] = 300;
  s->knee_command_ff[RIGHT] = 300;
  s->ankle_command_ff[LEFT] = 0;
  s->ankle_command_ff[RIGHT] = 0;

  s->hip_integrated_error[LEFT] = 0;
  s->hip_integrated_error[RIGHT] = 0;
  s->knee_integrated_error[LEFT] = 0;
  s->knee_integrated_error[RIGHT] = 0;
  s->ankle_integrated_error[LEFT] = 0;
  s->ankle_integrated_error[RIGHT] = 0;

  init_trajectory( desired_trajectory, N_BIPED_DOFS );
  init_trajectory( current_desireds, N_BIPED_DOFS );

  /* add actual position as knot to current desireds */
  add_knot_point( current_desireds, L_HIP, QUINTIC_SPLINE, 0.0, 
		  s->hip_angle_d[LEFT], s->hip_angled_d[LEFT], 0.0 );
  add_knot_point( current_desireds, L_KNEE, QUINTIC_SPLINE, 0.0,
		  s->knee_length_d[LEFT], s->knee_lengthd_d[LEFT], 0.0 );
  add_knot_point( current_desireds, L_ANKLE, QUINTIC_SPLINE, 0.0,
		  s->ankle_angle_d[LEFT], s->ankle_angled_d[LEFT], 0.0 );
  add_knot_point( current_desireds, R_HIP, QUINTIC_SPLINE, 0.0, 
		  s->hip_angle_d[RIGHT], s->hip_angled_d[RIGHT], 0.0 );
  add_knot_point( current_desireds, R_KNEE, QUINTIC_SPLINE, 0.0, 
		  s->knee_length_d[RIGHT], s->knee_lengthd_d[RIGHT], 0.0 );
  add_knot_point( current_desireds, R_ANKLE, QUINTIC_SPLINE, 0.0, 
		  s->ankle_angle_d[RIGHT], s->ankle_angled_d[RIGHT], 0.0 );

  reset_score( s );

  for ( i = 0; i < N_CSTATES; i++ )
    last_state_valid[i] = FALSE;

  // These should be negative to indicate not valid yet.
  s->ds_start_time = -1;
  s->ss_start_time = -1;

  return 0;
}

/*****************************************************************************/
/* call this many times to restart a simulation */

// int reinit_sim_original( SIM *s )
int reinit_sim( SIM *s )
{
  int i;
  double min;

  srand( s->rand_seed );

  s->time = 0.0;

  /* zero out X references */
  s->torso[XX] = 0;
  s->groin[XX] = 0;
  s->hip[LEFT][XX] = -s->pelvis_width/2;
  s->hip[RIGHT][XX] = s->pelvis_width/2;
  s->foot[LEFT][XX] = -s->pelvis_width/2;
  s->foot[RIGHT][XX] = s->pelvis_width/2;

  /* zero out velocities */
  for( i = 0; i < 3; i++ )
    s->torsod[i] = 0;

  /* zero out angles */
  s->roll = 0;
  s->rolld = 0;
  s->hip_angle[LEFT] = 0;
  s->hip_angle[RIGHT] = 0;
  s->knee_length[LEFT] = 0;
  s->knee_length[RIGHT] = 0;
  s->ankle_angle[LEFT] = 0;
  s->ankle_angle[RIGHT] = 0;
  s->hip_angled[LEFT] = 0;
  s->hip_angled[RIGHT] = 0;
  s->knee_lengthd[LEFT] = 0;
  s->knee_lengthd[RIGHT] = 0;
  s->ankle_angled[LEFT] = 0;
  s->ankle_angled[RIGHT] = 0;

  /* Make it fall */
  /*
  s->rolld = 1.0;
  s->hip_angled[LEFT] = s->hip_angled[RIGHT] = 1.0;
  s->knee_lengthd[LEFT] = s->knee_lengthd[RIGHT] = -1.0;
  */

  /* Indicate foot zero positions not set */
  s->foot_zero[LEFT][ZZ] = 1.0;
  s->foot_zero[RIGHT][ZZ] = 1.0;

  init_state_two_feet_on_ground( s );

  reinit_controller( s );
}

/*****************************************************************************/
/* Call this once to do one time operations like memory allocation */

int init_sim( SIM *s )
{
  init_dynamics( s );
  reinit_sim( s );
}

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

void
run_trajectory( SIM *s, double time )
{
  double  position, velocity, acceleration;

  if ( !lookup_spline3( desired_trajectory, L_HIP, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }
  /* update current desired position */
  set_first_knot_to( current_desireds, L_HIP, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->hip_angle_d[LEFT] = position;
  s->hip_angled_d[LEFT] = velocity;

  if ( !lookup_spline3( desired_trajectory, L_KNEE, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }
  /* update current desired position */
  set_first_knot_to( current_desireds, L_KNEE, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->knee_length_d[LEFT] = position;
  s->knee_lengthd_d[LEFT] = velocity;

  if ( !lookup_spline3( desired_trajectory, L_ANKLE, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }
  /* update current desired position */
  set_first_knot_to( current_desireds, L_ANKLE, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->ankle_angle_d[LEFT] = position;
  s->ankle_angled_d[LEFT] = velocity;

  if ( !lookup_spline3( desired_trajectory, R_HIP, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }

  /* update current desired position */
  set_first_knot_to( current_desireds, R_HIP, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->hip_angle_d[RIGHT] = position;
  s->hip_angled_d[RIGHT] = velocity;

  if ( !lookup_spline3( desired_trajectory, R_KNEE, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }
  /* update current desired position */
  set_first_knot_to( current_desireds, R_KNEE, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->knee_length_d[RIGHT] = position;
  s->knee_lengthd_d[RIGHT] = velocity;

  if ( !lookup_spline3( desired_trajectory, R_ANKLE, time,
			current_desireds,
			&position,
			&velocity,
			&acceleration ) )
    {
      printf( "Lookup failure.\n" );
      exit( -1 );
    }
  /* update current desired position */
  set_first_knot_to( current_desireds, R_ANKLE, QUINTIC_SPLINE, time,
		     position, velocity, acceleration );
  s->ankle_angle_d[RIGHT] = position;
  s->ankle_angled_d[RIGHT] = velocity;
}

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

void
run_servos( SIM *s )
{
  int i;

  s->hip_integrated_error[LEFT] 
    += s->time_step*(s->hip_angle[LEFT] - s->hip_angle_d[LEFT]);
  s->hip_integrated_error[RIGHT] 
    += s->time_step*(s->hip_angle[RIGHT] - s->hip_angle_d[RIGHT]);
  s->knee_integrated_error[LEFT] 
    += s->time_step*(s->knee_length[LEFT] - s->knee_length_d[LEFT]);
  s->knee_integrated_error[RIGHT] 
    += s->time_step*(s->knee_length[RIGHT] - s->knee_length_d[RIGHT]);
  s->ankle_integrated_error[LEFT] 
    += s->time_step*(s->ankle_angle[LEFT] - s->ankle_angle_d[LEFT]);
  s->ankle_integrated_error[RIGHT] 
    += s->time_step*(s->ankle_angle[RIGHT] - s->ankle_angle_d[RIGHT]);

  if ( s->hip_servo_mode[LEFT] == PD_MODE )
    {
      s->hip_command[LEFT] = -s->pos_gains[L_HIP][L_HIP]
	*( s->hip_angle[LEFT] - s->hip_angle_d[LEFT] )
	- s->vel_gains[L_HIP][L_HIP]
	*( s->hip_angled[LEFT] - s->hip_angled_d[LEFT] )
	- s->int_gains[L_HIP][L_HIP]*s->hip_integrated_error[LEFT]
	+ s->hip_command_ff[LEFT];
    }

  if ( s->hip_servo_mode[RIGHT] == PD_MODE )
    {
      s->hip_command[RIGHT] = -s->pos_gains[R_HIP][R_HIP]
	*( s->hip_angle[RIGHT] - s->hip_angle_d[RIGHT] )
	- s->vel_gains[R_HIP][R_HIP]
	*( s->hip_angled[RIGHT] - s->hip_angled_d[RIGHT] )
	- s->int_gains[R_HIP][R_HIP]*s->hip_integrated_error[RIGHT]
	+ s->hip_command_ff[RIGHT];
    }

  if ( s->knee_servo_mode[LEFT] == PD_MODE )
    {
      s->knee_command[LEFT] = -s->pos_gains[L_KNEE][L_KNEE]
	*( s->knee_length[LEFT] - s->knee_length_d[LEFT] )
	- s->vel_gains[L_KNEE][L_KNEE]
	*( s->knee_lengthd[LEFT] - s->knee_lengthd_d[LEFT] )
	- s->int_gains[L_KNEE][L_KNEE]*s->knee_integrated_error[LEFT]
	+ s->knee_command_ff[LEFT];
    }

  if ( s->knee_servo_mode[RIGHT] == PD_MODE )
    {
      s->knee_command[RIGHT] = -s->pos_gains[R_KNEE][R_KNEE]
	*( s->knee_length[RIGHT] - s->knee_length_d[RIGHT] )
	- s->vel_gains[R_KNEE][R_KNEE]
	*( s->knee_lengthd[RIGHT] - s->knee_lengthd_d[RIGHT] )
	- s->int_gains[R_KNEE][R_KNEE]*s->knee_integrated_error[RIGHT]
	+ s->knee_command_ff[RIGHT];
    }

  if ( s->ankle_servo_mode[LEFT] == PD_MODE )
    {
      s->ankle_command[LEFT] = -s->pos_gains[L_ANKLE][L_ANKLE]
	*( s->ankle_angle[LEFT] - s->ankle_angle_d[LEFT] )
	- s->vel_gains[L_ANKLE][L_ANKLE]
	*( s->ankle_angled[LEFT] - s->ankle_angled_d[LEFT] )
	- s->int_gains[L_ANKLE][L_ANKLE]*s->ankle_integrated_error[LEFT]
	+ s->ankle_command_ff[LEFT];
    }

  if ( s->ankle_servo_mode[RIGHT] == PD_MODE )
    {
      s->ankle_command[RIGHT] = -s->pos_gains[R_ANKLE][R_ANKLE]
	*( s->ankle_angle[RIGHT] - s->ankle_angle_d[RIGHT] )
	- s->vel_gains[R_ANKLE][R_ANKLE]
	*( s->ankle_angled[RIGHT] - s->ankle_angled_d[RIGHT] )
	- s->int_gains[R_ANKLE][R_ANKLE]*s->ankle_integrated_error[RIGHT]
	+ s->ankle_command_ff[RIGHT];
      /*
      printf( "%g: %g %g %g\n",
	      s->ankle_command[RIGHT], s->pos_gains[R_ANKLE][R_ANKLE],
	      s->ankle_angle[RIGHT], s->ankle_angle_d[RIGHT] );
      */
    }

  /* Limit ankle torques */
  for ( i = LEFT; i <= RIGHT; i++ )
    {
      if ( s->foot_status[i] == FOOT_IN_AIR 
	   || s->ground_force[i][ZZ] == 0.0 /* avoid divide by zero */ )
	{
	  s->ankle_command[i] = 0;
	}
      else if ( s->ankle_command[i]/s->ground_force[i][ZZ] < s->zmp_x_min )
	s->ankle_command[i] = s->zmp_x_min*s->ground_force[i][ZZ];
      else if ( s->ankle_command[i]/s->ground_force[i][ZZ] > s->zmp_x_max )
	s->ankle_command[i] = s->zmp_x_max*s->ground_force[i][ZZ];
    }

  s->hip_torque[LEFT] = s->hip_command[LEFT];
  s->knee_force[LEFT] = s->knee_command[LEFT];
  s->ankle_torque[LEFT] = s->ankle_command[LEFT];
  s->hip_torque[RIGHT] = s->hip_command[RIGHT]; 
  s->knee_force[RIGHT] = s->knee_command[RIGHT];
  s->ankle_torque[RIGHT] = s->ankle_command[RIGHT];

  /* Perturbation */
  /*
  if ( s->time < 5.0 )
    s->torso_perturbation = -60.0;
  else
    s->torso_perturbation = 0.0;
  */

  /* Wipe out -60 perturbation
  s->hip_torque[LEFT] = s->hip_command[LEFT] + 17/2;
  s->knee_force[LEFT] = s->knee_command[LEFT] + 41/2;
  s->ankle_torque[LEFT] = s->ankle_command[LEFT] + 63/2;
  s->hip_torque[RIGHT] = s->hip_command[RIGHT] + 17/2; 
  s->knee_force[RIGHT] = s->knee_command[RIGHT] + 41/2;
  s->ankle_torque[RIGHT] = s->ankle_command[RIGHT] + 63/2;
  */

  /*
  s->hip_torque[LEFT] = s->hip_torque[RIGHT] = 0.0;
  s->knee_force[LEFT] = s->knee_force[RIGHT] = 0.0;
  s->ankle_torque[LEFT] = s->ankle_torque[RIGHT] = 1.0;
  */

  /*
  s->hip_torque[LEFT] = s->hip_command[LEFT]
    + s->rand_scale*((2.0*rand())/RAND_MAX - 1.0);
  s->knee_force[LEFT] = s->knee_command[LEFT]
    + s->rand_scale*((2.0*rand())/RAND_MAX - 1.0);
  s->hip_torque[RIGHT] = s->hip_command[RIGHT] 
    + s->rand_scale*((2.0*rand())/RAND_MAX - 1.0);
  s->knee_force[RIGHT] = s->knee_command[RIGHT]
    + s->rand_scale*((2.0*rand())/RAND_MAX - 1.0);
  */
}

/*****************************************************************************/
/*
Predict how far the system will roll to, given no more input energy.
Both estimates are approximate: 
Energy assumes pure inverted pendulum dynamics with maximum height being zero.
Compression of leg is a source of error.
This estimate seems to lose energy as we go.
Predicted_x assumes COM moves horizontally only.
Works well in single support, but does something screwy in double support.
*/

double predict_x( SIM *s )
{
  double x, z, angled, length_2, value;
  int side;

  /*
    need to reference COM to something. use last state to figure out which foot just landed.
    plot(pos - 0.161*c - 0.0855,0.14*vel - 0.002,'b',0.02*cos(z),0.02*sin(z),'r')
  */
  if ( s->controller_state == DS1 
       || s->controller_state == SSL1 
       || s->controller_state == SSL2 
       || s->controller_state == WAITING
       || s->controller_state == LAUNCH
       || s->controller_state == LIFTOFF1 )
    {
      side = LEFT;
      // x = (s->com[XX] - (s->foot[LEFT][XX] + s->foot_width/2));
      x = s->com[XX] - s->foot[LEFT][XX];
      // printf( "left: %g %g %g %g %g\n", p, s->com[XX], (s->foot[LEFT][XX] + s->foot_width/2), s->foot[LEFT][XX], s->foot_width/2 );
    }
  else if ( s->controller_state == DS2 
	    || s->controller_state == SSR1 
	    || s->controller_state == SSR2
	    || s->controller_state == LIFTOFF2 )
    {
      side = RIGHT;
      // x = s->com[XX] - (s->foot[RIGHT][XX] - s->foot_width/2);
      x = s->com[XX] - s->foot[RIGHT][XX];
      // printf( "right: %g %g %g %g %g\n", p, s->com[XX], (s->foot[RIGHT][XX] + s->foot_width/2), s->foot[RIGHT][XX], s->foot_width/2 );
    }
  else
    {
      fprintf( stderr, "Unknown state in predict_x, %d\n",
	       s->controller_state );
      exit( -1 );
    }
  s->com_alt[XX] = x;
  s->com_alt[YY] = s->com[YY];
  s->com_alt[ZZ] = s->com[ZZ];
  z = s->com[ZZ] - s->ground_level;
  length_2 = x*x + z*z;
  angled = (x*s->comd[ZZ] - z*s->comd[XX])/length_2;

  // X version
  value = x*x - s->comd[XX]*s->comd[XX]/GRAVITY;
  if ( value >= 0 )
    s->predicted_x = sqrt( value );
  else
    s->predicted_x = -sqrt( -value );

  // Z version
  s->energy = GRAVITY*(z-s->com_height_zero) + length_2*angled*angled;

  return s->energy;
}

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

void transition_to( SIM *s, int cstate, char *message )
{
  int i;

  s->controller_state = cstate;
  s->cstate_start_time = s->time;
  s->cstate_elapsed_time = 0.0;
  if ( s->controller_print )
    {
      printf( message, s->time );
      // print_sdfast_state( s );
      print_standard_state( s );
    }
  for( i = 0; i < MAX_N_SDFAST_STATE; i++ )
    last_states[cstate][i] =  s->standard_state[i];
  last_state_valid[cstate] = TRUE;
}

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

void do_ds( SIM *s, int side1 )
{
  int side2;

  side2 = other_side( side1 );
  s->hip_command_ff[side1] = 0.0;
  s->hip_command_ff[side2] = 0.0;
  s->hip_servo_mode[side1] = TORQUE_MODE;
  s->hip_servo_mode[side2] = TORQUE_MODE;
  s->ankle_servo_mode[side1] = TORQUE_MODE;
  s->ankle_servo_mode[side2] = TORQUE_MODE;
  if ( side1 == LEFT )
    {
      s->hip_command[side1] = s->ds_torque;
      s->hip_command[side2] = s->ds_torque;
      s->ankle_command[side1] = -s->ds_torque;
      s->ankle_command[side2] = -s->ds_torque;
    }
  else
    {
      s->hip_command[side1] = -s->ds_torque;
      s->hip_command[side2] = -s->ds_torque;
      s->ankle_command[side1] = s->ds_torque;
      s->ankle_command[side2] = s->ds_torque;
    }
  // knee gains
  s->pos_gains[L_KNEE][L_KNEE] = s->stance_knee_pos_gain;
  s->vel_gains[L_KNEE][L_KNEE] = s->stance_knee_vel_gain;
  s->int_gains[L_KNEE][L_KNEE] = s->stance_knee_int_gain;
  s->pos_gains[R_KNEE][R_KNEE] = s->stance_knee_pos_gain;
  s->vel_gains[R_KNEE][R_KNEE] = s->stance_knee_vel_gain;
  s->int_gains[R_KNEE][R_KNEE] = s->stance_knee_int_gain;

  // Keep these updated.
  s->hip_angle_d[side1] = s->hip_angle[side1];
  s->hip_angle_d[side2] = s->hip_angle[side2];
  s->hip_angled_d[side1] = s->hip_angled[side1];
  s->hip_angled_d[side2] = s->hip_angled[side2];
}

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

void do_liftoff( SIM *s, int side1 )
{
  int side2;

  side2 = other_side( side1 );
  s->hip_servo_mode[side1] = PD_MODE;
  // s->hip_angle_d[side1] = -s->ankle_angle[side1];
  s->hip_angle_d[side1] 
    += 0.1*(-s->ankle_angle[side1] - s->hip_angle_d[side1]);
  s->hip_angled_d[side1] 
    += 0.1*(-s->ankle_angled[side1] - s->hip_angled_d[side1]);
  if ( side1 == LEFT )
    s->hip_command_ff[side1] = -s->hip_ff1;
  else
    s->hip_command_ff[side1] = s->hip_ff1;
  // Didn't make much difference slowing this down.
  // s->hip_command_ff[side1] += 0.01*(-s->hip_ff1 - s->hip_command_ff[side1]);

  s->hip_servo_mode[side2] = PD_MODE;
  // s->hip_angle_d[side2] = -s->ankle_angle[side1];
  s->hip_angle_d[side2]
    += 0.1*(-s->ankle_angle[side1] - s->hip_angle_d[side2]);
  s->hip_angled_d[side2]
    += 0.1*(-s->ankle_angled[side1] - s->hip_angled_d[side2]);

  s->knee_length_d[side2] = -s->knee_length_d1;
  if ( side2 == LEFT )
    {
      s->pos_gains[L_KNEE][L_KNEE] = s->swing_knee_pos_gain;
      s->vel_gains[L_KNEE][L_KNEE] = s->swing_knee_vel_gain;
      s->int_gains[L_KNEE][L_KNEE] = s->swing_knee_int_gain;
    }
  else
    {
      s->pos_gains[R_KNEE][R_KNEE] = s->swing_knee_pos_gain;
      s->vel_gains[R_KNEE][R_KNEE] = s->swing_knee_vel_gain;
      s->int_gains[R_KNEE][R_KNEE] = s->swing_knee_int_gain;
    }

  s->ankle_servo_mode[side1] = TORQUE_MODE;
  if ( side1 == LEFT )
    s->ankle_command[side1] 
      // = s->ss_torque_gain*(s->energy - s->energy_d);
      = -s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);
  else
    s->ankle_command[side1] 
      // = -s->ss_torque_gain*(s->energy - s->energy_d);
      = s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);

  s->ankle_servo_mode[side2] = TORQUE_MODE;
  s->ankle_command[side2] = 0;
}

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

void do_ss1( SIM *s, int side1 )
{
  int side2;

  side2 = other_side( side1 );
  s->hip_servo_mode[side1] = PD_MODE;
  // s->hip_angle_d[side1] = -s->ankle_angle[side1];
  s->hip_angle_d[side1] 
    += 0.1*(-s->ankle_angle[side1] - s->hip_angle_d[side1]);
  s->hip_angled_d[side1] 
    += 0.1*(-s->ankle_angled[side1] - s->hip_angled_d[side1]);
  if ( side1 == LEFT )
    s->hip_command_ff[side1] = -s->hip_ff1;
  else
    s->hip_command_ff[side1] = s->hip_ff1;
  
  s->hip_servo_mode[side2] = PD_MODE;
  // s->hip_angle_d[side2] = -s->ankle_angle[side1];
  s->hip_angle_d[side2]
    += 0.1*(-s->ankle_angle[side1] - s->hip_angle_d[side2]);
  s->hip_angled_d[side2]
    += 0.1*(-s->ankle_angled[side1] - s->hip_angled_d[side2]);
  
  s->knee_length_d[side2] = -s->knee_length_d1;
  if ( side2 == LEFT )
    {
      s->pos_gains[L_KNEE][L_KNEE] = s->swing_knee_pos_gain;
      s->vel_gains[L_KNEE][L_KNEE] = s->swing_knee_vel_gain;
      s->int_gains[L_KNEE][L_KNEE] = s->swing_knee_int_gain;
    }
  else
    {
      s->pos_gains[R_KNEE][R_KNEE] = s->swing_knee_pos_gain;
      s->vel_gains[R_KNEE][R_KNEE] = s->swing_knee_vel_gain;
      s->int_gains[R_KNEE][R_KNEE] = s->swing_knee_int_gain;
    }
  
  s->ankle_servo_mode[side1] = TORQUE_MODE;
  if ( side1 == LEFT )
    s->ankle_command[side1] 
      // = s->ss_torque_gain*(s->energy - s->energy_d);
      = -s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);
  else
    s->ankle_command[side1] 
      // = -s->ss_torque_gain*(s->energy - s->energy_d);
      = s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);

  s->ankle_servo_mode[side2] = TORQUE_MODE;
  s->ankle_command[side2] = 0;
}

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

void do_ss2( SIM *s, int side1 )
{
  int side2;

  side2 = other_side( side1 );
  s->hip_servo_mode[side1] = PD_MODE;
  // s->hip_angle_d[side1] = -s->ankle_angle[side1];
  /*
  s->hip_angle_d[side1] 
    += 0.1*(-s->ankle_angle[side1] - s->hip_angle_d[side1]);
  s->hip_angled_d[side1] 
    += 0.1*(0 - s->hip_angled_d[side1]);
  */
  // only this worked?
  s->hip_angle_d[side1] = -s->ankle_angle[side1];

  s->hip_servo_mode[side2] = PD_MODE;
  s->hip_angle_d[side2] = -s->ankle_angle[side1];
  /* This didn't work
  s->hip_angled_d[side2] = -s->ankle_angled[side1];
  */

  if ( side1 == LEFT )
    {
      if ( s->com[XX] - s->foot[side1][XX] > s->knee_extend1 )
	{
	  s->knee_length_d[side2] = 0.0;
	  s->pos_gains[R_KNEE][R_KNEE] = s->touchdown_knee_pos_gain;
	  s->vel_gains[R_KNEE][R_KNEE] = s->touchdown_knee_vel_gain;
	  s->int_gains[R_KNEE][R_KNEE] = s->touchdown_knee_int_gain;
	}
    }
  else
    {
      if ( s->foot[side1][XX] - s->com[XX] > s->knee_extend1 )
	{
	  s->knee_length_d[side2] = 0.0;
	  s->pos_gains[L_KNEE][L_KNEE] = s->touchdown_knee_pos_gain;
	  s->vel_gains[L_KNEE][L_KNEE] = s->touchdown_knee_vel_gain;
	  s->int_gains[L_KNEE][L_KNEE] = s->touchdown_knee_int_gain;
	}
    }

  s->ankle_servo_mode[side1] = TORQUE_MODE;
  // Here we are using the estimate of how far we would get
  // if we reversed velocity. This assumes symmetry
  if ( side1 == LEFT )
    s->ankle_command[side1] 
      // = -s->ss_torque_gain*(s->energy - s->energy_d);
      = s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);
  else
    s->ankle_command[side1] 
      // = s->ss_torque_gain*(s->energy - s->energy_d);
      = -s->ss_torque_gain*(s->predicted_x - s->predicted_x_d);

  s->ankle_servo_mode[side2] = TORQUE_MODE;
  s->ankle_command[side2] = 0;
}

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

void
run_state_machine( SIM *s )
{
  int i;

  // Calculate additional state
  predict_x( s );

  // Debounce state transitions
  // if ( s->time - s->cstate_start_time < 0.001 )
  //    return;
  s->cstate_elapsed_time = s->time - s->cstate_start_time;

  /* Do actions and transitions */
  switch( s->controller_state )
    {
    case WAITING:
      if ( s->cstate_elapsed_time >= s->wait_time )
	{
	  s->ds_torque = s->torque1;
	  transition_to( s, LAUNCH, "LAUNCH %g\n" );
	}
      break;

    case LAUNCH:
      do_ds( s, LEFT );
      if ( s->cstate_elapsed_time > s->launch_time )
	{
	  s->knee_integrated_error[RIGHT] = 0.0;
	  transition_to( s, LIFTOFF1, "LIFTOFF1 %g\n" );
	}
      break;

    case DS1:
      do_ds( s, LEFT );
      if ( s->predicted_x < s->predicted_x_d )
	{
	  s->knee_integrated_error[RIGHT] = 0.0;
	  transition_to( s, LIFTOFF1, "LIFTOFF1 %g\n" );
	}
      // need to check if wrong leg leaves the ground.
      if ( s->foot_status[RIGHT] == FOOT_IN_AIR )
	{
	  s->knee_integrated_error[RIGHT] = 0.0;
	  do_liftoff( s, LEFT );
	  s->ss_start_time = s->time;
	  transition_to( s, SSL1, "SSL1 %g\n" );
	}
      break;

    case LIFTOFF1: // lifting right leg off the ground.
      do_liftoff( s, LEFT );
      // need to check if wrong leg leaves the ground.
      if ( s->foot_status[RIGHT] == FOOT_IN_AIR )
	{
	  s->ss_start_time = s->time;
	  transition_to( s, SSL1, "SSL1 %g\n" );
	}
      break;

    case SSL1: // balancing on left leg, upswing
      do_ss1( s, LEFT );
      if ( s->comd[XX] > 0.0 )
	{
	  transition_to( s, SSL2, "SSL2 %g\n" );
	}
      break;

    case SSL2: // balancing on left leg, downswing
      do_ss2( s, LEFT );
      if ( s->foot_status[RIGHT] == FOOT_ON_GROUND )
	{
	  // need to debounce settling into DS
	  if ( 1 ) // s->foot_status[LEFT] == FOOT_ON_GROUND )
	    {
	      s->ds_start_time = s->time;
	      s->ds_torque = s->torque1;
	      // s->ds_torque += -s->ds_torque_gain*(s->energy - s->energy_d)
	      transition_to( s, DS2, "DS2 %g\n" );
	    }
	  else // no DS period
	    {
	      s->knee_integrated_error[LEFT] = 0.0;
	      do_liftoff( s, RIGHT );
	      s->ss_start_time = s->time;
	      transition_to( s, SSR1, "SSR1 %g\n" );
	    }
	}
      break;

    case DS2:
      do_ds( s, RIGHT );
      if ( s->predicted_x < s->predicted_x_d )
	{
	  s->knee_integrated_error[LEFT] = 0.0;
	  transition_to( s, LIFTOFF2, "LIFTOFF2 %g\n" );
	}
      // need to check if wrong leg leaves the ground.
      if ( s->foot_status[LEFT] == FOOT_IN_AIR )
	{
	  s->knee_integrated_error[LEFT] = 0.0;
	  do_liftoff( s, RIGHT );
	  s->ss_start_time = s->time;
	  transition_to( s, SSR1, "SSR1 %g\n" );
	}
      break;

    case LIFTOFF2: // lifting left leg off the ground.
      do_liftoff( s, RIGHT );
      // need to check if wrong leg leaves the ground.
      if ( s->foot_status[LEFT] == FOOT_IN_AIR )
	{
	  s->ss_start_time = s->time;
	  transition_to( s, SSR1, "SSR1 %g\n" );
	}
      break;

    case SSR1: // balancing on right leg, upswing
      do_ss1( s, RIGHT );
      if ( s->comd[XX] < 0.0 )
	{
	  transition_to( s, SSR2, "SSR2 %g\n" );
	}
      break;

    case SSR2: // balancing on right leg, downswing
      do_ss2( s, RIGHT );
      if ( s->foot_status[LEFT] == FOOT_ON_GROUND )
	{
	  // need to debounce settling into DS
	  if ( 1 ) // s->foot_status[RIGHT] == FOOT_ON_GROUND )
	    {
	      s->ds_start_time = s->time;
	      s->ds_torque = s->torque1;
	      // s->ds_torque += -s->ds_torque_gain*(s->energy - s->energy_d)
	      transition_to( s, DS1, "DS1 %g\n" );
	    }
	  else // no DS period
	    {
	      s->knee_integrated_error[RIGHT] = 0.0;
	      do_liftoff( s, LEFT );
	      s->ss_start_time = s->time;
	      transition_to( s, SSL1, "SSL1 %g\n" );
	    }
	}
      break;

    default:
      fprintf( stderr, "Unknown state %d in run_state_machine.\n",
	       s->controller_state );
      exit( -1 );
    }
}

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

int controller( SIM *s )
{
  int i;

  /* Most kinematics already computed by SDFAST */
  compute_standard_state( s );

  /* Forces already computed by SDFAST */

  run_state_machine( s );

  // run_trajectory( s, s->time );

  run_servos( s );

  update_score( s );

  return 0;
}

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