/*******************************************************************************
*
* 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