/************************************************************************/
/*
compass.c: This is where all the plant specific details are handled:
   dynamics (SDFAST interface), integration, cost function, state book keeping
*/
/************************************************************************/

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

#include "compass1.h"

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

double stance_d_desired = 1.34*LEG_LENGTH; /* 3 miles/hour */

double u_penalty = 1.0;
double stance_d_penalty = 1.0;

double time_step = 0.001;
double time = 0.0;

int switch_dynamics_enabled = 1;

int print_stuff = 0;

int traj_too_long = 10000;

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

/* These are handy constants. XX rather than X to avoid math.h conflicts */
#define XX 0
#define YY 1
#define ZZ 2

/* SDFAST constants and variables */
#define NQ 4
#define NU 4
#define NSTATE (NQ+NU)
#define CTOL	1e-4	/* constraint tolerance */
#define GND -1
const double baumgarte=100.0; /* could be 10.0 - 1000.0 */

/* Here we define the state ( x, y, l, lr, xd, yd, ld, lrd ) */
double xt[NSTATE];
/* Here we define the derivative of the state */
static double dxt[NSTATE];
/* Here we define the actuator command signal */
static double hip_torque = 0.0;
/* Here we define additional variables derived from the state */
double left_foot_position[3], right_foot_position[3];
double left_foot_velocity[3], right_foot_velocity[3];

/* Leg state machines */
const double GROUND_LEVEL = 0.0;
/* x_foot_active: is x an active constraint? */
int r_foot_active = 0, l_foot_active = 0;

/* define contact points. */
static double left_l[3]={0.0,-0.9,0.0};
static double right_l[3]={0.0,-0.9,0.0};
static double trunk_p[3]={0.0,0.1,0.0};

static int n_steps = 0;

static int my_counter = 0;

/************************************************************************/
/* Initialize the state vector */

void init_dynamics_state( double *full_state )
{
  int i, j;

  for( i = 0; i < NSTATE; i++ )
    {
      xt[i] = full_state[i];
      dxt[i] = 0;
    }

  for( i = 0, j = NQ; i < NU; i++, j++ )
    { 
      dxt[i] = xt[j];
    }

  l_foot_active = rint(full_state[L_FOOT_DOWN]);
  /* Enforce one or the other foot down */
  if ( l_foot_active )
    r_foot_active = 0;
  else
    r_foot_active = 1;

  time = 0.0;
  sdstate( 0.0, xt, dxt );
  sdpos(0, left_l, left_foot_position);
  sdvel(0, left_l, left_foot_velocity);
  sdpos(1, right_l, right_foot_position); 
  sdvel(1, right_l, right_foot_velocity);
}

/************************************************************************/
/* Initialize sdfast  */

void init_dynamics( double *full_state )
{

  sdinit(); /* initialize SDFAST model */

  sdstab(2.0*baumgarte, baumgarte*baumgarte); 

  init_dynamics_state( full_state );

  sdprinterr(stderr);
}

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

int dynamics_type( double *x_full, double *u_full )
{
  return rint( x_full[L_FOOT_DOWN] );
}

/************************************************************************/
/* Leg state machine */

void switch_dynamics()
{
  if ( !switch_dynamics_enabled )
    return;

  if( r_foot_active
      && left_foot_position[YY] < GROUND_LEVEL
      && left_foot_position[YY] - time_step*left_foot_velocity[YY] 
      >= GROUND_LEVEL
      && xt[L_LEG] >= 0.0 )
    {
      n_steps++; 
      if ( print_stuff )
	printf("%7.3f: l-# of steps: %d\n", time, n_steps);
      l_foot_active = 1;
      r_foot_active = 0;
    }
  else if( l_foot_active
	   && right_foot_position[YY] < GROUND_LEVEL
	   && right_foot_position[YY] - time_step*right_foot_velocity[YY] 
	   >= GROUND_LEVEL
	   && xt[L_LEG] + xt[LR_LEG] >= 0.0 )
    {
      n_steps++; 
      if ( print_stuff )
	printf("%7.3f: r-# of steps: %d\n", time, n_steps);
      r_foot_active = 1;
      l_foot_active = 0;
    }
}

/************************************************************************/
/* This is what is called on each integration step */

void integrate_one_time_step( double *action, double *next_full_state  )
{ 
  int i;
  int  flag; /* set to 1 when things are changed or first call */
  int err; 
    /* { OK, DERIVATIVE_DISCONTINUITY, SYSTEM_LOCKED, CONSTRAINTS_ERR } */
  double errest;

  hip_torque = action[0];

  flag=1;
  /* TOL = 1.0d-4
  sdmotion(&time,xt,dxt,TIME_STEP,CTOL,TOL,&flag,&err);
  */
  sdfmotion(&time,xt,dxt,time_step,CTOL,&flag,&errest,&err);

  sdpos(0, left_l, left_foot_position);
  sdvel(0, left_l, left_foot_velocity);
  sdpos(1, right_l, right_foot_position); 
  sdvel(1, right_l, right_foot_velocity);

  switch_dynamics();

  for( i = 0; i < NSTATE; i++ )
    next_full_state[i] = xt[i];

  next_full_state[L_FOOT_DOWN] = l_foot_active;
  next_full_state[R_FOOT_DOWN] = r_foot_active;

  sdprinterr(stderr);
}

/************************************************************************/
/************************************************************************/
/* SDFAST stuff */
/************************************************************************/
/* This is where the control (hip_torque) is applied. May be called many
times per integration step at any time or state. */

void sduforce(double t, double *q, double *u)
{
  sdhinget( 1, 0, hip_torque );
}

/************************************************************************/
/* SDFAST constraint stuff */

void sduperr(double t, double *q, double *errs) 
{

    errs[0]=errs[1]=errs[2]=errs[3]=errs[4]=0;

    /* Don't seem to need to recalculate these
    sdpos(0, left_l, left_foot_position);
    sdpos(1, right_l, right_foot_position); 
    */

    if( right_foot_position[YY] < GROUND_LEVEL && r_foot_active )
      {
	errs[0]=right_foot_position[YY];
      }
    if( left_foot_position[YY] < GROUND_LEVEL && l_foot_active )
      {
	errs[1]=left_foot_position[YY];
      }
    if( xt[YY] < GROUND_LEVEL )
      {
	errs[2]=q[YY];
      }
}

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

void sduverr(double t, double *q, double *u, double *errs)
{

    errs[0]=errs[1]=errs[2]=errs[3]=errs[4]=0;

    /* Don't know why we need to recalculate these, but we do */
    sdvel(0, left_l, left_foot_velocity);
    sdvel(1, right_l, right_foot_velocity);

    if( right_foot_position[YY] < GROUND_LEVEL && r_foot_active )
      {
	errs[0]=right_foot_velocity[YY];
	errs[3]=right_foot_velocity[XX];
      }
    if( left_foot_position[YY] < GROUND_LEVEL && l_foot_active )
      {
	errs[1]=left_foot_velocity[YY];
	errs[4]=left_foot_velocity[XX];
      }
    if( xt[YY] < GROUND_LEVEL )
      {
	errs[2]=u[YY];
      }
}   

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

void sduaerr(double t, double *q, double *u, double *udot, double *errs)
{
    double left_foot_acceleration[3], right_foot_acceleration[3];

    errs[0]=errs[1]=errs[2]=errs[3]=errs[4]=0;

    sdacc(0, left_l, left_foot_acceleration);
    sdacc(1, right_l, right_foot_acceleration);

    if( right_foot_position[YY] < GROUND_LEVEL && r_foot_active )
      {
	errs[0]=right_foot_acceleration[YY];
	errs[3]=right_foot_acceleration[XX];
      }
    if( left_foot_position[YY] < GROUND_LEVEL && l_foot_active )
      {
	errs[1]=left_foot_acceleration[YY];
	errs[4]=left_foot_acceleration[XX];
      }
    if( xt[YY] < GROUND_LEVEL )
      {
	errs[2]=udot[YY];
      }
}

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

void sduconsfrc(double t, double *q, double *u, double *m)
{
  FILE *fp;
  double frc1[3]={0,0,0}, frc2[3]={0,0,0}, frc3[3]={0,0,0};
  double frc4[3]={0,0,0}, frc5[3]={0,0,0};

  if( right_foot_position[YY] < GROUND_LEVEL && m[0] >= 0 && r_foot_active)
    {
       frc1[YY]=m[0];
       sdtrans(GND, frc1, 0, frc1);
       sdpointf(1, right_l, frc1);

       frc4[XX]=m[3];
       sdtrans(GND, frc4, 0, frc4);
       sdpointf(1, right_l, frc4);
    }

  if( left_foot_position[YY] < GROUND_LEVEL && m[1] >= 0 && l_foot_active )
    {  
       frc2[YY]=m[1];
       sdtrans(GND, frc2, 1, frc2);
       sdpointf(0, left_l, frc2);

       frc5[XX]=m[4];
       sdtrans(GND, frc5, 1, frc5);
       sdpointf(0, left_l, frc5);
    }
  if( xt[YY] < GROUND_LEVEL && m[2] >= 0 )
    {
       frc3[YY]=m[2];
       sdtrans(GND, frc3, 0, frc3);
       sdpointf(0, trunk_p, frc3);
    }
}

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

double one_step_cost( double *x_full, double *u_full )
{
  double cost = 0;
  double x[MAX_N_X];
  int n_x;
  double leg_product;
  double xdot, ydot;

  full_x_to_state_x( x_full, x, &n_x );

  /* What we really want */
  cost += u_penalty*u_full[0]*u_full[0];

  xdot = -LEG_LENGTH*cos(x[S1_STANCE_LEG])*x[S1_STANCE_LEG_D];
  ydot = -LEG_LENGTH*sin(x[S1_STANCE_LEG])*x[S1_STANCE_LEG_D];

  cost += stance_d_penalty*(xdot - stance_d_desired)*(xdot - stance_d_desired)
    + stance_d_penalty*ydot*ydot;

  return cost;
}

/************************************************************************/
/************************************************************************/
/************************************************************************/
/* Set up state that optimizer sees. */

void full_x_to_state_x( double *x_full, double *x, int *n_x )
{
  if ( x_full[L_FOOT_DOWN] )
    {
      x[0] = x_full[LR_LEG];
      x[1] = x_full[LR_LEG_D];
      x[2] = x_full[L_LEG];
      x[3] = x_full[L_LEG_D];
    }
  else
    {
      x[0] = -x_full[LR_LEG];
      x[1] = -x_full[LR_LEG_D];
      x[2] = x_full[LR_LEG] + x_full[L_LEG];
      x[3] = x_full[LR_LEG_D] + x_full[L_LEG_D];
    }
  if ( n_x != NULL )
    *n_x = 4;
}

/************************************************************************/
/* Set up command that optimizer sees. */

void full_u_to_state_u( double *u_full, double *u, int *n_u )
{
  u[0] = u_full[0];
  if ( n_u != NULL )
    *n_u = 1;
}

/************************************************************************/
/* Move command back to full command */

void state_u_to_full_u( double *u, double *u_full )
{

  u_full[0] = u[0];
}

/************************************************************************/
/* Set up state that DYNAMICS sees. May violate constraints. */

void state_x_to_full_x( double *x, double *x_full_context, 
			  double *x_full_result )
{
  int i;
  double foot_x, foot_y, foot_x_d, foot_y_d;

  for( i = 0; i < N_X_FULL; i++ )
    x_full_result[i] = x_full_context[i];

  if ( x_full_context[L_FOOT_DOWN] > 0.1 )
    {
      foot_x = x_full_context[HIP_X] 
	+ sin( x_full_context[L_LEG] )*LEG_LENGTH;
      foot_y = x_full_context[HIP_Y] 
	- cos( x_full_context[L_LEG] )*LEG_LENGTH;
      foot_x_d = x_full_context[HIP_X_D]
        + cos( x_full_context[L_LEG] )*LEG_LENGTH*x_full_context[L_LEG_D];
      foot_y_d = x_full_context[HIP_Y_D]
        + sin( x_full_context[L_LEG] )*LEG_LENGTH*x_full_context[L_LEG_D];
      x_full_result[L_LEG] = x[2];
      x_full_result[LR_LEG] = x[0];
      x_full_result[L_LEG_D] = x[3];
      x_full_result[LR_LEG_D] = x[1];
      x_full_result[HIP_X] = foot_x 
	- sin( x_full_result[L_LEG] )*LEG_LENGTH;
      x_full_result[HIP_Y] = foot_y 
	+ cos( x_full_result[L_LEG] )*LEG_LENGTH;
      x_full_result[HIP_X_D] = foot_x_d
        - cos( x_full_result[L_LEG] )*LEG_LENGTH*x_full_result[L_LEG_D];
      x_full_result[HIP_Y_D] = foot_y_d
        - sin( x_full_result[L_LEG] )*LEG_LENGTH*x_full_result[L_LEG_D];
    }
  else
    {
      foot_x = x_full_context[HIP_X]
        + sin( x_full_context[LR_LEG] + x_full_context[L_LEG] )*LEG_LENGTH;
      foot_y = x_full_context[HIP_Y]
        - cos( x_full_context[LR_LEG] + x_full_context[L_LEG] )*LEG_LENGTH;
      foot_x_d = x_full_context[HIP_X_D]
        + cos( x_full_context[LR_LEG] + x_full_context[L_LEG] )*LEG_LENGTH
        *( x_full_context[LR_LEG_D] + x_full_context[L_LEG_D] );
      foot_y_d = x_full_context[HIP_Y_D]
        + sin( x_full_context[LR_LEG] + x_full_context[L_LEG] )*LEG_LENGTH
        *( x_full_context[LR_LEG_D] + x_full_context[L_LEG_D] );
      x_full_result[L_LEG] = x[2] + x[0];
      x_full_result[LR_LEG] = -x[0];
      x_full_result[L_LEG_D] = x[3] + x[1];
      x_full_result[LR_LEG_D] = -x[1];
      x_full_result[HIP_X] = foot_x
        - sin( x_full_result[LR_LEG] + x_full_result[L_LEG] )*LEG_LENGTH;
      x_full_result[HIP_Y] = foot_y
        + cos( x_full_result[LR_LEG] + x_full_result[L_LEG] )*LEG_LENGTH;
      x_full_result[HIP_X_D] = foot_x_d
        - cos( x_full_result[LR_LEG] + x_full_result[L_LEG] )*LEG_LENGTH
        *( x_full_result[LR_LEG_D] + x_full_result[L_LEG_D] );
      x_full_result[HIP_Y_D] = foot_y_d
        - sin( x_full_result[LR_LEG] + x_full_result[L_LEG] )*LEG_LENGTH
        *( x_full_result[LR_LEG_D] + x_full_result[L_LEG_D] );
    }
}

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