// ****************************************************************************
//
//          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"

#include <algorithm>

namespace aevol {

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

  // Vector of reproducer IDs to be returned
  auto next_generation_reproducers = std::vector<reproducers_info_t>(population.size());

  // Collect fitness values and compute sum thereof
  auto [fitnesses, sum_fitnesses] = collect_fitness_values_global(population);

  // Deal with special case: SelectionScheme::NO_SELECTION
  if (exp_setup->selection().selection_scheme() == SelectionScheme::NO_SELECTION) {
    for (auto& reproducer : next_generation_reproducers) {
      reproducer[0] = prng.random(population.size());
      if (not exp_setup->selection().selfing_control()) {
        reproducer[1] = prng.random(population.size());
      } else {
        if (prng.random() < exp_setup->selection().selfing_rate()) {
          // Force selfing
          reproducer[1] = reproducer[0];
        } else {
          // Forbid selfing (draw second parent until it's != first parent)
          do reproducer[1] = prng.random(population.size());
          while (reproducer[1] == reproducer[0]);
        }
      }
    }
    return next_generation_reproducers;
  }

  // Deal with special case: SelectionScheme::FITTEST
  if (exp_setup->selection().selection_scheme() == SelectionScheme::FITTEST) {
    auto index_of_best = indiv_id_type(std::ranges::distance(fitnesses.begin(), std::ranges::max_element(fitnesses)));
    return std::vector<reproducers_info_t>(population.size(), {index_of_best, index_of_best});
  }

  // Compute corresponding reproduction probabilities
  auto probs = compute_reproduction_probabilities(population.size(), fitnesses, 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());
  }

  auto index    = indiv_id_type{0};
  auto index_p2 = indiv_id_type{0};

  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;
}


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);

  auto selection_patch_size = exp_setup->selection().selection_patch_size();
  auto neighborhood_size = selection_patch_size * selection_patch_size;
  auto& prng = grid[indiv_id].selection_prng();

  // Deal with special case: SelectionScheme::NO_SELECTION
  if (exp_setup->selection().selection_scheme() == SelectionScheme::NO_SELECTION) {
    auto parent_1_local_idx = indiv_id_type(prng.random(neighborhood_size));
    auto parent_2_local_idx = decltype(parent_1_local_idx){};
    if (not exp_setup->selection().selfing_control()) {
      parent_2_local_idx = prng.random(population.size());
    } else {
      if (prng.random() < exp_setup->selection().selfing_rate()) {
        // Force selfing
        parent_2_local_idx = parent_1_local_idx;
      } else {
        // Forbid selfing (draw second parent until it's != first parent)
        do parent_2_local_idx = prng.random(neighborhood_size);
        while (parent_2_local_idx == parent_1_local_idx);
      }
    }

    return {to_global_index(indiv_id, grid, parent_1_local_idx, selection_patch_size),
            to_global_index(indiv_id, grid, parent_2_local_idx, selection_patch_size)};
  }

  // Collect fitness values and compute sum thereof
  auto [local_fit_array, sum_local_fit] = collect_fitness_values_local(population, grid, indiv_id);

  // Deal with special case: SelectionScheme::FITTEST
  if (exp_setup->selection().selection_scheme() == SelectionScheme::FITTEST) {
    auto parent_local_idx = std::ranges::distance(local_fit_array.begin(), std::ranges::max_element(local_fit_array));
    auto parent_global_idx = to_global_index(indiv_id, grid, parent_local_idx, selection_patch_size);
    return {parent_global_idx, parent_global_idx};
  }

  // Compute corresponding reproduction probabilities
  auto probs = compute_reproduction_probabilities(neighborhood_size, local_fit_array, sum_local_fit);

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

  auto parent_global_idx = to_global_index(indiv_id, grid, parent_local_idx, selection_patch_size);

  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 that the same parent id is drawn for parent2 as it was for parent 1
    parent2_local_idx = prng.roulette_random(probs.data(), neighborhood_size);
  }

  auto parent2_global_idx = to_global_index(indiv_id, grid, parent2_local_idx, selection_patch_size);

  return {parent_global_idx, parent2_global_idx};
}

}  // namespace aevol
