/*******************************************************************************
*
* McXtrace, X-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
* Component: Molecule_2state
*
* %Identification
*
* Written by: Erik B Knudsen
* Date: October 2012
* Version: 1.0
* Release: McXtrace 1.1
* Origin: DTU Physics
*
* Disordered optical-excitable molecule sample.
*
* %Description
* A sample model for pump probe experiments which models disordered molecules in a volume (rectangular,
* cylindrical, or spherical). Molecules can be in one of two states (0 and 1).
* Scattering is either specified through F vs. q scattering curves or as a set of atom positions from which
* F vs. q is computed.
* At t=-delta_t, a fraction of the molecules are put in state 1, from which they decay exponentially,
* with time constant t_relax, into state 0. For t<-delta_t
* all of the molecules are in the state specified by <i>initial_state</i>.
* To improve statistics, scattering may be limited to a "forward" cone with opening angle in [psimin, psimax].
* Furthermore, scattering may be restricted to the azimuthal segment between [etamin,etamax].
* 
* Example: Molecule_2state(
*  nq=512,state_0_file="Fe_bpy_GS_DFT.txt",state_1_file="Fe_bpy_ES_DFT.txt",radius=0.01,
*  psimin=0, psimax=15*DEG2RAD, etamin=-1*DEG2RAD,etamax=1*DEG2RAD,
*  t_relax=600e-12, delta_t=100e-9, excitation_yield=0.2)
*
* %Parameters
* Input parameters:
* initial_state:  [0/1] Which state is Molecule_2state in for t<delta_t? Useful for modelling something that changes state slowly.
* form_factors:   [str] File from which to read atomic form factors. Defualt amounts to use the one shipped with McXtrace.
* state_0_file:   [str] Isotropic scattering factors (parameterized by q), or atom positions are specified for state 0.
* state_1_file:   [str] Isotropic scattering factors (parameterized by q), or atom positions are specified for state 1.
* nq:             [1]   Number of q-bins if F is to be computed from atom positions (Debye formalism).
* material_datafile: [str] Where to read f1 and f2 factors from in order to handle absorption.
* delta_t:        [s]   Delay between the exciting event t=0. delay is negative, i.e. delta_t>0 means the exciting event happens before t=0.
* excitation_yield: [1] Mean fraction of molecules that get excited.
* t_relax:        [s]   Mean relaxation time (into state 0) of excited molecules.
* psimin:         [rad] Minimum scattering angle off the optical axis.
* psimax:         [rad] Maximum scattering angle off the optical axis.
* etamin:         [rad] Minimum scattering angle around the optical axis.
* etamax:         [rad] Maximum scattering angle around the optical axis.
* radius:         [m]   Radius of cylindrical of spherical sample.
* xwidth:         [m]   Width of rectangular sample.
* yheight:        [m]   Height of rectangular or cylindrical sample.
* zdepth:         [m]   Depth (thickness) of rectangular sample.
* concentration:  [m]   Concentration or packing factor of sample.
* p_transmit:     [m]   Fraction of statistics devoted to sample direct (unscattered) beam.
* q_parametric:   [0/1] When 0: Assume that datafiles contains atom positions. 1: datafiles contains F vs. q data.
* Emax:           [keV] Maximal energy for which scattering factors are computed. Must be larger than the maximal impinging energy.
*  
* %End
*******************************************************************************/

DEFINE COMPONENT Molecule_2state

SETTING PARAMETERS (delta_t=100e-9,excitation_yield=0.2,t_relax=100e-9,initial_state=0,
        psimin=0,psimax=M_PI_2, etamin=-M_PI, etamax=M_PI,radius=0, yheight=0, xwidth=0, zdepth=0,
        concentration=1,p_transmit=0.1,string form_factors="FormFactors.txt",string state_0_file=NULL,string state_1_file=NULL, nq=512,
        string material_datafile="Be.txt",q_parametric=0, Emax=80)

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

SHARE
%{
  %include "read_table-lib"
  %include "form_factor-lib"
#ifndef MXMOLECULE_2STATE
#define MXMOLECULE_2STATE
  enum SHAPES {NONE,CYLINDER,SPHERE,BOX};
#endif
%}


DECLARE
%{
  double *q[2];
  double *f2[2];
  int lines[2];
  double f2sum[2];
  double f2mcsum[2];
  double qnmin[2];
  double qnmax[2];
  int Z;
  double At;
  double rho;
  DArray1d E;
  DArray1d Mu;
  t_Table matT;
  int shape;
%}


INITIALIZE
%{
  t_Table T[2],f0;
  int status;

  if( psimin>psimax || psimin<0 || psimax>2*M_PI ||etamin>etamax || etamin<-M_PI ||etamax>M_PI ){
      fprintf(stderr,"Error (%s): Nonsensical angle defs psi=[%g,%g] eta=[%g,%g]\n",NAME_CURRENT_COMP,psimin,psimax,etamin,etamax);
      exit(-1); 
  }

  if ( state_0_file && (status=Table_Read(&(T[0]),state_0_file,0))==-1){
      fprintf(stderr,"Error (%s): Could not parse file \"%s\"\n",NAME_CURRENT_COMP, state_0_file);
      exit(-1);
  }
  if ( state_1_file && (status=Table_Read(&(T[1]),state_1_file,0))==-1){
      fprintf(stderr,"Error (%s): Could not parse file \"%s\"\n",NAME_CURRENT_COMP,state_1_file);
      fprintf(stderr,"Proceeding as a single state sample\n");
      T[1]=T[0];
  }

  if(form_factors){
      if (status=Table_Read(&(f0),form_factors,0)==-1){
          fprintf(stderr,"Error(%s): Could not parse file \"%s\"\n",NAME_CURRENT_COMP,form_factors);
          exit(-1);
      }
  }

  if (!q_parametric){
    /*input files are atom positions - use that to compute f2s*/
    f2[0]=malloc(sizeof(double)*nq);/*could be done by read_table-lib*/
    q[0]=malloc(sizeof(double)*nq);
    f2[1]=malloc(sizeof(double)*nq);
    q[1]=malloc(sizeof(double)*nq);

    if (!psimax) psimax=M_PI;
    /*compute q-limits*/;
    if (psimin!=0) qnmin[0]=qnmin[1]=Emax*E2K*M_SQRT2*sqrt(1-cos(psimin));
    qnmax[0]=qnmax[1]=Emax*E2K*M_SQRT2*sqrt(1-cos(psimax));

    printf("%s: Computing F2(q) for two states:\n", NAME_CURRENT_COMP);
    int r,n,m,states;/*do this for two states*/
    double dq= qnmax[0]/(nq-1);
    for (states=0;states<2;states++){
      for (r=0;r<nq;r++){
        double f2_single=0;
        double q_single=qnmin[states] + r*dq;//Table_Index(f0,r,0);
        for (n=0;n<T[states].rows;n++){
          int Zn=(int) Table_Index(T[states],n,0);
          /*figure out Z_n*/
          double f0n=atomic_form_factor(Zn,0,q_single);//Table_Index(f0,r,Zn);
          double nx,ny,nz;
          nx=Table_Index(T[states],n,1);
          ny=Table_Index(T[states],n,2);
          nz=Table_Index(T[states],n,3);
          f2_single+=fabs(f0n*f0n);
          //printf("debug %d %g\n",Zn,f2_single);
          for (m=n+1;m<T[states].rows;m++){
            int Zm=(int) Table_Index(T[states],m,0);
            double f0m=atomic_form_factor(Zm,0,q_single);//Table_Index(f0,r,Zm);
            double mx,my,mz,dr;
            mx=Table_Index(T[states],m,1);
            my=Table_Index(T[states],m,2);
            mz=Table_Index(T[states],m,3);
            dr=sqrt((nx-mx)*(nx-mx)+(ny-my)*(ny-my)+(nz-mz)*(nz-mz));
            if (q_single){
              f2_single+=2 * f0m*f0n*sin(q_single*dr)/(q_single*dr);
            }else{
              f2_single+=2* f0m*f0n;
            }
          }
        }
        f2[states][r]=f2_single;
        q[states][r]=q_single;
        if (r>0){
          //double dq=q-Table_Index(f0,r-1,0);
          /*integrate using linear interpolation*/
          f2sum[states]+=0.5*(f2[states][r-1]+f2_single)*dq;
          if (q_single>qnmin[states] && q_single<qnmax[states]){
            double q1,q2,alpha;
            q1=(qnmin[states]>q_single-dq?qnmin[states]:q_single-dq);
            q2=(qnmax[states]<q_single?qnmax[states]:q_single);
            alpha=((q1+q2)*0.5-(q_single-dq))/q_single;
            f2mcsum[states]+=(q2-q1)*(alpha*f2[states][r-1]+ (1-alpha)*f2_single);
          }
        }
      }
      lines[states]=T[states].rows;
      printf("Integrated f2 for state %d= %g, Mu_s=%g\n",states,f2sum[states],RE*RE*f2sum[states]);
    }
  }else{
    /*input files contain f^2 parametrized by q*/
    double dq=0.1;
    int r,states;/*do this for two states*/
    f2[0]=malloc(sizeof(double)*(T[0].rows+1));/*could be done by read_table-lib*/
    q[0]=malloc(sizeof(double)*(T[0].rows+1));
    f2[1]=malloc(sizeof(double)*(T[1].rows+1));
    q[1]=malloc(sizeof(double)*(T[1].rows+1));

    if (!psimax) psimax=M_PI;
    /*compute q-limits*/;
    qnmin[0]=qnmin[1]=M_SQRT2*sqrt(1-cos(psimin));
    qnmax[0]=qnmax[1]=M_SQRT2*sqrt(1-cos(psimax));

    printf("%s: Computing F2(q) for two states:\n", NAME_CURRENT_COMP);
    for (states=0;states<2;states++){
      for (r=0;r<T[states].rows;r++){
        q[states][r]=T[states].data[r*2];
        f2[states][r]=T[states].data[r*2+1];
        //printf("%d %d %g %g\n",states,r,q[states][r], f2[states][r]);
        f2sum[states]+=f2[states][r]*dq;
      }
      lines[states]=T[states].rows;
      printf("Integrated f2 for state %d= %g, Mu_s=%g\n",states,f2sum[states],RE*RE*f2sum[states]);
      f2[states][r]=f2[states][r-1];
      q[states][r]=FLT_MAX;
    }
  }
  if (radius){
    if (yheight) shape=CYLINDER;
    else shape=SPHERE;
  }else if (xwidth && yheight && zdepth){
    shape=BOX;
  }
  if (shape==NONE){
    fprintf(stderr,"Error (%s): could not understand which shape the thing is\n",NAME_CURRENT_COMP);exit(1);
  }

  t_Table A;
  /*Read absorption table*/
  if ( (status=Table_Read(&A,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(A.header,"Z","A[r]","rho","Z/A","sigma[a]",NULL);
  //Prms=calloc(1,sizeof(struct mat_prms));
  E=malloc(sizeof(double)*(A.rows+1));
  Mu=malloc(sizeof(double)*(A.rows+1));
  if(header_parsed[2]){rho=strtod(header_parsed[2],NULL);}
  else{fprintf(stderr,"Warning(%s): %s not found in header of %s, set to 1\n",NAME_CURRENT_COMP,"rho",material_datafile);rho=1;}
  /*which columns holds the mus*/
  int mu_c=5;
  if (A.columns==3){
    /*three column format*/
    mu_c=1;
  }

  int i;
  for (i=0;i<A.rows;i++){
    E[i]=A.data[i*A.columns];
    Mu[i]=A.data[mu_c+i*A.columns]*rho*1e2;     /*mu is now in SI, [m^-1]*/
  }

  E[A.rows]=-1.0;
  Mu[A.rows]=-FLT_MAX;

  Table_Free(&A);

%}

TRACE
%{
  int hit,state;
  double alpha,e,k,mu;
  double l0,l1;
  int i;
  if (shape==CYLINDER){
    hit=cylinder_intersect(&l0,&l1,x,y,z,kx,ky,kz,radius,yheight);
    /*sample is a cylinder*/
  }else if (shape==SPHERE){
    /*sample is a sphere - unlikely*/
    hit=sphere_intersect(&l0,&l1,x,y,z,kx,ky,kz,radius);
  }else if (shape==BOX){
    /*sample is a box*/
    hit=box_intersect(&l0,&l1,x,y,z,kx,ky,kz,xwidth,yheight,zdepth);
  }
  if (hit){
    /*if we've intersected with the sample, propagate to intersection*/
    PROP_DL(l0);

    /*Absorption table interpolation*/
    k=sqrt(kx*kx+ky*ky+kz*kz);
    e=k*K2E;
    i=0;
    while (e>E[i]){
      i++;
      if (E[i]==-1){
        fprintf(stderr,"Photon energy (%g keV) is outside the filter's material data\n",e); ABSORB;
      }
    }
    alpha=(e-E[i-1])/(E[i]-E[i-1]);
    mu=(1-alpha)*Mu[i-1]+alpha*Mu[i];
    /*which state is the molecule in?*/
    /*delta_T is positive for optical pulse coming before x-ray pulse (t). Assuming the optical pulse to be
     *short, relaxation has progressed since t+delta_t. Then the probability
     *the molecule is in an excited state is: excitation_yield* t_relax * exp(-t_relax *(t+delta_t)*/
    if(delta_t<t){
      /*photon arrives before pump pulse - molecule cannot be excited*/
      state=initial_state;
    }else {
      double r=rand01();
      if( r< excitation_yield * exp(-(t+delta_t)/t_relax) ){
        /*excited state*/
        //printf("I'm excited %g %g %e\n",r, excitation_yield * exp(-(t+delta_t)/t_relax),t);
        state=1;
      }else{
        //printf("I'm bored %g %g %e\n",r, excitation_yield * exp(-(t+delta_t)/t_relax),t);
        state=0;
      }
    }
    /*now figure out if we scatter at all*/
    double dl=l1-l0;
    double mu_s=RE*RE*f2sum[state];
    double l_conc=pow(concentration,0.3333333333333333333333333333333333333);
    double pmul=1,p_s;
    double pr;
    p_s=1-exp(-mu_s*dl);

    pr=rand01();
    if (p_transmit<pr){
      /*scattering branch*/
      /*find scattering pt*/
      dl=rand01()*dl;

      PROP_DL(dl);
      SCATTER;

      /*Absorption before scattering*/
      p*=exp(-mu*dl);


      double qq,alpha,ff2,rr;
      /*choose a random scattering direction*/
      double kfx,kfy,kfz,solid_angle;
      randvec_target_circle(&kfx, &kfy, &kfz, &solid_angle, 0, 0, 1, tan(psimax) );
      if( (etamax-etamin)!=2*M_PI){
          kfx=sqrt(kfx*kfx+kfy*kfy);
          kfy=0;
          double eta=rand01()*(etamax-etamin)+etamin;
          /*rotate kf round 0,0,1 by eta, and reassign to kf*/
          rotate(kfx,kfy,kfz, kfx,kfy,kfz,eta,0,0,1);
          /*downweight since we're not using the full eta range*/
          p*=(etamax-etamin)/(2*M_PI);
      }

      NORM(kfx,kfy,kfz);
      kfx*=k;kfy*=k;kfz*=k;
      qq=sqrt( scalar_prod(kx-kfx,ky-kfy,kz-kfz,kx-kfx,ky-kfy,kz-kfz));

      /*apply new vector*/
      kx=kfx;ky=kfy;kz=kfz;

      /*find f2 for this q by interpolation*/
      int r;
      if (!q_parametric){
          for(r=1;r<nq-1;r++){
              if (q[state][r]>qq) break;
          }
      }else{
        for(r=1;r<lines[state]-1;r++){
              if (q[state][r]>qq) break;
          }
      }
      alpha=(qq-q[state][r-1])/(q[state][r]-q[state][r-1]);
      ff2=(1-alpha)*f2[state][r-1] + alpha*f2[state][r];
      if (ff2<0) ff2=0;

      /*using the new kf-vector, recompute the intersections to find length to go to correct for multiple scattering and absorption*/
      if (shape==CYLINDER){
        hit=cylinder_intersect(&l0,&l1,x,y,z,kx,ky,kz,radius,yheight);
        /*sample is a cylinder*/
      }else if (shape==SPHERE){
        /*sample is a sphere - unlikely*/
        hit=sphere_intersect(&l0,&l1,x,y,z,kx,ky,kz,radius);
      }else if (shape==BOX){
        /*sample is a box*/
        hit=box_intersect(&l0,&l1,x,y,z,kx,ky,kz,xwidth,yheight,zdepth);
      }

      /*scale p according to F2 and \int_0_inf F2*/
      //p*=(ff2/f2sum[state]) * (psimax-psimin)/M_PI * (etamax-etamin)/(2*M_PI) *  p_s/(1-p_transmit);
      p*=(ff2/f2sum[state]) * solid_angle * p_s/(1-p_transmit);
      //p*=f2mcsum[state]/f2sum[state] * (etamax-etamin)/(2*M_PI) *  p_s/(1-p_transmit);

      /*Absorption after scattering*/
      p*=exp(-mu*l1);
      //printf("ATT: %g %g %g\n",l1,mu,exp(-mu*l1));

    }else{
      /*tunneling branch*/
      /*downscale p by the total amount of scattering while going straight through ,weighted */
      p*=(1-p_s)/(p_transmit);
      /*also downscale for absorption effects*/
      l1-=l0;
      p*=exp(-mu*l1);
      /*here we should also take into account flourescence*/
    }
  }
%}

MCDISPLAY
%{
  
  if (shape==CYLINDER){
    /*sample is a cylinder*/
    circle("xz", 0,  yheight/2.0, 0, radius);
    circle("xz", 0, -yheight/2.0, 0, radius);
    line(-radius, -yheight/2.0, 0, -radius, +yheight/2.0, 0);
    line(+radius, -yheight/2.0, 0, +radius, +yheight/2.0, 0);
    line(0, -yheight/2.0, -radius, 0, +yheight/2.0, -radius);
    line(0, -yheight/2.0, +radius, 0, +yheight/2.0, +radius);
  }else if (shape==SPHERE){
    /*sample is a sphere*/
    circle("xy",0,0,0,radius);
    circle("xz",0,0,0,radius);
    circle("yz",0,0,0,radius);
  }else if (shape==BOX){
   /*sample is a box*/
    box(0,0,0,xwidth,yheight,zdepth,0, 0, 1, 0);
  }


  line(0,0,0,0.2,0,0);
  line(0,0,0,0,0.2,0);
  line(0,0,0,0,0,0.2);
%}

END