/* Variational Monte Carlo for the hydrogen and helium atoms */
/* Note that numerical derivation is employed here           */
/* The cusp condition is not implemented                     */
/* For the hydrogen atom choose 1 for number of particles    */

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

/* declaraton of functions */

/* Function to calculate random numbers (from lib.c) */
double ran1(long *);

/* Function to declear matrix (from lib.c) */
void **matrix(int, int, int);

/* Function to read in data from screen  */
void initialise(/*int *,*/ int *, int *, int *, int *, int *, double *, 
		double *);

/* The Mc sampling for the variational Monte Carlo */
void mc_sampling(long, int, int, int, int, int, int, double, double, double, 
		  double *, double *, double *, double *);

/* Function to find distance between particle i and particle j */
double distance(double **, double *, double *, double, int, int);

/* The local energy */
double local_energy(double **, double *, double *, double, double, double, double, int, int);

/* prints to screen the results of the calculations  */
void output(double *, double *, double *, double *, int, int, int, double, double);

/* Begin of main program   */

int main(int argc, char** argv){
  
  long idum;
  int number_blocks, number_cycles, max_variations_correlation, thermalization;
  int i, dimension, number_particles, my_rank, size; 
  double initial_correlation, delta_correlation, lowest_correlation;
  double *sum_e, *sum_e2, *sum_variance, *time, local_time;
  double *cumulative_e, *cumulative_e2, *cumulative_variance;
  
  MPI_Init (&argc, &argv);
  MPI_Comm_size (MPI_COMM_WORLD, &size);
  MPI_Comm_rank (MPI_COMM_WORLD, &my_rank);
  dimension = 3;
  /*   Read in data */
  /*initialise(&number_blocks, &number_particles, &max_variations_correlation, 
  **	     &number_cycles, &thermalization, &initial_correlation, 
  **	     &delta_correlation) ;
  */
  number_blocks = 10;
  number_particles = 32;
  max_variations_correlation = 1;
  number_cycles = 100000;
  thermalization = 10;
  initial_correlation = 1.16;
  delta_correlation = 0;
    
  cumulative_e = (double *) malloc((max_variations_correlation)*sizeof(double));
  cumulative_e2 = (double *) malloc((max_variations_correlation)*sizeof(double));
  cumulative_variance = (double *) malloc((max_variations_correlation)*sizeof(double));
  sum_e = (double *) malloc((max_variations_correlation)*sizeof(double));
  sum_e2 = (double *) malloc((max_variations_correlation)*sizeof(double));
  sum_variance = (double *) malloc((max_variations_correlation)*sizeof(double));
  time = (double *) malloc(size*sizeof(double));

  idum=-42*my_rank+size;

  local_time = MPI_Wtime();
  
  /*  Do the mc sampling  */
  mc_sampling(idum, dimension, number_blocks, number_particles, max_variations_correlation, 
	      thermalization, number_cycles, initial_correlation, size,
	      delta_correlation, &lowest_correlation, cumulative_e, cumulative_e2,
	      cumulative_variance);

  local_time = fabs(MPI_Wtime()-local_time);
  
  MPI_Reduce(cumulative_e,sum_e,max_variations_correlation,MPI_DOUBLE,MPI_SUM,0,
	     MPI_COMM_WORLD);
  MPI_Reduce(cumulative_e2,sum_e2,max_variations_correlation,MPI_DOUBLE,MPI_SUM,0,
	     MPI_COMM_WORLD);
  MPI_Reduce(cumulative_variance,sum_variance,max_variations_correlation,MPI_DOUBLE,
	     MPI_SUM,0, MPI_COMM_WORLD);

  if(my_rank==0){
    MPI_Status status;
    time[my_rank]=local_time;
    for(i=1; i<size; i++){
      MPI_Probe (MPI_ANY_SOURCE, 500, MPI_COMM_WORLD, &status);
      MPI_Recv (&local_time, 1, MPI_DOUBLE, status.MPI_SOURCE, 
		500, MPI_COMM_WORLD, &status);
      time[i]=local_time;
    }

    for(i=0; i<max_variations_correlation; i++){
      sum_e[i] /= size;
      sum_e2[i] /= size*size;
      sum_variance[i] /= size;
      printf("\n\ncorrelation %12.5E\nenergy %12.5E\nvariance %12.5E\nstd dev %12.5E\n", 
	     i*delta_correlation+initial_correlation,sum_e[i],
	     sum_variance[i],sum_variance[i]/size/(number_blocks-1));
    }
    output(sum_e, sum_e2, sum_variance, time, size, number_blocks, 
	   max_variations_correlation, delta_correlation, initial_correlation);
  }
  else
    MPI_Send(&local_time, 1, MPI_DOUBLE, 0, 500, MPI_COMM_WORLD);

  MPI_Finalize;

  return 0;
}


/* Monte Carlo sampling with the Metropolis algorithm  */

void mc_sampling(long idum, int dimension, int number_blocks, int number_particles, 
		 int max_variations_correlation, int thermalization, int number_cycles, 
		 double initial_correlation, double size,
		 double delta_correlation, double *lowest_correlation, 
		 double *cumulative_e, double *cumulative_e2, 
		 double *cumulative_variance){
  
  int blocks, cycles, variate, accept, i, j, ji, factor;
  double sum, correlation5, energy, energy2, delta_e, lowest_energy;
  double correlation, **r, *r_old, *r_new, *r_distance, *r_temp, r2_old, r2_new;
  
  /* constants needed in calculations: */
  double l, lh2, kinetic_factor, density, scale, four_epsilon;
  density = 0.3648*0.9;  //Schiff-Verlet in sigma units
  l = pow((number_particles/density),0.3333333);  //Simulation cube
  scale = 1;
  lh2 = l*l/4;
  four_epsilon = 40.88;
  correlation = initial_correlation-delta_correlation;
  kinetic_factor = 9.27647;  //5*(hbar**2/m)/sigma**2
  factor = number_blocks*number_particles;
  printf("density=%12.5E, l=%12.5E, scale=%12.5E\n", density, l, scale);
  /* allocate matrices which contain the position of the particles */ 
  r = (double **) matrix( number_particles, dimension, sizeof(**r));
  r_old = (double *) malloc((dimension)*sizeof(double));
  r_new = (double *) malloc((dimension)*sizeof(double));
  r_distance = (double *) malloc((dimension)*sizeof(double));
  r_temp = (double *) malloc((dimension)*sizeof(double));

  for (i = 0; i < max_variations_correlation; i++){
    cumulative_e[i] = cumulative_e2[i] = 0;
  }
  lowest_energy = r2_old = r2_new = 0.;
  /* loop over variational parameters */ 
  for (variate=0; variate < max_variations_correlation; variate++){
    /* initialisations of variational parameters and energies */
    correlation += delta_correlation;
    correlation5 = pow(correlation,5);
    energy = energy2 = 0; accept =0; delta_e=0;
    /*  initial trial position */
    for (i = 0; i < number_particles; i++) { 
      for ( j=0; j < dimension; j++) {
	r[i][j] = l*ran1(&idum);
      }
    }

    /* loop over blocks */
    for (blocks = 1; blocks <= number_blocks; blocks++){
      /* loop over monte carlo cycles */
      energy=0;
      for (cycles = 1; cycles <= number_cycles+thermalization; cycles++){ 
	/* new position */
	for (i = 0; i < number_particles; i++) { 
	  for ( j=0; j < dimension; j++) {
	    /* attempted move */
	    r_old[j] = r[i][j];
	    r_new[j] = r_old[j]+scale*(ran1(&idum)-0.5);
	    /* move particle to th fundamental cell */
	    if (r_new[j]<0.)
	      r_new[j] += l;
	    if (r_new[j]>l)
	      r_new[j] -= l;
	  }
	  sum = 0.;
	  for (ji=0; ji<number_particles; ji++){
	    if(ji != i){
	      r2_old = distance(r, r_old, r_distance, l, ji, dimension);
	      r2_new = distance(r, r_new, r_distance, l, ji, dimension);
	      if(r2_old > lh2)
		r2_old = lh2;
	      if(r2_new > lh2)
		r2_new = lh2;
	      sum += 1./(r2_old*r2_old*sqrt(r2_old))-1./(r2_new*r2_new*sqrt(r2_new));
	    }
	  }
	  sum *= correlation5;
	  /* Metropolis test */
	  if(ran1(&idum) <= exp(sum)){
	    for(j=0; j<dimension; j++){
	      r[i][j] = r_new[j];
	    }
	    accept ++;
	  }
	}// end for(number_particles)
	if ( cycles > thermalization ) {
	  /* compute local energy */ 
	  delta_e = local_energy(r, r_distance, r_temp, correlation5, l, kinetic_factor, 
				 four_epsilon, dimension, number_particles);
	  /* update energies */ 
	  energy += delta_e;
	  energy2 += delta_e*delta_e;
	  //energy2 += delta_e*delta_e;
	} /* end if() */
      }   /* end of loop over MC trials  */ 
      if(blocks!=0){ // Neglect first block
	energy /= (double)number_cycles;
	energy2 /= (double)number_cycles;	
	cumulative_e[variate] += energy;
	cumulative_e2[variate] += energy2;
      }
    } /* end of loop over blocks */
    printf("variational parameter %12.5E accepted steps %12.5E\n",
	   correlation, (double)accept/(double)number_blocks/(double)number_particles);
    /* update the energy average and its squared */
    cumulative_e[variate] /= (double)factor;
    cumulative_e2[variate] /= (double)factor*(double)number_particles;
    cumulative_variance[variate] = (cumulative_e2[variate]-
				    cumulative_e[variate]*cumulative_e[variate]);
    printf("energy %12.5E\nvariance %12.5E\nstd dev %12.5E\n", cumulative_e[variate],
	   cumulative_variance[variate],cumulative_variance[variate]/(number_blocks-1));
    
  }    /* end of loop over variational  steps */
}   /* end mc_sampling function  */
  
double distance(double **r, double *r_temp, double *r_distance, double l, int ji, int dimension){

  int dim;
  double dist;
  dist = 0;
  for(dim=0; dim<dimension; dim++){
    r_distance[dim] = fabs(r_temp[dim]-r[ji][dim]); // always positive
    if(r_distance[dim]>l/2){
      r_distance[dim] = l-r_distance[dim];
    }
    dist += r_distance[dim]*r_distance[dim];
  }
  return dist; //dist^2, that is
}

/* Function to calculate the local energy */

double local_energy(double **r, double *r_distance, double *r_temp, double correlation5, 
		    double l, double kinetic_factor, double four_epsilon, int dimension, 
		    int number_particles){

  int particle_i, particle_j, i, j, k, dim;
  double local_energy, kinetic_e, potential_e;
  double dist, dist6, lh, z1, z2, z3;
  local_energy = dist = 0;
  lh = l*l/4;
  kinetic_e = potential_e = 0;
  for(particle_i=0; particle_i<number_particles-1; particle_i++){
    for(particle_j=particle_i+1; particle_j<number_particles; particle_j++){
      for(dim=0; dim<dimension; dim++){
	r_temp[dim]=r[particle_j][dim];
      }
      dist = distance(r, r_temp, r_distance, l, particle_i, dimension);
      if(dist<=lh) // cutoff correlation
	kinetic_e += kinetic_factor/(dist*dist*dist*sqrt(dist)); //kin/(dist^7)
      // printf("le %12.5E ", local_energy);

      /* check image of x, y and z axis */
      for(i=-1; i<=0; i++){
	z1=(r_distance[0]+i*l)*(r_distance[0]+i*l);
	for(j=-1; j<=0; j++){ 
	  z2=(r_distance[1]+j*l)*(r_distance[1]+j*l);
	  for(k=-1; k<=0; k++){
	    z3=(r_distance[2]+k*l)*(r_distance[2]+k*l);
	    dist=z1+z2+z3;
	    dist6=dist*dist*dist; //dist^6
	    if(dist<=l*l) // cutoff potential at L
	      potential_e += (1./dist6-1)/dist6;
	  }
	}
      }
    }
  }
  kinetic_e *= correlation5;
  potential_e *= four_epsilon;
  local_energy = kinetic_e+potential_e;
  return local_energy;
}

void output(double *sum_e, double *sum_e2, double *sum_variance, double *time, int size, 
	    int number_blocks, int max_variations_correlation, double delta_correlation, 
	    double initial_correlation){
  
  int i;
  double sum_time;
  char filename_1[42];
  char filename_2[42];
  char filename_3[42];
  FILE *output_file1, *output_file2, *output_file3;
  sprintf(filename_1, "plots/time_plot%02d.dat", size);
  sprintf(filename_2, "data/variationalTest_MPI_Reduce%02d.dat", size);
  sprintf(filename_3, "plots/time_plot_all.dat");
  output_file1 = fopen(filename_1, "w");
  output_file2 = fopen(filename_2, "w");
  output_file3 = fopen(filename_3, "a");
  
  sum_time=0;
  for(i=0; i<size; i++){
    sum_time+=time[i];
    //my_rank my_rank's time
    fprintf(output_file1, "%d %12.5E\n", i, time[i]);
  }
  fprintf(output_file3, "%d %12.5E %12.5E\n", size, sum_time, sum_time/size);
  
  for(i=0; i<max_variations_correlation; i++){
    sum_e[i] /= size;
    sum_e2[i] /= size*size;
    sum_variance[i] /= size;
    fprintf(output_file2, "%12.5E %12.5E %12.5E %12.5E",
	    i*delta_correlation+initial_correlation,sum_e[i],
	    sum_variance[i],sum_variance[i]/size/(number_blocks-1));
  }
  fclose(output_file1);
  fclose(output_file2);
  fclose(output_file3);
   
}