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

// Include Bpplib
#include <Bpp/Phyl/Tree/PhyloTree.h>

// Include containers
#include <memory>
#include <treerecs/algorithms/NeighborJoining.h>
#include <treerecs/tools/Statistics.h>
#include <treerecs/containers/Grid.h>
#include <treerecs/containers/Table.h>
#include <treerecs/containers/GeneMap.h>

// Include handlers, tools
#include <treerecs/tools/SpeciesGeneMapper.h>
#include <treerecs/tools/IO/IO.h>
#include <treerecs/tools/utils.h>

// Include examples paths
#include "examples_paths.h"

using namespace treerecs;

/*!
 * @class PhyloTreeToolBox_tests
 * @brief Test PhyloTreeToolBox methods
 */
class PhyloTreeToolBox_tests: public testing::Test {
protected:
  virtual void SetUp() {

    speciestree = IO::readTreeFile(
        polytomy_examples_path + "species_polytomy.txt",
        TextFormat ::newick, false, false, false);
    genetree = IO::readTreeFile(ensembl_examples_path + "phyml_tree.txt",
                                TextFormat ::newick, true, false);

    numbered_tree = IO::readTreeFile(
        numbered_tree_path + "numbered_tree.xml", TextFormat::phyloxml);

    fake_genetree = IO::readTreesFile(multipolytomy_examples_path
                                      + "gene_tree.txt",
                                      TextFormat::newick, 1).front();
    fake_speciestree = IO::readTreeFile(multipolytomy_examples_path
                                         + "species_tree.txt");
    fake_map = SpeciesGeneMapper::mapWithTrees(
        *fake_genetree, *fake_speciestree);
    SpeciesGeneMapper::updateInnerNodeMapping(fake_map, *fake_genetree,
                                              *fake_speciestree);

  }

  virtual void TearDown() {}

  std::shared_ptr<bpp::PhyloTree> speciestree;
  std::shared_ptr<bpp::PhyloTree> genetree;
  std::shared_ptr<bpp::PhyloTree> fake_genetree;
  std::shared_ptr<bpp::PhyloTree> fake_speciestree;
  SpeciesGeneMap fake_map;
  std::shared_ptr<bpp::PhyloTree> numbered_tree;
};


TEST_F(PhyloTreeToolBox_tests, testDistanceMatrixComputation) {
  bool verbose = false;

  if (verbose) std::cout << "Tree: " << *genetree << std::endl;

  // Test distance matrix computations.
  // this method of PhyloTreeToolBox is unoptimized: it computes lengths between
  // each node by creating a path and then make the sum of the branches in the
  // path. This process is repeated for each pair. We are re-computing sub-path
  // lengths a lot of time as consequence. That would otherwise be quick, this
  // method is sure.
  auto dmatrix_ref_unoptimized =
      PhyloTreeToolBox::distanceMatrixBetweenNodes_unoptimized(
          *genetree, genetree->getAllNodes(), false);
  // this method is faster because we are computing distances between nodes by a
  // bottom-up reconstruction: we are
  // visiting each branch only one time.
  auto dmatrix_optimized = PhyloTreeToolBox::computeDistanceMatrix(
      *genetree, nullptr, false);

  // each distance matrix is expected to be identical. But we are using doubles,
  // so we use utils::double_equivalence to
  // check equality.
  bool identical_dmatrices = dmatrix_optimized.equalTo(
      dmatrix_ref_unoptimized,
      [](const double a,
          const double b) { return utils::double_equivalence(a, b); });
  ASSERT_TRUE(identical_dmatrices);
}

TEST_F(PhyloTreeToolBox_tests, testPhyloTreeToolBoxAlgorithms) {
  // Test compareTrees_Days
  ASSERT_FALSE(
      PhyloTreeToolBox::compare_trees_with_days_algorithm(
          *genetree, *speciestree, false));
  ASSERT_TRUE(
      PhyloTreeToolBox::compare_trees_with_days_algorithm(
          *genetree, *genetree, false));

  // Test PhyloTreeToolBox getters
  // - get all leaves
  std::vector<Node> leavesPO = {
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "1"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "2"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "3"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "6"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "7")
  };
  auto bpp_leaves = numbered_tree->getAllLeaves();
  auto treerecs_leaves = PhyloTreeToolBox::getLeaves(*numbered_tree);

  ASSERT_EQ(numbered_tree->getNumberOfLeaves(),
            treerecs_leaves.size());

  ASSERT_TRUE(utils::contains(treerecs_leaves.begin(), treerecs_leaves.end(),
            [&bpp_leaves](const Node& node) {
              return std::find(bpp_leaves.begin(), bpp_leaves.end(),
                               node) != bpp_leaves.end(); }));

  ASSERT_TRUE(utils::comp_all(leavesPO.begin(), leavesPO.end(),
                              treerecs_leaves.begin(), treerecs_leaves.end()));

  // - get all internal nodes
  std::vector<Node> inodesPO = {
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "4", false, false),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "5", false, false),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "8", false, false),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "9", false, false)
  };
  auto bpp_inodes = numbered_tree->getAllNodes();
  auto treerecs_inodes = PhyloTreeToolBox::getInternalNodes(*numbered_tree);

  ASSERT_EQ(numbered_tree->getNumberOfNodes()
            - numbered_tree->getNumberOfLeaves(),
            PhyloTreeToolBox::getInternalNodes(*numbered_tree).size());

  ASSERT_TRUE(utils::contains(treerecs_inodes.begin(), treerecs_inodes.end(),
                              [&bpp_inodes](const Node& node) {
                                return std::find(bpp_inodes.begin(),
                                                 bpp_inodes.end(),
                                                 node) != bpp_inodes.end();
                              }));

  // - get all leaves under a specific one
  ASSERT_EQ(genetree->getLeavesUnderNode(genetree->getRoot()).size(),
            PhyloTreeToolBox::getLeaves(*genetree, genetree->getRoot()).size());

  Node five = PhyloTreeToolBox::getNodeFromName(*numbered_tree, "5",
                                                false, false);
  std::vector<Node> under_five = {
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "1"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "2"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "3")
  };
  auto treerecs_under_five = PhyloTreeToolBox::getLeaves(
      *numbered_tree, five);

  ASSERT_TRUE(utils::comp_all(under_five.begin(), under_five.end(),
                              treerecs_under_five.begin(),
                              treerecs_under_five.end()));

  // - get all nodes under a specific one
  under_five = {
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "1"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "2"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "3"),
      PhyloTreeToolBox::getNodeFromName(*numbered_tree, "4", false, false),
      five
  };
  treerecs_under_five
      = PhyloTreeToolBox::getNodesInPostOrderTraversalIterative_list(
      *numbered_tree,
      five);

  ASSERT_TRUE(utils::comp_all(under_five.begin(), under_five.end(),
                              treerecs_under_five.begin(),
                              treerecs_under_five.end()));

  // Test In-Order Traversal Recursive.
  //   Use examples/numbered_tree.xml with node names as numbers.

  std::vector<std::string> IO_recursive_true
      {"1", "5", "2", "4", "3", "9", "6", "8", "7"};
  auto IO_recursive = PhyloTreeToolBox::getNodesInOrderTraversalRecursive_list(
      *numbered_tree, numbered_tree->getRoot());
  ASSERT_EQ(numbered_tree->getNumberOfNodes(), IO_recursive.size());

  std::vector<std::string> IO_recursive_test;
  IO_recursive_test.reserve(IO_recursive.size());
  for (const auto& IO_recursive_i : IO_recursive)
    IO_recursive_test.push_back(IO_recursive_i->getName());
  ASSERT_TRUE(
      utils::comp_all(IO_recursive_true.begin(), IO_recursive_true.end(),
      IO_recursive_test.begin(), IO_recursive_test.end()));

  // Test Post-Order Traversal Recursive.
  std::vector<std::string> PO_recursive_true {
      "1", "2", "3", "4", "5", "6", "7", "8", "9"};
  auto PO_recursive =
      PhyloTreeToolBox::getNodesInPostOrderTraversalRecursive_list(
      *numbered_tree, numbered_tree->getRoot());
  ASSERT_EQ(numbered_tree->getNumberOfNodes(), PO_recursive.size());

  std::vector<std::string> PO_recursive_test;
  PO_recursive_test.reserve(PO_recursive.size());
  for (const auto& PO_recursive_i : PO_recursive)
    PO_recursive_test.push_back(PO_recursive_i->getName());
  ASSERT_TRUE(
      utils::comp_all(PO_recursive_true.begin(), PO_recursive_true.end(),
                      PO_recursive_test.begin(), PO_recursive_test.end()));

  // Test Pre-Order Traversal Recursive.
  std::vector<std::string> PreO_recursive_true
      {"9", "5", "1", "4", "2", "3", "8", "6", "7"};
  auto PreO_recursive =
      PhyloTreeToolBox::getNodesInPreOrderTraversalRecursive_list(
      *numbered_tree, numbered_tree->getRoot());
  ASSERT_EQ(numbered_tree->getNumberOfNodes(), PreO_recursive.size());

  std::vector<std::string> PreO_recursive_test;
  PreO_recursive_test.reserve(PreO_recursive.size());
  for (const auto& PreO_recursive_i : PreO_recursive)
    PreO_recursive_test.push_back(PreO_recursive_i->getName());
  ASSERT_TRUE(
      utils::comp_all(PreO_recursive_true.begin(), PreO_recursive_true.end(),
                      PreO_recursive_test.begin(), PreO_recursive_test.end()));

  // Test Post-Order Traversal Recursive.
  PO_recursive.clear();
  PO_recursive = PhyloTreeToolBox::getNodesInPostOrderTraversalRecursive_list(
      *genetree, genetree->getRoot());
  ASSERT_EQ(genetree->getNumberOfNodes(), PO_recursive.size());

  // Test Post-Order Traversal Iterative.
  auto PO_iterative =
      PhyloTreeToolBox::getNodesInPostOrderTraversalIterative_list(
      *genetree, genetree->getRoot());
  ASSERT_EQ(genetree->getNumberOfNodes(), PO_iterative.size());

  // Check if these lists are similar.
  ASSERT_TRUE(utils::comp_all(PO_recursive.begin(), PO_recursive.end(),
                              PO_iterative.begin(), PO_iterative.end()));

  // Test last common ancestor getters.
  ASSERT_EQ(
  PhyloTreeToolBox::getLastCommonAncestor(numbered_tree->getAllLeaves(),
                                          *numbered_tree)
  , numbered_tree->getRoot());

  Node node1 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "1", false, false);
  Node node2 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "2", false, false);
  Node node3 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "3", false, false);
  Node node4 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "4", false, false);
  Node node5 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "5", false, false);
  Node node6 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "6", false, false);
  Node node7 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "7", false, false);
  Node node8 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "8", false, false);
  Node node9 = PhyloTreeToolBox::getNodeFromName(*numbered_tree,
                                                 "9", false, false);

  ASSERT_EQ(
  PhyloTreeToolBox::getLastCommonAncestor(
      {node6, node7}, *numbered_tree), node8);
  ASSERT_EQ(
  PhyloTreeToolBox::getLastCommonAncestor(
      {node8, node4}, *numbered_tree), node9);
  ASSERT_EQ(
  PhyloTreeToolBox::getLastCommonAncestor(
      {node1, node2, node3}, *numbered_tree), node5);
  ASSERT_EQ(
  PhyloTreeToolBox::getCommonAncestor(
      *numbered_tree, node3, node1), node5);
}

TEST_F(PhyloTreeToolBox_tests, testEvents) {
  ASSERT_NE(Event::duplication, Event::loss);

  // speciation
  Node b1_b = PhyloTreeToolBox::getNodeFromName(*fake_genetree, "b1_b");
  Node b1_b_father = fake_genetree->getFather(b1_b);
  Event b1_b_father_event =
      PhyloTreeToolBox::getEvent(*fake_genetree, fake_map, b1_b_father);
  ASSERT_EQ(Event::speciation, b1_b_father_event);

  // duplication
  Node a1_a = PhyloTreeToolBox::getNodeFromName(*fake_genetree, "a1_a");
  Node a1_a_father = fake_genetree->getFather(a1_a);
  Event a1_a_father_event =
      PhyloTreeToolBox::getEvent(*fake_genetree, fake_map, a1_a_father);
  ASSERT_EQ(Event::duplication, a1_a_father_event);

  // extant
  Event a1_a_event = PhyloTreeToolBox::getEvent(*fake_genetree, fake_map, a1_a);
  ASSERT_EQ(Event::extant, a1_a_event);

  // loss
  // Initialy, the fake_genetree has no loss/ArtificialGene:
  ASSERT_FALSE(PhyloTreeToolBox::isArtificalGene(a1_a));
  // We will add an "artificial gene"
  a1_a->setProperty("isArtificialGene", IsArtificialGene(true));
  a1_a_event = PhyloTreeToolBox::getEvent(*fake_genetree, fake_map, a1_a);
  ASSERT_EQ(Event::loss, a1_a_event);
  ASSERT_TRUE(PhyloTreeToolBox::isArtificalGene(a1_a));
  ASSERT_TRUE(PhyloTreeToolBox::hasArtificialGenes(*fake_genetree));

  a1_a->setProperty("isArtificialGene", IsArtificialGene(false));
  ASSERT_FALSE(PhyloTreeToolBox::isArtificalGene(a1_a));
  ASSERT_FALSE(PhyloTreeToolBox::hasArtificialGenes(*fake_genetree));

  // Compute events
  auto b1_b_grand_father_events
      = PhyloTreeToolBox::getGeneEventsInBifurcation(
          fake_genetree->getFather(b1_b_father),
          *fake_genetree, *fake_speciestree, fake_map);
  ASSERT_EQ(1, b1_b_grand_father_events.at(Event::duplication));
  ASSERT_EQ(1, b1_b_grand_father_events.at(Event::loss));
}

TEST_F(PhyloTreeToolBox_tests, testTreeCloning) {
  // Test tree cloning, we clone the tree before.
  auto genetree_copy = PhyloTreeToolBox::cloneTree(*genetree);
  // the tree is expected to be identical
  ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree, *genetree_copy, false));

  // Get all edges of the tree. With edges we can get the number of nodes and
  // branch lengths in the same time.
  ASSERT_TRUE(genetree->isValid());
  auto original_edges = genetree->getAllEdges();
  auto cloned_edges = genetree_copy->getAllEdges();

  // Compare.
  ASSERT_EQ(original_edges.size(), cloned_edges.size());
  ASSERT_FALSE(utils::comp_all(original_edges.begin(), original_edges.end(),
                               cloned_edges.begin(), cloned_edges.end()));
  ASSERT_TRUE(
      utils::comp_all(
          original_edges.begin(), original_edges.end(),
          cloned_edges.begin(), cloned_edges.end(),
          [](const Branch& a, const Branch& b){
            if(a->hasLength() == b->hasLength()) {
              if(a->hasLength())
                return utils::double_equivalence(a->getLength(),b->getLength());
              return true;
            } else {
              return false;
            }
          }
  ));

}

TEST_F(PhyloTreeToolBox_tests, testTreeContraction) {
  // Test tree contraction.
  auto genetree_copy = PhyloTreeToolBox::cloneTree(*genetree);

  // Check if there is no multi-furcation before contraction
  ASSERT_FALSE(PhyloTreeToolBox::hasNode(
      *genetree_copy,
      [genetree_copy](const Node& node){
        if(genetree_copy->isLeaf(node)) return false;
        return genetree_copy->getSons(node).size() > 2;
      }
  ));

  // Full tree contraction.
  double support_threshold = 100.0;
  auto removed_branches = PhyloTreeToolBox::mergeWeakBranches(
      *genetree_copy, support_threshold,
      DEFAULT_STRICT_SUPPORT_THRESHOLDS);
  ASSERT_TRUE(removed_branches.size() > 0);
  // by using a threshold equal to 100 we expect a total contraction of the tree
  // with only one polytomy at the root.

  // Check if PhyloTreeToolBox::hasNode as multi-furcation works
  ASSERT_TRUE(PhyloTreeToolBox::hasNode(
      *genetree_copy,
      [genetree_copy](const Node& node){
        if(genetree_copy->isLeaf(node)) return false;
        return genetree_copy->getSons(node).size() > 2;
      }
  ));

  // Check if the number of polytomies is equal to 1.
  ASSERT_EQ(1, PhyloTreeToolBox::getNodesInPostOrderTraversalIterative_list(
      *genetree_copy,
      [genetree_copy](const Node& node){
        if(genetree_copy->isLeaf(node)) return false;
        return genetree_copy->getSons(node).size() > 2;
      }
  ).size());

  // Check if the resulting tree is not binary.
  ASSERT_FALSE(PhyloTreeToolBox::isBinary(*genetree_copy));

  // Check if the resulting tree is different to the original one.
  ASSERT_FALSE(PhyloTreeToolBox::compare_trees_with_days_algorithm(
      *genetree, *genetree_copy, false));

  ASSERT_EQ(1,  // Check if the number of internal node is equal to 1.
  PhyloTreeToolBox::getNodesInPostOrderTraversalIterative_list(*genetree_copy,
  [genetree_copy](const Node &node) {
  return not genetree_copy->isLeaf(node);
  }).size()
  );

  // Check for a partial contraction
  support_threshold = 0.25;
  genetree_copy = PhyloTreeToolBox::cloneTree(*genetree);

  // get number of branches with the given threshold.
  std::size_t number_of_nodes_to_remove = 0;
  for(auto node : genetree_copy->getAllNodes()) {
    if((not genetree_copy->isLeaf(node))
       and (genetree_copy->getRoot() != node)) {
      auto edge = genetree_copy->getEdgeToFather(node);
      if (edge
          and edge->hasBootstrapValue()
          and (edge->getBootstrapValue() < support_threshold or
               utils::double_equivalence(edge->getBootstrapValue(),
               support_threshold)))
        number_of_nodes_to_remove++;
      else if ((not edge->hasBootstrapValue())
               and CONTRACT_BRANCHES_WITHOUT_SUPPORT)
        number_of_nodes_to_remove++;
    }
  }

  //  partial tree contraction.
  auto nodes_removed = PhyloTreeToolBox::mergeWeakBranches(
      *genetree_copy, support_threshold,
      DEFAULT_STRICT_SUPPORT_THRESHOLDS);
  auto number_of_nodes_removed = nodes_removed.size();
  ASSERT_EQ(number_of_nodes_to_remove, number_of_nodes_removed);

  genetree_copy = nullptr;
}

TEST_F(PhyloTreeToolBox_tests, testNJ) {
  // Test Neighbor joining execution.
  std::shared_ptr<bpp::PhyloTree> genetree = nullptr;
  genetree = IO::readTreeFile(
      polytomy_examples_path + "gene_polytomy.txt");

  treerecs::NeighborJoining nj;
  auto genetree_copy = PhyloTreeToolBox::cloneTree(*genetree);
  nj(*genetree_copy, PhyloTreeToolBox::distanceMatrix(*genetree_copy, true));

  auto original_distances = PhyloTreeToolBox::distanceMatrix(*genetree, true);
  auto solved_distances = PhyloTreeToolBox::distanceMatrix(*genetree_copy,
                                                           true);

  /*
  bool distances_are_similar =
      original_distances.equalTo(
        solved_distances,
        [](const double a, const double b){
          return utils::double_equivalence(a, b);
        }
      );

  if(not distances_are_similar){
    std::cout << "Solved distances:" << std::endl
              << solved_distances << std::endl;
    std::cout << "original distances:" << std::endl
              << original_distances << std::endl;
    ASSERT_TRUE(original_distances.equalTo(
      solved_distances,
      [](const double a, const double b){
        return utils::double_equivalence(a, b);
      }
    ));
  }
  */
}

TEST_F(PhyloTreeToolBox_tests, otherTests) {
  ASSERT_TRUE(PhyloTreeToolBox::node_name_match_with_pattern(
      "par_wel", "par_*", true));
  ASSERT_FALSE(PhyloTreeToolBox::node_name_match_with_pattern(
      "pAr_wel", "par_*", true));
  ASSERT_TRUE(PhyloTreeToolBox::node_name_match_with_pattern(
      "pAr_wel", "par_*", false));
  ASSERT_FALSE(PhyloTreeToolBox::node_name_match_with_pattern(
      "pAr_wel", "*_par", false));
  PhyloTreeToolBox::removeBranchLengths(*genetree);
  for(auto edge : genetree->getAllEdges()) {
    ASSERT_FALSE(edge->hasLength());
  }
}
