/*****************************************************************************/
/* Includes */

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

/*****************************************************************************/
/* Defines */

#define OK 1
#define CRASHED 0
#define ERROR -1

#define NO_MOUSE 0
#define USE_MOUSE 1

/* Generic good stuff */
#define TRUE 1
#define FALSE 0

#define SCREEN_UPDATE_INTERVAL 0.03

/*****************************************************************************/
/* Declarations */

Communication_Buffer *tcb = NULL;

float lms_k_pole = 0.0;
float lms_b_pole = 0.0;
float lms_k_cart = 0.0;
float lms_b_cart = 0.0;
float lms_gain = 0.025;

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

char *default_iv_filename( Communication_Buffer *cb )
{
  return "cart-pole.iv";
}

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

float
default_time_step( Communication_Buffer *cb )
{
  return SCREEN_UPDATE_INTERVAL;
}

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

void
dump_state( Communication_Buffer *cb )
{
  int i;
}

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

void
save_data( Communication_Buffer *cb )
{
  printf( "save_data() not implemented.\n" );
}

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

void
handle_goto( Communication_Buffer *cb )
{
  int i;

  printf( "handle_goto()\n" );
}

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

void
handle_step( Communication_Buffer *cb )
{
  if ( !cb->playing )
    {
      advance_one_step( cb );
    }
  else
    cb->playing = FALSE;
}

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

void
handle_play( Communication_Buffer *cb )
{
  cb->playing = TRUE;
}

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

void
handle_mouse_left( float mouse_x, float mouse_y,
		    Communication_Buffer *cb )
{
  handle_step( cb );
}

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

void
handle_mouse_middle( float mouse_x, float mouse_y,
		    Communication_Buffer *cb )
{
  if ( cb->playing )
    handle_step( cb );
  else
    handle_play( cb );
}

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

void
handle_mouse_right( float mouse_x, float mouse_y,
		    Communication_Buffer *cb )
{
  handle_play( cb );
}

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

int index = 0;

void
set_desireds( Communication_Buffer *cb )
{

  if ( !cb->playing )
    printf( "\n%d\n", index );

  index++;
}

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

void
restart_simulation( Communication_Buffer *cb )
{
  cb->ps->cart_x = 0;
  cb->ps->pole_angle = 0;
  cb->ps->cart_xd = 0;	
  cb->ps->pole_angled = 0;
  cb->ps->cart_xdd = 0;		
  cb->ps->pole_angledd = 0;	
  cb->ps->cart_force = 0;	
  cb->ps->pole_torque = 0;	
  cb->ps->cart_x_desired = 0;	
  cb->ps->pole_angle_desired = 0;
  cb->ps->sin_pole_angle = sin( cb->ps->pole_angle );
  cb->ps->cos_pole_angle = cos( cb->ps->pole_angle );
}

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

void
init_simulation( Communication_Buffer *cb )
{
  pSimulator init_simulator();

  printf( "init_simulation...\n" );
  tcb = cb;
  cb->playing = FALSE;
  cb->time = 0.0;

  cb->ps = init_simulator();
}

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

advance_one_step( Communication_Buffer *cb )
{
  int i;

  for( i = SCREEN_UPDATE_INTERVAL/cb->ps->time_step; i >= 0; i-- )
    simulate( cb->ps );
}

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

int count = 0;

void
simulate_up_to_time( float stop_time, float mouse_x, float mouse_y,
		    Communication_Buffer *cb )
{

  if ( !cb->playing )
    return;

  if ( count % 100 == 0 )
    {
      printf( "%d: %f %f %f\n", count, stop_time, cb->ps->time,
	      stop_time - cb->ps->time );
      printf( "LMS gains:\n %f %f %f %f\n",
	      lms_k_pole, lms_b_pole, lms_k_cart, lms_b_cart );
    }
  count++;

  cb->ps->cart_x_desired = (mouse_x - 600.0)/300;

  advance_one_step( cb );
}

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

pSimulator
init_simulator()
{
  int i, index;
  pSimulator ps;

  ps = (pSimulator) malloc( sizeof( Simulator ) );
  if ( ps == NULL )
    {
      fprintf( stderr, "Failed to allocate simulator in init_simulator()\n" );
      exit( -1 );
    }

  /* Parameters. */
  ps->time_step = 0.001;   		/* s */
  ps->gravity = 9.8;  			/* m/s^2 */
  ps->pole_length = 2.0;                /* m */
  ps->pole_cm = ps->pole_length/2;	/* height of center of mass on pole. */
  ps->cart_mass = 1.0;  		/* kg */
  ps->pole_mass = 1.0;  		/* kg */
  ps->total_mass = ps->cart_mass + ps->pole_mass;
  ps->pole_I =  			/* principal moment in kg-m^2 */
    ps->pole_mass * ps->pole_length * ps->pole_length / 12.0;
  ps->pole_I_base =                     /* moment of inertia of pole at base */
    ps->pole_I + ps->pole_mass * ps->pole_cm * ps->pole_cm;

  /* State variables */
  /* Positions */
  ps->cart_x = 0;			/* m */
  ps->cart_y = 0.0;			/* m */
  ps->pole_angle = 0.0;		/* radians */
  /* pole_angle is angle with respect to vertical. */

  /* Velocities */
  ps->cart_xd = 0;			/* m/sec */
  ps->pole_angled = 0;			/* radians/sec */

  /* Accelerations */
  ps->cart_xdd = 0;			/* m/sec^2 */
  ps->pole_angledd = 0;			/* radians/sec^2 */

  /* Forces */
  ps->cart_force = 0;			/* Newtons */
  ps->pole_torque = 0;			/* Newton-meters */

  /* Controller variables. */
  ps->cart_x_desired = 0;		/* m */
  ps->pole_angle_desired = 0;		/* radians */

  /* For Manual control of cart position. */
  /* design_manual_controller( ps ); */

  /* For Automatic control of cart position. */
  /* design_controller1( ps ); */
  /* design_controller2( ps ); */
  /* design_controller3( ps ); */
  /* design_controller4( ps ); */

  /* To use gains in file "cart-pole-gains" */
  use_gains_from_file( ps, "cart-pole-gains" );

  /* Other useful variables. */
  ps->sin_pole_angle = sin( ps->pole_angle );
  ps->cos_pole_angle = cos( ps->pole_angle );

  ps->min_x_display = -2.0;
  ps->max_x_display = 2.0;
  ps->min_y_display = -1.0;
  ps->max_y_display = 1.0;

  ps->max_excursion = ps->max_x_display;
  ps->max_angle = 1.5;

  return ps;
}

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

#define MANUAL_CONTROL_NAT_FREQ (10*2*3.14)

design_manual_controller( pSimulator ps )
{
  ps->k_pole = 0;
  ps->b_pole = 0;
  ps->k_cart =
    -ps->total_mass * MANUAL_CONTROL_NAT_FREQ * MANUAL_CONTROL_NAT_FREQ;
  ps->b_cart = -ps->total_mass * 2 * MANUAL_CONTROL_NAT_FREQ;
  printf( "Control gains set for manual control:\n %f %f %f %f\n",
	 ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}

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

use_gains_from_file( ps, file_name )
     pSimulator ps;
     char *file_name;
{
  FILE *fopen(), *pstream;

  pstream = fopen( file_name, "r" );
  fscanf( pstream, "%f %f %f %f",
	 &(ps->k_pole), &(ps->b_pole), &(ps->k_cart), &(ps->b_cart) );
  fclose( pstream );
  printf( "Using gains from file %s:\n %f %f %f %f\n", file_name,
	 ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}


#define POLE_NAT_FREQ (2*3.14)
#define POLE_DAMPING_RATIO 0.7

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

design_controller1( ps )
     pSimulator ps;
{
  ps->k_pole =
    -(POLE_NAT_FREQ * POLE_NAT_FREQ +
     ps->pole_mass * ps->pole_cm * ps->gravity / ps->pole_I_base)*
       (ps->pole_I_base * ps->total_mass) /
	 (ps->pole_mass * ps->pole_cm);
  ps->b_pole =
    -(ps->total_mass * ps->pole_I_base * 2 *
     POLE_NAT_FREQ * POLE_DAMPING_RATIO) / 
      (ps->pole_mass * ps->pole_cm);
  /* 10,25 okay, a bit slow. */
  /* 40,50 okay */
  /* 160,100 unstable, 160,50 looks underdamped. */
  ps->k_cart = 40.0;
  ps->b_cart = 50.0; 
  printf( "Control gains kludged by hand:\n %f %f %f %f\n",
	 ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}

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

#define NAT_FREQ (1*3.14)

design_controller2( ps )
     pSimulator ps;
{
  float a, b, c;
  a = ( ps->pole_mass * ps->pole_cm * ps->gravity ) / ps->pole_I_base;
  b = ( ps->pole_mass * ps->pole_cm ) /
    ( ps->total_mass * ps->pole_I_base );
  c = 1 / ps->total_mass;
  
  ps->k_cart = (NAT_FREQ * NAT_FREQ * NAT_FREQ * NAT_FREQ)/(a*c);
  ps->b_cart = (4 * NAT_FREQ * NAT_FREQ * NAT_FREQ)/(a*c);
  ps->k_pole = -(6 * NAT_FREQ * NAT_FREQ + a + c*ps->k_cart)/b;
  ps->b_pole = -(4 * NAT_FREQ + c*ps->b_cart)/b;
  printf( "All poles placed at %f rad/sec:\n %f %f %f %f\n",
	 NAT_FREQ, ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}

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

#define DAMPING_RATIO 0.7

design_controller3( ps )
     pSimulator ps;
{
  float a, b, c;
  a = ( ps->pole_mass * ps->pole_cm * ps->gravity ) / ps->pole_I_base;
  b = ( ps->pole_mass * ps->pole_cm ) /
    ( ps->total_mass * ps->pole_I_base );
  c = 1 / ps->total_mass;
  
  ps->k_cart = (NAT_FREQ * NAT_FREQ * NAT_FREQ * NAT_FREQ)/(a*c);
  ps->b_cart = (4 * DAMPING_RATIO * NAT_FREQ * NAT_FREQ * NAT_FREQ)/(a*c);
  ps->k_pole = -(2 * NAT_FREQ * NAT_FREQ + a + c*ps->k_cart +
		4 * DAMPING_RATIO * DAMPING_RATIO * NAT_FREQ * NAT_FREQ)/b;
  ps->b_pole = -(4 * DAMPING_RATIO * NAT_FREQ + c*ps->b_cart)/b;
  printf(
 "All poles placed at %f rad/sec with damping ratio %f:\n %f %f %f %f\n",
	 NAT_FREQ, DAMPING_RATIO,
	 ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}

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

#define W1 (2*3.14)
#define W2 (0.1*2*3.14)
#define D1 0.7
#define D2 0.7

design_controller4( ps )
     pSimulator ps;
{
  float a, b, c;
  a = ( ps->pole_mass * ps->pole_cm * ps->gravity ) / ps->pole_I_base;
  b = ( ps->pole_mass * ps->pole_cm ) /
    ( ps->total_mass * ps->pole_I_base );
  c = 1 / ps->total_mass;
  
  ps->k_cart = (W1*W1*W2*W2)/(a*c);
  ps->b_cart = (2*D1*W1*W2*W2 + 2*D2*W2*W1*W1)/(a*c);
  ps->k_pole = -(W1*W1 + W2*W2 + 4*D1*D2*W1*W2 + a + c*ps->k_cart )/b;
  ps->b_pole = -(2*D1*W1 + 2*D2*W2 + c*ps->b_cart)/b;
  printf( "2 poles placed at %f rad/sec with damping ratio %f\n",
	 W1, D1 );
  printf( "2 poles placed at %f rad/sec with damping ratio %f\n",
	 W2, D2 );
  printf( " %f %f %f %f\n",
	 ps->k_pole, ps->b_pole, ps->k_cart, ps->b_cart );
}

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

/* Simulate the cart-pole for one time step. */
simulate( pSimulator ps )
{
  int return_value;
  float screen_width;

  if ( ( fabs( ps->cart_x ) > ps->max_excursion ) ||
      ( fabs( ps->pole_angle ) > ps->max_angle ) )
    {
      return( CRASHED );
    }

  control( ps );
  lms_learn( ps );
  dynamics( ps );

  if ( ( fabs( ps->cart_x ) > ps->max_excursion ) ||
      ( fabs( ps->pole_angle ) > ps->max_angle ) )
    {
      printf( "CRASHED: %f %f\n", ps->cart_x, ps->pole_angle );
      return( CRASHED );
    }
  else
    return_value = OK;

  return( return_value );
}

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

control( pSimulator ps )
{
  ps->cart_force =
    ps->k_pole * (ps->pole_angle - ps->pole_angle_desired)
    + ps->b_pole * ps->pole_angled
    + ps->k_cart * (ps->cart_x - ps->cart_x_desired)
    + ps->b_cart * ps->cart_xd;
}

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

dynamics( pSimulator ps )
{
  ps->pole_torque = 0;

  compute_accelerations( ps );

  integrate_accelerations( ps );

  compute_useful_stuff( ps );
}

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

compute_accelerations( pSimulator ps )
{
  /* Full dynamics. */
  float a, b, d, e, f, determinant;

  a = ps->total_mass;
  b = ps->pole_mass * ps->pole_cm * ps->cos_pole_angle;
  d = ps->pole_I_base;
  determinant = a*d - b*b;
  e = ps->cart_force
    - ps->pole_mass * ps->pole_cm * ps->pole_angled * ps->pole_angled
      * ps->sin_pole_angle;
  f = ps->pole_torque
    - ps->pole_mass * ps->gravity * ps->pole_cm * ps->sin_pole_angle;
  ps->cart_xdd = (d*e - b*f)/determinant;
  ps->pole_angledd = -(-b*e + a*f)/determinant;
}

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

integrate_accelerations( pSimulator ps )
{
  ps->cart_xd_next = ps->cart_xdd * ps->time_step + ps->cart_xd;
  ps->pole_angled_next = ps->pole_angledd * ps->time_step +
    ps->pole_angled;

  ps->cart_x += (ps->cart_xd + ps->cart_xd_next) * 0.5 * ps->time_step;
  ps->pole_angle += (ps->pole_angled + ps->pole_angled_next) *
    0.5 * ps->time_step;

  ps->cart_xd = ps->cart_xd_next;
  ps->pole_angled = ps->pole_angled_next;
  ps->time += ps->time_step;
}

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

compute_useful_stuff( pSimulator ps )
{
  ps->sin_pole_angle = sin( ps->pole_angle );
  ps->cos_pole_angle = cos( ps->pole_angle );
}

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

/* To do lms learning use manual control and always attempt to have
 the cart at the center position. */

lms_learn( pSimulator ps )
{
  float output, error;
  output =
    lms_k_pole * ps->pole_angle
    + lms_b_pole * ps->pole_angled
    + lms_k_cart * ps->cart_x
    + lms_b_cart * ps->cart_xd;
  error = output - ps->cart_force;
  lms_k_pole -= lms_gain * error * ps->pole_angle;
  lms_b_pole -= lms_gain * error * ps->pole_angled;
  lms_k_cart -= lms_gain * error * ps->cart_x;
  lms_b_cart -= lms_gain * error * ps->cart_xd;
}

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