// ****************************************************************************
//
//          Aevol - An in silico experimental evolution platform
//
// ****************************************************************************
//
// Copyright: See the AUTHORS file provided with the package or <www.aevol.fr>
// Web: http://www.aevol.fr/
// E-mail: See <http://www.aevol.fr/contact/>
// Original Authors : Guillaume Beslon, Carole Knibbe, David Parsons
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//
// ****************************************************************************

#include "Selection.h"

namespace aevol {

// TODO<dpa> Refactor
#ifndef SEX
auto global_selection(const Population& population, JumpingMT& prng) -> std::vector<reproducers_info_t> {
  assert(exp_setup->selection().selection_scope() == SelectionScope::GLOBAL);

  // Collect fitness values and compute sum thereof
  auto fitnesses     = std::vector<double>(population.size());
  auto sum_fitnesses = 0.0;
  for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
    fitnesses[indiv_id] = population[indiv_id].fitness();
    sum_fitnesses += fitnesses[indiv_id];
  }

  // Compute corresponding reproduction probabilities
  auto probs = std::vector<double>(population.size());
  if (sum_fitnesses == 0.) {
    std::ranges::fill(probs, 1. / probs.size());
  } else {
    for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
      probs[indiv_id] = fitnesses[indiv_id] / sum_fitnesses;
    }
  }

  // Compute actual number of offspring for each individual
  auto nb_offsprings = std::vector<int32_t>(population.size());
  prng.multinomial_drawing(nb_offsprings.data(), probs.data(), population.size(), population.size());

  // Generate vector of reproducer IDs
  Grid::xy_type index = 0;
  auto next_generation_reproducers = std::vector<reproducers_info_t>(population.size());
  for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
    for (int32_t j = 0; j < nb_offsprings[indiv_id]; ++j) {
      next_generation_reproducers[index] = indiv_id;
      ++index;
    }
  }

  return next_generation_reproducers;
}
#else
auto global_selection(const Population& population, JumpingMT& prng) -> std::vector<reproducers_info_t> {
  assert(exp_setup->selection().selection_scope() == SelectionScope::GLOBAL);

  // Collect fitness values and compute sum thereof
  auto fitnesses     = std::vector<double>(population.size());
  auto sum_fitnesses = 0.0;

  // Compute corresponding reproduction probabilities
  for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
    fitnesses[indiv_id] = population[indiv_id].fitness();
    sum_fitnesses += fitnesses[indiv_id];
  }

  // Compute corresponding reproduction probabilities
  auto probs = std::vector<double>(population.size());
  if (sum_fitnesses == 0.) {
    std::ranges::fill(probs, 1. / probs.size());
  } else {
    for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
      probs[indiv_id] = fitnesses[indiv_id] / sum_fitnesses;
    }
  }

  // Generate vector of reproducer info depending on selfing options
  auto nb_offsprings = std::vector<int32_t>(population.size());
  prng.multinomial_drawing(nb_offsprings.data(), probs.data(), population.size(), population.size());
  auto nb_offsprings_p2 = std::vector<int32_t>{};
  if (not exp_setup->selection().selfing_control()) {
    nb_offsprings_p2.resize(population.size());
    prng.multinomial_drawing(nb_offsprings_p2.data(), probs.data(), population.size(), population.size());
  }

  Grid::xy_type index              = 0;
  Grid::xy_type index_p2           = 0;
  auto next_generation_reproducers = std::vector<reproducers_info_t>(population.size());

  for (indiv_id_type indiv_id = 0; indiv_id < population.size(); ++indiv_id) {
    for (int32_t j = 0; j < nb_offsprings[indiv_id]; ++j) {
      next_generation_reproducers[index][0] = indiv_id;
      ++index;
    }

    if (not exp_setup->selection().selfing_control()) {
      for (int32_t j = 0; j < nb_offsprings_p2[indiv_id]; ++j) {
        next_generation_reproducers[index_p2][1] = indiv_id;
        ++index_p2;
      }
    }
  }

  if (exp_setup->selection().selfing_control()) {
    for (indiv_id_type offspring_index = 0; offspring_index < population.size(); ++offspring_index) {
      auto parent1_idx = next_generation_reproducers[offspring_index][0];
      if (prng.random() < exp_setup->selection().selfing_rate()) {
        // Force selfing
        next_generation_reproducers[offspring_index][1] = parent1_idx;
      } else {
        // Forbid selfing (draw second parent until it's != first parent)
        auto parent2_idx = decltype(parent1_idx){};
        do {
          parent2_idx = prng.roulette_random(probs.data(), population.size());
        } while (parent2_idx == parent1_idx);

        next_generation_reproducers[offspring_index][1] = parent2_idx;
      }
    }
  }

  return next_generation_reproducers;
}
#endif  // SEX

auto local_selection(const Population& population, const Grid& grid, indiv_id_type indiv_id) -> reproducers_info_t {
  assert(exp_setup->selection().selection_scope() == SelectionScope::LOCAL);

  int32_t selection_scope_x = exp_setup->selection().selection_scope_x();
  int32_t selection_scope_y = exp_setup->selection().selection_scope_y();
  int16_t neighborhood_size = selection_scope_x * selection_scope_y;

  FitnessFunction fitness_function = exp_setup->selection().fitness_func();

  auto local_fit_array     = std::vector<double>(neighborhood_size);
  auto probs               = std::vector<double>(neighborhood_size);
  int16_t count            = 0;
  double sum_local_fit     = 0.0;

  auto x = grid[indiv_id].x();
  auto y = grid[indiv_id].y();
  int cur_x, cur_y;

  for (int8_t i = -(selection_scope_x / 2); i <= (selection_scope_x / 2); i++) {
    for (int8_t j = -(selection_scope_y / 2); j <= (selection_scope_y / 2); j++) {
      cur_x = (x + i + grid.width()) % grid.width();
      cur_y = (y + j + grid.height()) % grid.height();

      if (fitness_function == FitnessFunction::EXP) {
        local_fit_array[count] = population[cur_x * grid.height() + cur_y].fitness();
      }

      sum_local_fit += local_fit_array[count];
      count++;
    }
  }

  if (sum_local_fit == 0.) {
    std::ranges::fill(probs, 1. / probs.size());
  } else {
    for (int16_t i = 0; i < neighborhood_size; i++) {
      probs[i] = local_fit_array[i] / sum_local_fit;
    }
  }

  auto& prng = grid[indiv_id].selection_prng();

  auto parent_local_idx = indiv_id_type(prng.roulette_random(probs.data(), neighborhood_size));

  auto x_offset = indiv_id_difference_type((parent_local_idx / selection_scope_x) - 1);
  auto y_offset = indiv_id_difference_type((parent_local_idx % selection_scope_x) - 1);

  auto parent_global_idx = indiv_id_type(((x + x_offset + grid.width()) % grid.width()) * grid.height() +
                                         ((y + y_offset + grid.height()) % grid.height()));

  #ifdef SEX
  auto parent2_local_idx = decltype(parent_local_idx){};
  if (exp_setup->selection().selfing_control()) {
    if (prng.random() < exp_setup->selection().selfing_rate()) {
      // Force selfing
      parent2_local_idx = parent_local_idx;
    } else {
      // Forbid selfing (draw second parent until it's != first parent)
      do {
        parent2_local_idx = prng.roulette_random(probs.data(), neighborhood_size);
      } while (parent2_local_idx == parent_local_idx);
    }
  } else {
    // Selfing is not controlled, it may happen if the same parent id is drawn for parent2 as it was for parent 1
    parent2_local_idx = prng.roulette_random(probs.data(), neighborhood_size);
  }

  x_offset = (parent2_local_idx / selection_scope_x) - 1;
  y_offset = (parent2_local_idx % selection_scope_x) - 1;

  auto parent2_global_idx = indiv_id_type(((x + x_offset + grid.width()) % grid.width()) * grid.height() +
                                          ((y + y_offset + grid.height()) % grid.height()));
  #endif

  #ifndef SEX
  return parent_global_idx;
  #else
  return {parent_global_idx, parent2_global_idx};
  #endif
}

}  // namespace aevol
