#include <iostream>
#include <cstring>
#include <cmath>

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#include "particle.hh"
#include "constants.hh"
#include "file_stream.hh"

using namespace std;

#define GRADIENT	(40 * megavolt / m)
#define FREQUENCY	(12 * gigahertz)
#define PHASE		(0.0 * degree)

#define LENGTH		(1 * m)

static inline void accelerate(Particle &particle, double e0, double frequency, double phase)
{
  particle.energy += e0 * sin(2*M_PI * particle.z*um * frequency / c_light - phase);
}

int main(int argc, char **argv )
{
  if (argc<2) {
    printf("usage:  make_beam NUM_OF_PARTICLES [PHASE/deg]\n\n");
    exit(1);
  }
  
  if (gsl_rng *rng = gsl_rng_alloc(gsl_rng_mt19937)) {
    
    size_t num_of_particles = size_t(atoi(argv[1]));
    double phase = argc==3?(atof(argv[2])*degree):PHASE;
        
    double energy0 = 9 * GeV;
    double energy_spread = 0.32 * per_cent;
    
    struct {
      double charge;
      double sigma_z;
      struct {
	double beta;
	double alpha;
	double gamma;
	double emittance;
      } h, v;
      double sE_correlation;
    } bunch;
    
    bunch.charge = 0.67e-9 * coulomb;
    bunch.sigma_z = 175 * um;
    
    bunch.h.beta = 200 * m;
    bunch.h.alpha = 4.8;
    bunch.h.emittance = 570 * nm * rad * electron_mass_c2 / energy0;
    
    bunch.v.beta = 200 * m;
    bunch.v.alpha = 4.8;
    bunch.v.emittance = 4 * nm * rad * electron_mass_c2 / energy0;
    
    bunch.h.gamma = (1+bunch.h.alpha*bunch.h.alpha)/bunch.h.beta;
    bunch.v.gamma = (1+bunch.v.alpha*bunch.v.alpha)/bunch.v.beta;
    
    bunch.sE_correlation = 69.5 / m;
    
    struct {
      double x;
      double y;
      double xp;
      double yp;
    } sigma;
    
    sigma.x = sqrt(bunch.h.beta * bunch.h.emittance)/2;
    sigma.y = sqrt(bunch.v.beta * bunch.v.emittance)/2;
    sigma.xp = sqrt(bunch.h.gamma * bunch.h.emittance)/2;
    sigma.yp = sqrt(bunch.v.gamma * bunch.v.emittance)/2;
    
    cerr << "generating " << num_of_particles << " particles...\n";
    cerr << "\tsigma.x = " << sigma.x / um << " um\n";
    cerr << "\tsigma.y = " << sigma.y / um << " um\n";
    cerr << "\tsigma.xp = " << sigma.xp / urad << " urad\n";
    cerr << "\tsigma.yp = " << sigma.yp / urad << " urad\n";
    
    double rr = bunch.h.alpha;
    rr *= bunch.h.alpha;
    rr /= rr + 1;
    
    cerr << "\temittance x = " << 4 * sigma.x * sigma.xp * sqrt(1 - rr) / nm / rad * energy0 / electron_mass_c2 << " nm * rad\n";
    cerr << "\temittance y = " << 4 * sigma.y * sigma.yp * sqrt(1 - rr) / nm / rad * energy0 / electron_mass_c2 << " nm * rad\n";

    int nCavities = 1;
    double energy_gain;
    double voltage;
    for (int i=0; i<2; i++) {
      energy_gain = bunch.sE_correlation * energy0 / (2 * M_PI * FREQUENCY / c_light) / nCavities;
      voltage = nCavities * energy_gain / eplus;
      if (i==0) {
        nCavities = int(ceil(voltage / (GRADIENT * LENGTH)));
      }
    }
    
    std::cerr << "\tintegrated voltage " << voltage / megavolt << " MV\n";
    std::cerr << "\tgradient per cavity " << (voltage / megavolt) / (LENGTH / m) / nCavities << " MV/m\n";
    std::cerr << "\tin " << nCavities << " total cavites" << std::endl;
    
    for (size_t i=0; i<num_of_particles; i++) {
      Particle particle;
      particle.x  = gsl_ran_gaussian(rng, sigma.x) / um;
      particle.xp = gsl_ran_gaussian(rng, sigma.xp) / urad;
      particle.y  = gsl_ran_gaussian(rng, sigma.y) / um;
      particle.yp = gsl_ran_gaussian(rng, sigma.yp) / urad;
      particle.z  = gsl_ran_gaussian(rng, bunch.sigma_z) / um;
//      particle.energy = energy0 * (1 + gsl_ran_gaussian(rng, energy_spread) + bunch.sE_correlation * particle.z * um);
      particle.energy = energy0 * (1 + gsl_ran_gaussian(rng, energy_spread));
      particle.weight = bunch.charge / eplus / num_of_particles;

      accelerate(particle, energy_gain, FREQUENCY, phase);

      fprintf(stdout,"%.15g %g %g %g %g %g\n",
              particle.energy,
	      particle.x,
	      particle.y,
	      particle.z,
	      particle.xp,
	      particle.yp);
    }
    gsl_rng_free(rng);
  }
  return 0;
}