// ****************************************************************************
//
//          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 <getopt.h>
#include <cstdlib>
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <sstream>
#include <string>
#include <assert.h>

// =================================================================
//                            Project Files
// =================================================================
#include "aevol.h"

using namespace aevol;

using std::string;

// Command-line option variables
static string triangles_file_name;
static string sequence_file_name;
static bool all_indiv = false;
static bool by_index = false;
static bool x_axis = false;
static bool y_axis = false;
static int16_t gu = -1;
static int32_t timestep = -1;
static int32_t ind = -1;
static int16_t x_pos = -1;
static int16_t y_pos = -1;

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

void analyse_indiv(Individual* indiv,
                   std::ostream& triangles_file,
                   std::ostream& sequence_file,
                   int16_t gu,
                   const LegacyPhenotypicTarget& phenotypicTarget);
void analyse_gu(GeneticUnit* gen_unit,
                std::ostream& triangles_file,
                const LegacyPhenotypicTarget& phenotypicTarget);




int main(int argc, char* argv[]) {
  interpret_cmd_line_options(argc, argv);

  assert(x_axis == y_axis);
  //make sure that no more than one option is selected
  int nb_option = 0;
  if(all_indiv){nb_option++;}
  if(by_index){nb_option++;}
  if(x_axis){nb_option++;}
  assert(nb_option <= 1);

  // Open the files
  std::ofstream triangles_file;
  std::ofstream sequence_file;

  if (triangles_file_name.size() > 0) {
    triangles_file.open(triangles_file_name, std::ios_base::out);

    // Write file headers
    int key = 0;
    triangles_file << "# " << std::setw(2) << ++key << " individual's identifier (id)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " strand" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " protein position (pos)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " length (len)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " position of last translated nucleotide (lpos)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " primary sequence (sequence)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " mean (m)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " width (w)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " height (h)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " concentration (c)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " feature (f)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " promoter position (prom_pos)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " RNA length (rna_len)" << std::endl;
    triangles_file << "# " << std::setw(2) << ++key << " basal level (basal_level)" << std::endl;
    triangles_file << std::endl;
    triangles_file << "id strand pos len lpos sequence m w h c f prom_pos rna_len basal_level" << std::endl;
  }
  if (sequence_file_name.size() > 0) {
    sequence_file.open(sequence_file_name, std::ios_base::out);
  }

  exp_manager = new ExpManager();
  exp_manager->load(timestep, false);

  if(by_index){
      assert(ind >= 0 && ind < exp_manager->grid_width() * exp_manager->grid_height());
  }
  if(x_axis){
      //we already know that x_axis = y_axis
      assert(x_pos >= 0 && x_pos < exp_manager->grid_width());
      assert(y_pos >= 0 && y_pos < exp_manager->grid_height());
  }

  // The best individual is already known because it is the last in the list
  // Thus we do not need to know anything about the environment and to evaluate
  // the individuals

  // Parse the individuals
  if(by_index) {
      int32_t orig_x = ind / exp_manager->world()->height();
      int32_t orig_y = ind % exp_manager->world()->height();
      Individual* indiv = exp_manager->indiv_by_position(orig_x,orig_y);
      
      analyse_indiv(indiv, triangles_file, sequence_file, gu, indiv->habitat().phenotypic_target());
  }
  else if(x_axis) {
      Individual* indiv = exp_manager->indiv_by_position(x_pos, y_pos);

      analyse_indiv(indiv, triangles_file, sequence_file, gu, indiv->habitat().phenotypic_target());
  }
  else if (all_indiv) {
    for (const auto& indiv: exp_manager->indivs()) {
      indiv->do_transcription_translation_folding(); // We need to recompute proteins if not already done (ie if using a population file and not a full backup)

      analyse_indiv(indiv, triangles_file, sequence_file, gu, indiv->habitat().phenotypic_target());
    }
  }
  else { // best_indiv
    Individual* best = exp_manager->best_indiv();
    best->do_transcription_translation_folding(); // We need to recompute proteins if not already done (ie if using a population file and not a full backup)
    analyse_indiv(best, triangles_file, sequence_file, gu, best->habitat().phenotypic_target()); // list of GU of the individual
  }

  sequence_file.close();
  triangles_file.close();

  delete exp_manager;

  return EXIT_SUCCESS;
}

// Parsing an individual
void analyse_indiv(Individual* indiv,
                   std::ostream& triangles_file,
                   std::ostream& sequence_file,
                   int16_t gu,
                   const LegacyPhenotypicTarget& phenotypicTarget) {
  // Output FASTA definition line
  sequence_file << ">" << "gen_" << time() << "_id_" << indiv->id() << " [organism=Aevol]"
                << " [strain=" << indiv->strain_name() << "]" << std::endl;

  if (gu == -1) { // We want to treat all genetic units
    int32_t gen_unit_number = 0;
    for (auto& gen_unit: indiv->genetic_unit_list_nonconst()) {
      if(triangles_file) {
        analyse_gu(&gen_unit, triangles_file,
                   phenotypicTarget);
      }
      if (sequence_file) {
        const char* dna = gen_unit.dna()->data();
        // The sequences of different GUs are separated by a space
        if (gen_unit_number > 0) sequence_file << " ";
 #ifdef BASE_2
        sequence_file << dna;
#else
	int i;
	for (i = 0 ; i < gen_unit.dna()->length() ; ++i) {
	    if (dna[i] == '0') sequence_file << "A";
	    if (dna[i] == '1') sequence_file << "T";
	    if (dna[i] == '2') sequence_file << "C";
	    if (dna[i] == '3') sequence_file << "G";
	  }
	sequence_file << std::endl;
#endif
      }

      gen_unit_number++;
    }
  }
  else { // User has specified a genetic unit
    GeneticUnit* gen_unit = &indiv->genetic_unit_nonconst(gu);

    if(triangles_file) {
        analyse_gu(gen_unit, triangles_file, phenotypicTarget);
    }
    if (sequence_file) {
      const char* dna = gen_unit->dna()->data();
      sequence_file << dna;
    }
  }
  // We go to next line in each file
  if (triangles_file) {
    triangles_file << std::endl;
  }
  if (sequence_file) {
    sequence_file << std::endl;
  }
}

// Parsing a GU
inline void analyse_gu(GeneticUnit* gen_unit,
                       std::ostream& triangles_file,
                       const LegacyPhenotypicTarget& phenotypicTarget) {
  // Construct the list of all rnas
  auto llrnas = gen_unit->rna_list();
  auto lrnas = llrnas[std::to_underlying(Strand::LEADING)];
  lrnas.splice(lrnas.end(), llrnas[std::to_underlying(Strand::LAGGING)]);
  // Parse this list
  for (const auto& rna: lrnas) {
    for (const auto& protein: rna.transcribed_proteins()) {
      double mean = protein->mean();
      int nfeat = -1;
      for (size_t i = 0 ;
           i <= static_cast<size_t>(phenotypicTarget.nb_segments()) - 1 ;
           ++i) {
        if ((mean > phenotypicTarget.segments()[i]->start) and
            (mean < phenotypicTarget.segments()[i]->stop)) {
          nfeat = phenotypicTarget.segments()[i]->feature;
          break;
        }
      }

      triangles_file << gen_unit->indiv()->id() << ' '
          << to_string(protein->strand()) << ' '
          << protein->first_translated_pos() << ' '
          << protein->length() << ' '
          << protein->last_translated_pos() << ' '
          << protein->AA_sequence('_') << ' '
          << std::fixed << protein->mean() << ' '
          << std::fixed << protein->width() << ' '
          << std::fixed << protein->height() << ' '
          << std::fixed << protein->concentration() << ' '
          << nfeat << ' '
          << rna.promoter_pos() << ' '
          << rna.transcript_length() << ' '
          << std::fixed << rna.basal_level() << std::endl;
    }
  }
}


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("\tExtracts the genotype and/or data about the phenotype of individuals\n");
    printf("\tin the provided population and write them into text files easy to parse\n");
    printf("\twith e.g. matlab.\n");
    printf("\n");
    printf("Usage : %s -h\n", prog_name);
    printf("   or : %s -V or --version\n", prog_name);
    printf("   or : %s [-t TIMESTEP] [-S SEQ_FILE] [-T TRIANGLE_FILE] [-U NUM_GU] [-a] [-x X -y Y]\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("  -t TIMESTEP\n");
    printf("\tspecify timestep of the individual(s) of interest\n");
    printf("  -S SEQ_FILE\n");
    printf("\textract sequences into file SEQ_FILE\n");
    printf("  -T TRIANGLE_FILE\n");
    printf("\textract phenotypic data into file TRIANGLE_FILE\n");
    printf("  -U NUM_GU\n");
    printf("\tonly treat genetic unit #NUM_GU (default: treat all genetic units)\n");
    printf("  -a\n");
    printf("\ttreat all the individuals (default: treat only the best)\n");
    printf("  -i IND\n");
    printf("\tonly treat individual #IND (default: treat only the best)\n");
    printf("  -x X -y Y\n");
    printf("\tonly treat individual at position X, Y on the grid\n");


    printf("\n\
This program extracts some data about the individuals at a given generation\n\
and write them into text files that are easy to parse with e.g. matlab.\n\
\n\
Two kinds of data can be extracted:\n\
\n\
 * data about the phenotype (option -T or --triangles):\n\
   write information about the proteins in a text file.\n\
   A space delimits two proteins, a new line delimits two individuals.\n\
   For each protein, the output is \"m_h_w_c_r_s_f_l_z_g\" where:\n\
     * m, h, w and c are the mean, height, width and concentration of the protein\n\
     * r is an identifier of the rna it belongs (useful to\n\
         know if several proteins are on the same rna)\n\
     * s indicates the strand (LEADING/LAGGING)\n\
     * f and l are the first and last translated bases\n\
     * z indicates the feature (at the center of the protein)\n\
\n\
 * sequences of the individuals (option -S or --sequence): write the sequences\n\
   in a text file. A new line delimits two individuals. In case\n\
   there are several genetic units (GUs), they are separated by whitespaces.\n\
\n\
By default, only the best individual is treated. You may use one of the\n\
following options to modify this behaviour:\n\
  * With option -a, all the individuals in the population are treated.\n\
  * With option -i, only the individual with the provided ID is treated.\n\
  * With options -x and -y, only the individual at the provided position on\n\
    the grid is treated.\n\
\n\
Use option -t (or --timestep) to specify timestep of the individual(s) of\n\
interest. A full backup tree must exist for the provided timestep.\n\
\n\
Examples :\n\
\n\
For generation 20000, write info about the phenotypes of all the\n\
individuals in phe_020000 and the sequences of all the\n\
individuals in seq_020000 :\n\
\n\
   extract --timestep 20000 --all --triangles phe_020000 --sequence seq_020000\n\
\n\
For generation 20000, write the best individual's sequence in\n\
seq_020000_best :\n\
\n\
   extract --timestep 20000 --sequence seq_020000_best\n");
}

void interpret_cmd_line_options(int argc, char* argv[]) {
  // Define allowed options
  const char * options_list = "hVt:aU:S:T:i:x:y:";
  static struct option long_options_list[] = {
      {"help",      no_argument,        nullptr, 'h'},
      {"version",   no_argument,        nullptr, 'V'},
      {"timestep",  required_argument,  nullptr, 't'},
      {"all",       no_argument,        nullptr, 'a'},
      {"gu",        required_argument,  nullptr, 'U'},
      {"sequence",  required_argument,  nullptr, 'S'},
      {"triangles", required_argument,  nullptr, 'T'},
      {"index",     required_argument,  nullptr, 'i'},
      {"xaxis",     required_argument,  nullptr, 'x'},
      {"yaxis",     required_argument,  nullptr, 'y'},
      {0, 0, 0, 0}
  };

  // Get actual values of the command-line options
  int option;
  while ((option = getopt_long(argc, argv, options_list,
                               long_options_list, nullptr)) != -1) {
    switch (option) {
      case 'h' : {
        print_help(argv[0]);
        exit(EXIT_SUCCESS);
      }
      case 'V' : {
        Utils::PrintAevolVersion();
        exit(EXIT_SUCCESS);
      }
      case 't' : {
        timestep = atol(optarg);
        break;
      }
      case 'a' : {
        all_indiv = true;
        break;
      }
      case 'U' : {
        gu = atoi(optarg);
        break;
      }
      case 'S' : {
        std::ostringstream stream;
        stream << optarg;
        sequence_file_name = stream.str();
        break;
      }
      case 'T' : {
        std::ostringstream stream;
        stream << optarg;
        triangles_file_name = stream.str();
        break;
      }
      case 'i' : {
        ind = atoi(optarg);
        by_index = true;
        break;
      }
      case 'x' : {
        x_pos = atoi(optarg);
        x_axis = true;
        break;
      }
      case 'y' : {
        y_pos = atol(optarg);
        y_axis = true;
        break;
      }
    }
  }

  // If timestep wasn't provided, use default
  if (timestep < 0) {
    timestep = OutputManager::last_gener();
  }

  // If neither the sequence_file_name nor the triangles_file_name was provided,
  // we will output only the sequence in a default-named file
  if (sequence_file_name.size() == 0 && triangles_file_name.size() == 0) {
    sequence_file_name = "sequence";
  }
}
