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

// The input file is produced by the lineage post-treatment, please refer to it
// for e.g. the file format/content

// ============================================================================
//                                   Includes
// ============================================================================
#include <err.h>
#include <getopt.h>
#include <sys/stat.h>
#include <unistd.h>
#include <zlib.h>

#include <cinttypes>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cerrno>
#include <filesystem>
#include <format>
#include <fstream>
#include <iostream>
#include <limits>
#include <list>
#include <string>

#include "aevol.h"

#ifdef _OPENMP
#include <omp.h>
#else
typedef int omp_int_t;
inline omp_int_t omp_get_thread_num() { return 0;}
inline omp_int_t omp_get_max_threads() { return 1;}
#endif

#include "stats/SingleIndividualStatsHandler.h"
#include "stats/DefaultGlobalStatsData.h"
#include "stats/FitnessStatsData.h"
#include "stats/ChromosomeStatsData.h"
#include "stats/MutationsStatsData.h"
#include "stats/NonCodingStatsData.h"

using namespace aevol;

// Helper functions
auto interpret_cmd_line_options(int argc, char* argv[])
    -> std::tuple<std::filesystem::path, std::filesystem::path, bool, bool, bool>;
void print_help(char* prog_path);

#ifdef _OPENMP
static bool run_in_parallel = false;
#endif

int main(int argc, char* argv[]) {
  auto [lineage_file_name, param_file_name, verbose_sb, full_check_sb, trace_mutations_sb] = interpret_cmd_line_options(argc, argv);
  // Since capturing a structured binding is not yet supported in OpenMP, we must assign those values that will be
  // used in omp sections to plain old variables
  auto full_check = full_check_sb;
  auto verbose = verbose_sb;
  auto trace_mutations = trace_mutations_sb;

  // =======================
  //  Open the lineage file
  // =======================
  gzFile lineage_file = gzopen(lineage_file_name.c_str(), "r");
  if (lineage_file == Z_NULL) {
    std::cout << std::format("ERROR : Could not read the lineage file {}\n", lineage_file_name.string());
    exit(EXIT_FAILURE);
  }

  int64_t t0 = 0;
  int64_t t_end = 0;
  int32_t initial_indiv_index = 0;
  int32_t final_indiv_index = 0;

  gzread(lineage_file, &t0, sizeof(t0));
  gzread(lineage_file, &t_end, sizeof(t_end));
  gzread(lineage_file, &initial_indiv_index, sizeof(initial_indiv_index));
  gzread(lineage_file, &final_indiv_index, sizeof(final_indiv_index));

  if (verbose) {
    std::cout << std::format(
        "\n\n"
        "===============================================================================\n"
        " Statistics of the ancestors of indiv. {} from time {} to {}\n"
        "================================================================================\n",
        final_indiv_index, t0, t_end);
  }

  // ==============================
  // Initialize individual analyser
  // ==============================
  auto individual_analyser = IndividualAnalyser::make_from_param_file(param_file_name);

  auto backup_step    = individual_analyser->checkpoint_frequency();
  const auto pop_size = individual_analyser->population_size();
  const auto& target  = individual_analyser->target();

  // =========================
  //  Open the output file(s)
  // =========================
  // Create missing directories
  int status;
  status = mkdir("stats/ancestor_stats/", 0755);
  if ((status == -1) && (errno != EEXIST)) {
    err(EXIT_FAILURE, "stats/ancestor_stats/");
  }

  // Open main output file
  auto stats_handler = SingleIndividualStatsHandler<DefaultGlobalStatsData,
                                                    FitnessStatsData,
                                                    ChromosomeStatsData,
                                                    MutationsStatsData,
                                                    NonCodingStatsData>::
  make_empty("stats/ancestor_stats/stats_ancestor_best.csv");

  // Optional additional outputs

  // Open optional output files
  std::ofstream fixed_mutations_file;
  if (trace_mutations) {
    // Write readme file
    std::ofstream readme("stats/ancestor_stats/fixedmut-readme.txt");
    readme << "# #################################################################\n"
        << "#  Mutations in the lineage of the best indiv at generation " << t_end << "\n"
        << "# #################################################################\n"
        << "#  1.  Generation       (mut. occurred when producing the indiv. of this generation)\n"
        << "#  2.  Genetic unit     (which underwent the mutation, 0 = chromosome) \n"
        << "#  3.  Mutation type    (0: switch, 1: smallins, 2: smalldel, 3:dupl, 4: del, 5:trans, 6:inv, 7:insert, 8:ins_HT, 9:repl_HT) \n"
        << "#  4.  pos_0            (position for the small events, begin_segment for the rearrangements, begin_segment of the inserted segment for ins_HT, begin_segment of replaced segment for repl_HT) \n"
        << "#  5.  pos_1            (-1 for the small events, end_segment for the rearrangements, end_segment of the inserted segment for ins_HT, begin_segment of donor segment for repl_HT) \n"
        << "#  6.  pos_2            (reinsertion point for duplic., cutting point in segment for transloc., insertion point in the receiver for ins_HT, end_segment of the replaced segment for repl_HT, -1 for other events)\n"
        << "#  7.  pos_3            (reinsertion point for transloc., breakpoint in the donor for ins_HT, end_segment of the donor segment for repl_HT, -1 for other events)\n"
        << "#  8.  seq              (switched-to base (4b) or inserted sequence (small insertion))\n"
        << "#  9.  invert           (transloc, was the segment inverted (0/1)?, sense of insertion for ins_HT (0=DIRECT, 1=INDIRECT), sense of the donor segment for repl_HT (0=DIRECT, 1=INDIRECT),-1 for other events)\n"
        << "#  10. align_score      (score that was needed for the rearrangement to occur, score of the first alignment for ins_HT and repl_HT)\n"
        << "#  11. align_score2     (score for the reinsertion for transloc, score of the second alignment for ins_HT and repl_HT)\n"
        << "#  12. seg_len          (segment length for rearrangement, donor segment length for ins_HT and repl_HT)\n"
        << "#  13. repl_seg_len     (replaced segment length for repl_HT, -1 for the others)\n"
        << "#  14. GU_length        (before the event)\n"
        << "#  15. Impact of the mutation on the metabolic error (negative value = smaller gap after = beneficial mutation) \n"
        << "#  16. Impact of the mutation on fitness (positive value = higher fitness after = beneficial mutation) \n"
        << "#  17. Selection coefficient of the mutation (s > 1 = higher fitness after = beneficial mutation) \n"
        << "####################################################################################################################\n";

    auto fixed_mutations_file_name = std::filesystem::path(
        std::format("stats/ancestor_stats/fixedmut-b{:0>9}-e{:0>9}-i{}.csv", t0, t_end, final_indiv_index));
    fixed_mutations_file.open(fixed_mutations_file_name);
    if (not fixed_mutations_file) {
      exit_with_usr_msg(std::string("could not create output file ") + fixed_mutations_file_name.string());
    }

    // Write the header
    fixed_mutations_file << "gener,gen_unit,mut_type,pos_0,pos_1,pos_2,pos_3,seq,invert,align_score,align_score_2,"
                         << "seg_len,repl_seg_len,GU_len,impact_on_gap,impact_on_fitness,sel_coeff\n";
  }


  // Init Factory (Fuzzy/Dna)
  dna_factory.reset(new DnaFactory(DnaFactory_Policy::FIRSTFIT, 32, 5000));

  // Set aevol clock
  AeTime::set_time(t0);

  // ==================================================
  //  Prepare the initial ancestor and write its stats
  // ==================================================
  auto dna_size = int32_t{};
  gzread(lineage_file, &dna_size, sizeof(dna_size));
  auto dna_seq = std::string{};
  dna_seq.resize(dna_size);
  gzread(lineage_file, dna_seq.data(), dna_size);

  double w_max = exp_setup->w_max();
  double selection_pressure = exp_setup->selection_pressure();

  auto indiv = Individual::make_from_sequence(dna_seq);

  indiv->evaluate(w_max, selection_pressure, target);
  stats_handler->write_data(AeTime::time(), pop_size, initial_indiv_index, *indiv);

  if (verbose) {
    std::cout << std::format("Initial fitness     = {}\n"
                             "Initial genome size = {}\n",
                             indiv->fitness(),
                             indiv->dna().length());
  }

  // ==========================================================================
  //  Replay the mutations to get the successive ancestors and analyze them
  // ==========================================================================
  ReplicationReport* rep = nullptr;
  int32_t index;

  bool check_now = false;

  aevol::AeTime::plusplus();
  #pragma omp parallel
  {
  
  while (time() <= t_end)
  {
    #pragma omp single
    {
    rep = new ReplicationReport(lineage_file);
    index = rep->id(); // who we are building...

    // Check now?
    check_now = time() == t_end ||
        (full_check && mod(time(), backup_step) == 0);

    if (verbose) {
      std::cout << std::format("Rebuilding ancestor at generation {} (index {})...", time(), index);
    }

    indiv->evaluate(w_max, selection_pressure, target);

    // 2) Replay replication (create current individual's child)
    const auto& dnarep = rep->dna_replic_report();

    dnarep.iter_muts([&](const auto& mut) {
      int32_t unitlen_before = 0;
      double metabolic_error_before = 0.0;
      double fitness_before = 0.0;

      if (trace_mutations) {
        // Store initial values before the mutation
        metabolic_error_before = indiv->metabolic_error();
        fitness_before         = indiv->fitness();
        unitlen_before         = indiv->dna().length();
      }

      // Apply mutation
      indiv->mutable_dna().undergo_mutation(*mut);

      if (trace_mutations) {
        indiv->evaluate(w_max, selection_pressure, target);

        // Compute the metabolic impact of the mutation
        double impact_on_metabolic_error = indiv->metabolic_error() - metabolic_error_before;

        double impact_on_fitness = indiv->fitness() - fitness_before;

        double selection_coefficient = indiv->fitness() / fitness_before - 1.0;

        fixed_mutations_file << time() << ','
                             << 0 << ','
                             << *mut << ','
                             << std::format("{},{:.15e},{:.15e},{:.15e}\n",
                                            unitlen_before,
                                            impact_on_metabolic_error,
                                            impact_on_fitness,
                                            selection_coefficient);
      }
    });

    if (check_now) {
      if (verbose) {
        std::cout << "Checking the sequence of the unit..." << std::flush;
      }
      auto current_generation_checkpoint_explorer = aevol::CheckpointExplorer::make_from_checkpoint(time());
      const auto& stored_indiv = current_generation_checkpoint_explorer->individual(index);

      if (indiv->dna().length() == stored_indiv.dna().length() &&
          strncmp(indiv->dna().data(), stored_indiv.dna().data(), indiv->dna().length()) == 0) {
        if (verbose)
          std::cout << " OK\n";
      }
      else {
        if (verbose) std::cout << " ERROR !\n";
        std::cerr << std::format("Error: the rebuilt genetic unit is not the same as \n"
                                 "the one saved at generation {}... ", time());
        std::cerr << std::format("Rebuilt unit : {} bp\n {}\n", indiv->dna().length(), indiv->dna().data());
        std::cerr << std::format("Stored unit  : {} bp\n {}\n", stored_indiv.dna().length(), stored_indiv.dna().data());

        gzclose(lineage_file);
        exit(EXIT_FAILURE);
      }
    }

    // 3) All the mutations have been replayed, we can now evaluate the new individual
    auto cloned_generation = std::shared_ptr<Individual>{Individual::make_clone(*indiv)};
    int32_t current_generation = AeTime::time();

    #pragma omp task
    {
    cloned_generation->evaluate(w_max, selection_pressure, target);
    
    #pragma omp critical
    stats_handler->write_data(current_generation,
                              pop_size,
                              index,
                              *cloned_generation);
    }

    if (verbose) std::cout << " OK\n";

    delete rep;

    aevol::AeTime::plusplus();
  }
  }
  if (verbose) {
    std::cout << std::format("{} -- LOOP finished\n",omp_get_thread_num());
    std::cout << std::format("{} -- Finished\n",omp_get_thread_num());
  }
  }

  gzclose(lineage_file);

  return EXIT_SUCCESS;
}

auto interpret_cmd_line_options(int argc, char* argv[])
    -> std::tuple<std::filesystem::path, std::filesystem::path, bool, bool, bool> {
  // Command-line option variables
  auto lineage_file_name = std::filesystem::path{};
  auto param_file_name = std::filesystem::path{};
  auto verbose = false;
  auto full_check = false;
  auto trace_mutations = false;

  // =====================
  //  Parse command line
  // =====================
  const char * short_options = "p:hVFvM";
  static struct option long_options[] = {
    {"help",                no_argument, NULL, 'h'},
    {"version",             no_argument, NULL, 'V'},
    {"full-check",          no_argument, NULL, 'F'},
    {"trace-mutations",     no_argument, NULL, 'M'},
    {"verbose",             no_argument, NULL, 'v'},
    {"parallel",      required_argument, nullptr, 'p'},
    {0, 0, 0, 0}
  };

  int option;
  while((option = getopt_long(argc, argv, short_options,
                              long_options, nullptr)) != -1) {
    switch(option) {
      case 'h':
        print_help(argv[0]);
        exit(EXIT_SUCCESS);
      case 'V':
        print_aevol_version();
        exit(EXIT_SUCCESS);
      case 'v':
        verbose = true;
        break;
      case 'F':
        full_check = true;
        break;
      case 'M':
        trace_mutations = true;
        break;
      case 'p' : {
        #ifdef _OPENMP
        run_in_parallel = true;
          if (atoi(optarg) > 0) {
            omp_set_num_threads(atoi(optarg));
          } else
            omp_set_num_threads(1);
        #endif
        break;
      }
      default:
        // An error message is printed in getopt_long, we just need to exit
        exit(EXIT_FAILURE);
    }
  }
  
  #ifdef _OPENMP
  if (not run_in_parallel)
    omp_set_num_threads(1);
  #endif

  // There should be only exactly 2 remaining args: the lineage file and the param file
  if (optind != argc - 2) {
    exit_with_usr_msg("please specify both a lineage file and a parameter file");
  }

  lineage_file_name = argv[optind];
  param_file_name   = argv[optind + 1];

  return {lineage_file_name, param_file_name, verbose, full_check, trace_mutations};
}

void print_help(char* prog_path) {
  // Get the program file-name in prog_name (strip prog_path of the path)
  char* prog_name; // No new, it will point to somewhere inside prog_path
  if ((prog_name = strrchr(prog_path, '/'))) {
    prog_name++;
  }
  else {
    prog_name = prog_path;
  }

  std::cout << "******************************************************************************\n"
            << "*                                                                            *\n"
            << "*                        aevol - Artificial Evolution                        *\n"
            << "*                                                                            *\n"
            << "* Aevol is a simulation platform that allows one to let populations of       *\n"
            << "* digital organisms evolve in different conditions and study experimentally  *\n"
            << "* the mechanisms responsible for the structuration of the genome and the     *\n"
            << "* transcriptome.                                                             *\n"
            << "*                                                                            *\n"
            << "******************************************************************************\n"
            << "\n"
            << std::format("{}:\ncompute statistics on ancestry described in provided lineage file.\n", prog_name)
            << "\n"
            << std::format("Usage : {} -h or --help\n", prog_name)
            << std::format("   or : {} -V or --version\n", prog_name)
            << std::format("   or : {} [-FMv] [-p NB_THREADS] LINEAGE_FILE PARAM_FILE\n", prog_name)
            << "\nOptions\n"
            << "  -h, --help\n\tprint this help, then exit\n"
            << "  -V, --version\n\tprint version number, then exit\n"
            << "  -F, --full-check\n\tperform genome checks whenever possible\n"
            << "  -M, --trace-mutations\n\toutputs the fixed mutations (in a separate file)\n"
            << "  -v, --verbose\n\tbe verbose\n"
            << "  -p, --parallel NB_THREADS\n\trun on NB_THREADS threads (use -1 for system default)\n";
}
