/*
Copyright Emiliano Cristiani 2011
I thank Simone Cacace for the function "SaveBitmap" 

This code computes the basins of attraction of some attractors 1,2,...,num_attr.
Attractors attract particles following the Newton law (ma=F=Cm/r^2 + friction)

The final solution is a matrix. Each entry is associated to a point of the plane and represents the number 
of the attractor the point is attracted to (-1 means that trajectory did not end on any attractor after N iterations).
The matrix is converted in a bitmap image (colors should be defined by the user in the function "SaveBitmap")

It is also possible to compute a single trajectory of the particle defining by hand the initial point and the initial velocity.
In this case the solution is saved in the file "one_trajectory.txt"

Set the parameters you prefer below!
*/

#include<stdio.h>
#include<math.h>
#include<iostream>
using namespace std;

struct vettR2{
double x;
double y;
};

/*------------------FIRST SET OF PARAMETERS (see also below) ---------------*/
const int num_attr=4;	//number of attractors
const int pixel_x=200;	//number of pixels in the horizontal direction
const int pixel_y=200;	//number of pixels in the vertical direction
/*--------------------------------------------------------------------------*/

double C[num_attr], Dt, b, epsilon;
vettR2 M[num_attr];

//prototypes
void FIELD(double[], double[]);
int SaveBitmap(const char*,int,int,short int**);


///////// MAIN /////////
int main (void){
short int **sol;  
bool go;
int i,j,k,n,steps,close[num_attr];
long int N;
double left,right,bottom,top,Dx,Dy,Rsquare,dsquare[num_attr];
double xold[5],xnew[5],vx,vy;
double uapp[5],app1[5],app2[5],app3[5],app4[5];
FILE *pf;

sol=new short int *[pixel_x];
for(i=0;i<pixel_x;i++){
  sol[i]=new short int[pixel_y];
}

/*---------------------- SECOND SET OF PARAMETERS ------------------------*/
int choice_main=1;       //1: regular computation, 2:compute just one trajectory
int choice_ODE_solver=1; //1: Euler explicit I order, 2:Runge-Kutta IV order

M[0].x=2.0; M[0].y=2.0;    //position of attractor #1
M[1].x=2.0;  M[1].y=-2.0;  //position of attractor #2
M[2].x=-2.0;  M[2].y=2.0;  //position of attractor #3
M[3].x=-2.0;  M[3].y=-2.0; //position of attractor #4
//...
C[0]=16.0; //strenght of attractor #1 (16.0)
C[1]=16.0; //strenght of attractor #2 (16.0)
C[2]=16.0; //strenght of attractor #3 (16.0)
C[3]=16.0; //strenght of attractor #4 (16.0)
//...
N=100000;             //max number of iterations for the ODE solver (100000)
epsilon=0.1;          //size of attractors (0.1)
steps=600;            //stop criterion for the ODE solver: min number of steps "around" final attractor to stop (600)
Rsquare=1.0;          //how much "around" is large (1.0)
Dt=0.003;             //time step for ODE solver (0.003)
b=0.2;                //friction (0.2)
left=-20.0;           //bounds of the plane
right=20.0;
bottom=-20.0;
top=20.0;
vx=1.0; vy=1.0;       //initial velocity of any trajectory (1.0,1.0)
/*---------------------------------------------------------*/


if (choice_main==2){ //save a particular trajectory
int end;
pf=fopen("one_trajectory.txt","w");
xold[1]=-8.0; //initial position x
xold[2]=-2.0;  //initial position y 
xold[3]=vx; 
xold[4]=vy;    
for (k=0;k<num_attr;k++) {close[k]=0;}  
n=0;
go=true;
while (go && n<N-1){
  if (choice_ODE_solver==1){   //Euler
    FIELD(xold,app1);
    for (k=1;k<5;k++) {xnew[k]=xold[k]+Dt*app1[k];}
  }
  else if (choice_ODE_solver==2){   //Runge-Kutta IV order
    FIELD(xold,app1);
    for (k=1;k<5;k++) {uapp[k]=xold[k]+(Dt/2.0)*app1[k];}
    FIELD(uapp,app2);
    for (k=1;k<5;k++) {uapp[k]=xold[k]+(Dt/2.0)*app2[k];}
    FIELD(uapp,app3);
    for (k=1;k<5;k++) {uapp[k]=xold[k]+Dt*app3[k];}
    FIELD(uapp,app4);
    for (k=1;k<5;k++) {xnew[k]=xold[k]+Dt*(1.0/6.0)*(app1[k] + 2.0*app2[k] + 2.0*app3[k] + app4[k]);}
  }
  for (k=0;k<num_attr;k++) {dsquare[k]=(M[k].x-xnew[1])*(M[k].x-xnew[1]) + (M[k].y-xnew[2])*(M[k].y-xnew[2]);}
  for (k=0;k<num_attr;k++) {if (dsquare[k]<Rsquare) {close[k]++;} else {close[k]=0;}}
  for (k=0;k<num_attr;k++) {if (close[k]>steps) {end=k; go=false;}}
  for (k=1;k<5;k++) {xold[k]=xnew[k];}
  fprintf(pf,"%lf %lf\n",xnew[1],xnew[2]);
  n++;
}  
cout << "number of steps for the ODE solver= " << n << endl << "final attractor= #" << end << endl;
fclose(pf);
cout << "Trajectory saved in 'one_trajectory.txt'" << endl;
}


if (choice_main==1){   //regular computation
  
//space steps
Dx=(right-left)/(pixel_x-1.0); //cout << "Dx= " << Dx << endl;
Dy=(top-bottom)/(pixel_y-1.0); //cout << "Dy= " << Dy << endl;

//initialization of solution
for (i=0;i<pixel_x;i++){ 
  for (j=0;j<pixel_y;j++){
    sol[i][j]=-1;
  }
}

//here we go
cout << "Computation started" << endl;
for (i=0;i<pixel_x;i++){
  for (j=0;j<pixel_y;j++){
    xold[1]=left+(i*Dx);
    xold[2]=bottom+(j*Dy); 
    xold[3]=vx; 
    xold[4]=vy;  
    for (k=0;k<num_attr;k++) {close[k]=0;}  
    n=0;
    go=true;
    while (go && n<N-1){
      if (choice_ODE_solver==1){ //Euler
        FIELD(xold,app1);
        for (k=1;k<5;k++) {xnew[k]=xold[k]+Dt*app1[k];}
      }
      else if (choice_ODE_solver==2){   //Runge-Kutta IV order
        FIELD(xold,app1);
        for (k=1;k<5;k++) {uapp[k]=xold[k]+(Dt/2.0)*app1[k];}
        FIELD(uapp,app2);
        for (k=1;k<5;k++) {uapp[k]=xold[k]+(Dt/2.0)*app2[k];}
        FIELD(uapp,app3);
        for (k=1;k<5;k++) {uapp[k]=xold[k]+Dt*app3[k];}
        FIELD(uapp,app4);
        for (k=1;k<5;k++) {xnew[k]=xold[k]+Dt*(1.0/6.0)*(app1[k] + 2.0*app2[k] + 2.0*app3[k] + app4[k]);}
      }
      for (k=0;k<num_attr;k++) {dsquare[k]=(M[k].x-xnew[1])*(M[k].x-xnew[1]) + (M[k].y-xnew[2])*(M[k].y-xnew[2]);}
      for (k=0;k<num_attr;k++) {if (dsquare[k]<Rsquare) {close[k]++;} else {close[k]=0;}}
      for (k=0;k<num_attr;k++) {if (close[k]>steps) {sol[i][j]=k; go=false;}}
      for (k=1;k<5;k++) {xold[k]=xnew[k];}
      n++;
    }  
    //cout << "n= " << n << endl;
  }
  {
  cout << "line " << i << "/" << pixel_x-1 << " complete" << endl;
  }
}

//write final solution on file "solution.txt"
pf=fopen("solution.txt","w");
for (i=0;i<pixel_x;i++){ 
  for (j=0;j<pixel_y;j++){
    fprintf(pf,"%d ",sol[i][j]);
  }
  fprintf(pf,"\n");
}
fclose(pf);

//check solution
int count=0;
for (i=0;i<pixel_x;i++){ 
  for (j=0;j<pixel_y;j++){
      if (sol[i][j]==-1) {count++;}
  }
}
if (count==0){ 
  cout << "Computation ended successfully!" << endl;}
else{
  cout << count << " trajectories did not reach any attractor. Try increasing N" << endl;
}

cout << "Solution saved in 'solution.txt' " << endl;

SaveBitmap("solution.bmp",pixel_x,pixel_y,sol);  //create bitmap
cout << "Solution converted in 'solution.bmp' " << endl;
cout << "end" << endl;
}

system("pause");
return 0;
}


///////////////////////////////////////////////////////////////
void FIELD(double u[], double f[]){
 double dsquare[num_attr];
 int k;
 
 f[1]=u[3];
 f[2]=u[4];
 
 for (k=0;k<num_attr;k++){dsquare[k]=(M[k].x-u[1])*(M[k].x-u[1]) + (M[k].y-u[2])*(M[k].y-u[2]);}
 f[3]=0.0; f[4]=0.0;
 for (k=0;k<num_attr;k++){
   f[3]=f[3]+C[k]*(M[k].x-u[1])/(pow(dsquare[k],1.5)+epsilon);
   f[4]=f[4]+C[k]*(M[k].y-u[2])/(pow(dsquare[k],1.5)+epsilon);
 }
 f[3]=f[3]-b*u[3];
 f[4]=f[4]-b*u[4];
}


///////////////////////////////////////////////////////////////
int SaveBitmap(const char *filename, int w, int h, short int **mat)
{
 // This function writes an uncompressed 24bit bitmap
 FILE *fp;          
 int shift=0;   if((3*w)%4) shift=4-((3*w)%4);
 unsigned char  bits_r, bits_g, bits_b, bits_null=0;
 unsigned short bfType=19778;               // Magic number for bitmap file (19778='MB')
 unsigned int   bfSize=54+(3*w+shift)*h;    // File size (40bytes for header + 15bytes for info header + Xbytes image data)   
 unsigned short bfReserved1=0;              // Unused
 unsigned short bfReserved2=0;              // Unused
 unsigned int   bfOffBits=54;               // Offset to image data (size of header + info header)
 unsigned int   biSize=40;                  // Header size
 int            biWidth=w;                  // Image width
 int            biHeight=h;                 // Image height
 unsigned short biPlanes=1;                 // Number of planes of the target device
 unsigned short biBitCount=24;              // Color depth
 unsigned int   biCompression=0;            // Compression
 unsigned int   biSizeImage=(3*w+shift)*h;  // Image size
 int            biXPelsPerMeter=0;          // Unused
 int            biYPelsPerMeter=0;          // Unused 
 unsigned int   biClrUsed=0;                // Unused
 unsigned int   biClrImportant=0;           // Unused

 if ((fp = fopen(filename, "wb")) == NULL) return -1;

 // Write bitmap header and bitmap info header
 fwrite(&bfType,          sizeof(unsigned short), 1, fp);
 fwrite(&bfSize,          sizeof(unsigned int),   1, fp);
 fwrite(&bfReserved1,     sizeof(unsigned short), 1, fp);
 fwrite(&bfReserved2,     sizeof(unsigned short), 1, fp);
 fwrite(&bfOffBits,       sizeof(unsigned int),   1, fp);
 fwrite(&biSize,          sizeof(unsigned int),   1, fp);
 fwrite(&biWidth,         sizeof(int),            1, fp);
 fwrite(&biHeight,        sizeof(int),            1, fp);
 fwrite(&biPlanes,        sizeof(unsigned short), 1, fp);
 fwrite(&biBitCount,      sizeof(unsigned short), 1, fp);
 fwrite(&biCompression,   sizeof(unsigned int),   1, fp);
 fwrite(&biSizeImage,     sizeof(unsigned int),   1, fp);
 fwrite(&biXPelsPerMeter, sizeof(int),            1, fp);
 fwrite(&biYPelsPerMeter, sizeof(int),            1, fp);
 fwrite(&biClrUsed,       sizeof(unsigned int),   1, fp);
 fwrite(&biClrImportant,  sizeof(unsigned int),   1, fp);
 
 // Convert indexed values of mat into RGB colors and write on file
 for(int j=h-1; j>=0; j--) // Bitmap format writes the image with a vertical flip
 {
  for(int i=0; i<w; i++)
  {
   if(mat[i][j]==-1)    {bits_r=255; bits_g=255; bits_b=255;}     // white (error)
   else if(mat[i][j]==0){bits_r=255; bits_g=0;   bits_b=0;  }     // red (attractor #1)
   else if(mat[i][j]==1){bits_r=0;   bits_g=0;   bits_b=255;}     // blue (attractor #2)
   else if(mat[i][j]==2){bits_r=0;   bits_g=0;   bits_b=0;  }     // black (attractor #3)
   else if(mat[i][j]==3){bits_r=0;   bits_g=230; bits_b=0;  }     // green (attractor #4)
   // ... add colors!!!
   else
   {
    bits_r=255;bits_g=255;bits_b=255;  // white (error)
   }
   // Bitmap format needs colors in BGR order
   fwrite(&bits_b, sizeof(unsigned char), 1, fp);
   fwrite(&bits_g, sizeof(unsigned char), 1, fp);
   fwrite(&bits_r, sizeof(unsigned char), 1, fp);
  }
  // If the image width is not a multiple of 4, each line shound be completed with zeros up to the nearest multiple of 4 
  for(int s=0;s<shift;s++){ fwrite(&bits_null, sizeof(unsigned char), 1, fp);}  
 }
 fclose(fp);
 return 0;
}