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

#include <format>
#include <fstream>
#include <iostream>

#include "nlohmann/json.hpp"

// Shorthand for nlohmann's json library
// We use the ordered_json specialization that maintains the insertion order
// of object keys instead of the default that uses alphabetical order for
// serialization
using json = nlohmann::ordered_json;

#include "AeTime.h"
#include "Individual.h"
#include "Protein.h"
#include "Rna.h"

namespace aevol {

UserInterfaceOutput::UserInterfaceOutput(const std::filesystem::path& outdir, time_type frequency)
    : outdir_(outdir), frequency_(frequency) {
  if (frequency_ != 0) {
    std::filesystem::create_directories(outdir_);
  }
}

bool UserInterfaceOutput::nothing_to_output() const {
  return (frequency_ == 0) || (AeTime::time() % frequency_ != 0);
}

void UserInterfaceOutput::write_indiv_output(const Individual& indiv, const std::string* seqid) {
  // Open output file
  std::ofstream indiv_out(
      outdir_ / (seqid ? std::format("indiv_{}.json", *seqid)
                       : std::format("indiv_{:0>9}.json", AeTime::time())),
      std::ios::out);

  // Fill json with indiv's data
  json indiv_json;
  if (exp_setup->diploid()) {
    indiv_json["indiv"]["genome_len"] = indiv.dna(Chrsm::A).length() + indiv.dna(Chrsm::B).length();
    indiv_json["indiv"]["chrsmA_len"] = indiv.dna(Chrsm::A).length();
    indiv_json["indiv"]["chrsmB_len"] = indiv.dna(Chrsm::B).length();
  } else {
    indiv_json["indiv"]["genome_len"] = indiv.dna().length();
  }
  indiv_json["indiv"]["fitness"] = indiv.fitness();

  // Rnas
  indiv_json["indiv"]["rnas"] = json::array();
  for (const auto& rna : indiv.rnas()) {
    indiv_json["indiv"]["rnas"].push_back(json::object({
        {"begin", rna.begin},
        {"end", rna.end},
        {"strand", to_string(rna.strand_)},
        {"is_coding", rna.is_coding_},
        {"expression_level", rna.e}}));
  }
  if (exp_setup->diploid()) {
    indiv_json["indiv"]["rnas_chrsmB"] = json::array();
    for (const auto& rna : indiv.rnas(Chrsm::B)) {
      indiv_json["indiv"]["rnas_chrsmB"].push_back(json::object({
          {"begin", rna.begin},
          {"end", rna.end},
          {"strand", to_string(rna.strand_)},
          {"is_coding", rna.is_coding_},
          {"expression_level", rna.e}}));
    }
  }

  // Proteins
  indiv_json["indiv"]["proteins"] = json::array();
  for (const auto* protein : indiv.proteins()) {
    indiv_json["indiv"]["proteins"].push_back(json::object({
        {"begin", protein->position_first_aa()},
        {"end", protein->position_last_base_stop_codon()},
        {"strand", to_string(protein->strand())},
        {"m", protein->m()},
        {"w", protein->w()},
        {"h", protein->h()},
        {"expression_level", protein->e()}}));
  }
  if (exp_setup->diploid()) {
    indiv_json["indiv"]["proteins_chrsmB"] = json::array();
    for (const auto* protein : indiv.proteins(Chrsm::B)) {
      indiv_json["indiv"]["proteins_chrsmB"].push_back(json::object({
          {"begin", protein->position_first_aa()},
          {"end", protein->position_last_base_stop_codon()},
          {"strand", to_string(protein->strand())},
          {"m", protein->m()},
          {"w", protein->w()},
          {"h", protein->h()},
          {"expression_level", protein->e()}}));
    }
  }

  // Phenotype
  indiv_json["indiv"]["phenotype"]["points"] = json::array();
  auto nb_points = indiv.phenotype().nb_points();
  for (decltype(nb_points) i = 0 ; i < nb_points ; ++i) {
    indiv_json["indiv"]["phenotype"]["points"].push_back(json::object({
      {"x", indiv.phenotype().point(i).x},
      {"y", indiv.phenotype().point(i).y}}));
  }

  // Write json to output file
  indiv_out << indiv_json.dump(2) << std::endl;
}

void UserInterfaceOutput::write_grid_output(size_t grid_width, size_t grid_height, const Population& population) {
  // Open output file
  std::ofstream grid_out(outdir_ / std::format("grid_{:0>9}.json", AeTime::time()), std::ios::out);

  // Fill json with grid data
  json grid_json;
  grid_json["grid_data"]["grid_width"] = grid_width;
  grid_json["grid_data"]["grid_height"] = grid_height;

  // Individuals' fitness
  std::vector<double> fitnesses;
  auto grid_size = grid_width * grid_height;
  fitnesses.reserve(grid_size);
  for (const auto& indiv : population.individuals()) {
    fitnesses.push_back(indiv.fitness());
  }
  grid_json["grid_data"]["fitness_grid"] = fitnesses;

  // Write json to output file
  grid_out << grid_json.dump(2) << std::endl;
}

void UserInterfaceOutput::write_environment_output(const PhenotypicTarget& target) {
  // Open output file
  std::ofstream env_out(outdir_ / std::format("environment_{:0>9}.json", AeTime::time()), std::ios::out);

  // Fill json with environmental target data
  json env_json;
  env_json["environment"]["target"]["points"] = json::array();
  auto nb_points = target.nb_points();
  for (decltype(nb_points) i = 0 ; i < nb_points ; ++i) {
    env_json["environment"]["target"]["points"].push_back(json::object({
        {"x", target.point(i).x},
        {"y", target.point(i).y}}));
  }

  // Write json to output file
  env_out << env_json.dump(2) << std::endl;
}

}
