/*******************************************************************************
*
* McXtrace, X-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
* Component: Source_div_quasi
*
* %Identification
* Written by: Mads Carlsen and Erik Bergbäck Knudsen (erkn@fysik.dtu.dk)
* Date: Apr 21
* Origin: DTU Physics
* Release: McXtrace 1.6
*
* Quasi-stochastic X-ray source with Gaussian or uniform divergence
*
* %Description
* A flat rectangular surface source with uniform or Gaussian divergence profile and focussing.
* If the parametere gauss is not set (the default) the divergence profile is flat
* in the range [-focus_ax,focus_ay]. If gauss is set, the focux_ax,focus_ay is considered
* the standard deviation of the gaussian profile.
* Currently focussing is only active for flat profile. The "focus window" is defined by focus_xw,focus_yh and dist.
* The spectral intensity profile is uniformly distributed in the energy interval defined by e0+-dE/2 or 
* by wavelength lambda0+-dlambda/2
* 
* The phase space sapnned by the generated X-rays is sampled by means of Halton-sequences, instead of regular
* pseudo random numbers. This ensures that samples are evenly distributed within the phase space region of interest.
*
* Example: Source_div_quasi(xwidth=0.1, yheight=0.1, focus_aw=2, focus_ah=2, E0=14, dE=2, gauss=0)
*
* %Parameters
* xwidth:   [m]   Width of source
* yheight:  [m]   Height of source
* focus_aw: [rad] Std. dev. (Gaussian) or maximal (uniform) horz. width divergence. focus_xw overrrides if it is more restrictive.
* focus_ah: [rad] Std. dev. (Gaussian) or maximal (uniform) vert. height divergence. focus_yh overrrides if it is more restrictive.
* focus_xw: [m]   Width of sampling window
* focus_yh: [m]   Height of sampling window
* dist:     [m]   Downstream distance to place sampling target window
* E0:       [keV] Mean energy of X-rays.
* dE:       [keV] Energy spread of X-rays.
* lambda0:  [AA]  Mean wavelength of X-rays (only relevant for E0=0)
* dlambda:  [AA]  Wavelength half spread of X-rays.
* gauss:    [1]   Criterion: 0: uniform, 1: Gaussian distribution of energy/wavelength
* gauss_a:  [1]   Criterion: 0: uniform, 1: Gaussian divergence distribution
* flux:     [1/(s*cm**2*st*energy unit)] Flux per energy unit, Angs or meV
* verbose:  [0/1] Generate more output on the console.
* spectrum_file: [string] File from which to read the spectral intensity profile
* phase:    [rad] Set to finite value to define X-ray phase (0:2 pi)
* randomphase: [0/1] When=1, the X-ray phase is randomised 
*
* %End
*******************************************************************************/

DEFINE COMPONENT Source_div_quasi

SETTING PARAMETERS (string spectrum_file="", xwidth=0, yheight=0, focus_xw=0, focus_yh=0, dist=0, focus_aw=0, focus_ah=0,
    E0=0, dE=0, lambda0=0, dlambda=0, flux=0, gauss=0, gauss_a=0, randomphase=1, phase=0, int verbose=1)

/* X-ray parameters: (x,y,z,kx,ky,kz,phi,t,Ex,Ey,Ez,p) */ 

SHARE
%{
  #ifndef MX_SOURCE_DIV_QUASI_H
  #define MX_SOURCE_DIV_QUASI_H 1
  %include "read_table-lib"

  #pragma acc routine
  double
  _quasi_rand01 (int axis, long long uid, double shift) {
    const int no_primes = 32;
    const int primes[] = { 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113, 127 };
    double f, hn;
    long long n0, n1, r;

    hn = 0.0;
    f = 1.0 / primes[axis];
    n0 = uid + 1;
    while (n0 > 0) {
      n1 = n0 / primes[axis];
      r = n0 - n1 * primes[axis];
      hn += f * r;
      f = f / primes[axis];
      n0 = n1;
    }
    hn = modf (hn + shift, &f);
    return hn;
  }

  #define quasi_rand01(a,b) _quasi_rand01((a),_particle->_uid,(b))

  #endif /*MX_SOURCE_DIV_H*/
%}

DECLARE
%{
  int ray_number;
  double xmin;
  double xmax;
  double focus_xw_2;
  double ymin;
  double ymax;
  double focus_yh_2;
  double pmul;
  double pint;
  t_Table T;
  int spectrum_from_file;
  double shifts[32];
%}

INITIALIZE
%{
  int j;
  /*random shifts to avoid correlation between runs*/
  for (j = 0; j < 32; j++) {
    shifts[j] = rand01 ();
  }

  ray_number = 1;

  /* define function for generating numbers in Halton sequences */

  focus_xw_2 = focus_xw / 2.0;
  focus_yh_2 = focus_yh / 2.0;
  xmin = -xwidth / 2.0;
  ymin = -yheight / 2.0;
  xmax = xwidth / 2.0;
  ymax = yheight / 2.0;

  /*flag if we are using a datafile*/
  spectrum_from_file = (spectrum_file && strcmp (spectrum_file, "NULL") && strlen (spectrum_file));

  if (spectrum_from_file) {
    /*read spectrum from file*/
    int status = 0;
    if ((status = Table_Read (&(T), spectrum_file, 0)) == -1) {
      fprintf (stderr, "ERROR (%s): Could not parse file \"%s\"\n", NAME_CURRENT_COMP, spectrum_file ? spectrum_file : "");
      exit (-1);
    }
    /*data is now in table t*/
    /*integrate to get total flux, assuming raw numbers have been corrected for measuring aperture*/
    int i;
    pint = 0;
    for (i = 0; i < T.rows - 1; i++) {
      pint += ((T.data[i * T.columns + 1] + T.data[(i + 1) * T.columns + 1]) / 2.0) * (T.data[(i + 1) * T.columns] - T.data[i * T.columns]);
    }
    if (verbose) {
      printf ("INFO (%s): Integrated intensity radiated is %g pht/s\n", NAME_CURRENT_COMP, pint);
      if (E0)
        printf ("INFO (%s):, E0!=0 -> assuming intensity spectrum is parametrized by energy [keV]\n", NAME_CURRENT_COMP);
    }
  } else if (!E0 && !lambda0) {
    fprintf (stderr, "ERROR (%s): Error: Must specify either wavelength or energy distribution\n", NAME_CURRENT_COMP);
    exit (-1);
  }

  /*calculate the X-ray weight from the flux*/
  if (flux) {
    pmul = flux;
  } else {
    pmul = 1;
  }
  pmul *= 1.0 / ((double)mcget_ncount ());

  if (dist == 0 && (focus_xw != 0 || focus_yh != 0)) {
    fprintf (stderr, "ERROR (%s): Cannot have focus sampling window (focus_xw x focus_yh) = (%g x %g) with dist=0.\n", NAME_CURRENT_COMP);
    exit (-1);
  }

  /*check if divergence limits are compatible with focus_xw, focus_yh*/
  if (focus_xw != 0) {
    double maxdivh = atan ((xwidth + focus_xw) / dist);
    if (focus_aw > maxdivh || focus_aw == 0) {
      focus_aw = maxdivh;
    }
  }
  if (focus_yh != 0) {
    double maxdivv = atan ((yheight + focus_yh) / dist);
    if (focus_ah > maxdivv || focus_ah == 0) {
      focus_ah = maxdivv;
    }
  }
%}

TRACE
%{
  double kk, theta_x, theta_y, l, e, k;
  p = pmul;
  if (!gauss_a) {
    theta_x = (quasi_rand01 (0, shifts[0]) - 0.5) * focus_aw;
    if (focus_xw != 0.0) {
      double x0, x1, pi_x;
      x0 = -focus_xw_2 - dist * tan (theta_x);
      x0 = MAX (xmin, x0);
      x1 = focus_xw_2 - dist * tan (theta_x);
      x1 = MIN (xmax, x1);
      pi_x = (x1 - x0) / (xwidth);
      x = x0 + (x1 - x0) * quasi_rand01 (2, shifts[2]);
      p *= pi_x;
    } else {
      x = xmin + quasi_rand01 (2, shifts[2]) * xwidth;
    }
    theta_y = (quasi_rand01 (1, shifts[1]) - 0.5) * focus_ah;
    if (focus_yh != 0.0) {
      double y0, y1, pi_y;
      y0 = -focus_yh_2 - dist * tan (theta_y);
      y0 = MAX (ymin, y0);
      y1 = focus_yh_2 - dist * tan (theta_y);
      y1 = MIN (ymax, y1);
      pi_y = (y1 - y0) / (yheight);
      y = y0 + (y1 - y0) * quasi_rand01 (3, shifts[3]);
      p *= pi_y;
    } else {
      y = ymin + quasi_rand01 (3, shifts[3]) * yheight;
    }
  } else {
    theta_x = randnorm () * focus_aw;
    theta_y = randnorm () * focus_ah;
    x = xmin + rand01 () * xwidth;
    y = ymin + rand01 () * yheight;
  }

  if (spectrum_from_file) {
    double pp = 0;
    while (pp <= 0) {
      l = T.data[0] + (T.data[(T.rows - 1) * T.columns] - T.data[0]) * quasi_rand01 (4, shifts[4]);
      pp = Table_Value (T, l, 1);
    }
    p *= pp;
    /*if E0!=0 the tabled value is assumed to be energy in keV*/
    if (E0 != 0) {
      k = E2K * l;
    } else {
      k = (2 * M_PI / l);
    }
  } else if (E0) {
    if (!dE) {
      e = E0;
    } else if (gauss) {
      e = E0 + dE * randnorm ();
    } else {
      e = randpm1 () * dE + E0;
    }
    k = E2K * e;
  } else if (lambda0) {
    if (!dlambda) {
      l = lambda0;
    } else if (gauss) {
      l = lambda0 + dlambda * randnorm ();
    } else {
      l = randpm1 () * dlambda + lambda0;
    }
    k = (2 * M_PI / l);
  }

  kx = tan (theta_x);
  ky = tan (theta_y);
  kz = 1;
  NORM (kx, ky, kz);

  kx *= k;
  ky *= k;
  kz *= k;

  /*set polarization and phase to something known*/
  Ex = 0;
  Ey = 0;
  Ez = 0;
  if (!randomphase) {
    phi = phase;
  } else {
    phi = rand01 () * M_2_PI;
  }

  /*set polarization vector*/
  Ex = 0;
  Ey = 0;
  Ez = 0;
%}

MCDISPLAY
%{
  magnify ("xy");
  multiline (5, -xwidth / 2.0, -yheight / 2.0, 0.0, xwidth / 2.0, -yheight / 2.0, 0.0, xwidth / 2.0, yheight / 2.0, 0.0, -xwidth / 2.0, yheight / 2.0, 0.0,
             -xwidth / 2.0, -yheight / 2.0, 0.0);
  if (focus_xw || focus_yh) {
    dashed_line (0, 0, 0, -focus_xw / 2, -focus_yh / 2, dist, 4);
    dashed_line (0, 0, 0, focus_xw / 2, -focus_yh / 2, dist, 4);
    dashed_line (0, 0, 0, focus_xw / 2, focus_yh / 2, dist, 4);
    dashed_line (0, 0, 0, -focus_xw / 2, focus_yh / 2, dist, 4);
  } else {
    dashed_line (0, 0, 0, tan (-focus_ah / 2.0) * dist * 0.1, tan (-focus_aw / 2.0) * dist * 0.1, dist * 0.1, 4);
    dashed_line (0, 0, 0, tan (focus_ah / 2.0) * dist * 0.1, tan (-focus_aw / 2.0) * dist * 0.1, dist * 0.1, 4);
    dashed_line (0, 0, 0, tan (focus_ah / 2.0) * dist * 0.1, tan (focus_aw / 2.0) * dist * 0.1, dist * 0.1, 4);
    dashed_line (0, 0, 0, tan (-focus_ah / 2.0) * dist * 0.1, tan (focus_aw / 2.0) * dist * 0.1, dist * 0.1, 4);
  }
%}

END