//                      //
                   // chaotic scattering   //
                   //                      //

            // notation                                            //
            // variables    x --> 0, y --> 1, vx --> 2, vy --> 3   //
            // sources      0 --> 1, 1 --> 2, 2 --> 3              //


                                  // Michel Vallieres, Fall 2003 //

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


                // dimension of system
#define ndim 4
                // model parameters
int    NS;
double k, b, xs[3], ys[3], mass;



                // total energy //
double total_energy( double f[] )
{
  double energy, x, y, vx, vy, temp;
  int    source;

  vx = f[2];
  vy = f[3];
  energy = 0.5*mass*( vx*vx + vy*vy );
  for ( source=0; source<NS ; source++ )
    {
      x = f[0] - xs[source];
      y = f[1] - ys[source];
      temp = k * exp( -( x*x + y*y )/(b*b) );
      energy = energy + 0.5*b*b*temp;
    }

  return energy;
}


                // generates RHS of diff eq //
void rhs_function( double t, double f[], double rhsf[] )
{
  int    source;
  double x, y, temp;


  rhsf[0] = f[2];                             // x eq
  rhsf[1] = f[3];                             // y eq
  rhsf[2] = 0.0;                              // vx eq
  rhsf[3] = 0.0;                              // vy eq
  for ( source=0; source<NS ; source++ )
    {
      x = f[0] - xs[source];
      y = f[1] - ys[source];
      temp = k * exp( -( x*x + y*y )/(b*b) );
      rhsf[2] = rhsf[2] + temp * x; 
      rhsf[3] = rhsf[3] + temp * y;
    }

}


                // Runge-Kutta 4th order time step     //
                // f_t[] is the function to advance    //
                // requires a global definition for    //
                // int ndim: dimension of system       //
                // void rhs_function (t, f, rhsf )     //
          
void rk4_step ( double t, double dt, double f_t[] )
{
  double k1[ndim], k2[ndim], k3[ndim],
            k4, yt[ndim], rhsf[ndim];
  int i;

  rhs_function( t, f_t, rhsf );
  for ( i=0 ; i<ndim ; i++ )
    {
      k1[i] = dt * rhsf[i];
      yt[i] = f_t[i] + 0.5*k1[i];
    }
  rhs_function( t+0.5*dt, yt, rhsf );
  for ( i=0 ; i<ndim ; i++ )
    {
      k2[i] = dt * rhsf[i];
      yt[i] = f_t[i] + 0.5*k2[i];
    }
  rhs_function( t+0.5*dt, yt, rhsf );
  for ( i=0 ; i<ndim ; i++ )
    {
      k3[i] = dt * rhsf[i];
      yt[i] = f_t[i] + k3[i];
    }
  rhs_function( t+dt, yt, rhsf );
  for ( i=0 ; i<ndim ; i++ )
    {
      k4 = dt * rhsf[i];
      f_t[i] = f_t[i] + ( 0.5*k1[i] + k2[i] + k3[i] + 0.5*k4 ) / 3.0;
    }
}


                           // define time grid
void define_time_grid( double *t_min, double *t_max, double *dt )
{
  double eta;

  *t_min = 0.0;
  *t_max = 425.0;
  eta = 0.01;
  *dt = eta/sqrt(k);
}


                           // set initial conditions
void set_initial_conditions(  double f[] )
{
  f[0] = -4.0;
  f[1] = 0.05;
  f[2] = 0.15;
  f[3] = 0.0;
}


void define_model()
{
  k = 1.0;
  b = 0.25;
  mass = 1.0;
  NS = 3;
  xs[0] = 0;
  ys[0] = 0.5;
  xs[1] = 0;
  ys[1] = -0.5;
  xs[2] = sqrt(3.0)/2.0;
  ys[2] = 0.0;
}


int main( int argc, char * argv[] )
{

  double dt, t, t_min, t_max;
  double x, v, energy, e0;
  double f_t[ndim];
  double dist_squared, max_dist_squared;
  double exit_angle, x_init, y_init, vx_init, vy_init;
  int    print_control, input_control, all_input_control, i, n;


                           // define model
  define_model();

                           // time grid
  define_time_grid( &t_min, &t_max, &dt );

                           // print trajectories (or not)
                           // input y (or not)
  print_control = 0;
  input_control = 0;
  all_input_control = 0;
  for ( n=1; n < argc ; n++ )
    if ( argv[n][0] == '-' )
      switch ( argv[n][1] )
        {
          case 't': 
             print_control = 1;
	     break;
          case 'i': 
             input_control = 1;
	     break;
	  case 'a':
             all_input_control = 1;
	     break;
	  case 'm':
             t_max = atof( argv[++n] ); 
             break;
          default: 
            fprintf( stderr, "Syntax: scattering <-p> <-i> <-a> <-m #>\n" );
	    fprintf( stderr, "-t --> print traj; -i --> input y_init\n" );
            fprintf( stderr, "-a --> input initial x, y, vx, vy\n" );
            fprintf( stderr, "-m --> t_max\n" );
	    exit(1);
         }

                           // loop over initial y values
  for ( ; ; )
    {
                           // initial values
      set_initial_conditions( f_t );

                           // input initial y value if needed
      if ( input_control == 1 )
        scanf( "%lf", &f_t[1] );

                           // or all variables...
      if ( all_input_control == 1 )
        scanf( "%lf %lf %lf %lf", 
                &f_t[0], &f_t[1], &f_t[2], &f_t[3] );

                           // store initial conditions for output later
      x_init = f_t[0];
      y_init = f_t[1];
      vx_init = f_t[2];
      vy_init = f_t[3];

                           // end execution signal
      if ( x_init < -20 || y_init < -20 )
        exit(0);

                           // energy check
      energy = total_energy( f_t );
      e0 = energy;

                           // initial time
      t = t_min;

                           // print results
      if ( print_control == 1 )
        printf( " %f %f %f %f %f %f %f\n", 
                t, f_t[0], f_t[1], f_t[2], f_t[3], energy, energy-e0 ); 

                           // loop over time until
                           // particle escapes
      dist_squared = 0.0;
      max_dist_squared = 4.1*4.1;
      while (  t <= t_max )
        {
                           // RK4 step
          rk4_step ( t, dt, f_t );
          t = t+dt;
                           // energy check
          energy = total_energy( f_t );

                           // square of distance
          dist_squared = f_t[0]*f_t[0] + f_t[1]*f_t[1];

                           // print results
          if ( print_control == 1 )
            printf( " %f %f %f %f %f %f %f\n", 
                   t, f_t[0], f_t[1], f_t[2], f_t[3], energy, energy-e0 ); 

	                   // safety check on time
          if ( dist_squared >= max_dist_squared  ) break;

        } // end while over distance

                           // exit time & scattering angle
      exit_angle = atan2( f_t[3], f_t[2] );
      fprintf( stderr, " %f %f %f %f (%f %f %f %f -- %f)\n", 
         y_init, exit_angle, t, energy, x_init,y_init,vx_init,vy_init, e0  );

      if ( input_control == 0 && all_input_control == 0 ) 
         break;

    }  // end for over y initial

}