/*******************************************************************************
*
* McXtrace, x-ray tracing package
*         Copyright, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*         Synchrotron SOLEIL, Saint-Aubin, France
*
* Component: MCPL_input
*
* %Identification
* Written by: Erik B Knudsen
* Date: Aug 2016
* Origin: DTU Physics
*
* Source-like component that reads photon state parameters from an mcpl-file.
*
* %Description
* Source-like component that reads photon state parameters from a binary mcpl-file.
*
* MCPL is short for Monte Carlo Particle List, and is a new format for sharing events
* between e.g. MCNP(X), Geant-4 and McXtrace .
*
* When used with MPI, the --ncount given on the commandline is overwritten by 
* #MPI nodes x #events in the file.
*
* Example: MCPL_input(filename=voutput.mcpl,verbose=1,repeat_count=1,E_smear=0.1,pos_smear=0.001,dir_smear=0.01)
*
* %Parameters
* INPUT PARAMETERS
*
* filename: [str]       Name of photon mcpl file to read.
* verbose: [ ]          Print debugging information for first 10 particles read.
* polarisationuse: [ ]  If !=0 read polarisation vectors from file.
* Emin: [keV]           Lower energy bound. Particles found in the MCPL-file below the limit are skipped.
* Emax: [keV]           Upper energy bound. Particles found in the MCPL-file above the limit are skipped.
* repeat_count: [1]     Repeat contents of the MCPL file this number of times. NB: When running MPI, repeating is implicit and is taken into account by integer division. Should be combined sith the _smear options!
* E_smear: [1]          When repeating events, make a Gaussian MC choice within E_smear*E around particle energy E
* pos_smear: [m]        When repeating events, make a flat MC choice of position within pos_smear around particle starting position
* dir_smear: [deg]      When repeating events, make a Gaussian MC choice of direction within dir_smear around particle direction
*
* %End
*******************************************************************************/

DEFINE COMPONENT MCPL_input

SETTING PARAMETERS (string filename=0, polarisationuse=1,verbose=1, Emin=0, Emax=FLT_MAX, int repeat_count=1, E_smear=0, pos_smear=0, dir_smear=0)

DEPENDENCY "-Wl,-rpath,CMD(mcpl-config --show libdir) -LCMD(mcpl-config --show libdir) -lmcpl -ICMD(mcpl-config --show includedir)"

SHARE
%{
#include <mcpl.h>
%}

DECLARE
%{
mcpl_file_t inputfile;
long long nparticles;
long long read_photons;
long long used_photons;
int repeat_cnt;
int repeating;
int ismpislave;
#ifdef OPENACC
DArray1d X;
DArray1d Y;
DArray1d Z;
DArray1d KX;
DArray1d KY;
DArray1d KZ;
DArray1d EX;
DArray1d EY;
DArray1d EZ;
DArray1d T;
DArray1d P;
#endif
%}

INITIALIZE
%{
    char line[256];
    long long ncount;

    if(Emax<Emin){
        fprintf(stderr,"Warning(%s): Nonsensical energy interval: E=[%g,%g]. Aborting.\n",NAME_CURRENT_COMP,Emin,Emax);
        exit(-1);
    }
    /* No need to check if the file opens correctly since mcpl will
     * abort internally if it cannot open the file.*/
    inputfile = mcpl_open_file(filename);

  
    if ( !(nparticles=mcpl_hdr_nparticles(inputfile)) ) {
        fprintf(stderr,"Warning(%s): MCPL-file reports no present particles. Foolishly trying to go on.\n",NAME_CURRENT_COMP);
        #ifndef OPENACC
        nparticles=ncount;
        #endif
    }else{
        printf("Message(%s): MCPL file (%s) produced with %s.\n",NAME_CURRENT_COMP,filename,mcpl_hdr_srcname(inputfile));
        printf("Message(%s): MCPL file (%s) contains %lu particles.\n",NAME_CURRENT_COMP,filename,(long unsigned)nparticles); 
    }
    mcset_ncount(nparticles);

    if(repeat_count>1 && E_smear==0 && pos_smear==0 && dir_smear==0) {
      fprintf(stdout, "\n\n WARNING: You have requested a repeat_count=%i but have left all *_smear parameters at 0!\n --> This is not allowed! Resetting to repeat_count=1!\n",repeat_count);
      fprintf(stderr, "\n\n WARNING: You have requested a repeat_count=%i but have left all *_smear parameters at 0!\n --> This is not allowed! Resetting to repeat_count=1!\n",repeat_count);
      repeat_count=1;
      sleep(10);
    }

    if(repeat_count==0) repeat_count=1;
    repeat_cnt = repeat_count;
    ismpislave=0;
#if defined (USE_MPI)
    repeat_cnt = ceil(1.0*repeat_cnt/mpi_node_count);
    ismpislave = mpi_node_rank;
    MPI_MASTER(
#endif
	     fprintf(stdout, "\n\n Warning: You are using MCPL_input with a repeat_count of %lu:\n - Minimum photon count requested is %lu x %lu <= %lu",
	       (long unsigned)repeat_count,(long unsigned)nparticles,
	       (long unsigned)repeat_count,(long unsigned)repeat_cnt*nparticles); 
#if defined (USE_MPI)
  fprintf(stdout, " x %i MPI nodes = %lu photons total\n",
    mpi_node_count,(long unsigned)mpi_node_count*repeat_cnt*nparticles);
     );
  mcset_ncount(mpi_node_count*repeat_cnt*nparticles);
#else
  fprintf(stdout, " photons total\n\n");
  mcset_ncount(repeat_cnt*nparticles);
#endif

    ncount=mcget_ncount();
    fprintf(stdout,"Initialize ncount is %lu\n",(long unsigned)ncount);
    read_photons=0;
    used_photons=0;

#if defined (USE_MPI)   
    MPI_MASTER(
#endif

      if (verbose==1) {
        printf("MCPL_input verbose mode - outputting data on the 10 first read photons in MCPL units:\n");
      }

#if defined (USE_MPI) 
    );
#endif
    repeating = 0;
#ifdef OPENACC
      X = create_darr1d(nparticles);
      Y = create_darr1d(nparticles);
      Z = create_darr1d(nparticles);
      KX = create_darr1d(nparticles);
      KY = create_darr1d(nparticles);
      KZ = create_darr1d(nparticles);
      if(polarisationuse){
	EX = create_darr1d(nparticles);
      	EY = create_darr1d(nparticles);
      	EZ = create_darr1d(nparticles);
      }
      T = create_darr1d(nparticles);
      P = create_darr1d(nparticles);
      printf("Initiating file read...\n");
      int loop;
      for (loop=0; loop < nparticles ; loop++) {
	const mcpl_particle_t *particle;
	particle=mcpl_read(inputfile);
	if (particle) {
	  if (particle->pdgcode==22) {
	    if (verbose && read_photons<11) {
	      printf("id=%ld pdg=22\tekin=%g MeV\tx=%g cm\ty=%g cm\tz=%g cm\tux=%g\tuy=%g\tuz=%g\tt=%g ms\tweight=%g\tpolx=%g\tpoly=%g\tpolz=%g\n",
		     (long unsigned)read_photons, particle->ekin, particle->position[0], particle->position[1], particle->position[2],
		     particle->direction[0], particle->direction[1], particle->direction[2], particle->time, particle->weight,
		     particle->polarisation[0], particle->polarisation[1], particle->polarisation[2]);
	    }
	    /* check energy range*/
	    if ( particle->ekin>Emin*1e-9 || particle->ekin<Emax*1e-9 ) {
	      /* Particle energy in range */
	      /*positions are in cm*/
	      X[read_photons]=particle->position[0]/100;
	      Y[read_photons]=particle->position[1]/100;
	      Z[read_photons]=particle->position[2]/100;

	      if(polarisationuse){
		EX[read_photons]=(double)particle->polarisation[0];
		EY[read_photons]=(double)particle->polarisation[1];
		EZ[read_photons]=(double)particle->polarisation[2];
	      }else{
		EX[read_photons]=0;
		EY[read_photons]=0;
		EZ[read_photons]=0;
	      }
	      double nrm;
              nrm = particle->ekin *1e3*E2K;
              nrm = sqrt(nrm);

	      double d0=particle->direction[0],d1=particle->direction[1],d2=particle->direction[2];

	      KX[read_photons]=d0*nrm;
	      KY[read_photons]=d1*nrm;
	      KZ[read_photons]=d2*nrm;

	      /*time in ms:*/
	      T[read_photons] = particle->time*1e-3;
	      /*weight in unspecified units:*/
	      P[read_photons] = particle->weight;
	      read_photons++;
	    }
	  }
	}
      }
      printf("Done reading MCPL file, found %ld photons\n",(long unsigned)read_photons);
      mcpl_close_file(inputfile);
      fprintf(stdout, "\n\n Warning: You are using MCPL_input with a repeat_count of %lu:\n - Minimum photon count requested is %lu x %lu <= %lu",
	      (long unsigned)repeat_count,(long unsigned)read_photons,
	      (long unsigned)repeat_count,(long unsigned)repeat_cnt*read_photons);
      fprintf(stdout, " photons total\n\n");
      mcset_ncount(repeat_cnt*read_photons);

      ncount=mcget_ncount();
      fprintf(stdout,"Initialize ncount is %lu\n",(long unsigned)ncount);
#endif
%}

TRACE
%{
#ifndef OPENACC
    long long ncount;

    double nrm;
    const mcpl_particle_t *particle;// = (mcpl_particle_t *) calloc(sizeof(mcpl_particle_t),1);
    particle = mcpl_read(inputfile);

    ncount=mcget_ncount();
    // fprintf(stdout,"Trace ncount is %ld\n",ncount);      
    if (!particle) {
      if(repeat_cnt>1) {
	/* Trigger rewind of the file and ABSORB to get the first photon "again" */
	repeating++;
	mcpl_rewind(inputfile);
	particle = mcpl_read(inputfile);
#if defined (USE_MPI)   
	MPI_MASTER(
#endif
       printf("MCPL inputfile %s rewound %i time(s)\n",filename,repeating);
#if defined (USE_MPI) 
    );
#endif
      } else
	ABSORB;
    }
    if (particle->pdgcode!=22) {
        /*Either no particle read, particle is not a photon, or it has invalid energy - terminate to trigger next ray*/
       ABSORB;
    }
    read_photons++;

    /* check energy range*/
    if ( particle->ekin<Emin*1e-3 || particle->ekin>Emax*1e-3 ) {
        /*Particle energy out of range - terminate to trigger next ray*/
        ABSORB;
    }
    used_photons++;

#if defined (USE_MPI)   
  MPI_MASTER(
#endif

    if (verbose && used_photons<11) {
      printf("id=%ld\tpdg=22\tekin=%g MeV\tx=%g cm\ty=%g cm\tz=%g cm\tux=%g\tuy=%g\tuz=%g\tt=%g ms\tweight=%g\tpolx=%g\tpoly=%g\tpolz=%g\n",
	     read_photons, particle->ekin, particle->position[0], particle->position[1], particle->position[2],
	     particle->direction[0], particle->direction[1], particle->direction[2], particle->time, particle->weight,
	     particle->polarisation[0], particle->polarisation[1], particle->polarisation[2]);
    }

#if defined (USE_MPI)   
	    );
#endif    

    /*positions are in cm*/
    x=particle->position[0]/100;
    y=particle->position[1]/100;
    z=particle->position[2]/100;
    
    if (ismpislave || repeating) {
      double tmpx,tmpy,tmpz;
      // Position-MC:
      randvec_target_circle(&tmpx, &tmpy, &tmpz, NULL, 0, 0, 1, 0);
      NORM(tmpx,tmpy,tmpz);
      tmpx *= pos_smear*rand01(); tmpy *= pos_smear*rand01(); tmpz *= pos_smear*rand01();
      x+=tmpx; y+=tmpy; z+=tmpz;
    }
    
    if(polarisationuse){
        Ex=particle->polarisation[0];
        Ey=particle->polarisation[1];
        Ez=particle->polarisation[2];
    }else{
        Ex=Ey=Ez=0;
    }
    nrm = particle->ekin *1e3*E2K;
    nrm = sqrt(nrm);
    if (ismpislave || repeating) {
      // Energy-MC:
      double tmp=(1.0+E_smear*randpm1());
      //printf("Adjusting energy from %g to",nrm);
      nrm *= (1+E_smear*randpm1());
      //printf(" to %g\n",nrm);
    }
    double d0=particle->direction[0],d1=particle->direction[1],d2=particle->direction[2];
    
    if (ismpislave || repeating) {
      // Direction-MC:
      double tmpx,tmpy,tmpz;
      // Position-MC:
      randvec_target_circle(&d0, &d1, &d2, NULL, particle->direction[0], particle->direction[1], particle->direction[2], sin(dir_smear*DEG2RAD));
      NORM(d0,d1,d2);
    }
    
    kx=d0*nrm;
    ky=d1*nrm;
    kz=d2*nrm;   

    /*time in ms:*/
    t=particle->time*1e-3;
    /*weight in unspecified units:*/
    p=particle->weight;

    /* Correct for repetition, by repeat_count and/or MPI */
    p /= repeat_cnt;
#if defined (USE_MPI)   
    p /= mpi_node_count;
#endif
#else
  unsigned long long i=_particle->_uid;
  if (i>=nparticles) {
    repeating=1;
    i = i % nparticles;
  }
  x=X[i];
  y=Y[i];
  z=Z[i];
  kx=KX[i];
  ky=KY[i];
  kz=KZ[i];
  Ex=EX[i];
  Ey=EY[i];
  Ez=EZ[i];
  t=T[i];
  /*MCPL does not store photon phase information.*/
  phi=0;
  p=P[i];
  if (repeat_cnt>1) p = p /repeat_cnt;

  if (repeating) {
    /* Position smearing */
    double tmpx,tmpy,tmpz;
    // Position-MC:
    randvec_target_circle(&tmpx, &tmpy, &tmpz, NULL, 0, 0, 1, 0);
    NORM(tmpx,tmpy,tmpz);
    tmpx *= pos_smear*rand01(); tmpy *= pos_smear*rand01(); tmpz *= pos_smear*rand01();
    x+=tmpx; y+=tmpy; z+=tmpz;

    double d0=kx,d1=ky,d2=kz;
    /* Direction smearing: */
    randvec_target_circle(&d0, &d1, &d2, NULL, kx, ky, kz, sin(dir_smear*DEG2RAD));
    NORM(d0,d1,d2);

    /* Energy smearing: */
    double tmpe=(1.0+E_smear*randpm1());
    double nrm=K2E*sqrt(kx*kx+ky*ky+kz*kz);
    nrm *= (1.0+E_smear*randpm1());
    nrm = K2E*nrm;

    kx=nrm*d0;
    ky=nrm*d1;
    kz=nrm*d2;
  }
#endif
  SCATTER;
%}

SAVE
%{
    #ifndef OPENACC
    mcpl_close_file(inputfile);
    #endif
%}

FINALLY
%{
    long long ncount;

    ncount=mcget_ncount();

    if (used_photons!=read_photons){
        fprintf(stdout,"Message(%s): You have used %lu of %lu photons available in the MCPL file.\n",NAME_CURRENT_COMP,
          (long unsigned)used_photons,(long unsigned)read_photons);
    }
    if (ncount != used_photons){
        fprintf(stderr,"Warning (%s): You requested %ld photons from a file which contains %lu particles in general, of which only %lu are photons (within the wanted energy interval).\n"
                "Please examine the recorded intensities carefully.\n",
          NAME_CURRENT_COMP,(long unsigned)ncount,(long unsigned)nparticles,(long unsigned)used_photons);
    }
    
    destroy_darr1d(X);
    destroy_darr1d(Y);
    destroy_darr1d(Z);
    destroy_darr1d(KX);
    destroy_darr1d(KY);
    destroy_darr1d(KZ);
    if(polarisationuse){
      destroy_darr1d(EX);
      destroy_darr1d(EY);
      destroy_darr1d(EZ);
    }
    destroy_darr1d(T);
    destroy_darr1d(P);
%}

MCDISPLAY
%{
    multiline(5, 0.2,0.2,0.0, -0.2,0.2,0.0, -0.2,-0.2,0.0, 0.2,-0.2,0.0, 0.2,0.2,0.0);
    /*M*/
    multiline(5,-0.085,-0.085,0.0, -0.085,0.085,0.0, -0.045,-0.085,0.0, -0.005,0.085,0.0, -0.005,-0.085,0.0);
    /*I*/
    line(0.045,-0.085,0, 0.045, 0.085,0);    
    line(0.005, 0.085,0, 0.085, 0.085,0);    
    line(0.005,-0.085,0, 0.085,-0.085,0);    
%}

END