/*******************************************************************************
* Component: MCPL_input_once
*
* %I
* Written by: Gregory S Tucker
* Date: Sep 2024
* Origin: European Spallation Source ERIC
*
* Source-like component that reads neutron state parameters from a MCPL-file one time.
* %D
* Source-like component that reads neutron state parameters from a MCPL-file one time.
*
* MCPL is short for Monte Carlo Particle List, and is a format for sharing events
* between e.g. MCNP(X), Geant4 and McStas.
*
* When used with MPI, the file contents are shared between workers with each accessing
* approximately (#events in the file) / (#MPI nodes)
*
* %BUGS
*
* %P
* INPUT PARAMETERS
*
* filename: [str]       Name of neutron mcpl file to read.
* preload: [ ]          Load particles during INITIALIZE. On GPU preload is forced.
* polarisationuse: [ ]  If !=0 read polarisation vectors from file.
* Emin: [meV]           Lower energy bound. Particles found in the MCPL-file below the limit are skipped.
* Emax: [meV]           Upper energy bound. Particles found in the MCPL-file above the limit are skipped.
* v_smear: [1]          Relative fraction for randomness in replayed particle velocity v *= (1 + v_smear * randpm1())
* pos_smear: [m]        Maximum extent of random position for replayed particles
* dir_smear: [deg]      Maximum deviation of random direction for replayed particles
* always_smear: [ ]     Finite values force particle property smearing for all particles
* verbose: [ ]          Finite values may cause extra output (at some point in the future)
*
* %E
*******************************************************************************/

DEFINE COMPONENT MCPL_input_once

SETTING PARAMETERS (string filename=0, polarisationuse=1, Emin=0, Emax=FLT_MAX, v_smear=0, pos_smear=0, dir_smear=0, int always_smear=0, int preload=0, int verbose=0)

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

SHARE
%{
#include <mcpl.h>

typedef struct {
  // Prefixed names to avoid the _particle accessor macros 
  double _x, _y, _z;
  double _vx, _vy, _vz;
  double _sx, _sy, _sz;
  // insert [int, void*, int, double, double, double, unsigned long]
  // here in order to have matching memory layout with 
  // _struct_particle, e.g., _class_particle
  double _t, _p;
} mcpl_input_once_particle_t;

void mcpl_input_once_translator(const int use_polarisation, const mcpl_particle_t * input, mcpl_input_once_particle_t * output){
  // position in mm -> m
  output->_x = input->position[0] / 100;
  output->_y = input->position[1] / 100;
  output->_z = input->position[2] / 100;
  // ekin in MeV -> meV; then to velocity
  double nrm = sqrt(input->ekin  * 1e9 / VS2E); 
  output->_vx = input->direction[0] * nrm;
  output->_vy = input->direction[1] * nrm;
  output->_vz = input->direction[2] * nrm;
  // polarization is a direction, and we might ignore it
  output->_sx = (use_polarisation) ? input->polarisation[0] : 0;
  output->_sy = (use_polarisation) ? input->polarisation[1] : 0;
  output->_sz = (use_polarisation) ? input->polarisation[2] : 0;
  // time in msec -> sec
  output->_t = input->time  * 1e-3;
  // probability, unitless
  output->_p = input->weight;
}
%}

DECLARE
%{
mcpl_file_t inputfile;
uint64_t nparticles;
uint64_t read_neutrons;
uint64_t used_neutrons;
uint64_t emitted_neutrons;
uint64_t current_index;
uint64_t maximum_index;
uint64_t first_particle;
int times_replayed;
mcpl_input_once_particle_t * particles;
%}

INITIALIZE
%{
  uint64_t particles_per_node;
  uint64_t last_particle;
  if(Emax<Emin){
    fprintf(stderr, "Error(%s): Nonsensical energy interval: E=[%g,%g]. Aborting.\n", NAME_CURRENT_COMP, Emin, Emax);
    exit(-1);
  }
  inputfile = mcpl_open_file(filename);
  if (!(nparticles = mcpl_hdr_nparticles(inputfile))) {
    fprintf(stderr, "Error(%s): MCPL-file reports no present particles. Aborting.\n", NAME_CURRENT_COMP);
    exit(-1);
  } else {
#if defined (USE_MPI)   
    MPI_MASTER(
#endif
    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); 
#if defined (USE_MPI) 
    );
#endif
  }
  first_particle = 0;
  last_particle = nparticles;
#if defined (USE_MPI)
  // divy up the available particles between nodes
  particles_per_node = last_particle / mpi_node_count;
  // ensuring at least 1 particle per node (e.g., protecting against division by negative node count)
  if (particles_per_node < 1) particles_per_node = 1;
  // each node has first index given by how many particles each should do
  first_particle = particles_per_node * mpi_node_rank;
  // the last worker keeps 'nparticles' as its last particle index, to ensure the full range is covered
  if (mpi_node_rank != mpi_node_count - 1) last_particle = first_particle + particles_per_node;
#endif
  read_neutrons=0;
  used_neutrons=0;
#ifdef OPENACC
  preload=1;
  printf("OpenACC, preload implicit:\n");	    
#endif
  // Move this node's pointer into the file to its first particle:
  // in preparation for pre-loading or first pass through TRACE
  mcpl_seek(inputfile, first_particle);
  // Index will track how many particles this component has accessed
  current_index = 0;
  // Which we want to check against how large it can grow, to know if we've exhausted the available particles
  maximum_index = last_particle - first_particle;
  particles = NULL;
  if (preload) {
    printf("Preload requested, loading MCPLfile particles (%lu, %lu) in INITIALIZE\n", (long unsigned)first_particle, (long unsigned)last_particle);
    particles = (mcpl_input_once_particle_t *) calloc(last_particle - first_particle, sizeof(mcpl_input_once_particle_t));
    for (uint64_t loop=first_particle; loop < last_particle ; loop++) {
      const mcpl_particle_t *particle;
      particle=mcpl_read(inputfile);
      if (particle && particle->pdgcode==2112) {
        if (particle->ekin>Emin*1e-9 && particle->ekin<Emax*1e-9 ) {
          mcpl_input_once_translator(polarisationuse, particle, particles + used_neutrons++);
        }
        read_neutrons++;
      }
    }
    // keep track of how many particles are available to us in the `particles` array
    maximum_index = used_neutrons;
    printf("Done reading MCPL file (%lu, %lu), found %lu neutrons (and kept %lu)\n",
           (long unsigned)first_particle, (long unsigned)last_particle, (long unsigned)read_neutrons, (long unsigned)used_neutrons);
    mcpl_close_file(inputfile);
  }
  // keep track of how many times we had to replay the same MCPL input
  times_replayed = 0;
  // and the total number of neutrons emitted by this component
  emitted_neutrons = 0;
  // Determine the total number of read (or to be read) neutrons
  uint64_t total_index = maximum_index;
#if defined (USE_MPI)
  // add up the read neutrons (or to be read neutrons) from each node
  MPI_Reduce(&maximum_index, &total_index, 1, MPI_UINT64_T, MPI_SUM, 0, MPI_COMM_WORLD);
  // tell all nodes the value, so they can set ncount for themselves
  MPI_Bcast(&total_index, 1, MPI_UINT64_T, 0, MPI_COMM_WORLD);
#endif
  mcset_ncount(total_index); // will be divided by mpi_node_count before starting the raytrace
%}

TRACE
%{
  mcpl_input_once_particle_t * ptr = NULL;
  
  if (current_index >= maximum_index){
    // go back to the start of this worker's particles, rather than accessing out-of-range data
    #pragma acc atomic write
    current_index = 0;
#ifndef OPENACC
    if (!preload){
      // mcpl_seek is not available under openACC, and not used anyway
      mcpl_seek(inputfile, first_particle);
    }
#endif
    times_replayed++;
  }
#ifndef OPENACC
  // preload can only ever be false if OPENACC is not define (in which case it is the likely code path)
  if (!preload){
    ptr = (mcpl_input_once_particle_t *) calloc(1, sizeof(mcpl_input_once_particle_t));
    const mcpl_particle_t *particle;// = (mcpl_particle_t *) calloc(sizeof(mcpl_particle_t),1);
    particle = mcpl_read(inputfile);
    current_index++; // track how many particles have been accessed, to ensure we don't exceed our slice of the file
    if (!particle || particle->pdgcode != 2112) {
       ABSORB;
    }
    // keep track of how many neutrons have been read; but only on the first replay
    if (!times_replayed) read_neutrons++;
    if ( particle->ekin<Emin*1e-9 || particle->ekin>Emax*1e-9 ) {
        ABSORB;
    }
    mcpl_input_once_translator(polarisationuse, particle, ptr);
    // And how many of the read neutrons actually get used
    if (!times_replayed) used_neutrons++;
  } else {
#endif
    ptr = particles + (_particle->_uid % maximum_index);
#ifdef OPENACC
    #pragma acc atomic
    current_index++; // track how many particles have been accessed, to ensure we don't exceed our slice of the file
#else OPENACC
  }
#endif
  // copy from the component particle struct to the particle ray struct:
  if (ptr == NULL) {
    fprintf(stderr, "ERROR (%s): component particle struct pointer not set! Crash before out-of-bounds memory access.\n", NAME_CURRENT_COMP);
    exit(-1);
  }
  #pragma acc atomic
  emitted_neutrons++;
  // this could be done via memcpy if we ensure equal memory layout with the ray's struct
  x = ptr->_x;
  y = ptr->_y;
  z = ptr->_z;
  vx = ptr->_vx;
  vy = ptr->_vy;
  vz = ptr->_vz;
  sx = ptr->_sx;
  sy = ptr->_sy;
  sz = ptr->_sz;
  t = ptr->_t;
  p = ptr->_p;
  
  // done with the mcpl_input_once_particle_t pointer -- free it if not pointing into the particles list
  if (!preload && ptr != NULL) free(ptr);
  
  if (always_smear || times_replayed){
    // fuzz the input
    if (pos_smear) {
      double tmpx, tmpy, tmpz;
      randvec_target_circle(&tmpx, &tmpy, &tmpz, NULL, 0, 0, 1, 0);
      NORM(tmpx, tmpy, tmpz);
      x += tmpx * pos_smear * rand01();
      y += tmpy * pos_smear * rand01();
      z += tmpz * pos_smear * rand01();
    }
    if (v_smear) {
      double fraction;
      fraction = 1.0 + v_smear * randpm1();
      vx *= fraction;
      vy *= fraction;
      vz *= fraction;
    }
    if (dir_smear) {
      double vv, dx, dy, dz;
      vv = sqrt(vx * vx + vy * vy + vz * vz);
      dx = vx / vv;
      dy = vy / vv;
      dz = vz / vv;
      randvec_target_circle(&dx, &dy, &dz, NULL, dx, dy, dz, sin(dir_smear * DEG2RAD));
      NORM(dx, dy, dz);
      vx = dx * vv;
      vy = dy * vv;
      vz = dz * vv;
    }
  }
  SCATTER;
%}

SAVE
%{
  #ifndef OPENACC
  if (!preload) mcpl_close_file(inputfile);
  #endif
  if (particles != NULL) free(particles);
%}

FINALLY
%{
if (times_replayed && v_smear == 0 && pos_smear == 0 && dir_smear == 0){
  printf("Warning (%s): Forced to replay particle list %d time(s) without smearing\n", NAME_CURRENT_COMP, times_replayed);
}
  char mpi_nodes_message[256];
  mpi_nodes_message[0] = '\0';
  uint64_t requested_neutrons = (uint64_t) mcget_ncount();

#if defined (USE_MPI)
  uint64_t accumulated[4], distributed[4];
  distributed[0] = used_neutrons;
  distributed[1] = read_neutrons;
  distributed[2] = emitted_neutrons;
  distributed[3] = requested_neutrons;
  MPI_Reduce(&distributed, &accumulated, 4, MPI_UINT64_T, MPI_SUM, 0, MPI_COMM_WORLD);
if (mpi_node_rank == 0){
  used_neutrons = accumulated[0];
  read_neutrons = accumulated[1];
  emitted_neutrons = accumulated[2];
  requested_neutrons = accumulated[3];
  sprintf(mpi_nodes_message, "he %d MPI node copies of t", mpi_node_count);
#endif
  if (used_neutrons!=read_neutrons){
    fprintf(stdout,"Message(%s): You have used %llu of %llu neutrons available in the MCPL file.\n", NAME_CURRENT_COMP, used_neutrons, read_neutrons);
  }
  if (requested_neutrons != used_neutrons){
    char bad_particle_message[512];
    bad_particle_message[0] = '\0';
    if (used_neutrons != nparticles){
      sprintf(bad_particle_message, " (of which %llu are neutrons within the requested energy interval)", used_neutrons);
    }
    fprintf(stderr,
    "Warning (%s): You requested %llu neutrons from a file which contains %llu particles in general%s."
    " T%shis component emitted %llu neutrons in total. Please examine the recorded intensities carefully.\n", 
    NAME_CURRENT_COMP, requested_neutrons, nparticles, bad_particle_message, mpi_nodes_message, emitted_neutrons
    );
  }
#if defined (USE_MPI)
}
#endif
%}

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

