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

// ============================================================================
//                                   Includes
// ============================================================================
#include <cerrno>
#include <cinttypes>
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <format>
#include <fstream>

#include <zlib.h>
#include <sys/stat.h>
#include <getopt.h>

#include <list>
#include <queue>

#include "aevol.h"

enum check_type
{
    FULL_CHECK  = 0,
    LIGHT_CHECK = 1,
    NO_CHECK    = 2
};

using namespace aevol;

// Helper functions
void print_help(char* prog_path);
void interpret_cmd_line_options(int argc, char* argv[]);

// Command-line option variables
static bool full_check = false;
static bool best = true;
static bool verbose = false;
static int64_t t0 = 0;
static int64_t t_end = -1;
static int32_t final_indiv_index = -1;

int main(int argc, char** argv) {
  // The output file (lineage.ae) contains the following information:
  // You may check that this information is up-to-date by searching
  // "lineage_file" in this source file
  //
  // - t0
  // - t_end
  // - final individual index
  // - initial_ancestor information (including its genome)
  // - replication report of ancestor at generation t0+1
  // - replication report of ancestor at generation t0+2
  // - replication report of ancestor at generation t0+3
  // - ...
  // - replication report of ancestor at generation t_end

  interpret_cmd_line_options(argc, argv);

  printf("\n  WARNING : Parameter change in the middle of a simulation is not managed.\n");

  // Load end generation
  auto final_generation_checkpoint_explorer = aevol::CheckpointExplorer::make_from_checkpoint(t_end);

  // Check that the tree was recorded
  if (not final_generation_checkpoint_explorer->tree_output_frequency()) {
    exit_with_usr_msg("The phylogenetic tree wasn't recorded during evolution, could not reconstruct the lineage");
  }

  auto tree_step = final_generation_checkpoint_explorer->tree_output_frequency().value();

  // The tree
  Tree* tree = nullptr;

  // Indices and replication reports of the individuals in the lineage
  auto indices = new int32_t[t_end - t0 + 1];
  auto reports = new ReplicationReport*[t_end - t0];
  // NB: we do not need the report of the ancestor at time t0 since we have
  // retrieved the individual itself from the initial backup
  // (plus it might be the generation 0, for which we have no reports)
  // reports[0] = how ancestor at t0 + 1 was created
  // reports[i] = how ancestor at t0 + i + 1 was created
  // reports[t_end - t0 - 1] = how the final individual was created
  //
  //           ---------------------------------------------------------------
  //  reports |  t0 => t1   |  t1 => t2   |...| t_n-1 => t_n   | XXXXXXXXXXXX |
  //           ---------------------------------------------------------------
  //  indices | index at t0 | index at t1 |...| index at t_n-1 | index at t_n |
  //           ---------------------------------------------------------------



  // =========================
  //  Load the last tree file
  // =========================

  if (verbose) {
    printf("\n\n");
    printf("====================================\n");
    printf(" Loading the last tree file ... ");
    fflush(stdout);
  }


  // Example for ae_common::rec_params->tree_step() == 100 :
  //
  // tree_000100.ae ==>  timesteps 1 to 100.
  // tree_000200.ae ==>  timesteps 101 to 200.
  // tree_000300.ae ==>  timesteps 201 to 300.
  // etc.
  //

  auto tree_file_name = std::format("tree/tree_{:0>9}.ae", t_end);

  printf("Loading %" PRId64 "\n",t_end);

  tree = new Tree(final_generation_checkpoint_explorer->population_size(), tree_step, tree_file_name.c_str());

  if (verbose) {
    printf("OK\n");
    printf("====================================\n");
  }


  // ============================================================================
  //  Find the index of the final individual and retrieve its replication report
  // ============================================================================
  if (best) {
    final_indiv_index = final_generation_checkpoint_explorer->best_indiv_idx();
  } else {
    assert(final_indiv_index != -1);
  }

  reports[t_end - t0 - 1] = new ReplicationReport(*(tree->report_by_index(t_end, final_indiv_index)));
  indices[t_end - t0]  = final_indiv_index;

  if (verbose) {
    printf("The final individual has index %" PRId32 "\n", final_indiv_index);
  }


  // =======================
  //  Open the output file
  // =======================
  auto output_file_name = std::string("lineage");
  if (best) output_file_name += "-best";
  output_file_name += std::format("-b{:0>9}-e{:0>9}-i{}.ae", t0, t_end, final_indiv_index);

  gzFile lineage_file = gzopen(output_file_name.c_str(), "w");
  if (lineage_file == nullptr) {
    fprintf(stderr, "File %s could not be created.\n", output_file_name.c_str());
    fprintf(stderr, "Please check your permissions in this directory.\n");
    exit(EXIT_FAILURE);
  }




  // ===================================================
  //  Retrieve the replication reports of the ancestors
  // ===================================================

  if (verbose) {
    printf("\n\n\n");
    printf("======================================================================\n");
    printf(" Parsing tree files to retrieve the ancestors' replication reports... \n");
    printf("======================================================================\n");
  }


  // Retrieve the index of the first ancestor from the last replication report
  indices[t_end - t0 -1] = reports[t_end - t0 - 1]->parent_id();

  // For each generation (going backwards), retrieve the index of the parent and
  // the corresponding replication report
  for (int64_t i = t_end - t0 - 2 ; i >= 0 ; i--) {
    int64_t t = t0 + i + 1;

    // We want to fill reports[i], that is to say, how the ancestor
    // at generation begin_gener + i + 1  was created
    if (verbose)
      printf("Getting the replication report for the ancestor at generation %" PRId64 " : %d\n", t,indices[i+1]);


    // If we've exhausted the current tree file, load the next one
    if (mod(t, tree_step) == 0) {
      // Change the tree file
      delete tree;

      tree_file_name = std::format("tree/tree_{:0>9}.ae", t);

      tree = new Tree(final_generation_checkpoint_explorer->population_size(), tree_step, tree_file_name.c_str());
    }
    
    // Copy the replication report of the ancestor
    reports[i] = new ReplicationReport(*(tree->report_by_index(t, indices[i + 1])));
    // printf("Report for %d is %d => %d -- I %d\n",reports[i]->parent_id(),reports[i]->id(),indices[i+1]);

    // Retreive the index of the next ancestor from the report
    indices[i] = reports[i]->parent_id();
  }

  delete tree; // delete the last Tree


  if (verbose) printf("OK\n");


  // =============================================================================
  //  Get the initial genome from the backup file and write it in the output file
  // =============================================================================

  if (verbose) {
    printf("\n\n\n");
    printf("=============================================== \n");
    printf(" Getting the initial genome sequence... ");
    fflush(nullptr);
  }

  // Load the initial generation
  auto initial_generation_checkpoint_explorer = aevol::CheckpointExplorer::make_from_checkpoint(t0);

  // Copy the initial ancestor
  const auto ancestral_indiv = Individual::make_clone(initial_generation_checkpoint_explorer->individual(indices[0]));

  // Write file "header"
  // printf("Write from %d to %d\n",t0,t_end);
  auto initial_indiv_index = indices[0];
  gzwrite(lineage_file, &t0, sizeof(t0));
  gzwrite(lineage_file, &t_end, sizeof(t_end));
  gzwrite(lineage_file, &initial_indiv_index, sizeof(initial_indiv_index));
  gzwrite(lineage_file, &final_indiv_index, sizeof(final_indiv_index));

  auto dna_size = ancestral_indiv->dna().length();
  gzwrite(lineage_file, &dna_size, sizeof(dna_size));
  auto dna_seq = ancestral_indiv->dna().data();
  gzwrite(lineage_file, dna_seq, dna_size);

  if (verbose) {
    printf("OK\n");
    printf("=============================================== \n");
  }


  // ===============================================================================
  //  Write the replication reports of the successive ancestors in the output file
  //  (and, optionally, check that the rebuilt genome is correct each time a backup
  //  is available)
  // ===============================================================================

  if (verbose) {
    printf("\n\n\n");
    printf("============================================================ \n");
    printf(" Write the replication reports in the output file... \n");
    printf("============================================================ \n");
  }

  // NB: I must keep the genome encapsulated inside an Individual, because
  // replaying the mutations has side effects on the list of promoters,
  // which is stored in the individual
  bool check_genome_now = false;

  for (int64_t i = 0 ; i < t_end - t0 ; i++) {
    // Where are we in time...
    int64_t t = t0 + i + 1;

    // Do we need to check the genome now?
    check_genome_now =
        t == t_end || (full_check && mod(t, initial_generation_checkpoint_explorer->checkpoint_frequency()) == 0);

    // Write the replication report of the ancestor for current generation
    if (verbose) {
      printf("Writing the replication report for t= %" PRId64
             " (built from indiv %" PRId32 " %" PRId32" at t= %" PRId64 ") %llu => %llu\n",
             t, indices[i], indices[i+1], t-1, reports[i]->parent_id(), reports[i]->id());
    }
    reports[i]->write_to_tree_file(lineage_file);
    if (verbose) printf(" OK\n");

    // Warning: this portion of code won't work if the number of units changes
    // during the evolution

    // Replay the mutations stored in the current replication report on the current genome
    reports[i]->dna_replic_report().iter_muts([&](const auto& mut) {
      ancestral_indiv->mutable_dna().undergo_mutation(*mut);
    });

    if (check_genome_now) {
      if (verbose) {
        printf("%" PRId64 " -- Checking the sequence of the unit...",t);
        fflush(stdout);
      }

      // Load current generation
      auto cur_generation_checkpoint_explorer = aevol::CheckpointExplorer::make_from_checkpoint(t);

      // Copy the ancestor from the backup
      const auto& stored_indiv = cur_generation_checkpoint_explorer->individual(indices[i+1]);

      if (ancestral_indiv->dna().length() == stored_indiv.dna().length() &&
          strncmp(ancestral_indiv->dna().data(), stored_indiv.dna().data(), ancestral_indiv->dna().length()) == 0) {
        if (verbose) printf(" OK\n");
      }
      else {
        if (verbose) printf(" ERROR !\n");
        fprintf(stderr, "Error: the rebuilt unit is not the same as \n");
        fprintf(stderr, "the one stored in backup file at %" PRId64 "\n", t);
        fprintf(stderr,
                "Rebuilt unit : %" PRId32 " bp\n %s\n",
                ancestral_indiv->dna().length(),
                ancestral_indiv->dna().data());
        fprintf(stderr, "Stored unit  : %" PRId32 " bp\n %s\n", stored_indiv.dna().length(), stored_indiv.dna().data());

        gzclose(lineage_file);

        for (int64_t i_rep = 0 ; i_rep < t_end - t0 ; i_rep++) {
          delete reports[i_rep];
        }
        delete [] reports;
        delete [] indices;
        fflush(stdout);
        exit(EXIT_FAILURE);
      }
    }
  }


  gzclose(lineage_file);
  std::cout << "output written in " << output_file_name << std::endl;

  // Output the index of each individual in the lineage at each time step
  auto log_file_name = std::string("lineage");
  if (best) log_file_name += "-best";
  log_file_name += std::format("-b{:0>9}-e{:0>9}-i{}.log", t0, t_end, final_indiv_index);
  std::ofstream log_file(log_file_name);
  for (auto t = t0 ; t <= t_end ; ++t) {
    log_file << t << " " << indices[t - t0] << "\n";
  }
  log_file.close();
  std::cout << "log written in " << log_file_name << std::endl;

  for (int64_t i_rep = 0 ; i_rep < t_end - t0 ; i_rep++) {
    delete reports[i_rep];
  }
  delete [] reports;
  delete [] indices;

  return EXIT_SUCCESS;
}

/**
 * \brief print help and exist
 */
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;
  }

  printf("******************************************************************************\n");
  printf("*                                                                            *\n");
  printf("*                        aevol - Artificial Evolution                        *\n");
  printf("*                                                                            *\n");
  printf("* Aevol is a simulation platform that allows one to let populations of       *\n");
  printf("* digital organisms evolve in different conditions and study experimentally  *\n");
  printf("* the mechanisms responsible for the structuration of the genome and the     *\n");
  printf("* transcriptome.                                                             *\n");
  printf("*                                                                            *\n");
  printf("******************************************************************************\n");
  printf("\n");
  printf("%s:\n", prog_name);
  printf("\tReconstruct the lineage of a given individual from the tree files\n");
  printf("\n");
  printf("Usage : %s -h or --help\n", prog_name);
  printf("   or : %s -V or --version\n", prog_name);
  printf("   or : %s [-b TIMESTEP] [-e TIMESTEP] [-I INDEX] [-F] [-v]\n",
         prog_name);
  printf("\nOptions\n");
  printf("  -h, --help\n\tprint this help, then exit\n");
  printf("  -V, --version\n\tprint version number, then exit\n");
  printf("  -b, --begin TIMESTEP\n");
  printf("\tspecify time t0 up to which to reconstruct the lineage\n");
  printf("  -e, --end TIMESTEP\n");
  printf("\tspecify time t_end of the indiv whose lineage is to be reconstructed\n");
  printf("  -I, --index INDEX\n");
  printf("\tspecify the index of the indiv whose lineage is to be reconstructed\n");
  printf("\t(default: treat only the best)\n");
  printf("  -F, --full-check\n");
  printf("\tperform genome checks whenever possible\n");
  printf("  -v, --verbose\n\tbe verbose\n");
}

void interpret_cmd_line_options(int argc, char* argv[]) {
  // Define allowed options
  const char * short_options = "hVb:e:FI:R:vB";
  static struct option long_options[] = {
      {"help",      no_argument,       nullptr, 'h'},
      {"version",   no_argument,       nullptr, 'V'},
      {"begin",     required_argument, nullptr, 'b'},
      {"end",       required_argument, nullptr, 'e'},
      {"fullcheck", no_argument,       nullptr, 'F'},
      {"index",     required_argument, nullptr, 'I'},
      {"verbose",   no_argument,       nullptr, 'v'},
      {0, 0, 0, 0}
  };

  // Get actual values of the command-line options
  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 'b' : {
        if (strcmp(optarg, "") == 0) {
          printf("%s: error: Option -b or --begin : missing argument.\n",
                 argv[0]);
          exit(EXIT_FAILURE);
        }
        t0  = atol(optarg);
        break;
      }
      case 'e' : {
        if (strcmp(optarg, "") == 0) {
          printf("%s: error: Option -e or --end : missing argument.\n",
                 argv[0]);
          exit(EXIT_FAILURE);
        }
        t_end = atol(optarg);
        break;
      }
      case 'F' : {
        full_check = true;
        break;
      }
      case 'I' : {
        final_indiv_index  = atoi(optarg);
        best = false;
        break;
      }
      case 'v' : {
        verbose = true;
        break;
      }
      default : {
        // An error message is printed in getopt_long, we just need to exit
        exit(EXIT_FAILURE);
      }
    }
  }

  // If t_end wasn't provided, use default
  if (t_end < 0) {
    t_end = aevol::get_last_gener_from_file();
  }
}
