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

#include "AeTime.h"
#include "ExpSetup.h"
#include "macros.h"

namespace aevol {

DnaMutator::DnaMutator(Dna::size_type initial_genome_size,
                       std::shared_ptr<MutationParams> mutation_params,
                       const GenomeLimits& genome_limits)
    : initial_genome_size_(initial_genome_size),
      mutation_params_(mutation_params),
      genome_limits_(genome_limits) {
}

void DnaMutator::compute_nb_mutations_per_type(JumpingMT& prng) {
  nb_large_duplications_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->duplication_rate());
  nb_large_deletions_to_generate_    =
      prng.binomial_random(initial_genome_size_, mutation_params_->deletion_rate());
  nb_large_translocations_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->translocation_rate());
  nb_large_inversions_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->inversion_rate());
  nb_rears_to_generate_ = nb_large_duplications_to_generate_ + nb_large_deletions_to_generate_ +
                          nb_large_translocations_to_generate_ + nb_large_inversions_to_generate_;

  nb_switch_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->point_mutation_rate());
  nb_small_insertions_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->small_insertion_rate());
  nb_small_deletions_to_generate_ =
      prng.binomial_random(initial_genome_size_, mutation_params_->small_deletion_rate());
  nb_local_muts_to_generate_ =
      nb_switch_to_generate_ + nb_small_insertions_to_generate_ + nb_small_deletions_to_generate_;
}

void DnaMutator::generate_all_mutations(JumpingMT& prng) {
  while (mutations_to_generate()) {
    generate_next_mutation(prng);
  }
}

auto DnaMutator::generate_mutation_event(
    MutationEvent::Type mutation_type,
    JumpingMT& prng,
    const Dna::size_type genome_size_before,
    Dna::size_type* genome_size_after) const -> std::unique_ptr<MutationEvent> {
  switch (mutation_type) {
    case MutationEvent::Type::SWITCH:
      return generate_switch_event(prng, genome_size_before);
    case MutationEvent::Type::SMALL_INSERTION:
      return generate_small_insertion_event(prng, genome_size_before, genome_size_after);
    case MutationEvent::Type::SMALL_DELETION:
      return generate_small_deletion_event(prng, genome_size_before, genome_size_after);
    case MutationEvent::Type::DUPLICATION:
      return generate_large_duplication_event(prng, genome_size_before, genome_size_after);
    case MutationEvent::Type::DELETION:
      return generate_large_deletion_event(prng, genome_size_before, genome_size_after);
    case MutationEvent::Type::TRANSLOCATION:
      return generate_large_translocation_event(prng, genome_size_before);
    case MutationEvent::Type::INVERSION:
      return generate_large_inversion_event(prng, genome_size_before);
    default:
      exit_with_dev_msg("unexpectedly reached switch default label", __FILE__, __LINE__);
      exit(EXIT_FAILURE);
  }
}

auto DnaMutator::generate_large_duplication_event(
    JumpingMT& prng,
    const Dna::size_type genome_size_before,
    Dna::size_type* genome_size_after) const -> std::unique_ptr<MutationEvent> {
  if (genome_size_before == 1) {
    printf("*** genome of size 1 ; duplication not done *** \n");
    return nullptr;
  }

  Dna::size_type pos_1, pos_2, pos_3;
  pos_1 = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));
  do { // TODO: why are WGD forbidden ?
    pos_2 = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));
  } while (pos_2 == pos_1);

  // Enforce pos_1 < pos_2 if we have a linear chrsm
  if (exp_setup->linear_chrsm() and pos_1 > pos_2) {
    std::swap(pos_1, pos_2);
  }

  pos_3 = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));

  Dna::size_type seqlen = 0;
  if (pos_1 < pos_2) {
    seqlen = pos_2 - pos_1;
  } else {
    auto tmp1_len = genome_size_before - pos_1;
    auto tmp2_len = pos_2;
    seqlen        = tmp1_len + tmp2_len;
  }

  auto tentative_genome_size_after = genome_size_before + seqlen;
  if (tentative_genome_size_after > genome_limits_.max_genome_size()) {
    return nullptr;
  }
   if (genome_size_after) {
     *genome_size_after = tentative_genome_size_after;
   }

  return MutationEvent::make_large_duplication(pos_1, pos_2, pos_3);
}

auto DnaMutator::generate_large_deletion_event(
    JumpingMT& prng,
    const Dna::size_type genome_size_before,
    Dna::size_type* genome_size_after) const -> std::unique_ptr<MutationEvent> {
  if (genome_size_before == 1) {
    printf("*** genome of size 1 ; deletion not done *** \n");
    return nullptr;
  }

  Dna::size_type pos_1, pos_2;
  pos_1 = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));
  do {
    pos_2 = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));
  } while (pos_2 == pos_1);

  // Enforce pos_1 < pos_2 if we have a linear chrsm
  if (exp_setup->linear_chrsm() and pos_1 > pos_2) {
    std::swap(pos_1, pos_2);
  }

  auto tentative_genome_size_after = -1;

  if (pos_1 < pos_2) {
    tentative_genome_size_after = genome_size_before - (pos_2 - pos_1);
  } else {
    tentative_genome_size_after = genome_size_before - (genome_size_before - pos_1);
    tentative_genome_size_after = tentative_genome_size_after - pos_2;
  }

  if (tentative_genome_size_after < genome_limits_.min_genome_size()) {
    return nullptr;
  }

    if (genome_size_after) {
      *genome_size_after = tentative_genome_size_after;
    }

  return MutationEvent::make_large_deletion(pos_1, pos_2);
}

auto DnaMutator::generate_large_translocation_event(
    JumpingMT& prng,
    Dna::size_type genome_size) const -> std::unique_ptr<MutationEvent> {
  if (genome_size == 1) {
    printf("*** genome of size 1 ; translocation not done *** \n");
    return nullptr;
  }

  Dna::size_type pos_1, pos_2, pos_3, pos_4;
  Dna::size_type segment_length;
  bool invert;

  pos_1 = prng.random(genome_size + (exp_setup->linear_chrsm() ? 1 : 0));
  do {
    pos_2 = prng.random(genome_size + (exp_setup->linear_chrsm() ? 1 : 0));
  } while (pos_2 == pos_1);

  //if (pos_1 == pos_2) return nullptr;

  // As it is commented in do_translocation(int32_t pos_1, int32_t pos_2,
  // int32_t pos_3, int32_t pos_4, bool invert), translocating segment
  // [pos_1, pos_2] is the same as translocating segment [pos_2, pos_1]
  // Since OriC must be at position 0, we will always translocate segment
  // [pos_1, pos_2] with pos_1 < pos_2
  if (pos_1 > pos_2) {
    std::swap(pos_1, pos_2);
  }

  segment_length = pos_2 - pos_1;

  // Generate a position in [pos_1, pos_2).
  // Note that this range excludes pos_2 regardless of whether the chrsm is circular or linear.
  pos_3 = pos_1 + prng.random(segment_length);

  // Generate a position that is NOT between pos_1 and pos_2
  pos_4 = prng.random(genome_size - segment_length + (exp_setup->linear_chrsm() ? 1 : 0));
  if (pos_4 >= pos_1) {
    pos_4 += segment_length;
  }

  invert = (prng.random(2) == 0);

  return MutationEvent::make_translocation(pos_1, pos_2, pos_3, pos_4, invert);
}

auto DnaMutator::generate_large_inversion_event(
    JumpingMT& prng,
    Dna::size_type genome_size) const -> std::unique_ptr<MutationEvent> {
  if (genome_size == 1) {
    printf("*** genome of size 1 ; inversion not done *** \n");
    return nullptr;
  }

  Dna::size_type pos_1, pos_2;
  pos_1 = prng.random(genome_size + (exp_setup->linear_chrsm() ? 1 : 0));
  do {
    pos_2 = prng.random(genome_size + (exp_setup->linear_chrsm() ? 1 : 0));
  } while (pos_2 == pos_1);

  // if (pos_1 == pos_2) return nullptr; // Invert everything <=> Invert nothing!

  // Invert the segment that don't contain OriC
  // Convenient for prokaryotes, absolutely necessary for eukaryotes
  if (pos_1 > pos_2)
    std::swap(pos_1, pos_2);

  return MutationEvent::make_inversion(pos_1, pos_2);
}

auto DnaMutator::generate_switch_event(JumpingMT& prng,
                                       Dna::size_type genome_size) const -> std::unique_ptr<MutationEvent> {
  auto pos = prng.random(genome_size);

  #ifdef BASE_4
  return MutationEvent::make_switch(pos, 1 + prng.random(NB_BASE - 1));
  #else
  return MutationEvent::make_switch(pos);
  #endif
}

auto DnaMutator::generate_small_insertion_event(
    JumpingMT& prng,
    const Dna::size_type genome_size_before,
    Dna::size_type* genome_size_after) const -> std::unique_ptr<MutationEvent> {
  auto pos = prng.random(genome_size_before + (exp_setup->linear_chrsm() ? 1 : 0));
  int16_t nb_insert;

  if (mutation_params_->max_indel_size() == 1) {
    nb_insert = 1;
  } else {
    nb_insert = 1 + prng.random(mutation_params_->max_indel_size());
  }

  if (genome_size_before + nb_insert > genome_limits_.max_genome_size()) {
    return nullptr;
  }

  // Prepare the sequence to be inserted
  char* inserted_seq = new char[nb_insert + 1];
  char inserted_char;
  for (int16_t j = 0; j < nb_insert; j++) {
    inserted_char   = static_cast<char>('0' + prng.random(NB_BASE));
    inserted_seq[j] = inserted_char;
  }
  inserted_seq[nb_insert] = '\0';

    if (genome_size_after) {
      *genome_size_after = genome_size_before + nb_insert;
    }

  return MutationEvent::make_small_insertion(pos, nb_insert, inserted_seq);
}

auto DnaMutator::generate_small_deletion_event(
    JumpingMT& prng,
    const Dna::size_type genome_size_before,
    Dna::size_type* genome_size_after) const -> std::unique_ptr<MutationEvent> {
  auto pos = prng.random(genome_size_before);
  int16_t nb_pos_to_del;

  if (mutation_params_->max_indel_size() == 1) {
    nb_pos_to_del = 1;
  } else {
    nb_pos_to_del = 1 + prng.random(mutation_params_->max_indel_size());
  }

  if (genome_size_before - nb_pos_to_del < genome_limits_.min_genome_size()) {
    return nullptr;
  }

  if (exp_setup->linear_chrsm()) {
    if (pos + nb_pos_to_del > genome_size_before){
      nb_pos_to_del = genome_size_before - pos; // Doing the maximal possible deletion
    }
  }

  if (genome_size_after) {
    *genome_size_after = genome_size_before - nb_pos_to_del;
  }

  return MutationEvent::make_small_deletion(pos, nb_pos_to_del);
}

void DnaMutator::generate_next_mutation(JumpingMT& prng) {
  std::unique_ptr<MutationEvent> generated_mutation;
  auto genome_size = genome_size_after_mutations_.empty() ? initial_genome_size_ : genome_size_after_mutations_.back();

  int32_t random_value;
  auto next_generation_type = MutationEvent::Type{};

  // Determine the type of mutation to generate
  if (nb_rears_to_generate_ > 0) {
    random_value = prng.random(nb_rears_to_generate_);
    --nb_rears_to_generate_;
    if (random_value < nb_large_duplications_to_generate_) {
      --nb_large_duplications_to_generate_;  // Updating the urn (no replacement!)...
      next_generation_type = MutationEvent::Type::DUPLICATION;
    } else if (random_value < nb_large_duplications_to_generate_ + nb_large_deletions_to_generate_) {
      --nb_large_deletions_to_generate_;
      next_generation_type = MutationEvent::Type::DELETION;
    } else if (random_value < nb_large_duplications_to_generate_ + nb_large_deletions_to_generate_ +
                                  nb_large_translocations_to_generate_) {
      --nb_large_translocations_to_generate_;
      next_generation_type = MutationEvent::Type::TRANSLOCATION;
    } else {
      --nb_large_inversions_to_generate_;
      next_generation_type = MutationEvent::Type::INVERSION;
    }
  } else if (nb_local_muts_to_generate_ > 0) {
    random_value = prng.random(nb_local_muts_to_generate_);
    --nb_local_muts_to_generate_;
    if (random_value < nb_switch_to_generate_) {
      --nb_switch_to_generate_;
      next_generation_type = MutationEvent::Type::SWITCH;
    } else if (random_value < nb_switch_to_generate_ + nb_small_insertions_to_generate_) {
      --nb_small_insertions_to_generate_;
      next_generation_type = MutationEvent::Type::SMALL_INSERTION;
    } else {  // (random_value >= nb_swi + nb_ins) => del
      --nb_small_deletions_to_generate_;
      next_generation_type = MutationEvent::Type::SMALL_DELETION;
    }
  }

  // Generate the mutation event according to the mutation type we've just determined
  generated_mutation = generate_mutation_event(next_generation_type, prng, genome_size, &genome_size);

  if (generated_mutation) {
    mutation_list_.push_back(std::move(generated_mutation));
    genome_size_after_mutations_.push_back(genome_size);
  }
}

}  // namespace aevol
