/**********************************************************************************
 * 2D crystal model with Verlet algorithm molecular dynamics                      *
 **********************************************************************************
 * Lattice size: L x L (with periodic boundary condition)                         *
 * lattice constant: a = 1;                                                       *
 * atomic mass: m = 1;                                                            *
 * with stressed springs between nearest neigbors (spring                         *
 * constant: k, unstressed length: l0).                                           *
 *                                                                                *
 * Verlet algorithm                                                               *
 *   see: http://en.wikipedia.org/wiki/Verlet_integration#St.C3.B6rmer.27s_method *
 *   or:  http://newton.phy.bme.hu/%7Ekertesz/Ea9.pdf pp. 7-8.                    *
 *                                                                                *
 * Output: (dispersion.dat)                                                       *
 *   k     omega     amplitude                                                    *
 *   .       .           .                                                        *
 *   .       .           .                                                        *
 *   .       .           .                                                        *
 * Each k-point is separated by blank line. This data file It can be plotted with *
 * gnuplot splot immediately.                                                     *
 **********************************************************************************
 * Budapest, 2011/11/14                                                           *
 * author: Laszlo BALOGH, blaci947 (a) yahoo.co.uk                                *
 **********************************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_fft_complex.h>
#include <gsl/gsl_rng.h>

#define PI 3.1415926

// Macros for array indexing:
//  * P.b.c. is implemented here
//  * Valid index range: i € [-1, L+1]; j € [-1, L+1]
#define X0(i,j)  (x0[ ((i)+L)%L * L + ((j)+L)%L ] )
#define X1(i,j)  (x1[ ((i)+L)%L * L + ((j)+L)%L ] )
#define X2(i,j)  (x2[ ((i)+L)%L * L + ((j)+L)%L ] )
#define Y0(i,j)  (y0[ ((i)+L)%L * L + ((j)+L)%L ] )
#define Y1(i,j)  (y1[ ((i)+L)%L * L + ((j)+L)%L ] )
#define Y2(i,j)  (y2[ ((i)+L)%L * L + ((j)+L)%L ] )

// Macros for the FFT complex:
#define REAL(z,i) ((z)[2*(i)])
#define IMAG(z,i) ((z)[2*(i)+1])


int main(void) {
  // parameters:
  const int    L  = 100;        // sample size: L x L
  const double T  = 129.0;       // max. simulation time
  const double dt = 2e-2;       // time step in the O.D.E. solver
  const int    N  = T / dt + 1; // maximum number of output data
  const int    fftN = 1024;      // maximum number of points in the FFT (not the best solution...)
  const double k  = 1.0;        // spring constant
  const double l0 = 0.5;        // stressless length of the springs (latt. const. = 1) 
  const double A  = 0.01;       // initial amplitude

  // data:
  double *x0 = NULL, *x1 = NULL, *x2 = NULL, // x(t-dt), x(t), x(t+dt), ...
         *y0 = NULL, *y1 = NULL, *y2 = NULL; // all will be LxL arrays
  double fftdata[2*fftN];
  int    fftsample;
  
  // what to simulate:
  double kx = 2*PI/L, ky = 0.0; // wave vector (kx,ky)
  const double omega = 1.0;            // phonon frequency (only for initial excitaion)
  const int mesi = 0, mesj = 0;        // mesure the (x(t),y(t)) function of the (mesi,mesj) atom
  // polarization vector for longitudinal wave:
  const double u = A * kx / sqrt(kx*kx + ky*ky),  v = A * ky / sqrt(kx*kx + ky*ky);
  // polarization vector for transvesal wave:
  //const double u = -A * ky / sqrt(kx*kx + ky*ky),  v = A * kx / sqrt(kx*kx + ky*ky);
  
  // variables:
  double m = 1.0;        // mass 
  double dx, dy, l, F;   // auxiliary quantities
  double xdd, ydd;       // acceleration; "x-dot-dot", "y-dot-dot"
  int i,j,t;
  double *tmpptr = NULL;
  gsl_rng *rng = gsl_rng_alloc(gsl_rng_mt19937);
  
  // set stdout as output (or open a file if you want)
  FILE* outf = stdout;
  
  // We need some memory:
  fprintf(stderr, "memory...");
  x0 = (double*) calloc(L*L, sizeof(double));
  x1 = (double*) calloc(L*L, sizeof(double));
  x2 = (double*) calloc(L*L, sizeof(double));
  y0 = (double*) calloc(L*L, sizeof(double));
  y1 = (double*) calloc(L*L, sizeof(double));
  y2 = (double*) calloc(L*L, sizeof(double));
  
  fprintf(stderr, "ok\n");
  
  FILE* fftoutf = NULL;
  fftoutf = fopen("dispersion.dat", "wt");

  // here we have to define the curve (here: straight line) in the Brilouin zone
  kx = PI; ky = PI;
  for ( ; ky>1e-5; ky -= 2.0*PI/L , kx -= 2.0*PI/L )
  {
    fprintf(stderr, "kx,ky = (%e,%e) ", kx, ky);
  
  // reset the array for the FFT:
  for (t = 0; t < fftN; t++){
                    REAL(fftdata,t) = 0.0;
                    IMAG(fftdata,t) = 0.0;
  }
  
  // Initialization:
  for (i=0; i<L; i++) {
    for (j=0; j<L; j++) {
      X0(i,j) = u * sin(kx*i + ky*j);
      Y0(i,j) = v * sin(kx*i + ky*j);
      X1(i,j) = X0(i,j) - u * omega * cos(kx*i + ky*j) * dt;
      Y1(i,j) = Y0(i,j) - v * omega * cos(kx*i + ky*j) * dt;
    }
  }
  
  
  // Simulation:
  fprintf(stderr, "simulation...");
  for (t=0; t<N; t++) {
    gsl_rng_set(rng, 986754987); // Not really nice, but the most simple solution. It works.
    for (i=0; i<L; i++) {
      for (j=0; j<L; j++) {
	
	xdd = ydd = 0.0;
	
        // Problem 1: evry atom is the same:
        //m = 1.0;
        
        // Problem 2: sublattice:
	//if ((i+j)%2) m = 14.0/12.0; else m = 1.0;
	
	// problem 3: 50-50% alloy:
	if (gsl_rng_uniform_int(rng, 2)) m = 1.0; else m = 14.0/12.0;
	
	// left:
	dx = X1(i,j) - X1(i-1,j);
	dy = Y1(i-1,j) - Y1(i,j);
	l = sqrt( (1.0 + dx)*(1.0 + dx) + dy*dy );
	F = k*(l-l0);
	xdd += -(1.0 + dx)/l * F;
	ydd += dy / l * F;
	
	// right:
	dx = X1(i+1,j) - X1(i,j);
	dy = Y1(i+1,j) - Y1(i,j);
	l = sqrt( (1.0 + dx)*(1.0 + dx) + dy*dy );
	F = k*(l-l0);
	xdd += (1.0 + dx)/l * F;
	ydd += dy / l * F;
	
	// up:
	dx = X1(i,j+1) - X1(i,j);
	dy = Y1(i,j+1) - Y1(i,j);
	l = sqrt( dx*dx + (1.0 + dy)*(1.0 + dy) );
	F = k*(l-l0);
	xdd += dx / l * F;
	ydd += (1.0 + dy) / l * F;
	
	// down:
	dx = X1(i,j-1) - X1(i,j);
	dy = Y1(i,j) - Y1(i,j-1);
	l = sqrt( dx*dx + (1.0 + dy)*(1.0 + dy) );
	F = k*(l-l0);
	xdd += dx / l * F;
	ydd += -(1.0 + dy) / l * F;

        // these 2 lines are the Verlet algorithm:
	X2(i,j) = 2.0 * X1(i,j) - X0(i,j) + dt*dt*xdd/m;
	Y2(i,j) = 2.0 * Y1(i,j) - Y0(i,j) + dt*dt*ydd/m;
      }
    }
    //fprintf(outf, "%12.8lf   %10.6lf   %10.6lf\n", t*dt, X2(mesi,mesj), X2(mesi+1,mesj));

    // We add every fftsample^th data to the FFT dataset:
    fftsample = 50;
    if ( (t % fftsample == 0) && (t / fftsample < fftN) ) {
      REAL(fftdata,t/fftsample) = X2(mesi,mesj);
      IMAG(fftdata,t/fftsample) = 0.0;
    }
    
    // prepare for the next ireration:
    // need to copy: X1 --> X0; X2 --> X1.
    // TRICK: there is no real copying :)
    tmpptr = x0;
    x0 = x1;
    x1 = x2;
    x2 = tmpptr;
    tmpptr = y0;
    y0 = y1;
    y1 = y2;
    y2 = tmpptr;
  } // next t step
  fprintf(stderr, "OK\n");
  
  // FFT:
  fprintf(stderr, "fft...");

  // this line does the job:
  gsl_fft_complex_radix2_forward (fftdata, 1, fftN);

  /*
  for (t = 0; t < fftN; t++){
    fprintf(fftoutf,"%d %e %e\n", t,
                    REAL(fftdata,t)/sqrt(fftN),
                    IMAG(fftdata,t)/sqrt(fftN));
  }
  */
  // here is the output of the spectrum in this k-point:
  for (t = 0; t < fftN; t++){
    fprintf(fftoutf,"%e %e %e\n",
	            PI + PI + (PI - ky)*sqrt(2.0),
	            (double)t,
                    sqrt( REAL(fftdata,t)*REAL(fftdata,t) + IMAG(fftdata,t)*IMAG(fftdata,t) ) / sqrt(fftN) );
  }
  fprintf(fftoutf,"\n");
  fflush(fftoutf); // write out buffer
  
  fprintf(stderr, "OK\n");

  }// kx
  fclose(fftoutf);
  return 0;
}