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

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

            // setup for surveying   E  &  y                       //

                                               // Michel Vallieres //


// USE
// gcc survey_scattering.c -lm -o survey_scattering
// ./survey_scattering 120 120 > see
// cat see | python pixelplot.py 0.0 0.7 0.005 0.07 120 120

#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 )
{
  *t_min = 0.0;
  *t_max = 425.0;
  *dt = 0.01;
}


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;
  int    yinit, einit, Nyinit, Neinit;
  double emin, emax, eng, ymin, ymax, de, dy;
  int    print_theta, print_final;

                           // print control
  print_control = 0;       // - trajectory
  print_theta = 1;         // - theta (image)
  print_final = 0;         // - summary

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

                           // energy & y initial ranges
  if ( argc != 3  )
    {
      fprintf( stderr, "Use ./survey_scattering Ne Ny\n");
      exit( 1 );
    }
  Neinit = atoi( argv[1] );
  Nyinit = atoi( argv[2] );
  emin = 0.005;
  emax = 0.07;
  ymin = 0.0;
  ymax = 0.7;

  de = ( emax-emin ) / (double)(Neinit-1);
  dy = ( ymax-ymin ) / (double)(Nyinit-1);

                           // loop over initial energy
  for ( einit=Neinit-1 ; einit>=0 ; einit-- )
    {
      eng = emin + (double)einit * de;
                           // loop over initial y values
      for ( yinit=0 ; yinit<Nyinit ; yinit++ )
	{
                           // initial values
	  f_t[0] = -4.0;
	  f_t[1] = ymin + (double)yinit * dy;
	  f_t[2] = sqrt( 2.0*eng/mass );
	  f_t[3] = 0.0;
                           // 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];

                           // 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] );

          if ( print_final == 1 )
	    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 ( print_theta == 1 )
	    printf( "%f\n", exit_angle );


	}  // end loop over y initial

    }  // end loop over E initial

}