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

#include <stdio.h>
#include <math.h>
/* X11 graphics stuff */
#include "pixels.h"
#include "views.h"

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

/* How many things can we have in the simulator? */
#define MAX_N_ARMS 1
#define MAX_N_BLOCKS 10

/* Return values: */
#define OK 1
#define ERROR -1 

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

typedef struct arm
{
  /* Link lengths in meters */
  float link1;
  float link2;
  float link3;
  float finger;

  /* Joint values */
  float base;                 /* meters, zero is middle */
  float shoulder;             /* radians, zero is straight down. */
  float elbow;                /* radians, zero is straight down. */
  float wrist;                /* radians, zero is straight down. */
  float gripper;              /* meters, zero is all the way down */

  /* Joint velocities */
  float base_velocity;
  float shoulder_velocity;
  float elbow_velocity;
  float wrist_velocity;
  float gripper_velocity;

  /* Stuff calculated by forward kinematics */
  float base_x;
  float base_y;
  float elbow_x;
  float elbow_y;
  float wrist_x;
  float wrist_y;
  float gripper_x;
  float gripper_y;
  float lgripper_x;
  float lgripper_y;
  float rgripper_x;
  float rgripper_y;
  float ltip_x;
  float ltip_y;
  float rtip_x;
  float rtip_y;
  float c_shoulder;
  float s_shoulder;
  float c_elbow;
  float s_elbow;
  float c_wrist;
  float s_wrist;

  /* Stuff for quick erasing. */
  float base_x_display;
  float base_y_display;
  float elbow_x_display;
  float elbow_y_display;
  float wrist_x_display;
  float wrist_y_display;
  float gripper_x_display;
  float gripper_y_display;
  float lgripper_x_display;
  float lgripper_y_display;
  float rgripper_x_display;
  float rgripper_y_display;
  float ltip_x_display;
  float ltip_y_display;
  float rtip_x_display;
  float rtip_y_display;
} Arm, *pArm;

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

typedef struct block
{
  /* size */
  float width;
  float height;

  /* position */
  float x; /* location of corner (left lower corner when angle = 0) */
  float y;
  float angle;

  /* Stuff calculated by forward kinematics */
  float c2_x;
  float c2_y;
  float c3_x;
  float c3_y;
  float c4_x;
  float c4_y;
  float c_angle;
  float s_angle;

  /* Stuff for quick erasing. */
  float x_display;
  float y_display;
  float c2_x_display;
  float c2_y_display;
  float c3_x_display;
  float c3_y_display;
  float c4_x_display;
  float c4_y_display;

} Block, *pBlock;

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

typedef struct simulator
{
  float time;  		        /* seconds */
  float time_step;   		/* seconds */

  pArm arms[MAX_N_ARMS];

  pBlock blocks[MAX_N_BLOCKS];

  float min_x_display;
  float max_x_display;
  float min_y_display;
  float max_y_display;

  float display_clock;
  float display_interval;

  pView pv;   /* View of simulated object. */

} Simulator, *pSimulator;

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

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 );
    }

  /* Initialize simulator structure. */
  ps->time = 0.0;
  ps->time_step = 0.03;

  /* Start with no arms */
  for( i = 0; i < MAX_N_ARMS; i++ )
    ps->arms[i] = NULL;

  /* Start with no blocks */
  for( i = 0; i < MAX_N_BLOCKS; i++ )
    ps->blocks[i] = NULL;

  ps->min_x_display = -2;
  ps->max_x_display = 2;
  ps->min_y_display = 0;
  ps->max_y_display = 1;

  ps->display_clock = 0;
  ps->display_interval = 0.03;
  if ( ps->display_interval < ps->time_step )
    ps->display_interval = ps->time_step;

  ps->pv = (pView) malloc( sizeof( View ) );
  if ( ps->pv == NULL )
    {
      fprintf( stderr, "Failed to allocate view in init_simulator()\n" );
      exit( -1 );
    }

  set_up_view( ps->pv, ps->min_x_display, ps->min_y_display,
	      ps->max_x_display, ps->max_y_display );

  /* Add an arm */
  if ( MAX_N_ARMS > 0 )
    {
      ps->arms[0] = (pArm) malloc( sizeof( Arm ) );
      if ( ps->arms[0] == NULL )
	{
	  fprintf( stderr, "Failed to allocate arm in init_simulator()\n" );
	  exit( -1 );
	}
      ps->arms[0]->link1 = 0.4;
      ps->arms[0]->link2 = 0.4;
      ps->arms[0]->link3 = 0.1;
      ps->arms[0]->finger = 0.05;
      ps->arms[0]->base = 0;
      ps->arms[0]->shoulder = 0;
      ps->arms[0]->elbow = 0;
      ps->arms[0]->wrist = 0;
      ps->arms[0]->gripper = 0.1;
    }

  /* Add a block */
  if ( MAX_N_BLOCKS > 0 )
    {
      index = 0;
      ps->blocks[index] = (pBlock) malloc( sizeof( Block ) );
      if ( ps->blocks[index] == NULL )
	{
	  fprintf( stderr, "Failed to allocate block in init_simulator()\n" );
	  exit( -1 );
	}
      ps->blocks[index]->width = 0.1;
      ps->blocks[index]->height = 0.2;
      ps->blocks[index]->x = 1.0;
      ps->blocks[index]->y = 0.0;
      ps->blocks[index]->angle = 0;
    }

  /* Add a block */
  if ( MAX_N_BLOCKS > 1 )
    {
      index = 1;
      ps->blocks[index] = (pBlock) malloc( sizeof( Block ) );
      if ( ps->blocks[index] == NULL )
	{
	  fprintf( stderr, "Failed to allocate block in init_simulator()\n" );
	  exit( -1 );
	}
      ps->blocks[index]->width = 0.1;
      ps->blocks[index]->height = 0.1;
      ps->blocks[index]->x = 0.7;
      ps->blocks[index]->y = 0.0;
      ps->blocks[index]->angle = 0;
    }

  /* Add a block */
  if ( MAX_N_BLOCKS > 2 )
    {
      index = 2;
      ps->blocks[index] = (pBlock) malloc( sizeof( Block ) );
      if ( ps->blocks[index] == NULL )
	{
	  fprintf( stderr, "Failed to allocate block in init_simulator()\n" );
	  exit( -1 );
	}
      ps->blocks[index]->width = 0.1;
      ps->blocks[index]->height = 0.1;
      ps->blocks[index]->x = -1.0;
      ps->blocks[index]->y = 0.0;
      ps->blocks[index]->angle = 0;
    }

  for( i = 0; i < MAX_N_ARMS; i++ )
    {
      arm_forward_kinematics( i, ps );
    }

  for( i = 0; i < MAX_N_BLOCKS; i++ )
    {
      block_forward_kinematics( i, ps );
    }

  return ps;
}

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

main()
{
  pSimulator ps; /* default simulated object. */
  
  init_window( 4.0 );

  ps = init_simulator();

  init_display( ps );

  printf( "Press return to start\n" );
  getchar();

  run( 100.0, ps );

  printf( "GAME OVER. Press return to quit\n" );
  getchar();

  done_with_window();
}

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

run( n_seconds, ps )
     float n_seconds;
     pSimulator ps;
{
  float time_to_stop;

  for( time_to_stop = ps->time + n_seconds; time_to_stop > ps->time; )
    if ( simulate( ps ) != OK )
      break;
}

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

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

  your_planner_goes_here( ps );
  dynamics( ps );

  ps->display_clock -= ps->time_step;
  if ( ps->display_clock <= 0.0 )
    {
      erase_display( ps );
      draw_display( ps );
      ps->display_clock += ps->display_interval;
      flush_graphics();
      for( i = 0; i < 100000; i++ )
	/* waste time to slow this down */;
    }
  
  /* Check for CRASH goes here */
  return_value = OK;

  return( return_value );
}

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

your_planner_goes_here( ps )
     register pSimulator ps;
{
  /* simulate sensing here. */
  /* then do planning */
  /* then generate robot commands (in this case the velocities) */

  /* A cheap example. */
  ps->arms[0]->base_velocity = cos( (double) (ps->time) );
  ps->arms[0]->shoulder_velocity = cos( (double) (ps->time) );
  ps->arms[0]->elbow_velocity = -2*cos( (double) (ps->time) );
  ps->arms[0]->wrist_velocity = cos( (double) (ps->time) );
  ps->arms[0]->gripper_velocity = 0.1*cos( (double) (ps->time) );
}

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

dynamics( ps )
     register pSimulator ps;
{
  int i;

  ps->time += ps->time_step;

  for( i = 0; i < MAX_N_ARMS; i++ )
    {
      arm_forward_dynamics( i, ps );
      arm_forward_kinematics( i, ps );
    }

  for( i = 0; i < MAX_N_BLOCKS; i++ )
    {
      block_forward_kinematics( i, ps );
    }
}

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

arm_forward_dynamics( index, ps )
     int index;
     register pSimulator ps;
{
  register pArm pa;

  pa = ps->arms[index];

  /* Check if there is an arm there */
  if ( pa == NULL )
    return;

  pa->base += pa->base_velocity*ps->time_step;
  pa->shoulder += pa->shoulder_velocity*ps->time_step;
  pa->elbow += pa->elbow_velocity*ps->time_step;
  pa->wrist += pa->wrist_velocity*ps->time_step;
  pa->gripper += pa->gripper_velocity*ps->time_step;
}

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

arm_forward_kinematics( index, ps )
     int index;
     register pSimulator ps;
{
  register pArm pa;

  pa = ps->arms[index];

  /* Check if there is an arm there */
  if ( pa == NULL )
    return;

  pa->c_shoulder = cos( (double) (pa->shoulder) );
  pa->s_shoulder = sin( (double) (pa->shoulder) );
  pa->c_elbow = cos( (double) (pa->shoulder + pa->elbow) );
  pa->s_elbow = sin( (double) (pa->shoulder + pa->elbow) );
  pa->c_wrist = cos( (double) (pa->shoulder + pa->elbow + pa->wrist) );
  pa->s_wrist = sin( (double) (pa->shoulder + pa->elbow + pa->wrist) );

  pa->base_x = pa->base;
  pa->base_y = 1.0;
  pa->elbow_x = pa->base_x + pa->link1*pa->s_shoulder;
  pa->elbow_y = pa->base_y - pa->link1*pa->c_shoulder;
  pa->wrist_x = pa->elbow_x + pa->link2*pa->s_elbow;
  pa->wrist_y = pa->elbow_y - pa->link2*pa->c_elbow;
  pa->gripper_x = pa->wrist_x + pa->link3*pa->s_wrist;
  pa->gripper_y = pa->wrist_y - pa->link3*pa->c_wrist;
  pa->lgripper_x = pa->gripper_x - 0.5*pa->gripper*pa->c_wrist;
  pa->lgripper_y = pa->gripper_y - 0.5*pa->gripper*pa->s_wrist;
  pa->rgripper_x = pa->gripper_x + 0.5*pa->gripper*pa->c_wrist;
  pa->rgripper_y = pa->gripper_y + 0.5*pa->gripper*pa->s_wrist;
  pa->ltip_x = pa->lgripper_x + pa->finger*pa->s_wrist;
  pa->ltip_y = pa->lgripper_y - pa->finger*pa->c_wrist;
  pa->rtip_x = pa->rgripper_x + pa->finger*pa->s_wrist;
  pa->rtip_y = pa->rgripper_y - pa->finger*pa->c_wrist;
}

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

block_forward_kinematics( index, ps )
     int index;
     register pSimulator ps;
{
  register pBlock pb;

  pb = ps->blocks[index];

  /* Check if there is a block there */
  if ( pb == NULL )
    return;

  pb->c_angle = cos( (double) (pb->angle) );
  pb->s_angle = sin( (double) (pb->angle) );

  pb->c2_x = pb->x + pb->width*pb->c_angle;
  pb->c2_y = pb->y + pb->width*pb->s_angle;
  pb->c3_x = pb->c2_x - pb->height*pb->s_angle;
  pb->c3_y = pb->c2_y + pb->height*pb->c_angle;
  pb->c4_x = pb->x - pb->height*pb->s_angle;
  pb->c4_y = pb->y + pb->height*pb->c_angle;
}

/*****************************************************************************/
/*****************************************************************************/
/* Graphics stuff, change for your machine and OS */
/*****************************************************************************/

init_display( ps )
     pSimulator ps;
{
  clear_window();
  draw_display( ps );
  flush_graphics();
}

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

draw_display( ps )
     register pSimulator ps;
{
  int i;

  for( i = 0; i < MAX_N_ARMS; i++ )
    {
      draw_arm( i, ps );
    }

  for( i = 0; i < MAX_N_BLOCKS; i++ )
    {
      draw_block( i, ps );
    }
}

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

erase_display( ps )
     register pSimulator ps;
{
  int i;

  for( i = 0; i < MAX_N_ARMS; i++ )
    {
      erase_arm( i, ps );
    }

  for( i = 0; i < MAX_N_BLOCKS; i++ )
    {
      erase_block( i, ps );
    }
}

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

/* These come from views.c */
extern float last_x1, last_y1, last_x2, last_y2;

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

draw_arm( index, ps )
     int index;
     register pSimulator ps;
{
  register pArm pa;

  pa = ps->arms[index];

  /* Check if there is an arm there */
  if ( pa == NULL )
    return;

  vdraw_line( pa->base_x, pa->base_y, pa->elbow_x, pa->elbow_y,
	     XOR, ps->pv );
  pa->base_x_display = last_x1;
  pa->base_y_display = last_y1;
  pa->elbow_x_display = last_x2;
  pa->elbow_y_display = last_y2;

  vdraw_line( pa->elbow_x, pa->elbow_y, pa->wrist_x, pa->wrist_y,
	     XOR, ps->pv );
  pa->wrist_x_display = last_x2;
  pa->wrist_y_display = last_y2;

  vdraw_line( pa->wrist_x, pa->wrist_y, pa->gripper_x, pa->gripper_y,
	     XOR, ps->pv );
  pa->gripper_x_display = last_x2;
  pa->gripper_y_display = last_y2;

  vdraw_line( pa->lgripper_x, pa->lgripper_y, pa->rgripper_x, pa->rgripper_y,
	     XOR, ps->pv );
  pa->lgripper_x_display = last_x1;
  pa->lgripper_y_display = last_y1;
  pa->rgripper_x_display = last_x2;
  pa->rgripper_y_display = last_y2;

  vdraw_line( pa->lgripper_x, pa->lgripper_y, pa->ltip_x, pa->ltip_y,
	     XOR, ps->pv );
  pa->ltip_x_display = last_x2;
  pa->ltip_y_display = last_y2;

  vdraw_line( pa->rgripper_x, pa->rgripper_y, pa->rtip_x, pa->rtip_y,
	     XOR, ps->pv );
  pa->rtip_x_display = last_x2;
  pa->rtip_y_display = last_y2;
}

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

erase_arm( index, ps )
     int index;
     register pSimulator ps;
{
  register pArm pa;

  pa = ps->arms[index];

  /* Check if there is an arm there */
  if ( pa == NULL )
    return;

  _vdraw_line( (double) pa->rgripper_x_display,
	      (double) pa->rgripper_y_display,
	      (double) pa->rtip_x_display,
	      (double) pa->rtip_y_display, XOR );

  _vdraw_line( (double) pa->lgripper_x_display,
	      (double) pa->lgripper_y_display,
	      (double) pa->ltip_x_display,
	      (double) pa->ltip_y_display, XOR );

  _vdraw_line( (double) pa->lgripper_x_display,
	      (double) pa->lgripper_y_display,
	      (double) pa->rgripper_x_display,
	      (double) pa->rgripper_y_display, XOR );

  _vdraw_line( (double) pa->wrist_x_display,
	      (double) pa->wrist_y_display,
	      (double) pa->gripper_x_display,
	      (double) pa->gripper_y_display, XOR );

  _vdraw_line( (double) pa->elbow_x_display,
	       (double) pa->elbow_y_display,
	      (double) pa->wrist_x_display,
	      (double) pa->wrist_y_display, XOR );

  _vdraw_line( (double) pa->base_x_display,
	      (double) pa->base_y_display,
	      (double) pa->elbow_x_display,
	      (double) pa->elbow_y_display, XOR );
}

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

draw_block( index, ps )
     int index;
     register pSimulator ps;
{
  register pBlock pb;

  pb = ps->blocks[index];

  /* Check if there is a block there */
  if ( pb == NULL )
    return;

  vdraw_line( pb->x, pb->y, pb->c2_x, pb->c2_y, XOR, ps->pv );
  pb->x_display = last_x1;
  pb->y_display = last_y1;
  pb->c2_x_display = last_x2;
  pb->c2_y_display = last_y2;

  vdraw_line( pb->c2_x, pb->c2_y, pb->c3_x, pb->c3_y, XOR, ps->pv );
  pb->c3_x_display = last_x2;
  pb->c3_y_display = last_y2;

  vdraw_line( pb->c3_x, pb->c3_y, pb->c4_x, pb->c4_y, XOR, ps->pv );
  pb->c4_x_display = last_x2;
  pb->c4_y_display = last_y2;

  vdraw_line( pb->c4_x, pb->c4_y, pb->x, pb->y, XOR, ps->pv );
}

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

erase_block( index, ps )
     int index;
     register pSimulator ps;
{
  register pBlock pb;

  pb = ps->blocks[index];

  /* Check if there is a block there */
  if ( pb == NULL )
    return;

  _vdraw_line( (double) pb->c4_x_display,
	      (double) pb->c4_y_display,
	      (double) pb->x_display,
	      (double) pb->y_display, XOR );

  _vdraw_line( (double) pb->c3_x_display,
	      (double) pb->c3_y_display,
	      (double) pb->c4_x_display,
	      (double) pb->c4_y_display, XOR );

  _vdraw_line( (double) pb->c2_x_display,
	      (double) pb->c2_y_display,
	      (double) pb->c3_x_display,
	      (double) pb->c3_y_display, XOR );

  _vdraw_line( (double) pb->x_display,
	      (double) pb->y_display,
	      (double) pb->c2_x_display,
	      (double) pb->c2_y_display, XOR );
}

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