// Copyright (C) 2018  INRIA
//
// Treerecs is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// Treerecs 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 Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

// Include Googletest
#include <gtest/gtest.h>

// Incude standard libs
#include <unistd.h>

// Include Treerecs constants
#include <treerecs/Constants.h>
#include "examples_paths.h"

// Include Treerecs handlers, tools
#include <treerecs/tools/IO/IO.h>
#include <treerecs/tools/utils.h>
#include <treerecs/tools/PhyloTreeToolBox.h>
#include <treerecs/tools/SpeciesGeneMapper.h>
#include <treerecs/tools/treeReconciliation/TreeReconciliationConductor.h>

using namespace treerecs;

class IO_tests: public testing::Test {
public:
  std::string speciestree_file = multipolytomy_examples_path + "species_tree.txt";
  std::string genetree_newick_file = multipolytomy_examples_path + "gene_tree.txt";
  std::string genetree_nhx_file = multipolytomy_examples_path + "gene_tree.nhx";
  std::string genetree_phyloxml_file = multipolytomy_examples_path + "gene_tree.xml";
};

/// Compare edges (branch length and bootstraps) returns true if identicals.
/// Trees must have the same topology, leaves share the same name.
static bool compare_edges_between_two_trees(const bpp::PhyloTree& a
                                            , const bpp::PhyloTree& b) {
  auto anodes = PhyloTreeToolBox::getLeaves(a);
  auto bleaves = b.getAllLeaves();
  decltype(anodes) bnodes;
  for(auto& anode: anodes) {
    bnodes.push_back(PhyloTreeToolBox::getNodeFromName(
        bleaves, anode->getName()));
  }

  std::shared_ptr<bpp::PhyloNode> anode = nullptr;
  std::shared_ptr<bpp::PhyloNode> bnode = nullptr;
  while(not anodes.empty()) {
    anode = anodes.front();
    anodes.pop_front();
    bnode = bnodes.front();
    bnodes.pop_front();

    if(anode != a.getRoot()) {
      auto abranch = a.getEdgeToFather(anode);
      auto bbranch = b.getEdgeToFather(bnode);

      if(abranch->hasLength()) {
        if(not bbranch->hasLength()) return false;
        auto alength = abranch->getLength();
        auto blength = bbranch->getLength();
        if (not utils::double_equivalence(alength, blength)) {
          return false;
        }
      }

      if(abranch->hasBootstrapValue()) {
        if(not bbranch->hasBootstrapValue()) return false;
        auto asupport = abranch->getBootstrapValue();
        auto bsupport = bbranch->getBootstrapValue();
        if (not utils::double_equivalence(asupport, bsupport)) {
          return false;
        }
      }

      auto afather = a.getFather(anode);
      auto bfather = b.getFather(bnode);

      // If current nodes are son left we append the father into the list, so we
      // will visit only once the father next time.
      if((anode == a.getSons(afather).front()) and (afather != a.getRoot())) {
        anodes.push_back(afather);
        bnodes.push_back(bfather);
      }
    }
  }

  return true;
}

TEST_F(IO_tests, io_pre_process_tests) {
  ASSERT_TRUE(IO::exists(genetree_newick_file));

  // Open speciestree
  // We use a speciestree because it has no bootstrap and/or branch length.
  auto speciestree = IO::readTreeFile(speciestree_file);
  ASSERT_TRUE(speciestree != nullptr);

  // Check if branches are not associated with a bootstrap
  ASSERT_FALSE(PhyloTreeToolBox::hasNode(
      *speciestree,
      [&speciestree](const Node& node) {
        if ((not speciestree->isLeaf(node))
            and (node != speciestree->getRoot())) {
          auto edgeToFather = speciestree->getEdgeToFather(node);
          if(edgeToFather->hasBootstrapValue()) {
            return true;
          }
        }
        return false;
      })
  );

  // Check if all branches are associated with a branch length (with default
  // setting, we run IO::readTreeFile and ask to associate each branch with a
  // length).
  ASSERT_FALSE(PhyloTreeToolBox::hasNode(
      *speciestree,
      [&speciestree](const Node& node) {
        if (node != speciestree->getRoot()) {
          auto edgeToFather = speciestree->getEdgeToFather(node);
          if(not edgeToFather->hasLength()) {
            return true;
          }
        }
        return false;
      })
  );

  ASSERT_TRUE(speciestree != nullptr);
}

TEST_F(IO_tests, several_io_methods) {
  // IO::exists
  ASSERT_TRUE(IO::exists(genetree_newick_file));
  ASSERT_FALSE(IO::exists("\0"));
  ASSERT_FALSE(IO::exists(""));
  std::string random_filename = "test_file" + std::to_string(rand()%100);
  ASSERT_FALSE(IO::exists(random_filename));
  std::ofstream orandomfile(random_filename);
  orandomfile.close();
  ASSERT_TRUE(IO::exists(random_filename));
  ASSERT_TRUE(std::remove(random_filename.c_str()) == 0);
  ASSERT_FALSE(IO::exists(random_filename));

  // IO::nlines
  orandomfile = std::ofstream(random_filename);

  std::list<std::string> apollinaire {
  "Pauvre automne",
  "Meurs en blancheur et en richesse",
  "De neige et de fruits mûrs",
  "Au fond du ciel",
  "Des éperviers planent",
  "Sur les nixes nicettes aux cheveux verts et naines",
  "Qui n'ont jamais aimé"};

  for(auto verse: apollinaire)
    orandomfile << verse << std::endl;

  orandomfile.close();
  ASSERT_TRUE(IO::exists(random_filename));

  std::ifstream irandomfile(random_filename);
  ASSERT_EQ(IO::nlines(irandomfile), apollinaire.size());
  irandomfile.close();
  ASSERT_TRUE(std::remove(random_filename.c_str()) == 0);
}

TEST_F(IO_tests, io_read_newick_tests) {
  ASSERT_TRUE(IO::exists(genetree_newick_file));

  // Opening a single tree
  auto speciestree = IO::readTreeFile(speciestree_file);
  ASSERT_TRUE(speciestree != nullptr);
  ASSERT_EQ(4, speciestree->getNumberOfLeaves());

  // The file contains 3 trees.
  auto genetrees = IO::readTreesFile(genetree_newick_file);
  ASSERT_EQ(3, genetrees.size());
  for(auto genetree: genetrees) {
    ASSERT_TRUE(genetree->getNumberOfLeaves() > 0);
  }

  // Check if trees has bootstraps and branch lengths
  for(auto genetree: genetrees) {
    ASSERT_TRUE(
        PhyloTreeToolBox::hasNode(*genetree, [&genetree](const Node &node) {
          if (node != genetree->getRoot()) {
            if (genetree->hasFather(node)) {
              return genetree->getEdgeToFather(node)->hasBootstrapValue()
                  and genetree->getEdgeToFather(node)->hasLength();
            }
          }
          return false;
        })
    );
  }

  // Load only the second tree
  genetrees.clear();
  genetrees = IO::readTreesFile(genetree_newick_file, TextFormat::newick, 1);
  ASSERT_EQ(1, genetrees.size());

  // The second tree has 15 leaves
  ASSERT_EQ(15, genetrees.at(0)->getNumberOfLeaves());
}

TEST_F(IO_tests, io_write_newick_tests) {
  // Opening and loading gene trees
  auto genetrees = IO::readTreesFile(genetree_newick_file);

  auto genetree0 = genetrees.at(0);
  auto genetree1 = genetrees.at(1);
  auto genetree2 = genetrees.at(2);

  // Transform bpp::PhyloTree to Newick.
  auto genetree0_newick = IO::PhyloTreeToNewick(*genetree0);
  auto genetree1_newick = IO::PhyloTreeToNewick(*genetree1);
  auto genetree2_newick = IO::PhyloTreeToNewick(*genetree2);

  // Transform these newick trees in bpp::PhyloTree (again)
  auto genetree0_newick_phylotree = IO::newickToPhyloTree(genetree0_newick,
                                                          true);
  auto genetree1_newick_phylotree = IO::newickToPhyloTree(genetree1_newick,
                                                          true);
  auto genetree2_newick_phylotree = IO::newickToPhyloTree(genetree2_newick,
                                                          true);

  // Compare leaves between original and "copy", expected to be identical.
  auto genetree0_leaves = genetree0->getAllLeavesNames();
  auto genetree0_newick_phylotree_leaves = genetree0_newick_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(
      genetree0_leaves.begin(), genetree0_leaves.end(),
      genetree0_newick_phylotree_leaves.begin(), genetree0_newick_phylotree_leaves.end()));
  auto genetree1_leaves = genetree1->getAllLeavesNames();
  auto genetree1_newick_phylotree_leaves = genetree1_newick_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(
      genetree1_leaves.begin(), genetree1_leaves.end(),
      genetree1_newick_phylotree_leaves.begin(), genetree1_newick_phylotree_leaves.end()));
  auto genetree2_leaves = genetree2->getAllLeavesNames();
  auto genetree2_newick_phylotree_leaves = genetree2_newick_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(
      genetree2_leaves.begin(), genetree2_leaves.end(),
      genetree2_newick_phylotree_leaves.begin(), genetree2_newick_phylotree_leaves.end()));

  // Then compare edges
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree0, *genetree0_newick_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree1, *genetree1_newick_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree2, *genetree2_newick_phylotree));

}

TEST_F(IO_tests, io_read_nhx_tests) {
  ASSERT_TRUE(IO::exists(genetree_nhx_file));

  // Opening a single tree
  auto speciestree = IO::readTreeFile(speciestree_file, TextFormat::newick);
  ASSERT_TRUE(speciestree != nullptr);
  ASSERT_EQ(4, speciestree->getNumberOfLeaves());

  // The file contains 3 trees.
  auto genetrees = IO::readTreesFile(genetree_nhx_file, TextFormat::nhx);
  for(auto genetree: genetrees) {
    ASSERT_TRUE(genetree->getNumberOfLeaves() > 0);
  }
  ASSERT_EQ(3, genetrees.size());

  // Check if tree has bootstraps
  auto genetree = *genetrees.front();
  ASSERT_TRUE(
      PhyloTreeToolBox::hasNode(genetree, [&genetree](const Node& node) {
        if(node != genetree.getRoot()){
          if(genetree.hasFather(node)){
            if(genetree.getEdgeToFather(node)->hasBootstrapValue())
              return true;
          }
        }
        return false;
      })
  );

  // The second tree has 15 leaves
  ASSERT_EQ(15, genetrees.at(1)->getNumberOfLeaves());
}

TEST_F(IO_tests, io_write_nhx_tests) {
  auto genetrees = IO::readTreesFile(genetree_nhx_file, TextFormat::nhx);

  auto genetree0 = genetrees.at(0);
  auto genetree1 = genetrees.at(1);
  auto genetree2 = genetrees.at(2);

  auto genetree0_nhx = IO::PhyloTreeToNhx(*genetree0);
  auto genetree1_nhx = IO::PhyloTreeToNhx(*genetree1);
  auto genetree2_nhx = IO::PhyloTreeToNhx(*genetree2);

  auto genetree0_nhx_phylotree = IO::nhxToPhyloTree(genetree0_nhx);
  auto genetree1_nhx_phylotree = IO::nhxToPhyloTree(genetree1_nhx);
  auto genetree2_nhx_phylotree = IO::nhxToPhyloTree(genetree2_nhx);

  // Compare leaves between original and "copy".
  auto genetree0_leaves = genetree0->getAllLeavesNames();
  auto genetree0_nhx_phylotree_leaves = genetree0_nhx_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree0_leaves.begin(), genetree0_leaves.end(),
                              genetree0_nhx_phylotree_leaves.begin(), genetree0_nhx_phylotree_leaves.end()));
  auto genetree1_leaves = genetree1->getAllLeavesNames();
  auto genetree1_nhx_phylotree_leaves = genetree1_nhx_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree1_leaves.begin(), genetree1_leaves.end(),
                              genetree1_nhx_phylotree_leaves.begin(), genetree1_nhx_phylotree_leaves.end()));
  auto genetree2_leaves = genetree2->getAllLeavesNames();
  auto genetree2_nhx_phylotree_leaves = genetree2_nhx_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree2_leaves.begin(), genetree2_leaves.end(),
                              genetree2_nhx_phylotree_leaves.begin(), genetree2_nhx_phylotree_leaves.end()));

  // Then compare edges
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree0, *genetree0_nhx_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree1, *genetree1_nhx_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree2, *genetree2_nhx_phylotree));

  // Compare trees with days algorithm.
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(*genetree0, *genetree0_nhx_phylotree));
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(*genetree1, *genetree1_nhx_phylotree));
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(*genetree2, *genetree2_nhx_phylotree));
  ASSERT_FALSE(PhyloTreeToolBox::compare_trees_with_days_algorithm(*genetree1, *genetree2_nhx_phylotree));
}

TEST_F(IO_tests, io_read_phyloxml_tests) {
  ASSERT_TRUE(IO::exists(genetree_phyloxml_file));

  // Opening a single tree
  auto speciestree = IO::readTreeFile(speciestree_file, TextFormat::newick);
  ASSERT_TRUE(speciestree != nullptr);
  ASSERT_EQ(4, speciestree->getNumberOfLeaves());

  // The file contains 3 trees.
  auto genetrees = IO::readTreesFile(genetree_phyloxml_file, TextFormat::phyloxml);
  ASSERT_EQ(3, genetrees.size());

  // Check if tree has bootstraps
  auto genetree = *genetrees.front();
  ASSERT_TRUE(
      PhyloTreeToolBox::hasNode(genetree, [&genetree](const Node& node) {
        if(node != genetree.getRoot()){
          if(genetree.hasFather(node)){
            if(genetree.getEdgeToFather(node)->hasBootstrapValue()){
              return true;
            }
          }
        }
        return false;
      })
  );

  // The second tree has 15 leaves
  ASSERT_EQ(15, genetrees.at(1)->getNumberOfLeaves());
}

TEST_F(IO_tests, io_write_phyloxml_tests) {

  auto genetrees = IO::readTreesFile(genetree_phyloxml_file,
                                     TextFormat::phyloxml);
  ASSERT_EQ(3, genetrees.size());

  auto genetree0 = genetrees.at(0);
  auto genetree1 = genetrees.at(1);
  auto genetree2 = genetrees.at(2);

  auto genetree0_phyloxml = IO::PhyloTreeToPhyloXML(*genetree0);
  auto genetree1_phyloxml = IO::PhyloTreeToPhyloXML(*genetree1);
  auto genetree2_phyloxml = IO::PhyloTreeToPhyloXML(*genetree2);

  PhyloXML phyloxml;
  auto genetree0_phyloxml_phylotree
      = phyloxml.phyloXMLToPhyloTree(genetree0_phyloxml);
  auto genetree1_phyloxml_phylotree
      = phyloxml.phyloXMLToPhyloTree(genetree1_phyloxml);
  auto genetree2_phyloxml_phylotree
      = phyloxml.phyloXMLToPhyloTree(genetree2_phyloxml);

  // Compare leaves between original and "copy".
  auto genetree0_leaves = genetree0->getAllLeavesNames();
  auto genetree0_phyloxml_phylotree_leaves = genetree0_phyloxml_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree0_leaves.begin(), genetree0_leaves.end(),
                              genetree0_phyloxml_phylotree_leaves.begin(),
                              genetree0_phyloxml_phylotree_leaves.end()));
  auto genetree1_leaves = genetree1->getAllLeavesNames();
  auto genetree1_phyloxml_phylotree_leaves = genetree1_phyloxml_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree1_leaves.begin(), genetree1_leaves.end(),
                              genetree1_phyloxml_phylotree_leaves.begin(),
                              genetree1_phyloxml_phylotree_leaves.end()));
  auto genetree2_leaves = genetree2->getAllLeavesNames();
  auto genetree2_phyloxml_phylotree_leaves = genetree2_phyloxml_phylotree->getAllLeavesNames();
  ASSERT_TRUE(utils::comp_all(genetree2_leaves.begin(), genetree2_leaves.end(),
                              genetree2_phyloxml_phylotree_leaves.begin(),
                              genetree2_phyloxml_phylotree_leaves.end()));

  // Then compare edges
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree0,
                                              *genetree0_phyloxml_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree1,
                                              *genetree1_phyloxml_phylotree));
  ASSERT_TRUE(compare_edges_between_two_trees(*genetree2,
                                              *genetree2_phyloxml_phylotree));

  // Compare trees with days algorithm.
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree0, *genetree0_phyloxml_phylotree));
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree1, *genetree1_phyloxml_phylotree));
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree2, *genetree2_phyloxml_phylotree));
  ASSERT_FALSE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree1, *genetree2_phyloxml_phylotree));
}

TEST_F(IO_tests, io_write_one_tree_test) {
  int itree = 1; // We work with the second tree in genetree_newick_file
  auto genetrees = IO::readTreesFile(genetree_newick_file, TextFormat::newick);
  auto speciestree = IO::readTreeFile(speciestree_file, TextFormat::newick);
  auto map = SpeciesGeneMapper::map(genetrees.begin(), genetrees.end(),
                                    *speciestree, MappingMethod::TREES);

  TreeReconciliationConductor treeReconciliationConductor;
  auto solutions = treeReconciliationConductor(
      *genetrees[itree], *speciestree, map[itree]
      , nullptr
      , false
      , DEFAULT_DUPLICATION_COST
      , DEFAULT_LOSS_COST
      , 0.5
      , DEFAULT_STRICT_SUPPORT_THRESHOLDS
      , true
      , 1
      , false
      , false
      , true //add loss events
      , true //compute distances
      , true
      , false
  );

  // Keep only one tree to test and compute number of losses in leaves
  auto final_genetree = solutions.front().genetree();
  auto final_map = solutions.front().map();

  unsigned long ndup = solutions.front().ndup();
  unsigned long napparentdup = ndup;

  std::unordered_map<std::shared_ptr<bpp::PhyloTree>
      , std::map<double, std::vector<ReconciledRootedTree>>> results;

  results[genetrees[itree]][0.5] = std::move(solutions);

  // Test with newick format (without gene loss)
  std::stringstream sstream_newick;
  IO::writeReconciledTrees(sstream_newick, results,
                           DEFAULT_STRICT_SUPPORT_THRESHOLDS,
                           {genetrees[itree]},
                           *speciestree, true, TextFormat::newick, false);
  std::string sstream_newick_str = sstream_newick.str();

  ASSERT_EQ(0, utils::count(sstream_newick_str, "loss"));
  ASSERT_EQ(1, utils::count(sstream_newick_str, "\n"));

  // Test with NHX format (without loss but duplications annotated)
  std::stringstream sstream_nhx;
  IO::writeReconciledTrees(sstream_nhx, results,
                           DEFAULT_STRICT_SUPPORT_THRESHOLDS,
                           {genetrees[itree]},
                           *speciestree, true, TextFormat::nhx, false);
  std::string sstream_nhx_str = sstream_nhx.str();

  ASSERT_EQ(0, utils::count(sstream_nhx_str, "loss"));
  ASSERT_EQ(napparentdup, utils::count(sstream_nhx_str, "D=Y"));
  ASSERT_EQ(1, utils::count(sstream_nhx_str, "\n"));
  ASSERT_EQ(1, utils::count(sstream_nhx_str, ";"));
}