// ****************************************************************************
//
//          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 "EvolutionRunner.h"

#include "alignments/RecombAlignmentSearcher.h"
#include "io/last_gener.h"
#include "population/Selection.h"

#ifdef __OMP_LIST_SORT
  #if 1 == __OMP_LIST_SORT
    #include <random>       // std::default_random_engine
  #endif
#endif

namespace aevol {

#define __VECTORIZE_STRCMP

EvolutionRunner::EvolutionRunner(std::unique_ptr<CheckpointHandler> checkpoint_handler, CheckpointData&& chkpt_data):
    checkpoint_handler_(std::move(checkpoint_handler)) {
  // Set aevol clock
  AeTime::set_time(chkpt_data.time);

  auto nb_indivs = chkpt_data.grid.width() * chkpt_data.grid.height();

  // Init phenotypic target
  target_ = std::move(chkpt_data.phenotypic_target);

  // Init grid
  grid_ = std::move(chkpt_data.grid);

  // Init experiment setup
  aevol::exp_setup = std::move(chkpt_data.exp_setup);

  // Init checkpoint frequency
  checkpoint_frequency_ = chkpt_data.checkpoint_freq;

  // Init tree handler
  if (chkpt_data.tree_output_freq) {
    tree_output_frequency_ = chkpt_data.tree_output_freq;
    tree_handler_ = std::make_unique<PhylogeneticTreeHandler>("tree", nb_indivs, tree_output_frequency_.value());
  }

  // Init stats output frequencies
  stats_orchestrator_ = std::make_unique<StatsOrchestrator>(chkpt_data.best_indiv_stats_output_freq,
                                                            chkpt_data.whole_pop_stats_output_freq);

  // Evaluate the individuals
  for (auto& [seqid, indiv, modifiers] : chkpt_data.population) {
    indiv->evaluate(exp_setup->w_max(), exp_setup->selection_pressure(), *target_);
  }

  population_.reset_from_fasta_population(nb_indivs, chkpt_data.population);
  population_.update_best();

  // Initialize stats and write UI output
  init_stats();
  write_ui_output();
}

auto EvolutionRunner::make_from_checkpoint(time_type time) -> std::unique_ptr<EvolutionRunner> {
  auto checkpoint_handler = CheckpointHandler::make_default();
  auto checkpoint_data = checkpoint_handler->read_checkpoint(time);
  return std::unique_ptr<EvolutionRunner>(
      new EvolutionRunner(std::move(checkpoint_handler), std::move(checkpoint_data)));
}

auto EvolutionRunner::make_from_checkpoint_path(const std::filesystem::path& checkpoint_file)
    -> std::unique_ptr<EvolutionRunner> {
  auto checkpoint_handler = CheckpointHandler::make_from_ckpt_file(checkpoint_file);
  auto checkpoint_data = checkpoint_handler->read_checkpoint(checkpoint_file);
  return std::unique_ptr<EvolutionRunner>(
      new EvolutionRunner(std::move(checkpoint_handler), std::move(checkpoint_data)));
}

void EvolutionRunner::write_checkpoint() const {
  checkpoint_handler_->write_checkpoint(AeTime::time(),
                                        population_.to_indivs_idxs_sorted_map(),
                                        grid_,
                                        *target_,
                                        *exp_setup,
                                        checkpoint_frequency_,
                                        tree_output_frequency_,
                                        stats_orchestrator_->best_indiv_stats_output_frequency(),
                                        stats_orchestrator_->mean_stats_output_frequency());
  aevol::write_last_gener_file(AeTime::time());
}

auto EvolutionRunner::handle_clone(aevol::indiv_id_type indiv_id, aevol::indiv_id_type parent_id)
    -> std::shared_ptr<Individual> {
  auto cloned_indiv = population_.indiv_shared_ptr(parent_id);

  if (tree_handler_) {
    tree_handler_->report_new_indiv(cloned_indiv.get(), &population_[parent_id], indiv_id, parent_id);
  }

  return cloned_indiv;
}

#ifndef SEX
auto EvolutionRunner::create_offspring(indiv_id_type indiv_id, reproducers_info_t reproducers_info)
    -> std::unique_ptr<Individual> {
  auto offspring = Individual::make_clone(population_[reproducers_info]);

  if (tree_handler_) {
    tree_handler_->report_new_indiv(offspring.get(), &population_[reproducers_info], indiv_id, reproducers_info);
  }

  return offspring;
}
#else // SEX
auto EvolutionRunner::create_offspring(indiv_id_type indiv_id, reproducers_info_t reproducers_info, JumpingMT& prng)
    -> std::unique_ptr<Individual> {
  using recombination_info_t = struct {
    Chrsm chrsm_start;
    Dna::size_type recomb_pos_1;
    Dna::size_type recomb_pos_2;
    Chrsm chrsm_end;
  };

  std::array<recombination_info_t, 2> recombination_info;

  for (auto parent_idx = size_t{0}; parent_idx < 2; ++parent_idx) {
    // Pick which chromosome will start the recombinant chromosome (the other one will end it)
    recombination_info[parent_idx].chrsm_start = prng.random() < 0.5 ? Chrsm::A : Chrsm::B;
    recombination_info[parent_idx].chrsm_end =
        recombination_info[parent_idx].chrsm_start == Chrsm::B ? Chrsm::A : Chrsm::B;

    // Pick points for recombination, either alignment-based or at random
    auto alignment_params = exp_setup->alignment_params();
    if (alignment_params) {
      // Search for an alignement to determine the breakpoints for recombination
      const auto& alignment = align::RecombAlignmentSearcher::find_alignment(
          {&population_[reproducers_info[parent_idx]].dna(recombination_info[parent_idx].chrsm_start),
           &population_[reproducers_info[parent_idx]].dna(recombination_info[parent_idx].chrsm_end)},
          prng,
          *alignment_params);
      recombination_info[parent_idx].recomb_pos_1 = alignment.posA;
      recombination_info[parent_idx].recomb_pos_2 = alignment.posB;
    } else {
      // Pick random points for recombination
      recombination_info[parent_idx].recomb_pos_1 = prng.random(
          population_[reproducers_info[parent_idx]].dna(recombination_info[parent_idx].chrsm_start).length());
      recombination_info[parent_idx].recomb_pos_2 = prng.random(
          population_[reproducers_info[parent_idx]].dna(recombination_info[parent_idx].chrsm_end).length());
    }
  }

  // Create recombinant chromosomes
  auto parent_chrsms = std::array<std::unique_ptr<AnnotatedChromosome>, 2>{
      AnnotatedChromosome::make_recombinant(
          population_[reproducers_info[0]].annotated_chromosome(recombination_info[0].chrsm_start),
          recombination_info[0].recomb_pos_1,
          population_[reproducers_info[0]].annotated_chromosome(recombination_info[0].chrsm_end),
          recombination_info[0].recomb_pos_2),
      AnnotatedChromosome::make_recombinant(
          population_[reproducers_info[1]].annotated_chromosome(recombination_info[1].chrsm_start),
          recombination_info[1].recomb_pos_1,
          population_[reproducers_info[1]].annotated_chromosome(recombination_info[1].chrsm_end),
          recombination_info[1].recomb_pos_2)};

  // Randomly decide which inherited chromosome will become chromosome A resp. B
  if (prng.random() < 0.5) {
    parent_chrsms[0].swap(parent_chrsms[1]);
  }

  auto offspring = Individual::make_from_chromosomes(std::move(parent_chrsms));

  if (tree_handler_) {
    // TODO<dpa> store both parents
    tree_handler_->report_new_indiv(offspring.get(), &population_[reproducers_info[0]], indiv_id, reproducers_info[0]);
  }

  return offspring;
}
#endif // SEX

void EvolutionRunner::init_stats() {
  if (stats_orchestrator_) stats_orchestrator_->init_stats(AeTime::time());
  if (AeTime::time() == 0) {
    stats_orchestrator_->write_stats(AeTime::time(),
                                     population_.size(),
                                     population_.best_indiv_idx(),
                                     population_.best_indiv(),
                                     population_.individuals());
  }
}

void EvolutionRunner::run_evolution(time_type t_end, double selection_pressure) {
  bool finished = AeTime::time() >= t_end;

  // The individuals of the generation being produced
  auto next_generation_individuals = std::vector<std::shared_ptr<Individual>>(population_.size());
  // ID of the parent for each individual in the generation being produced
  auto next_generation_reproducers = std::vector<reproducers_info_t>(population_.size());
  // IDs of individuals that (will mutate/have mutated) during the generation being produced
  // along with their dna mutator
  auto mutant_list = std::vector<std::pair<indiv_id_type, std::unique_ptr<Individual::mutator_type>>>();
  // Individuals that have not mutated but that need to be reevaluated nonetheless
  auto non_mutants_to_reevaluate = std::set<std::shared_ptr<Individual>>();
  // Joined list of all the individuals to reevaluate
  // Each elt of this vector contains an individual to reevaluate and whether its phenotype needs to be recomputed
  auto indivs_to_reevaluate = std::vector<std::pair<std::shared_ptr<Individual>, bool>>();

  #pragma omp parallel default(shared) \
      shared(next_generation_individuals, next_generation_reproducers, mutant_list, non_mutants_to_reevaluate, \
             indivs_to_reevaluate)
  while (!finished) {
    // Take one step in the evolutionary loop
    run_a_step(exp_setup->w_max(),
               selection_pressure,
               next_generation_individuals,
               next_generation_reproducers,
               mutant_list,
               non_mutants_to_reevaluate,
               indivs_to_reevaluate);

    #pragma omp single
    {
      if (AeTime::time() % cli_output_frequency_ == 0) {
        std::cout << "============================== " << AeTime::time() << " ==============================\n"
                  << "  Best individual's fitness: " << best_indiv().fitness()
                  << std::endl;
      }
    }

    #pragma omp single
    {
      if (AeTime::time() >= t_end)
        finished = true;
    }

    // Write checkpoint
    #pragma omp single
    {
      if (AeTime::time() % checkpoint_frequency_ == 0 or finished) {
        std::cout << "  Writing checkpoint for generation " << AeTime::time() << "...";
        std::cout.flush();
        write_checkpoint();
        std::cout << " OK" << std::endl;
      }
    }

    // Write tree
    #pragma omp single
    {
      if (tree_output_frequency_ && (AeTime::time() % tree_output_frequency_.value() == 0)) {
        std::cout << "  Writing tree for generation " << AeTime::time() << "...";
        std::cout.flush();
        tree_handler_->write_tree(AeTime::time());
        std::cout << " OK" << std::endl;
      }
    }
  }

  std::cout << "================================================================\n"
            << "  The run is finished." << std::endl;
}

void EvolutionRunner::run_a_step(
    double w_max,
    double selection_pressure,
    std::vector<std::shared_ptr<Individual>>& next_generation_individuals,
    std::vector<reproducers_info_t>& next_generation_reproducers,
    std::vector<std::pair<indiv_id_type, std::unique_ptr<Individual::mutator_type>>>& mutant_list,
    std::set<std::shared_ptr<Individual>>& non_mutants_to_reevaluate,
    std::vector<std::pair<std::shared_ptr<Individual>, bool>>& indivs_to_reevaluate) {
  // ****************************************************************************************************
  // - Apply step-wide actions (e.g. environmental variation) if any
  // - Take a step in time
  #pragma omp single
  {
    // TODO<dpa> Apply environment variation

    // Take a step in time
    AeTime::plusplus();
  }

  // ****************************************************************************************************
  // - Reinitialize semantically step-local variables
  // - If selection scope is global, run global selection
  #pragma omp single
  {
    mutant_list.clear();
    mutant_list.reserve(population_.size());
    indivs_to_reevaluate.clear();
    indivs_to_reevaluate.reserve(population_.size());
    non_mutants_to_reevaluate.clear();

    for (auto& indiv : population_.mutable_individuals()) {
      indiv.mutable_dna().reset_stat();
    }

    if (exp_setup->selection().selection_scope() == SelectionScope::GLOBAL) {
      next_generation_reproducers = global_selection(population_, grid_[0].selection_prng());
    }
  }

  // ****************************************************************************************************
  // For each offspring
  //  - If selection scope is local, run local selection
  //  - Generate mutations (don't apply them yet)
  //  - If it has mutations to undergo, add to mutant_list
  //  - If it doesn't but still *do* need to be reevaluated, add to non_mutants_to_reevaluate
  #pragma omp for schedule(dynamic)
  for (indiv_id_type indiv_id = 0 ; indiv_id < population_.size() ; ++indiv_id) {
    if (exp_setup->selection().selection_scope() == SelectionScope::LOCAL) {
      next_generation_reproducers[indiv_id] = local_selection(population_, grid_, indiv_id);
    }
    #ifndef DIPLOID
    auto individual_mutator =
        std::make_unique<Individual::mutator_type>(population_[next_generation_reproducers[indiv_id]],
                                                   exp_setup->mut_params(),
                                                   exp_setup->genome_limits());
    individual_mutator->generate_mutations(grid_[indiv_id].mutation_prng());
    // Add offspring to the list of mutants if it has mutations to undergo
    // If it hasn't, make a clone of the parent and add it to the list of non-mutants to reevaluate if needed
    if (individual_mutator->has_mutations()) {
      next_generation_individuals[indiv_id] = create_offspring(indiv_id, next_generation_reproducers[indiv_id]);
      #pragma omp critical
      {
        mutant_list.emplace_back(indiv_id, std::move(individual_mutator));
      }
    } else {
      // Handle clone
      next_generation_individuals[indiv_id] = handle_clone(indiv_id, next_generation_reproducers[indiv_id]);
      if (whole_reevaluation_required()) {
        #pragma omp critical
        {
          non_mutants_to_reevaluate.insert(population_.indiv_shared_ptr(next_generation_reproducers[indiv_id]));
        }
      }
    }
    #else
    #ifndef SEX
    #error "Not implemented"
    #else // SEX
    // Sexual reproduction implies no clones. All the offspring are singular and will have to be reevaluated.
    // We can hence create all the offsprings regardless of whether they have mutations to undergo and
    // categorize them into either mutants or non_mutants_to_reevaluate.
    // Note that even non-mutants will require phenotypic recomputation before the actual reevaluation
    next_generation_individuals[indiv_id] =
        create_offspring(indiv_id, next_generation_reproducers[indiv_id], grid_[indiv_id].selection_prng());
    auto individual_mutator = std::make_unique<Individual::mutator_type>(*next_generation_individuals[indiv_id],
                                                                         exp_setup->mut_params(),
                                                                         exp_setup->genome_limits());
    individual_mutator->generate_mutations(grid_[indiv_id].mutation_prng());
    if (individual_mutator->has_mutations()) {
      #pragma omp critical
      {
        mutant_list.emplace_back(indiv_id, std::move(individual_mutator));
      }
    } else {
      #pragma omp critical
      {
        non_mutants_to_reevaluate.insert(next_generation_individuals[indiv_id]);
      }
    }
    #endif // SEX
    #endif
  }

  // ****************************************************************************************************
  // Add non mutants that need to be reevaluated (e.g. because of an environmental change) to indivs_to_reevaluate
  #pragma omp single
  {
    for (const auto& non_mutant_to_reevaluate : non_mutants_to_reevaluate) {
      // Add individual to indivs_to_reevaluate.
      // The individual's phenotype does not need to be recomputed when reproduction is asexual since the offspring is
      // a clone of its parent. It does however, when reproduction is sexual.
      indivs_to_reevaluate.emplace_back(non_mutant_to_reevaluate, exp_setup->sex());
    }
  }

  #ifdef __OMP_LIST_SORT
  #pragma omp single
  {
    #if 0 == __OMP_LIST_SORT
    sort(mutant_list.begin(), mutant_list.end(),
        [this](std::remove_reference_t<decltype(mutant_list)>::value_type& a,
               std::remove_reference_t<decltype(mutant_list)>::value_type& b) {
          #ifndef DIPLOID
          return a.second->final_genome_size() > b.second->final_genome_size();
          #else
          return a.second[std::to_underlying(Chrsm::A)]->final_genome_size() +
                 a.second[std::to_underlying(Chrsm::B)]->final_genome_size() >
                 b.second[std::to_underlying(Chrsm::A)]->final_genome_size() +
                 b.second[std::to_underlying(Chrsm::B)]->final_genome_size();
          #endif
        }
    );
    #elif 1 == __OMP_LIST_SORT
    unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();

    std::shuffle (mutant_list.begin(), mutant_list.end(), std::default_random_engine(seed));
    #endif
  }
  #endif

  // ****************************************************************************************************
  // For each mutant
  //  - Apply mutations
  //  - Add to indivs_to_reevaluate
  #ifdef __GNUC__
  #pragma omp for schedule(dynamic,1)
  #else
  #pragma omp for schedule(monotonic:dynamic,1)
  #endif
  for (size_t index = 0; index < mutant_list.size(); index++) {
    auto indiv_id     = mutant_list[index].first;
    auto& individual_mutator = mutant_list[index].second;

    assert(individual_mutator->has_mutations());

    individual_mutator->apply_mutations(*next_generation_individuals[indiv_id], indiv_id, tree_handler_.get());

    #pragma omp critical(indivs_to_reevaluate)
    {
      indivs_to_reevaluate.emplace_back(next_generation_individuals[indiv_id], true);
    }
  }

  #ifdef __OMP_LIST_SORT
  #pragma omp single
  {
    #if 0 == __OMP_LIST_SORT
    sort(indivs_to_reevaluate.begin(),
         indivs_to_reevaluate.end(),
         [this](std::remove_reference_t<decltype(indivs_to_reevaluate)>::value_type& a,
                std::remove_reference_t<decltype(indivs_to_reevaluate)>::value_type& b) {
           return a.first->dna()->length() > b.first->dna()->length();
         });
    #elif 1 == __OMP_LIST_SORT
    std::shuffle(indivs_to_reevaluate.begin(),
                 indivs_to_reevaluate.end(),
                 std::default_random_engine(std::chrono::system_clock::now().time_since_epoch().count()));
    #endif
  }
  #endif

  // ****************************************************************************************************
  // - Evaluate indivs in indivs_to_reevaluate (mutants + non-mutants to reevaluate)
  #ifdef __GNUC__
  #pragma omp for schedule(dynamic,1)
  #else
  #pragma omp for schedule(monotonic:dynamic,1)
  #endif
  for (auto& [indiv, recompute_phenotype] : indivs_to_reevaluate) {
    if (recompute_phenotype) {
      indiv->prom_compute_RNA();
      indiv->start_protein();
      indiv->compute_protein();
      indiv->translate_protein(w_max);
      indiv->compute_phenotype();
      indiv->normalize_phenotype(target_->fuzzy().get_geometric_area());
    }

    indiv->compute_metabolic_error(*target_);
    indiv->compute_fitness(selection_pressure);
  }

  // ****************************************************************************************************
  // - Handle tree (end replication for each indiv)
  #pragma omp for schedule(dynamic)
  for (indiv_id_type indiv_id = 0 ; indiv_id < population_.size() ; ++indiv_id) {
    if (tree_handler_) {
      tree_handler_->report_end_of_replication(next_generation_individuals[indiv_id].get(), indiv_id);
    }
  }

  // ****************************************************************************************************
  // - Replace individuals_ with the new generation (next_generation_individuals)
  population_.reset(next_generation_individuals);

  // ****************************************************************************************************
  // - Reduce dna_factory memory usage
  #pragma omp single
  {
    if (AeTime::time() % DnaFactory::cleanup_step == 0) {
      aevol::dna_factory->reduce_space(population_);
    }
  }

  // ****************************************************************************************************
  // - Update best individual
  #pragma omp single
  {
    population_.update_best();
  }

  #pragma omp single
  {
    stats_orchestrator_->write_stats(AeTime::time(),
                                     population_.size(),
                                     population_.best_indiv_idx(),
                                     population_.best_indiv(),
                                     population_.individuals());
  }

  // ****************************************************************************************************
  // - Write output for the User Interface (if any)
  #pragma omp single
  write_ui_output();
}

void EvolutionRunner::setup_ui_output(const std::filesystem::path& outdir, time_type frequency) {
  user_interface_output_ = std::make_unique<UserInterfaceOutput>(outdir, frequency);
}

void EvolutionRunner::setup_cli_output(time_type frequency) {
  cli_output_frequency_ = frequency;
}

void EvolutionRunner::write_ui_output() const {
  if (not user_interface_output_) return;
  if (user_interface_output_->nothing_to_output()) return;

  user_interface_output_->write_indiv_output(population_.best_indiv());
  user_interface_output_->write_grid_output(grid_.width(), grid_.height(), population_);
  user_interface_output_->write_environment_output(*target_);
}

}  // namespace aevol
