/************************************************************************
*
* McXtrace, X-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
*
* Component: Undulator
*
* %Identification
* Written by: Erik B. Knudsen
* Date: May, 2013.
* Version: 1.0
* Origin: DTU Physics
*
* Model of an undulator source
*
* %Description
* A undulator source model based on the derivation by K.J. Kim, AIP, conf. proc., 184, 1989. doi:10.1063/1.38046.
*
* SOLEIL_PX2a U24
* Example: Undulator( E0=12.65, dE=1, Ee=2.75, dEe=0.001, Ie=0.5, K=1.788, Nper=80, 
*   lu=24e-3, sigey=9.3e-6, sigex=215.7e-6, sigepx=29.3e-6, sigepy=4.2e-6, 
*   dist=29.5, E1st=12.400 )
*
* %Parameters
* Ee:       [GeV] Storage ring electron energy [typically a few GeV).
* dEe:      [percent] Relative electron energy beam spread (sigma/Ee).
* Ie:       [A]   Ring current.
* B:        [T]   Peak magnet field strength. Overrides K.
* Nper:     [int] Number of magnetic periods in the undulator.
* lu:       [m]   Magnetic period length of the undulator aka lambda_u.
* K:        [1]   Dimensionless deflection undulator parameter. When K >> 1 (ie B*lu is large) you get a wiggler.
* sigex:    [m]   Electron ring beam size in horizontal plane (rms).
* sigey:    [m]   Electron ring beam size in vertical plane (rms).
* sigepx:   [rad] Electron ring beam horizontal divergence (rms).
* sigepy:   [rad] Electron ring beam vertical divergence (rms).
* phase:    [rad] Initial phase of radiation.
* randomphase: [0/1] If !=0 phase will be random (I.e. the emitted radiation is completely incoherent).
* focus_xw: [m]   Width of target window.
* focus_yh: [m]   Height of target window.
* dist:     [m]   Distance from source plane to target window along the optical axis.
* E0:       [keV] Center of emitted energy spectrum.
* dE:       [keV] Half-width of emitted energy spectrum.
* E1st:     [keV] Energy of the fundmental (1st) undulator harmonic.
* verbose:  [0/1] If nonzero, output extra information.
* quick_integ: [0/1] If nonzero, use faster (but less accurate) integration scheme.
* Br:       [T] Remanent field (1.35T for Nd2Fe14B) for gap estimate
*
* %End
*****************************************************************/

DEFINE COMPONENT Undulator

SETTING PARAMETERS (E0=0, dE=0, phase=0, randomphase=1, Ee=2.4, dEe=0, Ie=0.4, B=0, K=0,
        int Nper=1, lu=16e-3, sigey=0, sigex=0, sigepx=0, sigepy=0, focus_xw=0, focus_yh=0, dist=1, quick_integ=0,
        E1st=0, int verbose=0, Br=1.35)

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

SHARE
%{

#include <gsl/gsl_sf_bessel.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_integration.h>

double mxundulator_Bsig_integrand(double x, void *params){
    double w_w1 = *((double *) params);
    double p = *((double *) params+1);
    double q = *((double *) params+2);
    double angle_term = *((double *) params+3);/* xi*gamma/K  horizontal angle term*/

    double f1 = angle_term - cos(x);
    double inner= w_w1*x - p*sin(x) + q*sin(2*x);
    double f2 =  cos(inner);

    return f1*f2;
}

double mxundulator_Bpi_integrand(double x, void *params){
    double w_w1 = *((double *) params);
    double p = *((double *) params+1);
    double q = *((double *) params+2);
    double angle_term = *((double *) params+3); /*phi*gamma/K vertical angle term*/

    double f1 = angle_term;
    double inner= w_w1*x - p*sin(x) + q*sin(2*x);
    double f2 =  cos(inner);

    return f1*f2;
}

double mxundulator_S_N(double w_w1, int N){
    return pow(sin(N*M_PI*w_w1)/(N*sin(M_PI*w_w1)),2.0);
}


%}

DECLARE
%{
  double gamma;
  double gamma2;
  double igamma;
  double s1x;
  double s1y; /*beam's size at dist (convolution of sigex/sigey and igamma)*/
  double length; /*undulator magnetic length*/
  double kc; /*undulator kritical wavenumber*/
  double pmul; /*initial photon weight*/
  double gap;     /* gap estimate */
  gsl_function Bsig;
  gsl_function Bpi;
  gsl_integration_workspace *gsl_int_ws;

  /*fine structure*/
  double alpha;
  /*electron mass*/
  double MELE;
  //double besselj[nharm],besselh[nharm];
%}


INITIALIZE
%{
  /*fine structure constant from CODATA*/
  alpha=7.2973525698e-3;
  /*electron mass from CODATA in kg*/
  MELE=9.10938291e-31;
  
  length=lu*Nper;

  if( (!E1st && !K && B<=0) || Ee<=0 || Ie<=0 ){
    fprintf(stderr, "Error (%s): E1st, K, B, Ee, and Ie do not have a sane set of values. Found (%g %g %g %g %g). Aborting.\n",NAME_CURRENT_COMP,E1st,K,B,Ee,Ie);
    exit(1);
  }

  /*compute gamma*/
  gamma=(Ee*1e9)/(MELE/CELE*M_C*M_C); /*the extra CELE is to convert to eV*/
  gamma2=gamma*gamma;
  igamma=1.0/gamma;

  
  if(E1st && lu && !K){
      /*compute K and B from desired target energy*/
      K=2*(4*M_PI*gamma2/(E2K*E1st*lu*1e10) -1);
      if (K>0) K=sqrt(K); else K=0;
  } else if (!lu && E1st && K) {
      lu=4*M_PI*gamma2/(K*K/2+1)/(E2K*E1st*1e10);
  }
  if(!K && B && lu) {
    K=CELE*B*lu/(2*M_PI*MELE*M_C);
  } else if (!B && K && lu) {
    B=2*M_PI*MELE*M_C*K/CELE/lu;
  }
  if (!E1st && K && lu) {
      E1st=4*M_PI*gamma2/(K*K/2+1)/(E2K*lu*1e10);
  }
  if (!E0 && E1st) E0=E1st;
  
  if( E1st<=0 || K<=0 || B<=0 || Ee<=0 || Ie<=0 || lu<=0 || E0<=0){
    fprintf(stderr, "Error (%s): (E1st, K, B, Ee, Ie, lu, E0) do not have a sane set of values. Found (%g %g %g %g %g %g %g). Aborting.\n",
    NAME_CURRENT_COMP,E1st,K,B,Ee,Ie,lu,E0);
    exit(1);
  }
  
  /* compute gap estimate */
  if (Br > B) { // remanent field, rather usual for Nd2Fe14B
    double a=0.55*Br +2.835;
    double b=-1.95*Br+7.22;
    double c=-1.3*Br +2.97;
    gap = -log(B/a)*lu/b;
  } else gap=3e-3;
  
  if (verbose) 
    printf("Undulator (%s) K=%g B=%g[T] lu=%g[m] E1st=%g[keV] E0=%g[keV] gap=%g[m]\n",
      NAME_CURRENT_COMP, K,B,lu,E1st,E0,gap);

  if (sigex <0 || sigey<0){
   fprintf(stderr, "Error (%s): sigex (= %g) and sigey (= %g) must both be >= 0. Negative beam size isn't meaningful. Aborting.\n",NAME_CURRENT_COMP,sigex,sigey);
    exit(1);
  }
  if (dist<=0 || focus_xw < 0 || focus_yh < 0){
    fprintf(stderr,"Error (%s): Target undefined. Set dist, focus_xw and focus_yh.\n",NAME_CURRENT_COMP);
    exit(1);
  }


  //printf("Undulator (%s): gamma=%g, divergence is 1/gamma=%g rad.\n",NAME_CURRENT_COMP,gamma,igamma);
  /*compute characteristic energy in keV*/
  double Ec=0.665*Ee*Ee*B;
  //double Ec=1.5*gamma2*HBAR*CELE*B/MELE *1e-3; /*check units on this one. The 1e-3 factor is because energy is assumed to be in keV*/
  /*We normally do computations in k so use that for transfer*/
  kc=E2K*Ec;

  /*allocate an integration workspace*/
  if(!quick_integ){
      gsl_int_ws = gsl_integration_workspace_alloc (1000);
  }
  Bsig.function= &mxundulator_Bsig_integrand;
  Bpi.function = &mxundulator_Bpi_integrand;
  gsl_set_error_handler_off();

  /*correct for number of rays*/
  pmul=1.0/( (double) mcget_ncount());

  /*correct for finite energy interval*/
  if(dE){
      pmul*=dE*2.0;
  }

%}


TRACE
%{

    double k,e,l,w_u,r,w;
    double xo,yo,zo,xi,psi,theta2,Omega;
    double bsigma_integral, bpi_integral, s_n2;
    double bsigma_error, bpi_error;

    /* pick an energy in the given interval */
    e=E0+randpm1()*dE;

    /* add electron beam energy spread to gamma parameters (if necessary).*/
    if( dEe){
        double deltaEe=(randnorm()*dEe*Ee)+ Ee;
        gamma=(deltaEe*1e9)/(MELE/CELE*M_C*M_C);/*the extra CELE is to convert to eV*/
        gamma2=gamma*gamma;
        igamma=1.0/gamma;
    }

    /* the undulator's fundamental angular frequency*/
    w_u=2*M_PI*M_C/lu;

    /* now pick angles (\xi and \psi => \theta) within the focus_window 
     * (point on window + point in electron beam)
     * ... we can now weight properly according to ang. flux density*/
    if(focus_xw && focus_yh){
        randvec_target_rect(&xo,&yo,&zo, &Omega, 0,0,dist, focus_xw, focus_yh, ROT_A_CURRENT_COMP);
    }else if (focus_yh!=0){
        xi=0;xo=0;yo=0;
        yo=randpm1()*0.5*focus_yh;
        Omega=1;
    }else if (focus_xw!=0){
        psi=0;xo=0;yo=0;
        xo=randpm1()*0.5*focus_xw;
        Omega=1;
    }else{
        xi=0;psi=0;xo=0;yo=0;
        Omega=1;
    }
    p=pmul*Omega/(4*M_PI);

    /*add emittance effects - note that doing it this way will shoot some rays outside the focus-window*/
    x=y=z=0;
    if(sigex){
        x=randnorm()*sigex;
    }
    if(sigey){
        y=randnorm()*sigey;
    }

    xi=fabs(atan2(xo,dist));
    psi=fabs(atan2(yo,dist));
    /* This has to be after (xi,psi), else it will be convoluted into the weight calculation.*/
    if(sigepx){
        xo+=randnorm()*sigepx*dist;
    }
    if(sigepy){
        yo+=randnorm()*sigepy*dist;
    }

    theta2=xi*xi+psi*psi;

    k=E2K*e;
    /*angular frequency w:*/
    w=M_C*k*1e10;

    kx=xo;ky=yo;kz=dist;
    NORM(kx,ky,kz);
    kx*=k;
    ky*=k;
    kz*=k;


    double w1theta=2*gamma2/(1+K*K/2.0 + gamma2 *theta2) * w_u;
    double w10=2*gamma2/(1+K*K/2.0) * w_u;
    double w_w1=w/w1theta;
    double w_w10=w/w10;

    double Bsig_integ_prms[4],Bpi_integ_prms[4];
    Bsig_integ_prms[0] = Bpi_integ_prms[0]=w_w1; /*relative frequency*/
    Bsig_integ_prms[1] = Bpi_integ_prms[1] = 2*w_w10*xi*gamma*K/(1+K*K/2.0); /*p*/
    Bsig_integ_prms[2] = Bpi_integ_prms[2] = 0.25 * w_w10*K*K/(1+K*K/2.0); /*q*/
    Bsig_integ_prms[3] = xi *gamma/K; /*angle terms*/
    Bpi_integ_prms[3]  = psi*gamma/K;

    /*pass the parameters to the integrand*/
    Bsig.params=Bsig_integ_prms;
    Bpi.params =Bpi_integ_prms;
    if (!quick_integ){
        gsl_integration_qags (&(Bsig),    0, M_PI,    0, 1e-6, 1000, gsl_int_ws, &bsigma_integral, &bsigma_error);
        gsl_integration_qags (&(Bpi),     0, M_PI,    0, 1e-6, 1000, gsl_int_ws, &bpi_integral, &bpi_error);
    }else{
        size_t neval;
        gsl_integration_qng (&(Bsig),    0, M_PI,    0, 1e-6, &bsigma_integral, &bsigma_error,&neval);
        gsl_integration_qng (&(Bpi),     0, M_PI,    0, 1e-6, &bpi_integral, &bpi_error,&neval);
    }
    bsigma_integral*=M_2_PI;/*correct for only integrating half interval and normalize by pi.*/
    bpi_integral*=M_2_PI;

    s_n2=mxundulator_S_N(w_w1, Nper);
    
    double prefactor=alpha*Ie/CELE*pow(K*gamma/(1+K*K/2.0),2.0)*Nper*Nper*w_w10*w_w10;

    p*=prefactor* ( pow(bsigma_integral,2.0)+pow(bpi_integral,2.0) ) *s_n2;

    /*randomly pick phase*/
    if (randomphase){
        phi=rand01()*2*M_PI;
    }else{
        phi=phase;
    }

    /*Set polarization vector. TODO: Do this right.*/
    Ex=0;Ey=0;Ez=0;
%}
FINALLY
%{
    if(!quick_integ){
        gsl_integration_workspace_free (gsl_int_ws);
    }
%}
MCDISPLAY
%{
  
  double zz,dz;
  const double xwidth=1e-2;
  const double D=dist;
  double x0,z0,x1,z1;

  zz=-(length+lu)/2.0;
  dz=lu/2.0;

  while (zz<=(length-lu)/2.0){
    box(0.0,gap/2.0+5e-4,zz,xwidth,1e-3,lu/2.0,0, 0, 1, 0);
    box(0.0,-gap/2.0-5e-4,zz,xwidth,1e-3,lu/2.0,0, 0, 1, 0);
    zz+=dz;
  }

  line(0.0,0.0,0.0, K*D*sin(igamma), 0.0, D);
  line(0.0,0.0,0.0,-K*D*sin(igamma), 0.0, D);
  line(0.0,0.0,0.0, 0.0, D*sin(igamma), D);
  line(0.0,0.0,0.0, 0.0,-D*sin(igamma), D);
  
  double phi,dphi;  
  phi =-igamma;
  dphi= 2.0*igamma/32;
  while(phi<igamma){
    x0=D*sin(phi);
    x1=D*sin(phi+dphi);
    z0=D*cos(phi);
    z1=D*cos(phi+dphi);
    line(K*x0,0.0,z0,K*x1,0.0,z1);
    line(0.0,x0,z0,0.0,x1,z1);
    phi+=dphi;
  }
%}

END