/*******************************************************************************
*
* McXtrace, x-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
* Component: Lens_parab
*
* %Identification
* Written by: Jana Baltser and Erik Knudsen
* Date: August 2010, modified July 2011
* Version: 1.0
* Release: McXtrace 0.1
* Origin: NBI
*
* X-ray compound refractive lens (CRL) with a profile of the parabola
*
* %Description
* A simple X-ray compound refractive lens (CRL) with a profile of the parabola in rotation simulates the photons' movement on passing through it. The CRL focuses in 2D
*
* Example: Lens_parab(material_datafile = "Be.txt", r=200e-6, r_ap=0.5e-3, d=50e-6, N=16)
*
* %Parameters
* Input parameters:
* r:        [m] Radius of curvature (circular approximation at the tip of the profile).
* r_ap:     [m] Radius of circular aperture, which also defines the depth of the lens profile.
* d:        [m] Distance between two surfaces of the lens along the propagation axis.
* N:        [m] Number of single lenses in a stack.
* rough_xy: [rad] RMS value of random slope error along x and y.
* rough_z:  [rad] RMS value of random slope error along z.
* material_datafile: [str]   Datafile containing f1 constants
*
* %End
*******************************************************************************/

DEFINE COMPONENT Lens_parab

SETTING PARAMETERS (string material_datafile="Be.txt", r=0.5e-3, r_ap=1.4e-3, d=.1e-3, int N=1, rough_z=0, rough_xy=0)

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

SHARE
%{
  %include "read_table-lib"
  struct incom_parab {
      double coord[3];
      double k[3];
  };

#pragma acc routine seq
  struct incom_parab intersection_lens_parab(struct incom_parab a, double *b, double roughness_xy,double roughness_z){
      struct incom_parab result;//={a.coord[0],a.coord[1],a.coord[2],a.k[0],a.k[1],a.k[2]};
      int i;
      double A,B,C,D,rr;
      double t[2],p[3],knorm[3],k[3],pos1_tmp[3],pos_tmp[3];

      double nxn,nyn,nzn,Nx,Ny,Nz,NORM,Knorm;
      double cos_theta,cos_theta1,Arg,Arg1,s,q;
      double k_new[3],k_new1[3],M,Sign,dd;
      double yh,xw;

      double tx,ty,tz,tnorm,txn,tyn,tzn;
      double v,w;

      for(i=0;i<=2;i++){
        result.k[i]=k[i]=a.k[i];
        result.coord[i]=p[i]=a.coord[i];
      }

      Knorm=sqrt(k[0]*k[0]+k[1]*k[1]+k[2]*k[2]);
      knorm[0]=k[0]/Knorm;
      knorm[1]=k[1]/Knorm;
      knorm[2]=k[2]/Knorm;

      rr=b[0];
      yh=b[1];
      xw=b[2];
      dd=b[3];
      M=b[4];
      Sign=b[5];

      A=knorm[0]*knorm[0]+knorm[1]*knorm[1];
      B= 2.0*(p[0]*knorm[0]+p[1]*knorm[1] - Sign*rr*knorm[2]);
      C=p[0]*p[0]+p[1]*p[1] - Sign*2*rr*p[2] + Sign*rr*2*dd;
      D=B*B-4.0*A*C;
    
    if (D<0) { /*ray does not intersect the parabola*/
        return result;
    }
    if (A==0){ /*incident k-vector is parallel (exactly) to the z-axis. Thus, the eq. becomes linear*/
      if(B==0){
        return result;
      }
      t[0]=-C/B;
      for(i=0;i<=2;i++){
	result.coord[i]=p[i]+t[0]*knorm[i];
      }
    } else {
      double qq;
      if (B<0){
	qq=-0.5*(B-sqrt(D));
      }else{
        qq=-0.5*(B+sqrt(D));
      }
      t[0]=qq/A;
      t[1]=C/qq;

      for(i=0;i<=2;i++){
        pos_tmp[i]=p[i]+t[0]*knorm[i];
        pos1_tmp[i]=p[i]+t[1]*knorm[i];
      }
      if ( fabs(pos_tmp[1])<=fabs(yh/2) && fabs(pos_tmp[0])<=fabs(xw/2) ){
        for(i=0;i<=2;i++){
          result.coord[i]=pos_tmp[i];
        }
      } else if ( fabs(pos1_tmp[1])<=fabs(yh/2) && fabs(pos1_tmp[0])<=fabs(xw/2) ){
        for(i=0;i<=2;i++){
          result.coord[i]=pos1_tmp[i];
        }
      } else return result;
    }
    
    /* Calculating tangential vector  */
    if (result.coord[0]==0 && result.coord[1]==0){ // incoming ray is along the axis, so it does not refract
        k_new[0]=k[0];
        k_new[1]=k[1];
        k_new[2]=k[2];
        for(i=0;i<3;i++) {
            result.k[i]=k_new[i];
        }
        return result;
    } else if (result.coord[0]!=0 && result.coord[1]!=0){
        Nx=-Sign*(result.coord[0]/rr); // surface normal
        Ny=-Sign*(result.coord[1]/rr);
        Nz=1;

        if (roughness_xy) {
            //Nx=Nx+roughness_xy*randnorm();
            //Ny=Ny+roughness_xy*randnorm();
        }
        if (roughness_z) {
            //Nz=Nz+roughness_z*randnorm();
        }

        NORM=sqrt(Nx*Nx+Ny*Ny+Nz*Nz);
        nxn=Nx/NORM;
        nyn=Ny/NORM;
        nzn=Nz/NORM;

        double cos_chi;
        cos_chi=knorm[0]*nxn+knorm[1]*nyn+knorm[2]*nzn;
        w=1/(sqrt(1-cos_chi*cos_chi)); // tangential vector
        v=-w*cos_chi;

        tx=v*nxn+w*knorm[0];
        ty=v*nyn+w*knorm[1];
        tz=v*nzn+w*knorm[2];
    }
    else if (result.coord[0]==0){
        tx=0;
        ty=Sign*(rr/result.coord[1]);
        tz=1;
    }
    else if (result.coord[1]==0){
        tx=Sign*(rr/result.coord[0]);
        ty=0;
        tz=1;
    }
	
    tnorm=sqrt(tx*tx+ty*ty+tz*tz);
    txn=tx/tnorm;
    tyn=ty/tnorm;
    tzn=tz/tnorm;
    
    cos_theta=txn*knorm[0]+tyn*knorm[1]+tzn*knorm[2];
    cos_theta1=M*cos_theta; /* Snell's law*/

    /* new k vector */
    if ((1.0-cos_theta*cos_theta)==0) {
       return result; 
    }
   
    Arg=(1.0-cos_theta1*cos_theta1)/(1.0-cos_theta*cos_theta);
    s=(1/M)*sqrt(Arg);
    q=(Knorm/tnorm)*((1/M)*cos_theta1-s*cos_theta);
    
    k_new[0]=q*tx+s*k[0];
    k_new[1]=q*ty+s*k[1];
    k_new[2]=q*tz+s*k[2];
    
    for(i=0;i<3;i++) { 
      result.k[i]=k_new[i]; 
    }
   return result;
  }
%}

DECLARE
%{
  int Z;
  double Ar;
  double rho;
  t_Table matT;
%}

INITIALIZE
%{
  int status=0;

  if ( (status=Table_Read(&(matT),material_datafile,0))==-1){
    fprintf(stderr,"Error: Could not parse file \"%s\" in COMP %s\n",material_datafile,NAME_CURRENT_COMP);
    exit(-1);
  }
  char **header_parsed;
  header_parsed=Table_ParseHeader(matT.header,"Z","A[r]","rho",NULL);
  if (!Z) Z=strtol(header_parsed[0],NULL,10);
  if (!Ar) Ar=strtod(header_parsed[1],NULL);
  if (!rho) rho=strtod(header_parsed[2],NULL);
%}

TRACE
%{
  double parab[6];
  double E,mu,f,rhoel,dl,d_total,e,k,delta,beta,Refractive_Index_Re,Refractive_Index_Im,w;
  struct incom_parab incid,refr,outg;

  int i=0,nr;
  
  parab[0]=r;
  parab[1]=r_ap*2;
  parab[2]=r_ap*2;
  
  w=(r_ap*r_ap)/(2.0*r);
  k=sqrt(kx*kx+ky*ky+kz*kz);
  e=K2E*k;  /*Energy in KeV, same unit as datafile */

  /*Interpolation of Table Values*/
  mu=Table_Value(matT,e,5)*rho*1e2;/*mu is now in SI, [m^-1]*/;

  f=Table_Value(matT,e,1);

  /*Calculation of Refractive Index */
  rhoel= f*NA*(rho*1e-24)/Ar; /*Material's Number Density of Electrons [e/A^3] incl f' scattering length correction*/
  delta= 2.0*M_PI*RE*rhoel/(k*k);
  beta= mu*1e-10/(2.0*k); /*mu and k in  A^-1*/

  Refractive_Index_Re = 1.0-delta; 
  Refractive_Index_Im = beta; 

  incid.k[0]=kx;
  incid.k[1]=ky; 
  incid.k[2]=kz;
  
  incid.coord[0]=x;
  incid.coord[1]=y;
  incid.coord[2]=z;

  d_total=0;
  for (nr=0;nr<=(N-1);nr++){
    parab[3]=nr*d+nr*2*w; // d constant
    parab[4]=1.0/Refractive_Index_Re; // M constant
    parab[5]=-1.0; //Sign constant
    
    refr=intersection_lens_parab(incid,parab,rough_xy,rough_z);
    
    if(refr.k[0]==0 && refr.k[1]==0 && refr.k[2]==0) continue;
    
    dl=sqrt( (refr.coord[0]-x)*(refr.coord[0]-x) + (refr.coord[1]-y)*(refr.coord[1]-y) + (refr.coord[2]-z)*(refr.coord[2]-z) );
    PROP_DL(dl);
    SCATTER;	
    
    kx=refr.k[0]; 
    ky=refr.k[1]; 
    kz=refr.k[2];

    //alter parabolic input to match second parabola
    parab[3]=(nr+1)*d+nr*2*w;
    parab[4]=Refractive_Index_Re;
    parab[5]=1.0; 

    outg=intersection_lens_parab(refr,parab,rough_xy,rough_z);
    
    dl=sqrt( (outg.coord[0]-x)*(outg.coord[0]-x) + (outg.coord[1]-y)*(outg.coord[1]-y) + (outg.coord[2]-z)*(outg.coord[2]-z) );
    PROP_DL(dl);
    d_total+=dl;
    SCATTER;

    kx=outg.k[0]; ky=outg.k[1]; kz=outg.k[2];
    incid=outg;
  }
  /*Add absorption according to the path length inside the lens material*/
  p*=exp(-mu*d_total);

%}

FINALLY
%{
    Table_Free(&(matT));
%}

MCDISPLAY
%{
  double z_c,zdepth,w;
  w=(r_ap*r_ap)/(2*r);
  zdepth=N*(2*w+d);
  z_c=zdepth/2.0-w;
  /*draw individiual lenses*/
  /*draw a circle at the maximal aperture and a parabola along x and y*/
  int i,j=0;
  for (j=0;j<(N<20?N:20);j++){
      circle("xy",0,0,-w+j*(d+2*w),r_ap);
      circle("xy",0,0,d+w+j*(d+2*w),r_ap);
      double zz0,zz1,yy0,yy1,dz,s;
      yy0=r_ap;
      zz0=zz1=-w;
      dz=w/(64.0-1.0);
      s=j*(d+2*w);
      /*first parabola*/
      while (zz1<=0){
          zz1+=dz;
          yy1=sqrt(2*r*fabs(zz1));
          line(0,yy0,s+zz0,0,yy1,s+zz1);
          line(0,-yy0,s+zz0,0,-yy1,s+zz1);
          line(yy0,0,s+zz0,yy1,0,s+zz1);
          line(-yy0,0,s+zz0,-yy1,0,s+zz1);
          zz0=zz1;yy0=yy1;
      }
      zz0=0;zz1=0;yy0=0;
      /*2nd parabola*/
      while (zz1<=w){
          zz1+=dz;
          yy1=sqrt(2*r*fabs(zz1));
          line(0,yy0,d+s+zz0,0,yy1,d+s+zz1);
          line(0,-yy0,d+s+zz0,0,-yy1,d+s+zz1);
          line(yy0,0,d+s+zz0,yy1,0,d+s+zz1);
          line(-yy0,0,d+s+zz0,-yy1,0,d+s+zz1);
          zz0=zz1;yy0=yy1;
      }
  }
  /*draw a circle at the last aperture of the lens*/
  circle("xy",0,0,zdepth-w,r_ap);
%}

END