<commit_before>#include "openmc/simulation.h"

#include "openmc/bank.h"
#include "openmc/capi.h"
#include "openmc/cell.h"
#include "openmc/container_util.h"
#include "openmc/eigenvalue.h"
#include "openmc/error.h"
#include "openmc/geometry.h"
#include "openmc/material.h"
#include "openmc/message_passing.h"
#include "openmc/mgxs_interface.h"
#include "openmc/nuclide.h"
#include "openmc/output.h"
#include "openmc/particle.h"
#include "openmc/photon.h"
#include "openmc/physics.h"
 #include "openmc/physics_mg.h"
#include "openmc/random_lcg.h"
#include "openmc/settings.h"
#include "openmc/source.h"
#include "openmc/state_point.h"
#include "openmc/thermal.h"
#include "openmc/timer.h"
#include "openmc/tallies/derivative.h"
#include "openmc/tallies/filter.h"
#include "openmc/tallies/tally.h"
#include "openmc/tallies/tally_scoring.h"
#include "openmc/tallies/trigger.h"
#include "openmc/track_output.h"

#ifdef _OPENMP
#include <omp.h>
#endif
#include "xtensor/xview.hpp"

#ifdef OPENMC_MPI
#include <mpi.h>
#endif

#include <algorithm>
#include <string>

namespace openmc {

	/*
extern std::vector<Particle*> calculate_fuel_xs_queue;
extern std::vector<Particle*> calculate_nonfuel_xs_queue;
extern std::vector<Particle*> advance_particle_queue;
extern std::vector<Particle*> surface_crossing_queue;
extern std::vector<Particle*> collision_queue;
#pragma omp threadprivate(calculate_fuel_xs_queue, calculate_nonfuel_xs_queue, advance_particle_queue, surface_crossing_queue, collision_queue)

std::vector<Particle*> calculate_fuel_xs_queue;
std::vector<Particle*> calculate_nonfuel_xs_queue;
std::vector<Particle*> advance_particle_queue;
std::vector<Particle*> surface_crossing_queue;
std::vector<Particle*> collision_queue;
*/

int * calculate_fuel_xs_queue;
int * calculate_nonfuel_xs_queue;
int * advance_particle_queue;
int * surface_crossing_queue;
int * collision_queue;
Particle * particles;

int calculate_fuel_xs_queue_length    = 0;
int calculate_nonfuel_xs_queue_length = 0;
int advance_particle_queue_length     = 0;
int surface_crossing_queue_length     = 0;
int collision_queue_length            = 0;

const int MAX_PARTICLES_IN_FLIGHT = 100;

void init_event_queues(void)
{
	int n_particles = MAX_PARTICLES_IN_FLIGHT;
	calculate_fuel_xs_queue =    new int[n_particles];
	calculate_nonfuel_xs_queue = new int[n_particles];
	advance_particle_queue =     new int[n_particles];
	surface_crossing_queue =     new int[n_particles];
	collision_queue =            new int[n_particles];
	particles =                  new Particle[n_particles];
}

void free_event_queues(void)
{
	delete calculate_fuel_xs_queue;
	delete calculate_nonfuel_xs_queue;
	delete advance_particle_queue;
	delete surface_crossing_queue;
	delete collision_queue;
	delete[] particles;
}

constexpr size_t MAX_PARTICLES_PER_THREAD {100};

Particle::Bank * shared_fission_bank;
int shared_fission_bank_length = 0;
int shared_fission_bank_max;

void init_shared_fission_bank(int max)
{
	shared_fission_bank_max = max;
	shared_fission_bank = new Particle::Bank[max];
}

void free_shared_fission_bank(void)
{
	delete[] shared_fission_bank;
	shared_fission_bank_length = 0;
}

// TODO: What is going on here?
void revive_particle_from_secondary(Particle* p)
{
  //p->from_source(&simulation::secondary_bank.back());
  p->from_source(&p->secondary_bank_.back());
  //simulation::secondary_bank.pop_back();
  p->secondary_bank_.pop_back();
  // n_event = 0;

  // Enter new particle in particle track file
  if (p->write_track_) add_particle_track();
}

void dispatch_xs_event(int i)
{
	Particle * p = particles + i;
	  int idx;
  if (p->material_ == MATERIAL_VOID) {

	#pragma omp atomic capture
	idx = calculate_nonfuel_xs_queue_length++;
	//std::cout << "Dispatching particle to non Fuel XS queue idx = " << idx << std::endl;
    calculate_nonfuel_xs_queue[idx] = i;
  } else {
    if (model::materials[p->material_]->fissionable_) {
	#pragma omp atomic capture
	idx = calculate_fuel_xs_queue_length++;
	//std::cout << "Dispatching particle to Fuel XS queue idx = " << idx << std::endl;
      calculate_fuel_xs_queue[idx] = i;
    } else {
	#pragma omp atomic capture
	idx = calculate_nonfuel_xs_queue_length++;
	//std::cout << "Dispatching particle to non Fuel XS queue idx = " << idx << std::endl;
      calculate_nonfuel_xs_queue[idx] = i;
    }
  }
}

void process_calculate_xs_events(int * queue, int n)
{
  // Save last_ members, find grid index
  int lost_particles = 0;
  #pragma omp parallel for reduction(+:lost_particles)
  for (int i = 0; i < n; i++) {
	  Particle *p = particles + queue[i]; 
	  //std::cout << "particle offset = " << queue[i] << std::endl;
    // Set the random number stream
	// TODO: Move RNG seeds to particle storage
    if (p->type_ == Particle::Type::neutron) {
		p->stream_ = STREAM_TRACKING;
    } else {
		p->stream_ = STREAM_PHOTON;
    }

    // Store pre-collision particle properties
    p->wgt_last_ = p->wgt_;
    p->E_last_ = p->E_;
    p->u_last_ = p->u();
    p->r_last_ = p->r();

    // If the cell hasn't been determined based on the particle's location,
    // initiate a search for the current cell. This generally happens at the
    // beginning of the history and again for any secondary particles
    if (p->coord_[p->n_coord_ - 1].cell == C_NONE) {
      if (!find_cell(p, false)) {
        p->mark_as_lost("Could not find the cell containing particle "
          + std::to_string(p->id_));
        //return;
		lost_particles += 1;
		continue;
      }

      // set birth cell attribute
      if (p->cell_born_ == C_NONE) p->cell_born_ = p->coord_[p->n_coord_ - 1].cell;
    }

    // Write particle track.
    if (p->write_track_) write_particle_track(*p);

    if (settings::check_overlaps) check_cell_overlap(p);

    if (settings::run_CE) {
      if (p->material_ == p->material_last_ && p->sqrtkT_ != p->sqrtkT_last_) {
        // Remove particle from queue
      }
    }

    // Find energy index on energy grid
    // TODO: Calculate this separately?
    //int neutron = static_cast<int>(Particle::Type::neutron);
    //p->macro_xs_.i_grid = std::log(p->E_/data::energy_min[neutron]) / simulation::log_spacing;
  }
  if( lost_particles > 0 )
	  exit(1);
  #pragma omp parallel for
  for( int i = 0; i < n; i++ )
  {
	  Particle * p = particles + queue[i];
	  // Calculate microscopic and macroscopic cross sections
	  if (p->material_ != MATERIAL_VOID) {
		  if (settings::run_CE) {
			  if (p->material_ != p->material_last_ || p->sqrtkT_ != p->sqrtkT_last_) {
				  // If the material is the same as the last material and the
				  // temperature hasn't changed, we don't need to lookup cross
				  // sections again.
				  model::materials[p->material_]->calculate_xs(*p);
			  }
		  } // else MG not supported
	  } else {
		  p->macro_xs_.total      = 0.0;
		  p->macro_xs_.absorption = 0.0;
		  p->macro_xs_.fission    = 0.0;
		  p->macro_xs_.nu_fission = 0.0;
	  }
	  /*
	  int idx;
	  #pragma omp atomic capture
	  idx = advance_particle_queue_length++;
	  advance_particle_queue[idx] = queue[i];
	  */
  }
  int start = advance_particle_queue_length;
  int end = start + n;
  int j = 0;
  for( int i = start; i < end; i++ )
  {
	  advance_particle_queue[i] = queue[j];
	  j++;
  }
  advance_particle_queue_length += n;

  /*
  // Calculate nuclide micros
  for (int i = 0; i < data::nuclides.size(); ++i) {
	  // loop over particles
    for (int j = 0; j < n; j++) {
		//Particle * p = particles + queue[i];
		Particle * p = particles + queue[j];
      if (p->material_ == MATERIAL_VOID) continue;

      // If material doesn't have this nuclide, skip it
      const auto& mat {model::materials[p->material_]};
      if (mat->mat_nuclide_index_[i] == -1) continue;

      // ======================================================================
      // CHECK FOR S(A,B) TABLE

      // Check if this nuclide matches one of the S(a,b) tables specified.
      // This relies on thermal_tables_ being sorted by .index_nuclide
      int i_sab = C_NONE;
      double sab_frac = 0.0;
      for (const auto& sab : mat->thermal_tables_) {
        if (i == sab.index_nuclide) {
          // Get index in sab_tables
          i_sab = sab.index_table;
          sab_frac = sab.fraction;

          // If particle energy is greater than the highest energy for the
          // S(a,b) table, then don't use the S(a,b) table
          //if (p->E_ > data::thermal_scatt[i_sab]->threshold()) i_sab = C_NONE;
			std::cout << "S(alpha, beta) max energy = " << data::thermal_scatt[i_sab]->energy_max_ <<std::endl;
          if (p->E_ > data::thermal_scatt[i_sab]->energy_max_) i_sab = C_NONE;
        }
      }

      // ======================================================================
      // CALCULATE MICROSCOPIC CROSS SECTION

      // Calculate microscopic cross section for this nuclide
      const auto& micro {p->neutron_xs_[i]};
      if (p->E_ != micro.last_E
          || p->sqrtkT_ != micro.last_sqrtkT
          || i_sab != micro.index_sab
          || sab_frac != micro.sab_frac) {
        data::nuclides[i]->calculate_xs(i_sab, p->macro_xs_.i_grid, sab_frac, *p);
      }
    }
  }

  for (int i = 0; i < n; i++) {
	  Particle * p = particles + queue[i];
    // Calculate microscopic and macroscopic cross sections
    if (p->material_ != MATERIAL_VOID) {
	  // Only works for CE, no MG support
      //if (settings::run_CE) {
        // If the material is the same as the last material and the
        // temperature hasn't changed, we don't need to lookup cross
        // sections again.
        model::materials[p->material_]->calculate_xs(*p);
      }
		//else {
      //}
     else {
      p->macro_xs_.total      = 0.0;
      p->macro_xs_.absorption = 0.0;
      p->macro_xs_.fission    = 0.0;
      p->macro_xs_.nu_fission = 0.0;
    }

	 int idx;
	 #pragma omp atomic capture
	 idx = advance_particle_queue_length++;
    advance_particle_queue[idx] = queue[i];
  }
	*/
}

void process_advance_particle_events()
{
  //for (auto& p : advance_particle_queue) {
  #pragma omp parallel for
  for (int i = 0; i < advance_particle_queue_length; i++) {
	  Particle * p = particles + advance_particle_queue[i];
    simulation::trace == (p->id_ == 0);

    // Sample a distance to collision
    double d_collision;
    if (p->type_ == Particle::Type::electron ||
        p->type_ == Particle::Type::positron) {
      d_collision = 0.0;
    } else if (p->macro_xs_.total == 0.0) {
      d_collision = INFINITY;
    } else {
      d_collision = -std::log(prn(p->prn_seeds_, p->stream_)) / p->macro_xs_.total;
    }

    // -------------- break here? -------------------

    // Find the distance to the nearest boundary
    p->boundary_ = distance_to_boundary(p);

    // Select smaller of the two distances
    double distance;
	int idx;
    if (p->boundary_.distance < d_collision) {
		#pragma omp atomic capture
		idx = surface_crossing_queue_length++;
      surface_crossing_queue[idx] = advance_particle_queue[i];
      distance = p->boundary_.distance;
    } else {
		#pragma omp atomic capture
		idx = collision_queue_length++;
		collision_queue[idx] = advance_particle_queue[i];
      distance = d_collision;
    }

    // -------------- break here? -------------------

    // Advance particle
    for (int j = 0; j < p->n_coord_; ++j) {
      p->coord_[j].r += distance * p->coord_[j].u;
    }

    // -------------- break here? -------------------

    // Score track-length tallies
    if (!model::active_tracklength_tallies.empty()) {
      score_tracklength_tally(p, distance);
    }

    // Score track-length estimate of k-eff
    if (settings::run_mode == RUN_MODE_EIGENVALUE &&
        p->type_ == Particle::Type::neutron) {
      global_tally_tracklength += p->wgt_ * distance * p->macro_xs_.nu_fission;
    }

    // Score flux derivative accumulators for differential tallies.
    if (!model::active_tallies.empty()) {
      score_track_derivative(p, distance);
    }
  }

  advance_particle_queue_length = 0;
}

void process_surface_crossing_events()
{
  //for (auto& p : surface_crossing_queue) {
  #pragma omp parallel for
  for (int i = 0; i < surface_crossing_queue_length; i++) {
	  Particle * p = particles + surface_crossing_queue[i];
    // Set surface that particle is on and adjust coordinate levels
    p->surface_ = p->boundary_.surface_index;
    p->n_coord_ = p->boundary_.coord_level;

    // Saving previous cell data
    for (int j = 0; j < p->n_coord_; ++j) {
      p->cell_last_[j] = p->coord_[j].cell;
    }
    p->n_coord_last_ = p->n_coord_;

    if (p->boundary_.lattice_translation[0] != 0 ||
        p->boundary_.lattice_translation[1] != 0 ||
        p->boundary_.lattice_translation[2] != 0) {
      // Particle crosses lattice boundary
      cross_lattice(p, p->boundary_);
      p->event_ = EVENT_LATTICE;
    } else {
      // Particle crosses surface
      p->cross_surface();
      p->event_ = EVENT_SURFACE;
    }
    // Score cell to cell partial currents
    if (!model::active_surface_tallies.empty()) {
      score_surface_tally(p, model::active_surface_tallies);
    }

    //if (!p->alive_ && !simulation::secondary_bank.empty()) {
    if (!p->alive_ && !p->secondary_bank_.empty()) {
      revive_particle_from_secondary(p);
    }

    if (p->alive_)
	{
		dispatch_xs_event(surface_crossing_queue[i]);
	}
  }

  surface_crossing_queue_length = 0;
}

void process_collision_events()
{
  //for (auto& p : collision_queue) {
  #pragma omp parallel for
  for (int i = 0; i < collision_queue_length; i++) {
	  Particle * p = particles + collision_queue[i];
	  //std::cout << "Beginning collision of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;
    // Score collision estimate of keff
    if (settings::run_mode == RUN_MODE_EIGENVALUE &&
        p->type_ == Particle::Type::neutron) {
      global_tally_collision += p->wgt_ * p->macro_xs_.nu_fission
        / p->macro_xs_.total;
    }

    // Score surface current tallies -- this has to be done before the collision
    // since the direction of the particle will change and we need to use the
    // pre-collision direction to figure out what mesh surfaces were crossed

    if (!model::active_meshsurf_tallies.empty())
      score_surface_tally(p, model::active_meshsurf_tallies);
	    
	//std::cout << "After surface tally of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Clear surface component
    p->surface_ = 0;

    if (settings::run_CE) {
      collision(p);
    } else {
      collision_mg(p);
    }

	//std::cout << "After collision() of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Score collision estimator tallies -- this is done after a collision
    // has occurred rather than before because we need information on the
    // outgoing energy for any tallies with an outgoing energy filter
    if (!model::active_collision_tallies.empty()) score_collision_tally(p);
    if (!model::active_analog_tallies.empty()) {
      if (settings::run_CE) {
        score_analog_tally_ce(p);
      } else {
        score_analog_tally_mg(p);
      }
    }
	//std::cout << "After analog tally of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Reset banked weight during collision
    p->n_bank_ = 0;
    p->n_bank_second_ = 0;
    p->wgt_bank_ = 0.0;
    for (int& v : p->n_delayed_bank_) v = 0;

    // Reset fission logical
    p->fission_ = false;

    // Save coordinates for tallying purposes
    p->r_last_current_ = p->r();

    // Set last material to none since cross sections will need to be
    // re-evaluated
    p->material_last_ = C_NONE;

    // Set all directions to base level -- right now, after a collision, only
    // the base level directions are changed
    for (int j = 0; j < p->n_coord_ - 1; ++j) {
      if (p->coord_[j + 1].rotated) {
        // If next level is rotated, apply rotation matrix
        const auto& m {model::cells[p->coord_[j].cell]->rotation_};
        const auto& u {p->coord_[j].u};
        p->coord_[j + 1].u.x = m[3]*u.x + m[4]*u.y + m[5]*u.z;
        p->coord_[j + 1].u.y = m[6]*u.x + m[7]*u.y + m[8]*u.z;
        p->coord_[j + 1].u.z = m[9]*u.x + m[10]*u.y + m[11]*u.z;
      } else {
        // Otherwise, copy this level's direction
        p->coord_[j+1].u = p->coord_[j].u;
      }
    }

    // Score flux derivative accumulators for differential tallies.
    if (!model::active_tallies.empty()) score_collision_derivative(p);

    //if (!p->alive_ && !simulation::secondary_bank.empty()) {
    if (!p->alive_ && !p->secondary_bank_.empty()) {
      revive_particle_from_secondary(p);
    }

    if (p->alive_)
	{
		dispatch_xs_event(collision_queue[i]);
	    //std::cout << "Ended collision of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;
		assert(std::isfinite(p->E_) );
	}
  }

  collision_queue_length = 0;
}

void check_energies(void)
{
	int * Q;
	int n;


	Q = calculate_fuel_xs_queue;
	n = calculate_fuel_xs_queue_length;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index xs FUEL " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = calculate_nonfuel_xs_queue;
	n = calculate_nonfuel_xs_queue_length ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index xs Non fuel " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = advance_particle_queue;
	n = advance_particle_queue_length     ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index advance particle " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = surface_crossing_queue;
	n = surface_crossing_queue_length     ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index surface crossing " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = collision_queue;
	n = collision_queue_length            ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index collision " << Q[i] << std::endl;
			assert(0);
		}
	}
}

void transport()
{
	int remaining_work = simulation::work_per_rank;
	int source_offset = 0;
		
	init_event_queues();

	double time_fuel_xs = 0;
	double time_nonfuel_xs = 0;
	double time_advance = 0;
	double time_collision = 0;
	double time_surf = 0;
	double stop, start;

	// Subiterations to complete sets of particles
	while (remaining_work > 0) {


		// Figure out work for this subiteration
		int n_particles = MAX_PARTICLES_IN_FLIGHT;
		if( n_particles > remaining_work)
			n_particles = remaining_work;

		//std::cout << "Initializing particle histories..." << std::endl;
		// Initialize all histories
		// TODO: Parallelize

		/*
		for( int p = 0; p < mpi::n_procs; p++)
		{
			MPI_Barrier(mpi::intracomm);
			if( p == mpi::rank )
			{
			*/
				for (int i = 0; i < n_particles; i++) {
					initialize_history(particles + i, source_offset + i + 1);
				}
				/*
			}
		}
		*/

		//std::cout << "Enqueing particles for XS Lookups..." << std::endl;
		// Add all particles to advance particle queue
		// TODO: Parallelize
		for (int i = 0; i < n_particles; i++) {
			dispatch_xs_event(i);
		}

		int event_kernel_executions = 0;
		while (true) {
			event_kernel_executions++;
			/*
			std::cout << "Fuel XS Lookups = " << calculate_fuel_xs_queue_length << std::endl;
			std::cout << "Non Fuel XS Lookups = " << calculate_nonfuel_xs_queue_length << std::endl;
			std::cout << "Advance Particles = " << advance_particle_queue_length << std::endl;
			std::cout << "Surface Crossings = " << surface_crossing_queue_length << std::endl;
			std::cout << "Collisions = " << collision_queue_length << std::endl;
			*/
			/*
			Particle * p = particles +  1;
		   std::cout << "E = " << p->E_ << " and Position {" <<
			   p->r().x << ", " <<
			   p->r().y << ", " <<
			   p->r().z << "}" << std::endl;
			   */
			int max = std::max({calculate_fuel_xs_queue_length, calculate_nonfuel_xs_queue_length, advance_particle_queue_length, surface_crossing_queue_length, collision_queue_length});
			//check_energies();
			if (max == 0) {
				break;
			} else if (max == calculate_fuel_xs_queue_length) {
				//std::cout << "pre fuel XS check..." << std::endl;
				//check_energies(calculate_fuel_xs_queue, calculate_fuel_xs_queue_length);
				//std::cout << "Performing Fuel XS Lookups..." << std::endl;
				start = omp_get_wtime();
				process_calculate_xs_events(calculate_fuel_xs_queue, calculate_fuel_xs_queue_length);
				stop = omp_get_wtime();
				time_fuel_xs += (stop-start);
				calculate_fuel_xs_queue_length = 0;
			} else if (max == calculate_nonfuel_xs_queue_length) {
				//std::cout << "pre non fuel XS check..." << std::endl;
				//check_energies(calculate_nonfuel_xs_queue, calculate_nonfuel_xs_queue_length);
			//	std::cout << "Performing Non Fuel XS Lookups..." << std::endl;
				start = omp_get_wtime();
				process_calculate_xs_events(calculate_nonfuel_xs_queue, calculate_nonfuel_xs_queue_length);
				stop = omp_get_wtime();
				time_nonfuel_xs += (stop-start);
				calculate_nonfuel_xs_queue_length = 0;
			} else if (max == advance_particle_queue_length) {
				//std::cout << "pre advancing check..." << std::endl;
				//check_energies(advance_particle_queue, advance_particle_queue_length);
				//std::cout << "Advancing Particles..." << std::endl;
				start = omp_get_wtime();
				process_advance_particle_events();
				stop = omp_get_wtime();
				time_advance += (stop-start);
			} else if (max == surface_crossing_queue_length) {
				//std::cout << "pre surface crossing check..." << std::endl;
				//check_energies(surface_crossing_queue, surface_crossing_queue_length);
				//std::cout << "Surface Crossings..." << std::endl;
				start = omp_get_wtime();
				process_surface_crossing_events();
				stop = omp_get_wtime();
				time_surf += (stop-start);
			} else if (max == collision_queue_length) {
				//std::cout << "pre Colliding check..." << std::endl;
				//check_energies(collision_queue, collision_queue_length);
				//std::cout << "Colliding..." << std::endl;
				start = omp_get_wtime();
				process_collision_events();
				stop = omp_get_wtime();
				time_collision += (stop-start);
			}
		}

		remaining_work -= n_particles;
		source_offset += n_particles;

		// Should all be zero
		/*
		calculate_fuel_xs_queue_length    = 0;
		calculate_nonfuel_xs_queue_length = 0;<commit_msg>changed to more flexible timers<commit_after>#include "openmc/simulation.h"

#include "openmc/bank.h"
#include "openmc/capi.h"
#include "openmc/cell.h"
#include "openmc/container_util.h"
#include "openmc/eigenvalue.h"
#include "openmc/error.h"
#include "openmc/geometry.h"
#include "openmc/material.h"
#include "openmc/message_passing.h"
#include "openmc/mgxs_interface.h"
#include "openmc/nuclide.h"
#include "openmc/output.h"
#include "openmc/particle.h"
#include "openmc/photon.h"
#include "openmc/physics.h"
 #include "openmc/physics_mg.h"
#include "openmc/random_lcg.h"
#include "openmc/settings.h"
#include "openmc/source.h"
#include "openmc/state_point.h"
#include "openmc/thermal.h"
#include "openmc/timer.h"
#include "openmc/tallies/derivative.h"
#include "openmc/tallies/filter.h"
#include "openmc/tallies/tally.h"
#include "openmc/tallies/tally_scoring.h"
#include "openmc/tallies/trigger.h"
#include "openmc/track_output.h"

#ifdef _OPENMP
#include <omp.h>
#endif
#include "xtensor/xview.hpp"

#ifdef OPENMC_MPI
#include <mpi.h>
#endif

#include <algorithm>
#include <string>
#include <chrono> 

namespace openmc {

	/*
extern std::vector<Particle*> calculate_fuel_xs_queue;
extern std::vector<Particle*> calculate_nonfuel_xs_queue;
extern std::vector<Particle*> advance_particle_queue;
extern std::vector<Particle*> surface_crossing_queue;
extern std::vector<Particle*> collision_queue;
#pragma omp threadprivate(calculate_fuel_xs_queue, calculate_nonfuel_xs_queue, advance_particle_queue, surface_crossing_queue, collision_queue)

std::vector<Particle*> calculate_fuel_xs_queue;
std::vector<Particle*> calculate_nonfuel_xs_queue;
std::vector<Particle*> advance_particle_queue;
std::vector<Particle*> surface_crossing_queue;
std::vector<Particle*> collision_queue;
*/

int * calculate_fuel_xs_queue;
int * calculate_nonfuel_xs_queue;
int * advance_particle_queue;
int * surface_crossing_queue;
int * collision_queue;
Particle * particles;

int calculate_fuel_xs_queue_length    = 0;
int calculate_nonfuel_xs_queue_length = 0;
int advance_particle_queue_length     = 0;
int surface_crossing_queue_length     = 0;
int collision_queue_length            = 0;

const int MAX_PARTICLES_IN_FLIGHT = 100;

void init_event_queues(void)
{
	int n_particles = MAX_PARTICLES_IN_FLIGHT;
	calculate_fuel_xs_queue =    new int[n_particles];
	calculate_nonfuel_xs_queue = new int[n_particles];
	advance_particle_queue =     new int[n_particles];
	surface_crossing_queue =     new int[n_particles];
	collision_queue =            new int[n_particles];
	particles =                  new Particle[n_particles];
}

void free_event_queues(void)
{
	delete calculate_fuel_xs_queue;
	delete calculate_nonfuel_xs_queue;
	delete advance_particle_queue;
	delete surface_crossing_queue;
	delete collision_queue;
	delete[] particles;
}

constexpr size_t MAX_PARTICLES_PER_THREAD {100};

Particle::Bank * shared_fission_bank;
int shared_fission_bank_length = 0;
int shared_fission_bank_max;

void init_shared_fission_bank(int max)
{
	shared_fission_bank_max = max;
	shared_fission_bank = new Particle::Bank[max];
}

void free_shared_fission_bank(void)
{
	delete[] shared_fission_bank;
	shared_fission_bank_length = 0;
}

// TODO: What is going on here?
void revive_particle_from_secondary(Particle* p)
{
  //p->from_source(&simulation::secondary_bank.back());
  p->from_source(&p->secondary_bank_.back());
  //simulation::secondary_bank.pop_back();
  p->secondary_bank_.pop_back();
  // n_event = 0;

  // Enter new particle in particle track file
  if (p->write_track_) add_particle_track();
}

void dispatch_xs_event(int i)
{
	Particle * p = particles + i;
	  int idx;
  if (p->material_ == MATERIAL_VOID) {

	#pragma omp atomic capture
	idx = calculate_nonfuel_xs_queue_length++;
	//std::cout << "Dispatching particle to non Fuel XS queue idx = " << idx << std::endl;
    calculate_nonfuel_xs_queue[idx] = i;
  } else {
    if (model::materials[p->material_]->fissionable_) {
	#pragma omp atomic capture
	idx = calculate_fuel_xs_queue_length++;
	//std::cout << "Dispatching particle to Fuel XS queue idx = " << idx << std::endl;
      calculate_fuel_xs_queue[idx] = i;
    } else {
	#pragma omp atomic capture
	idx = calculate_nonfuel_xs_queue_length++;
	//std::cout << "Dispatching particle to non Fuel XS queue idx = " << idx << std::endl;
      calculate_nonfuel_xs_queue[idx] = i;
    }
  }
}

void process_calculate_xs_events(int * queue, int n)
{
  // Save last_ members, find grid index
  int lost_particles = 0;
  #pragma omp parallel for reduction(+:lost_particles)
  for (int i = 0; i < n; i++) {
	  Particle *p = particles + queue[i]; 
	  //std::cout << "particle offset = " << queue[i] << std::endl;
    // Set the random number stream
	// TODO: Move RNG seeds to particle storage
    if (p->type_ == Particle::Type::neutron) {
		p->stream_ = STREAM_TRACKING;
    } else {
		p->stream_ = STREAM_PHOTON;
    }

    // Store pre-collision particle properties
    p->wgt_last_ = p->wgt_;
    p->E_last_ = p->E_;
    p->u_last_ = p->u();
    p->r_last_ = p->r();

    // If the cell hasn't been determined based on the particle's location,
    // initiate a search for the current cell. This generally happens at the
    // beginning of the history and again for any secondary particles
    if (p->coord_[p->n_coord_ - 1].cell == C_NONE) {
      if (!find_cell(p, false)) {
        p->mark_as_lost("Could not find the cell containing particle "
          + std::to_string(p->id_));
        //return;
		lost_particles += 1;
		continue;
      }

      // set birth cell attribute
      if (p->cell_born_ == C_NONE) p->cell_born_ = p->coord_[p->n_coord_ - 1].cell;
    }

    // Write particle track.
    if (p->write_track_) write_particle_track(*p);

    if (settings::check_overlaps) check_cell_overlap(p);

    if (settings::run_CE) {
      if (p->material_ == p->material_last_ && p->sqrtkT_ != p->sqrtkT_last_) {
        // Remove particle from queue
      }
    }

    // Find energy index on energy grid
    // TODO: Calculate this separately?
    //int neutron = static_cast<int>(Particle::Type::neutron);
    //p->macro_xs_.i_grid = std::log(p->E_/data::energy_min[neutron]) / simulation::log_spacing;
  }
  if( lost_particles > 0 )
	  exit(1);
  #pragma omp parallel for
  for( int i = 0; i < n; i++ )
  {
	  Particle * p = particles + queue[i];
	  // Calculate microscopic and macroscopic cross sections
	  if (p->material_ != MATERIAL_VOID) {
		  if (settings::run_CE) {
			  if (p->material_ != p->material_last_ || p->sqrtkT_ != p->sqrtkT_last_) {
				  // If the material is the same as the last material and the
				  // temperature hasn't changed, we don't need to lookup cross
				  // sections again.
				  model::materials[p->material_]->calculate_xs(*p);
			  }
		  } // else MG not supported
	  } else {
		  p->macro_xs_.total      = 0.0;
		  p->macro_xs_.absorption = 0.0;
		  p->macro_xs_.fission    = 0.0;
		  p->macro_xs_.nu_fission = 0.0;
	  }
	  /*
	  int idx;
	  #pragma omp atomic capture
	  idx = advance_particle_queue_length++;
	  advance_particle_queue[idx] = queue[i];
	  */
  }
  int start = advance_particle_queue_length;
  int end = start + n;
  int j = 0;
  for( int i = start; i < end; i++ )
  {
	  advance_particle_queue[i] = queue[j];
	  j++;
  }
  advance_particle_queue_length += n;

  /*
  // Calculate nuclide micros
  for (int i = 0; i < data::nuclides.size(); ++i) {
	  // loop over particles
    for (int j = 0; j < n; j++) {
		//Particle * p = particles + queue[i];
		Particle * p = particles + queue[j];
      if (p->material_ == MATERIAL_VOID) continue;

      // If material doesn't have this nuclide, skip it
      const auto& mat {model::materials[p->material_]};
      if (mat->mat_nuclide_index_[i] == -1) continue;

      // ======================================================================
      // CHECK FOR S(A,B) TABLE

      // Check if this nuclide matches one of the S(a,b) tables specified.
      // This relies on thermal_tables_ being sorted by .index_nuclide
      int i_sab = C_NONE;
      double sab_frac = 0.0;
      for (const auto& sab : mat->thermal_tables_) {
        if (i == sab.index_nuclide) {
          // Get index in sab_tables
          i_sab = sab.index_table;
          sab_frac = sab.fraction;

          // If particle energy is greater than the highest energy for the
          // S(a,b) table, then don't use the S(a,b) table
          //if (p->E_ > data::thermal_scatt[i_sab]->threshold()) i_sab = C_NONE;
			std::cout << "S(alpha, beta) max energy = " << data::thermal_scatt[i_sab]->energy_max_ <<std::endl;
          if (p->E_ > data::thermal_scatt[i_sab]->energy_max_) i_sab = C_NONE;
        }
      }

      // ======================================================================
      // CALCULATE MICROSCOPIC CROSS SECTION

      // Calculate microscopic cross section for this nuclide
      const auto& micro {p->neutron_xs_[i]};
      if (p->E_ != micro.last_E
          || p->sqrtkT_ != micro.last_sqrtkT
          || i_sab != micro.index_sab
          || sab_frac != micro.sab_frac) {
        data::nuclides[i]->calculate_xs(i_sab, p->macro_xs_.i_grid, sab_frac, *p);
      }
    }
  }

  for (int i = 0; i < n; i++) {
	  Particle * p = particles + queue[i];
    // Calculate microscopic and macroscopic cross sections
    if (p->material_ != MATERIAL_VOID) {
	  // Only works for CE, no MG support
      //if (settings::run_CE) {
        // If the material is the same as the last material and the
        // temperature hasn't changed, we don't need to lookup cross
        // sections again.
        model::materials[p->material_]->calculate_xs(*p);
      }
		//else {
      //}
     else {
      p->macro_xs_.total      = 0.0;
      p->macro_xs_.absorption = 0.0;
      p->macro_xs_.fission    = 0.0;
      p->macro_xs_.nu_fission = 0.0;
    }

	 int idx;
	 #pragma omp atomic capture
	 idx = advance_particle_queue_length++;
    advance_particle_queue[idx] = queue[i];
  }
	*/
}

void process_advance_particle_events()
{
  //for (auto& p : advance_particle_queue) {
  #pragma omp parallel for
  for (int i = 0; i < advance_particle_queue_length; i++) {
	  Particle * p = particles + advance_particle_queue[i];
    simulation::trace == (p->id_ == 0);

    // Sample a distance to collision
    double d_collision;
    if (p->type_ == Particle::Type::electron ||
        p->type_ == Particle::Type::positron) {
      d_collision = 0.0;
    } else if (p->macro_xs_.total == 0.0) {
      d_collision = INFINITY;
    } else {
      d_collision = -std::log(prn(p->prn_seeds_, p->stream_)) / p->macro_xs_.total;
    }

    // -------------- break here? -------------------

    // Find the distance to the nearest boundary
    p->boundary_ = distance_to_boundary(p);

    // Select smaller of the two distances
    double distance;
	int idx;
    if (p->boundary_.distance < d_collision) {
		#pragma omp atomic capture
		idx = surface_crossing_queue_length++;
      surface_crossing_queue[idx] = advance_particle_queue[i];
      distance = p->boundary_.distance;
    } else {
		#pragma omp atomic capture
		idx = collision_queue_length++;
		collision_queue[idx] = advance_particle_queue[i];
      distance = d_collision;
    }

    // -------------- break here? -------------------

    // Advance particle
    for (int j = 0; j < p->n_coord_; ++j) {
      p->coord_[j].r += distance * p->coord_[j].u;
    }

    // -------------- break here? -------------------

    // Score track-length tallies
    if (!model::active_tracklength_tallies.empty()) {
      score_tracklength_tally(p, distance);
    }

    // Score track-length estimate of k-eff
    if (settings::run_mode == RUN_MODE_EIGENVALUE &&
        p->type_ == Particle::Type::neutron) {
      global_tally_tracklength += p->wgt_ * distance * p->macro_xs_.nu_fission;
    }

    // Score flux derivative accumulators for differential tallies.
    if (!model::active_tallies.empty()) {
      score_track_derivative(p, distance);
    }
  }

  advance_particle_queue_length = 0;
}

void process_surface_crossing_events()
{
  //for (auto& p : surface_crossing_queue) {
  #pragma omp parallel for
  for (int i = 0; i < surface_crossing_queue_length; i++) {
	  Particle * p = particles + surface_crossing_queue[i];
    // Set surface that particle is on and adjust coordinate levels
    p->surface_ = p->boundary_.surface_index;
    p->n_coord_ = p->boundary_.coord_level;

    // Saving previous cell data
    for (int j = 0; j < p->n_coord_; ++j) {
      p->cell_last_[j] = p->coord_[j].cell;
    }
    p->n_coord_last_ = p->n_coord_;

    if (p->boundary_.lattice_translation[0] != 0 ||
        p->boundary_.lattice_translation[1] != 0 ||
        p->boundary_.lattice_translation[2] != 0) {
      // Particle crosses lattice boundary
      cross_lattice(p, p->boundary_);
      p->event_ = EVENT_LATTICE;
    } else {
      // Particle crosses surface
      p->cross_surface();
      p->event_ = EVENT_SURFACE;
    }
    // Score cell to cell partial currents
    if (!model::active_surface_tallies.empty()) {
      score_surface_tally(p, model::active_surface_tallies);
    }

    //if (!p->alive_ && !simulation::secondary_bank.empty()) {
    if (!p->alive_ && !p->secondary_bank_.empty()) {
      revive_particle_from_secondary(p);
    }

    if (p->alive_)
	{
		dispatch_xs_event(surface_crossing_queue[i]);
	}
  }

  surface_crossing_queue_length = 0;
}

void process_collision_events()
{
  //for (auto& p : collision_queue) {
  #pragma omp parallel for
  for (int i = 0; i < collision_queue_length; i++) {
	  Particle * p = particles + collision_queue[i];
	  //std::cout << "Beginning collision of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;
    // Score collision estimate of keff
    if (settings::run_mode == RUN_MODE_EIGENVALUE &&
        p->type_ == Particle::Type::neutron) {
      global_tally_collision += p->wgt_ * p->macro_xs_.nu_fission
        / p->macro_xs_.total;
    }

    // Score surface current tallies -- this has to be done before the collision
    // since the direction of the particle will change and we need to use the
    // pre-collision direction to figure out what mesh surfaces were crossed

    if (!model::active_meshsurf_tallies.empty())
      score_surface_tally(p, model::active_meshsurf_tallies);
	    
	//std::cout << "After surface tally of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Clear surface component
    p->surface_ = 0;

    if (settings::run_CE) {
      collision(p);
    } else {
      collision_mg(p);
    }

	//std::cout << "After collision() of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Score collision estimator tallies -- this is done after a collision
    // has occurred rather than before because we need information on the
    // outgoing energy for any tallies with an outgoing energy filter
    if (!model::active_collision_tallies.empty()) score_collision_tally(p);
    if (!model::active_analog_tallies.empty()) {
      if (settings::run_CE) {
        score_analog_tally_ce(p);
      } else {
        score_analog_tally_mg(p);
      }
    }
	//std::cout << "After analog tally of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;

    // Reset banked weight during collision
    p->n_bank_ = 0;
    p->n_bank_second_ = 0;
    p->wgt_bank_ = 0.0;
    for (int& v : p->n_delayed_bank_) v = 0;

    // Reset fission logical
    p->fission_ = false;

    // Save coordinates for tallying purposes
    p->r_last_current_ = p->r();

    // Set last material to none since cross sections will need to be
    // re-evaluated
    p->material_last_ = C_NONE;

    // Set all directions to base level -- right now, after a collision, only
    // the base level directions are changed
    for (int j = 0; j < p->n_coord_ - 1; ++j) {
      if (p->coord_[j + 1].rotated) {
        // If next level is rotated, apply rotation matrix
        const auto& m {model::cells[p->coord_[j].cell]->rotation_};
        const auto& u {p->coord_[j].u};
        p->coord_[j + 1].u.x = m[3]*u.x + m[4]*u.y + m[5]*u.z;
        p->coord_[j + 1].u.y = m[6]*u.x + m[7]*u.y + m[8]*u.z;
        p->coord_[j + 1].u.z = m[9]*u.x + m[10]*u.y + m[11]*u.z;
      } else {
        // Otherwise, copy this level's direction
        p->coord_[j+1].u = p->coord_[j].u;
      }
    }

    // Score flux derivative accumulators for differential tallies.
    if (!model::active_tallies.empty()) score_collision_derivative(p);

    //if (!p->alive_ && !simulation::secondary_bank.empty()) {
    if (!p->alive_ && !p->secondary_bank_.empty()) {
      revive_particle_from_secondary(p);
    }

    if (p->alive_)
	{
		dispatch_xs_event(collision_queue[i]);
	    //std::cout << "Ended collision of particle id " << collision_queue[i] << " with energy E = " << p->E_ << std::endl;
		assert(std::isfinite(p->E_) );
	}
  }

  collision_queue_length = 0;
}

void check_energies(void)
{
	int * Q;
	int n;


	Q = calculate_fuel_xs_queue;
	n = calculate_fuel_xs_queue_length;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index xs FUEL " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = calculate_nonfuel_xs_queue;
	n = calculate_nonfuel_xs_queue_length ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index xs Non fuel " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = advance_particle_queue;
	n = advance_particle_queue_length     ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index advance particle " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = surface_crossing_queue;
	n = surface_crossing_queue_length     ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index surface crossing " << Q[i] << std::endl;
			assert(0);
		}
	}
	Q = collision_queue;
	n = collision_queue_length            ;
	for( int i = 0; i < n; i++ )
	{
		if( !std::isfinite(particles[Q[i]].E_ ) )
		{
			std::cout << "NAN energy particle found at index collision " << Q[i] << std::endl;
			assert(0);
		}
	}
}

double get_time()
{
	#ifdef _OPENMP
	return omp_get_wtime();
	#endif

	#ifdef OPENMC_MPI
	return MPI_Wtime();
	#endif

	unsigned long us_since_epoch = std::chrono::high_resolution_clock::now().time_since_epoch() / std::chrono::microseconds(1);
	return (double) us_since_epoch / 1.0e6;
}

void transport()
{
	int remaining_work = simulation::work_per_rank;
	int source_offset = 0;
		
	init_event_queues();

	double time_fuel_xs = 0;
	double time_nonfuel_xs = 0;
	double time_advance = 0;
	double time_collision = 0;
	double time_surf = 0;
	double stop, start;

	// Subiterations to complete sets of particles
	while (remaining_work > 0) {


		// Figure out work for this subiteration
		int n_particles = MAX_PARTICLES_IN_FLIGHT;
		if( n_particles > remaining_work)
			n_particles = remaining_work;

		//std::cout << "Initializing particle histories..." << std::endl;
		// Initialize all histories
		// TODO: Parallelize

		/*
		for( int p = 0; p < mpi::n_procs; p++)
		{
			MPI_Barrier(mpi::intracomm);
			if( p == mpi::rank )
			{
			*/
				for (int i = 0; i < n_particles; i++) {
					initialize_history(particles + i, source_offset + i + 1);
				}
				/*
			}
		}
		*/

		//std::cout << "Enqueing particles for XS Lookups..." << std::endl;
		// Add all particles to advance particle queue
		// TODO: Parallelize
		for (int i = 0; i < n_particles; i++) {
			dispatch_xs_event(i);
		}

		int event_kernel_executions = 0;
		while (true) {
			event_kernel_executions++;
			/*
			std::cout << "Fuel XS Lookups = " << calculate_fuel_xs_queue_length << std::endl;
			std::cout << "Non Fuel XS Lookups = " << calculate_nonfuel_xs_queue_length << std::endl;
			std::cout << "Advance Particles = " << advance_particle_queue_length << std::endl;
			std::cout << "Surface Crossings = " << surface_crossing_queue_length << std::endl;
			std::cout << "Collisions = " << collision_queue_length << std::endl;
			*/
			/*
			Particle * p = particles +  1;
		   std::cout << "E = " << p->E_ << " and Position {" <<
			   p->r().x << ", " <<
			   p->r().y << ", " <<
			   p->r().z << "}" << std::endl;
			   */
			int max = std::max({calculate_fuel_xs_queue_length, calculate_nonfuel_xs_queue_length, advance_particle_queue_length, surface_crossing_queue_length, collision_queue_length});
			//check_energies();
			if (max == 0) {
				break;
			} else if (max == calculate_fuel_xs_queue_length) {
				//std::cout << "pre fuel XS check..." << std::endl;
				//check_energies(calculate_fuel_xs_queue, calculate_fuel_xs_queue_length);
				//std::cout << "Performing Fuel XS Lookups..." << std::endl;
				start = get_time();
				process_calculate_xs_events(calculate_fuel_xs_queue, calculate_fuel_xs_queue_length);
				stop = get_time();
				time_fuel_xs += (stop-start);
				calculate_fuel_xs_queue_length = 0;
			} else if (max == calculate_nonfuel_xs_queue_length) {
				//std::cout << "pre non fuel XS check..." << std::endl;
				//check_energies(calculate_nonfuel_xs_queue, calculate_nonfuel_xs_queue_length);
			//	std::cout << "Performing Non Fuel XS Lookups..." << std::endl;
				start = get_time();
				process_calculate_xs_events(calculate_nonfuel_xs_queue, calculate_nonfuel_xs_queue_length);
				stop = get_time();
				time_nonfuel_xs += (stop-start);
				calculate_nonfuel_xs_queue_length = 0;
			} else if (max == advance_particle_queue_length) {
				//std::cout << "pre advancing check..." << std::endl;
				//check_energies(advance_particle_queue, advance_particle_queue_length);
				//std::cout << "Advancing Particles..." << std::endl;
				start = get_time();
				process_advance_particle_events();
				stop = get_time();
				time_advance += (stop-start);
			} else if (max == surface_crossing_queue_length) {
				//std::cout << "pre surface crossing check..." << std::endl;
				//check_energies(surface_crossing_queue, surface_crossing_queue_length);
				//std::cout << "Surface Crossings..." << std::endl;
				start = get_time();
				process_surface_crossing_events();
				stop = get_time();
				time_surf += (stop-start);
			} else if (max == collision_queue_length) {
				//std::cout << "pre Colliding check..." << std::endl;
				//check_energies(collision_queue, collision_queue_length);
				//std::cout << "Colliding..." << std::endl;
				start = get_time();
				process_collision_events();
				stop = get_time();
				time_collision += (stop-start);
			}
		}

		remaining_work -= n_particles;
		source_offset += n_particles;

		// Should all be zero
		/*
		calculate_fuel_xs_queue_length    = 0;
		calculate_nonfuel_xs_queue_length = 0;
{'fraction_non_alphanumeric': 0.14968888107121656, 'fraction_numerical': 0.005188980013632066, 'mean_word_length': 3.206233237769352, 'pattern_counts': {'":': 0, '<': 492, '<?xml version=': 0, '>': 322, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6300087', 'n_tokens_mistral': 17077, 'n_tokens_neox': 15654, 'n_words': 3606}