/*******************************************************************************
*
* McStas, neutron ray-tracing package
*         Copyright (C) 1997-2026, All rights reserved
*         DTU Physics, Kgs. Lyngby, Denmark
*
* Component: MCPL_input
*
* %I
* Written by: Erik B Knudsen, Gregory S Tucker and Peter Willendrup
* Date: Mar 2016, consolidated with MCPL_input_once Oct 2026
* Origin: DTU Physics / European Spallation Source ERIC
*
* Source-like component that reads neutron state parameters from an mcpl-file.
* %D
* Source-like component that reads neutron state parameters from a binary mcpl-file.
*
* MCPL is short for Monte Carlo Particle List, and is a format for sharing events
* between e.g. MCNP(X), Geant4 and McStas.
*
* The --ncount given on the commandline is overwritten by
* (#particles in the file) x repeat_count. Particles that are not neutrons, or are
* outside [Emin, Emax], are counted but ABSORBed.
*
* The rays form a stream that runs through the file repeat_count times. As in KDSource,
* this stream is split into one contiguous chunk per MPI process: each process seeks once
* to its start and reads on, wrapping around to the start of the file (with preload=1 it
* loads only the part of the file its chunk covers). Every event is thus used
* repeat_count times across the whole run, and its weight is divided by repeat_count.
* All passes after the first are smeared using v_smear, pos_smear and dir_smear (the first
* pass is too if always_smear is set), so repeat_count > 1 MUST be combined with at least
* one *_smear option; if all are 0, repeat_count is reset to 1.
*
* With replicate_across_processes=0 (the default, as in MCPL_input_once) the file is shared
* between the MPI processes and --mpi=N does not change the number of rays.
* replicate_across_processes=1 restores the pre-Oct 2026 MPI behaviour: repeat_count is
* rounded up to a multiple of the number of MPI processes, so that every process runs
* through the whole file at least once, and all but the first pass on the master are
* smeared (or, with all *_smear=0, are identical copies).
*
* Example: MCPL_input(filename=voutput,verbose=1,repeat_count=1,v_smear=0.1,pos_smear=0.001,dir_smear=0.01)
*
* %P
* INPUT PARAMETERS
*
* filename: [str]       Name of neutron mcpl file to read
* verbose: [ ]          Print debugging information for the first 10 neutrons read (on the master process)
* 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
* repeat_count: [1]     Replay the contents of the MCPL file this number of times, dividing the event weights by the same number. MUST be combined with _smear options!
* v_smear: [1]          When replaying events, make a flat MC choice within v_smear*V around particle velocity V
* pos_smear: [m]        When replaying events, make a flat MC choice of position within pos_smear around particle starting position
* dir_smear: [deg]      When replaying events, make an MC choice of direction within dir_smear around particle direction
* always_smear: [ ]     If !=0, also smear the events on their first pass through the file
* replicate_across_processes: [ ] If !=0, with MPI round repeat_count up to a multiple of the number of processes, so that each process runs through the whole file
*
* %E
*******************************************************************************/

DEFINE COMPONENT MCPL_input

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

DEPENDENCY "@MCPLFLAGS@"

SHARE
%{
  #include <mcpl.h>
  #include <stdint.h>
  #include <sys/stat.h>

  #ifndef MCPLINPUT_SHARED
  #define MCPLINPUT_SHARED
  typedef struct {
    // Prefixed names to avoid the _particle accessor macros
    double _x, _y, _z;
    double _vx, _vy, _vz;
    double _sx, _sy, _sz;
    double _t, _p;
    int _keep; // 0: not a neutron, 1: outside [Emin, Emax], 2: used
  } mcplinput_particle_t;

  /* MCPL units -> McStas units */
  void
  mcplinput_translator (const int use_polarisation, double weight_scale, const mcpl_particle_t* input, mcplinput_particle_t* output) {
    // position in cm -> 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;
    // polarisation 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 ms -> s
    output->_t = input->time * 1e-3;
    // weight in unspecified units
    output->_p = input->weight * weight_scale;
  }

  void
  mcplinput_print (const char* compname, uint64_t index, const mcpl_particle_t* particle) {
    printf ("%s: id=%llu pdg=%ld\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", compname,
            (unsigned long long)index, (long)particle->pdgcode, 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]);
  }

  int
  mcplinput_file_exist (char* fn) {
    struct stat buffer;
    return (stat (fn, &buffer) == 0);
  }
  #endif
%}

DECLARE
%{
  mcpl_file_t inputfile;
  uint64_t nparticles;
  uint64_t first_ray;
  uint64_t node_ncount;
  uint64_t stream_index;
  uint64_t read_neutrons;
  uint64_t used_neutrons;
  uint64_t emitted_neutrons;
  int repeat_tot;
  int file_open;
  double weight_scale;
  mcplinput_particle_t* particles;
  char* resolved_filename;
%}

INITIALIZE
%{
  {
    if (!filename || !filename[0]) {
      fprintf (stderr, "ERROR(%s): Requires filename parameter.\n", NAME_CURRENT_COMP);
      exit (-1);
    }
    char* fn_mcpl = mcpl_name_helper (filename, 'M');
    char* fn_mcplgz = mcpl_name_helper (filename, 'G');
    if (mcplinput_file_exist (fn_mcpl)) {
      if (mcplinput_file_exist (fn_mcplgz)) {
        fprintf (stderr,
                 "ERROR(%s): Can not resolve input file unambiguously"
                 " since both %s and %s exist.\n",
                 NAME_CURRENT_COMP, fn_mcpl, fn_mcplgz);
        exit (-1);
      }
      resolved_filename = fn_mcpl;
      free (fn_mcplgz);
    } else {
      resolved_filename = fn_mcplgz;
      free (fn_mcpl);
    }
  }

  if (Emax < Emin) {
    fprintf (stderr, "ERROR(%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 (resolved_filename);
  file_open = 1;
  double mcpl_ray_count = mcpl_hdr_stat_sum (inputfile, "initial_ray_count");
  if (mcpl_ray_count == -2.0) {
    // legacy format without ray count:
    weight_scale = 1.0;
  } else if (!(mcpl_ray_count > 0.0)) {
    fprintf (stderr,
             "ERROR(%s): Input MCPL file has invalid initial_ray_count"
             " (%g). Unable to determine weight scale.\n",
             NAME_CURRENT_COMP, mcpl_ray_count);
    exit (-1);
  } else {
    weight_scale = 1.0 / mcpl_ray_count;
  }

  if (!(nparticles = mcpl_hdr_nparticles (inputfile))) {
    fprintf (stderr, "ERROR(%s): MCPL file (%s) reports no particles. Aborting.\n", NAME_CURRENT_COMP, resolved_filename);
    exit (-1);
  }
  MPI_MASTER (printf ("Message(%s): MCPL file (%s) produced with %s.\n", NAME_CURRENT_COMP, resolved_filename, mcpl_hdr_srcname (inputfile));
              printf ("Message(%s): MCPL file (%s) contains %llu particles.\n", NAME_CURRENT_COMP, resolved_filename, (unsigned long long)nparticles););

  if (repeat_count < 1)
    repeat_count = 1;
  if (repeat_count > 1 && v_smear == 0 && pos_smear == 0 && dir_smear == 0) {
    MPI_MASTER (fprintf (stderr,
                         "WARNING(%s): You have requested repeat_count=%i but have left all *_smear parameters at 0!\n"
                         " --> Replaying identical events is not allowed. Resetting to repeat_count=1!\n",
                         NAME_CURRENT_COMP, repeat_count););
    repeat_count = 1;
  }
  /* Number of passes through the file across all processes */
  repeat_tot = repeat_count;
  #if defined (USE_MPI)
  if (replicate_across_processes)
    repeat_tot = mpi_node_count * ((repeat_count + mpi_node_count - 1) / mpi_node_count);
  #endif

  /* The rays are a stream of nparticles x repeat_tot events, split into one contiguous chunk
     per process. Ray _uid = first_ray + i is event _uid % nparticles on pass _uid / nparticles */
  if (!mcdotrace) {
    mcset_ncount (nparticles * repeat_tot);
    MPI_MASTER (if (repeat_tot > 1) printf ("Message(%s): ncount set to %llu x %i = %llu.\n", NAME_CURRENT_COMP, (unsigned long long)nparticles,
                                            repeat_tot, (unsigned long long)(nparticles * repeat_tot)););
  }
  // the main loop gives each MPI process floor(ncount / #processes) rays
  node_ncount = (uint64_t)mcget_ncount ();
  first_ray = 0;
  #if defined (USE_MPI)
  node_ncount /= mpi_node_count;
  first_ray = (uint64_t)mpi_node_rank * node_ncount;
  #endif
  stream_index = first_ray;

  read_neutrons = 0;
  used_neutrons = 0;
  emitted_neutrons = 0;
  particles = NULL;

  #ifdef OPENACC
  preload = 1;
  MPI_MASTER (printf ("Message(%s): OpenACC, preload implicit.\n", NAME_CURRENT_COMP););
  #endif

  // Move this process' pointer into the file to the start of its chunk
  if (first_ray % nparticles)
    mcpl_seek (inputfile, first_ray % nparticles);

  if (preload) {
    // load only the part of the file this process' chunk covers, in stream order
    uint64_t nload = node_ncount < nparticles ? node_ncount : nparticles;
    MPI_MASTER (printf ("Message(%s): Preload requested, loading %llu particles per process in INITIALIZE.\n", NAME_CURRENT_COMP, (unsigned long long)nload););
    if (nload) {
      particles = (mcplinput_particle_t*)calloc (nload, sizeof (mcplinput_particle_t));
      if (!particles) {
        fprintf (stderr, "ERROR(%s): Could not allocate memory for %llu particles.\n", NAME_CURRENT_COMP, (unsigned long long)nload);
        exit (-1);
      }
    }
    for (uint64_t j = 0; j < nload; j++) {
      if (j && (first_ray + j) % nparticles == 0)
        mcpl_rewind (inputfile);
      const mcpl_particle_t* particle = mcpl_read (inputfile);
      if (particle && particle->pdgcode == 2112) {
        MPI_MASTER (if (verbose && j < 10) mcplinput_print (NAME_CURRENT_COMP, (first_ray + j) % nparticles, particle););
        mcplinput_translator (polarisationuse, weight_scale, particle, particles + j);
        particles[j]._keep = (particle->ekin > Emin * 1e-9 && particle->ekin < Emax * 1e-9) ? 2 : 1;
      }
    }
    mcpl_close_file (inputfile);
    file_open = 0;
  }
%}

TRACE
%{
  mcplinput_particle_t buffer;
  mcplinput_particle_t* ptr = NULL;
  uint64_t g;

  #ifndef OPENACC
  // preload can only be false without OpenACC
  if (!preload) {
    g = stream_index++;
    // wrap around to the start of the file
    if (g != first_ray && g % nparticles == 0)
      mcpl_rewind (inputfile);
    const mcpl_particle_t* particle = mcpl_read (inputfile);
    buffer._keep = 0;
    if (particle && particle->pdgcode == 2112) {
      MPI_MASTER (if (verbose && g - first_ray < 10) mcplinput_print (NAME_CURRENT_COMP, g % nparticles, particle););
      mcplinput_translator (polarisationuse, weight_scale, particle, &buffer);
      buffer._keep = (particle->ekin > Emin * 1e-9 && particle->ekin < Emax * 1e-9) ? 2 : 1;
    }
    ptr = &buffer;
  } else {
    #endif
    g = (uint64_t)_particle->_uid;
    ptr = particles + (g - first_ray) % nparticles;
    #ifndef OPENACC
  }
  #endif
  uint64_t pass = g / nparticles;

  if (pass == 0 && ptr->_keep) {
    #pragma acc atomic
    read_neutrons++;
  }
  if (ptr->_keep != 2)
    ABSORB;
  if (pass == 0) {
    #pragma acc atomic
    used_neutrons++;
  }
  #pragma acc atomic
  emitted_neutrons++;

  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;
  // each event is emitted repeat_tot times
  p = ptr->_p / repeat_tot;

  if (always_smear || pass > 0) {
    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 = 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;
%}

FINALLY
%{
  #ifndef OPENACC
  if (file_open) {
    mcpl_close_file (inputfile);
    file_open = 0;
  }
  #endif

  uint64_t requested_neutrons = (uint64_t)mcget_ncount ();
  uint64_t tolerance = 1;
  #if defined (USE_MPI)
  // the main loop rounds the number of rays per process down
  tolerance = mpi_node_count;
  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];
    #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, (unsigned long long)used_neutrons,
               (unsigned long long)read_neutrons);
    }
    uint64_t expected = nparticles * (uint64_t)repeat_tot;
    if (requested_neutrons + tolerance <= expected || requested_neutrons >= expected + tolerance) {
      fprintf (stderr,
               "Warning(%s): %llu rays were traced, but the MCPL file holds %llu particles x %i passes = %llu,\n"
               "so events were %s%s. %llu neutrons were emitted in total. Please examine the recorded intensities carefully.\n",
               NAME_CURRENT_COMP, (unsigned long long)requested_neutrons, (unsigned long long)nparticles, repeat_tot, (unsigned long long)expected,
               requested_neutrons > expected ? "replayed" : "left unused",
               (requested_neutrons > expected && v_smear == 0 && pos_smear == 0 && dir_smear == 0) ? " WITHOUT smearing" : "",
               (unsigned long long)emitted_neutrons);
    }
    #if defined (USE_MPI)
  }
  #endif

  if (particles != NULL) {
    free (particles);
    particles = NULL;
  }
  free (resolved_filename);
%}

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