#include "openmc/source.h"

#if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
#define HAS_DYNAMIC_LINKING
#endif

#include <utility> // for move

#ifdef HAS_DYNAMIC_LINKING
#include <dlfcn.h> // for dlopen, dlsym, dlclose, dlerror
#endif

#include "openmc/tensor.h"
#include <fmt/core.h>

#include "openmc/bank.h"
#include "openmc/capi.h"
#include "openmc/cell.h"
#include "openmc/container_util.h"
#include "openmc/error.h"
#include "openmc/file_utils.h"
#include "openmc/geometry.h"
#include "openmc/hdf5_interface.h"
#include "openmc/material.h"
#include "openmc/mcpl_interface.h"
#include "openmc/memory.h"
#include "openmc/message_passing.h"
#include "openmc/mgxs_interface.h"
#include "openmc/nuclide.h"
#include "openmc/random_lcg.h"
#include "openmc/search.h"
#include "openmc/settings.h"
#include "openmc/simulation.h"
#include "openmc/state_point.h"
#include "openmc/string_utils.h"
#include "openmc/xml_interface.h"

namespace openmc {

std::atomic<int64_t> source_n_accept {0};
std::atomic<int64_t> source_n_reject {0};

namespace {

void validate_particle_type(ParticleType type, const std::string& context)
{
  if (type.is_transportable())
    return;

  fatal_error(
    fmt::format("Unsupported source particle type '{}' (PDG {}) in {}.",
      type.str(), type.pdg_number(), context));
}

} // namespace

//==============================================================================
// Global variables
//==============================================================================

namespace model {

vector<unique_ptr<Source>> external_sources;

vector<unique_ptr<Source>> adjoint_sources;

DiscreteIndex external_sources_probability;

} // namespace model

//==============================================================================
// Source implementation
//==============================================================================

Source::Source(pugi::xml_node node)
{
  // Check for source strength
  if (check_for_node(node, "strength")) {
    strength_ = std::stod(get_node_value(node, "strength"));
    if (strength_ < 0.0) {
      fatal_error("Source strength is negative.");
    }
  }

  // Check for additional defined constraints
  read_constraints(node);
}

unique_ptr<Source> Source::create(pugi::xml_node node)
{
  // if the source type is present, use it to determine the type
  // of object to create
  if (check_for_node(node, "type")) {
    std::string source_type = get_node_value(node, "type");
    if (source_type == "independent") {
      return make_unique<IndependentSource>(node);
    } else if (source_type == "file") {
      return make_unique<FileSource>(node);
    } else if (source_type == "compiled") {
      return make_unique<CompiledSourceWrapper>(node);
    } else if (source_type == "mesh") {
      return make_unique<MeshSource>(node);
    } else {
      fatal_error(fmt::format("Invalid source type '{}' found.", source_type));
    }
  } else {
    // support legacy source format
    if (check_for_node(node, "file")) {
      return make_unique<FileSource>(node);
    } else if (check_for_node(node, "library")) {
      return make_unique<CompiledSourceWrapper>(node);
    } else {
      return make_unique<IndependentSource>(node);
    }
  }
}

void Source::read_constraints(pugi::xml_node node)
{
  // Check for constraints node. For backwards compatibility, if no constraints
  // node is given, still try searching for domain constraints from top-level
  // node.
  pugi::xml_node constraints_node = node.child("constraints");
  if (constraints_node) {
    node = constraints_node;
  }

  // Check for domains to reject from
  if (check_for_node(node, "domain_type")) {
    std::string domain_type = get_node_value(node, "domain_type");
    if (domain_type == "cell") {
      domain_type_ = DomainType::CELL;
    } else if (domain_type == "material") {
      domain_type_ = DomainType::MATERIAL;
    } else if (domain_type == "universe") {
      domain_type_ = DomainType::UNIVERSE;
    } else {
      fatal_error(
        std::string("Unrecognized domain type for constraint: " + domain_type));
    }

    auto ids = get_node_array<int>(node, "domain_ids");
    domain_ids_.insert(ids.begin(), ids.end());
  }

  if (check_for_node(node, "time_bounds")) {
    auto ids = get_node_array<double>(node, "time_bounds");
    if (ids.size() != 2) {
      fatal_error("Time bounds must be represented by two numbers.");
    }
    time_bounds_ = std::make_pair(ids[0], ids[1]);
  }
  if (check_for_node(node, "energy_bounds")) {
    auto ids = get_node_array<double>(node, "energy_bounds");
    if (ids.size() != 2) {
      fatal_error("Energy bounds must be represented by two numbers.");
    }
    energy_bounds_ = std::make_pair(ids[0], ids[1]);
  }

  if (check_for_node(node, "fissionable")) {
    only_fissionable_ = get_node_value_bool(node, "fissionable");
  }

  // Check for how to handle rejected particles
  if (check_for_node(node, "rejection_strategy")) {
    std::string rejection_strategy = get_node_value(node, "rejection_strategy");
    if (rejection_strategy == "kill") {
      rejection_strategy_ = RejectionStrategy::KILL;
    } else if (rejection_strategy == "resample") {
      rejection_strategy_ = RejectionStrategy::RESAMPLE;
    } else {
      fatal_error(std::string(
        "Unrecognized strategy source rejection: " + rejection_strategy));
    }
  }
}

void check_rejection_fraction(int64_t n_reject, int64_t n_accept)
{
  // Don't check unless we've hit a minimum number of total sites rejected
  if (n_reject < EXTSRC_REJECT_THRESHOLD)
    return;

  // Compute fraction of accepted sites and compare against minimum
  double fraction = static_cast<double>(n_accept) / n_reject;
  if (fraction <= settings::source_rejection_fraction) {
    fatal_error(fmt::format(
      "Too few source sites satisfied the constraints (minimum source "
      "rejection fraction = {}). Please check your source definition or "
      "set a lower value of Settings.source_rejection_fraction.",
      settings::source_rejection_fraction));
  }
}

SourceSite Source::sample_with_constraints(uint64_t* seed) const
{
  bool accepted = false;
  int64_t n_local_reject = 0;
  SourceSite site {};

  while (!accepted) {
    // Sample a source site without considering constraints yet
    site = this->sample(seed);

    if (constraints_applied()) {
      accepted = true;
    } else {
      // Check whether sampled site satisfies constraints
      accepted = satisfies_spatial_constraints(site.r) &&
                 satisfies_energy_constraints(site.E) &&
                 satisfies_time_constraints(site.time);
      if (!accepted) {
        ++n_local_reject;

        // Check per-particle rejection limit
        if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
          fatal_error("Exceeded maximum number of source rejections per "
                      "sample. Please check your source definition.");
        }

        // For the "kill" strategy, accept particle but set weight to 0 so that
        // it is terminated immediately
        if (rejection_strategy_ == RejectionStrategy::KILL) {
          accepted = true;
          site.wgt = 0.0;
        }
      }
    }
  }

  // Flush local rejection count, update accept counter, and check overall
  // rejection fraction
  if (n_local_reject > 0) {
    source_n_reject += n_local_reject;
  }
  ++source_n_accept;
  check_rejection_fraction(source_n_reject, source_n_accept);

  return site;
}

bool Source::satisfies_energy_constraints(double E) const
{
  return E > energy_bounds_.first && E < energy_bounds_.second;
}

bool Source::satisfies_time_constraints(double time) const
{
  return time > time_bounds_.first && time < time_bounds_.second;
}

bool Source::satisfies_spatial_constraints(Position r) const
{
  GeometryState geom_state;
  geom_state.r() = r;
  geom_state.u() = {0.0, 0.0, 1.0};

  // Reject particle if it's not in the geometry at all
  bool found = exhaustive_find_cell(geom_state);
  if (!found)
    return false;

  // Check the geometry state against specified domains
  bool accepted = true;
  if (!domain_ids_.empty()) {
    if (domain_type_ == DomainType::MATERIAL) {
      auto mat_index = geom_state.material();
      if (mat_index == MATERIAL_VOID) {
        accepted = false;
      } else {
        accepted = contains(domain_ids_, model::materials[mat_index]->id());
      }
    } else {
      for (int i = 0; i < geom_state.n_coord(); i++) {
        auto id =
          (domain_type_ == DomainType::CELL)
            ? model::cells[geom_state.coord(i).cell()].get()->id_
            : model::universes[geom_state.coord(i).universe()].get()->id_;
        if ((accepted = contains(domain_ids_, id)))
          break;
      }
    }
  }

  // Check if spatial site is in fissionable material
  if (accepted && only_fissionable_) {
    // Determine material
    auto mat_index = geom_state.material();
    if (mat_index == MATERIAL_VOID) {
      accepted = false;
    } else {
      accepted = model::materials[mat_index]->fissionable();
    }
  }

  return accepted;
}

//==============================================================================
// IndependentSource implementation
//==============================================================================

IndependentSource::IndependentSource(
  UPtrSpace space, UPtrAngle angle, UPtrDist energy, UPtrDist time)
  : space_ {std::move(space)}, angle_ {std::move(angle)},
    energy_ {std::move(energy)}, time_ {std::move(time)}
{}

IndependentSource::IndependentSource(pugi::xml_node node) : Source(node)
{
  // Check for particle type
  if (check_for_node(node, "particle")) {
    auto temp_str = get_node_value(node, "particle", false, true);
    particle_ = ParticleType(temp_str);
    if (particle_ == ParticleType::photon() ||
        particle_ == ParticleType::electron() ||
        particle_ == ParticleType::positron()) {
      settings::photon_transport = true;
    }
  }
  validate_particle_type(particle_, "IndependentSource");

  // Check for external source file
  if (check_for_node(node, "file")) {

  } else {

    // Spatial distribution for external source
    if (check_for_node(node, "space")) {
      space_ = SpatialDistribution::create(node.child("space"));
    } else {
      // If no spatial distribution specified, make it a point source
      space_ = UPtrSpace {new SpatialPoint()};
    }

    // For backwards compatibility, check for only fissionable setting on box
    // source
    auto space_box = dynamic_cast<SpatialBox*>(space_.get());
    if (space_box) {
      if (!only_fissionable_) {
        only_fissionable_ = space_box->only_fissionable();
      }
    }

    // Determine external source angular distribution
    if (check_for_node(node, "angle")) {
      angle_ = UnitSphereDistribution::create(node.child("angle"));
    } else {
      angle_ = UPtrAngle {new Isotropic()};
    }

    // Determine external source energy distribution
    if (check_for_node(node, "energy")) {
      pugi::xml_node node_dist = node.child("energy");
      energy_ = distribution_from_xml(node_dist);

      // For decay photon sources, use the absolute photon emission rate in
      // [photons/s] as the source strength
      if (dynamic_cast<DecaySpectrum*>(energy_.get())) {
        if (strength_ != 1.0) {
          warning(fmt::format(
            "Source strength of {} is ignored because the source uses a "
            "DecaySpectrum energy distribution. The source strength will be "
            "set from the DecaySpectrum emission rate.",
            strength_));
        }
        strength_ = energy_->integral();
      }
    } else {
      // Default to a Watt spectrum with parameters 0.988 MeV and 2.249 MeV^-1
      energy_ = UPtrDist {new Watt(0.988e6, 2.249e-6)};
    }

    // Determine external source time distribution
    if (check_for_node(node, "time")) {
      pugi::xml_node node_dist = node.child("time");
      time_ = distribution_from_xml(node_dist);
    } else {
      // Default to a Constant time T=0
      double T[] {0.0};
      double p[] {1.0};
      time_ = UPtrDist {new Discrete {T, p, 1}};
    }
  }
}

SourceSite IndependentSource::sample(uint64_t* seed) const
{
  SourceSite site {};
  site.particle = particle_;
  double r_wgt = 1.0;
  double E_wgt = 1.0;

  // Repeat sampling source location until a good site has been accepted
  bool accepted = false;
  int64_t n_local_reject = 0;

  while (!accepted) {

    // Sample spatial distribution
    auto [r, r_wgt_temp] = space_->sample(seed);
    site.r = r;
    r_wgt = r_wgt_temp;

    // Check if sampled position satisfies spatial constraints
    accepted = satisfies_spatial_constraints(site.r);

    // Check for rejection
    if (!accepted) {
      ++n_local_reject;
      if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
        fatal_error("Exceeded maximum number of source rejections per "
                    "sample. Please check your source definition.");
      }
    }
  }

  // Sample angle
  auto [u, u_wgt] = angle_->sample(seed);
  site.u = u;

  site.wgt = r_wgt * u_wgt;

  // Sample energy and time for neutron and photon sources
  if (settings::solver_type != SolverType::RANDOM_RAY) {
    // Check for monoenergetic source above maximum particle energy
    auto p = particle_.transport_index();
    auto energy_ptr = dynamic_cast<Discrete*>(energy_.get());
    auto decay_spectrum = dynamic_cast<DecaySpectrum*>(energy_.get());
    if (energy_ptr) {
      auto energies =
        tensor::Tensor<double>(energy_ptr->x().data(), energy_ptr->x().size());
      if ((energies > data::energy_max[p]).any()) {
        fatal_error("Source energy above range of energies of at least "
                    "one cross section table");
      }
    }

    while (true) {
      // Sample energy spectrum. For decay photon sources, also get the parent
      // nuclide index to store in the source site for tallying purposes.
      if (decay_spectrum) {
        auto sample = decay_spectrum->sample_with_parent(seed);
        site.E = sample.energy;
        E_wgt = sample.weight;
        site.parent_nuclide = sample.parent_nuclide;
      } else {
        auto [E, E_wgt_temp] = energy_->sample(seed);
        site.E = E;
        E_wgt = E_wgt_temp;
      }

      // Resample if energy falls above maximum particle energy
      if (site.E < data::energy_max[p] &&
          (satisfies_energy_constraints(site.E)))
        break;

      ++n_local_reject;
      if (n_local_reject >= MAX_SOURCE_REJECTIONS_PER_SAMPLE) {
        fatal_error("Exceeded maximum number of source rejections per "
                    "sample. Please check your source definition.");
      }
    }

    // Sample particle creation time
    auto [time, time_wgt] = time_->sample(seed);
    site.time = time;

    site.wgt *= (E_wgt * time_wgt);
  }

  // Flush local rejection count into global counter
  if (n_local_reject > 0) {
    source_n_reject += n_local_reject;
  }

  return site;
}

//==============================================================================
// FileSource implementation
//==============================================================================

FileSource::FileSource(pugi::xml_node node) : Source(node)
{
  auto path = get_node_value(node, "file", false, true);
  load_sites_from_file(path);
}

FileSource::FileSource(const std::string& path)
{
  load_sites_from_file(path);
}

void FileSource::load_sites_from_file(const std::string& path)
{
  // If MCPL file, use the dedicated file reader
  if (ends_with(path, ".mcpl") || ends_with(path, ".mcpl.gz")) {
    sites_ = mcpl_source_sites(path);
  } else {
    // Check if source file exists
    if (!file_exists(path)) {
      fatal_error(fmt::format("Source file '{}' does not exist.", path));
    }

    write_message(6, "Reading source file from {}...", path);

    // Open the binary file
    hid_t file_id = file_open(path, 'r', true);

    // Check to make sure this is a source file
    std::string filetype;
    read_attribute(file_id, "filetype", filetype);
    if (filetype != "source" && filetype != "statepoint") {
      fatal_error("Specified starting source file not a source file type.");
    }

    // Read in the source particles
    read_source_bank(file_id, sites_, false);

    // Close file
    file_close(file_id);
  }

  // Make sure particles in source file have valid types
  for (const auto& site : this->sites_) {
    validate_particle_type(site.particle, "FileSource");
  }
}

SourceSite FileSource::sample(uint64_t* seed) const
{
  // Sample a particle randomly from list
  size_t i_site = sites_.size() * prn(seed);
  return sites_[i_site];
}

//==============================================================================
// CompiledSourceWrapper implementation
//==============================================================================

CompiledSourceWrapper::CompiledSourceWrapper(pugi::xml_node node) : Source(node)
{
  // Get shared library path and parameters
  auto path = get_node_value(node, "library", false, true);
  std::string parameters;
  if (check_for_node(node, "parameters")) {
    parameters = get_node_value(node, "parameters", false, true);
  }
  setup(path, parameters);
}

void CompiledSourceWrapper::setup(
  const std::string& path, const std::string& parameters)
{
#ifdef HAS_DYNAMIC_LINKING
  // Open the library
  shared_library_ = dlopen(path.c_str(), RTLD_LAZY);
  if (!shared_library_) {
    fatal_error("Couldn't open source library " + path);
  }

  // reset errors
  dlerror();

  // get the function to create the custom source from the library
  auto create_compiled_source = reinterpret_cast<create_compiled_source_t*>(
    dlsym(shared_library_, "openmc_create_source"));

  // check for any dlsym errors
  auto dlsym_error = dlerror();
  if (dlsym_error) {
    std::string error_msg = fmt::format(
      "Couldn't open the openmc_create_source symbol: {}", dlsym_error);
    dlclose(shared_library_);
    fatal_error(error_msg);
  }

  // create a pointer to an instance of the custom source
  compiled_source_ = create_compiled_source(parameters);

#else
  fatal_error("Custom source libraries have not yet been implemented for "
              "non-POSIX systems");
#endif
}

CompiledSourceWrapper::~CompiledSourceWrapper()
{
  // Make sure custom source is cleared before closing shared library
  if (compiled_source_.get())
    compiled_source_.reset();

#ifdef HAS_DYNAMIC_LINKING
  dlclose(shared_library_);
#else
  fatal_error("Custom source libraries have not yet been implemented for "
              "non-POSIX systems");
#endif
}

//==============================================================================
// MeshElementSpatial implementation
//==============================================================================

std::pair<Position, double> MeshElementSpatial::sample(uint64_t* seed) const
{
  return {model::meshes[mesh_index_]->sample_element(elem_index_, seed), 1.0};
}

//==============================================================================
// MeshSource implementation
//==============================================================================

MeshSource::MeshSource(pugi::xml_node node) : Source(node)
{
  int32_t mesh_id = stoi(get_node_value(node, "mesh"));
  int32_t mesh_idx = model::mesh_map.at(mesh_id);
  const auto& mesh = model::meshes[mesh_idx];

  std::vector<double> strengths;
  // read all source distributions and populate strengths vector for MeshSpatial
  // object
  for (auto source_node : node.children("source")) {
    auto src = Source::create(source_node);
    if (auto ptr = dynamic_cast<IndependentSource*>(src.get())) {
      src.release();
      sources_.emplace_back(ptr);
    } else {
      fatal_error(
        "The source assigned to each element must be an IndependentSource.");
    }
    strengths.push_back(sources_.back()->strength());
  }

  // Set spatial distributions for each mesh element
  for (int elem_index = 0; elem_index < sources_.size(); ++elem_index) {
    sources_[elem_index]->set_space(
      std::make_unique<MeshElementSpatial>(mesh_idx, elem_index));
  }

  // Make sure sources use valid particle types
  for (const auto& src : sources_) {
    validate_particle_type(src->particle_type(), "MeshSource");
  }

  // the number of source distributions should either be one or equal to the
  // number of mesh elements
  if (sources_.size() > 1 && sources_.size() != mesh->n_bins()) {
    fatal_error(fmt::format("Incorrect number of source distributions ({}) for "
                            "mesh source with {} elements.",
      sources_.size(), mesh->n_bins()));
  }

  space_ = std::make_unique<MeshSpatial>(mesh_idx, strengths);
}

SourceSite MeshSource::sample(uint64_t* seed) const
{
  // Sample a mesh element based on the relative strengths
  int32_t element = space_->sample_element_index(seed);

  // Sample the distribution for the specific mesh element; note that the
  // spatial distribution has been set for each element using MeshElementSpatial
  return source(element)->sample_with_constraints(seed);
}

//==============================================================================
// Non-member functions
//==============================================================================

void initialize_source()
{
  write_message("Initializing source particles...", 5);

// Generation source sites from specified distribution in user input
#pragma omp parallel for
  for (int64_t i = 0; i < simulation::work_per_rank; ++i) {
    // initialize random number seed
    int64_t id = simulation::total_gen * settings::n_particles +
                 simulation::work_index[mpi::rank] + i + 1;
    uint64_t seed = init_seed(id, STREAM_SOURCE);

    // sample external source distribution
    simulation::source_bank[i] = sample_external_source(&seed);
  }

  // Write out initial source
  if (settings::write_initial_source) {
    write_message("Writing out initial source...", 5);
    std::string filename = settings::path_output + "initial_source.h5";
    hid_t file_id = file_open(filename, 'w', true);
    write_source_bank(file_id, simulation::source_bank, simulation::work_index);
    file_close(file_id);
  }
}

SourceSite sample_external_source(uint64_t* seed)
{
  // Sample from among multiple source distributions
  int i = 0;
  int n_sources = model::external_sources.size();
  if (n_sources > 1) {
    if (settings::uniform_source_sampling) {
      i = prn(seed) * n_sources;
    } else {
      i = model::external_sources_probability.sample(seed);
    }
  }

  // Sample source site from i-th source distribution
  SourceSite site {model::external_sources[i]->sample_with_constraints(seed)};

  // For uniform source sampling, multiply the weight by the ratio of the actual
  // probability of sampling source i to the biased probability of sampling
  // source i, which is (strength_i / total_strength) / (1 / n)
  if (n_sources > 1 && settings::uniform_source_sampling) {
    double total_strength = model::external_sources_probability.integral();
    site.wgt *=
      model::external_sources[i]->strength() * n_sources / total_strength;
  }

  // If running in MG, convert site.E to group
  if (!settings::run_CE) {
    site.E = lower_bound_index(data::mg.rev_energy_bins_.begin(),
      data::mg.rev_energy_bins_.end(), site.E);
    site.E = data::mg.num_energy_groups_ - site.E - 1.;
  }

  return site;
}

void free_memory_source()
{
  model::external_sources.clear();
  model::adjoint_sources.clear();
  reset_source_rejection_counters();
}

void reset_source_rejection_counters()
{
  source_n_accept = 0;
  source_n_reject = 0;
}

//==============================================================================
// C API
//==============================================================================

extern "C" int openmc_sample_external_source(
  size_t n, uint64_t* seed, void* sites)
{
  if (!sites || !seed) {
    set_errmsg("Received null pointer.");
    return OPENMC_E_INVALID_ARGUMENT;
  }

  if (model::external_sources.empty()) {
    set_errmsg("No external sources have been defined.");
    return OPENMC_E_OUT_OF_BOUNDS;
  }

  auto sites_array = static_cast<SourceSite*>(sites);

  // Derive independent per-particle seeds from the base seed so that
  // each iteration has its own RNG state for thread-safe parallel sampling.
  uint64_t base_seed = *seed;

#pragma omp parallel for schedule(static)
  for (size_t i = 0; i < n; ++i) {
    uint64_t particle_seed = init_seed(base_seed + i, STREAM_SOURCE);
    sites_array[i] = sample_external_source(&particle_seed);
  }
  return 0;
}

} // namespace openmc