// 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 "gtest/gtest.h"
#include <cassert>
#include <iostream>
#include <fstream>
#include <Bpp/Phyl/Io/Newick.h>
#include <treerecs/tools/PhyloTreeToolBox.h>
#include <treerecs/tools/SpeciesGeneMapper.h>
#include <treerecs/tools/IO/IO.h>
#include "examples_paths.h"

using namespace treerecs;

class treeMapping_test: public testing::Test {
protected:

virtual void SetUp() {}

virtual void TearDown() {}

  std::string genetree_str =
      "((a1_a, ((a2_a, a3_a), (a4_a, b1_b))), (c1_c, c2_c));";
  std::string genetree_nhx_str =
      "((a1_a[&&NHX:S=a],((a2_a[&&NHX:S=a],a3_a[&&NHX:S=a]),"
          "(a4_a[&&NHX:S=a],b1_b[&&NHX:S=b]))),"
          "(c1_c[&&NHX:S=c],c2_c[&&NHX:S=c]));";
  std::string speciestree_str = "((Aa, (a, b)), c);";

  //Open and load genes tree
  std::shared_ptr<bpp::PhyloTree> genetree
      = IO::readTreeFile(polytomy_examples_path
                         + "gene_polytomy.txt");

  std::shared_ptr<bpp::PhyloTree> genetree_nhx
      = IO::readTreeFile(polytomy_examples_path
                         + "gene_polytomy.nhx", TextFormat::nhx);

  //Open a load big gene tree
  std::shared_ptr<bpp::PhyloTree> big_genetree
      = IO::readTreeFile(ensembl_examples_path + "phyml_tree.txt");

  //Open and load species tree
  std::shared_ptr<bpp::PhyloTree> speciestree = IO::readTreeFile(
      polytomy_examples_path + "species_polytomy.txt");

  //Open and load big species tree
  std::shared_ptr<bpp::PhyloTree> big_speciestree
      = IO::readTreeFile(ensembl_examples_path + "species_tree.txt");

  std::string mapping_file = polytomy_examples_path + "map.txt";

};

TEST_F(treeMapping_test, testSpeciesGeneMap) {
  // Get species
  Node species_a = PhyloTreeToolBox::getNodeFromName(*speciestree, "a");
  Node species_b = PhyloTreeToolBox::getNodeFromName(*speciestree, "b");
  Node species_c = PhyloTreeToolBox::getNodeFromName(*speciestree, "c");

  // Get leaves
  Node a1_a = PhyloTreeToolBox::getNodeFromName(*genetree, "a1_a");
  Node a2_a = PhyloTreeToolBox::getNodeFromName(*genetree, "a2_a");
  Node b1_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b1_b");
  Node b2_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b2_b");
  Node b3_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b3_b");
  Node c1_c = PhyloTreeToolBox::getNodeFromName(*genetree, "c1_c");

  // We will build a map0 from scratch
  SpeciesGeneMap map0;
  ASSERT_EQ(0, map0.nspecies());

  // Add one gene from species a
  ASSERT_NE(nullptr, species_a);
  ASSERT_NE(nullptr, a1_a);
  map0.addGene(species_a, a1_a);
  ASSERT_EQ(1, map0.nspecies());
  ASSERT_TRUE(map0.speciesHasGene(species_a, a1_a));
  ASSERT_FALSE(map0.speciesHasGene(species_a, a2_a)); //a2_a not added yet
  ASSERT_TRUE(map0.hasGene(a1_a));
  ASSERT_FALSE(map0.hasGene(a2_a));
  ASSERT_EQ(1, map0.getGenes(species_a).size());
  ASSERT_EQ(1, map0.ngenes(species_a));
  ASSERT_EQ(a1_a, map0.getGenes(species_a).front());
  ASSERT_EQ(species_a, map0.getAssociatedSpecies(a1_a));

  // Add one gene from species c
  ASSERT_NE(nullptr, species_c);
  ASSERT_NE(nullptr, c1_c);
  map0.setGene(species_c, c1_c);
  ASSERT_EQ(2, map0.nspecies());
  ASSERT_TRUE(map0.speciesHasGene(species_c, c1_c));
  ASSERT_FALSE(map0.speciesHasGene(species_c, a1_a));
  ASSERT_TRUE(map0.hasGene(c1_c));
  ASSERT_EQ(1, map0.getGenes(species_a).size());
  ASSERT_EQ(1, map0.ngenes(species_c));
  ASSERT_EQ(species_c, map0.getAssociatedSpecies(c1_c));

  // We will build an other map to merge it with the previous one
  SpeciesGeneMap map1;

  std::vector<Node> species_b_genes {b1_b, b2_b, b3_b};
  std::sort(species_b_genes.begin(), species_b_genes.end());

  map1.addGenes(species_b, species_b_genes);
  ASSERT_EQ(1, map1.nspecies());
  ASSERT_EQ(3, map1.ngenes(species_b));

  auto species_b_mapping = map1.getGenes(species_b);
  std::sort(species_b_mapping.begin(), species_b_mapping.end());

  ASSERT_TRUE(utils::comp_all(species_b_genes.begin(), species_b_genes.end(),
                              species_b_mapping.begin(), species_b_mapping.end()
  ));
  //we will add the second gene of species a in this map
  map1.addGene(species_a, a2_a);

  // Then merge the two maps (map0 <- map1)
  map0.addMap(map1);
  ASSERT_EQ(3, map0.nspecies());
  ASSERT_TRUE(map0.hasGene(b2_b));
  ASSERT_TRUE(map0.speciesHasGene(species_b, b1_b));
  ASSERT_TRUE(map0.speciesHasGene(species_c, c1_c));
  ASSERT_EQ(2, map0.ngenes(species_a));
  ASSERT_EQ(3, map0.ngenes(species_b));
  ASSERT_EQ(1, map0.ngenes(species_c));

  std::vector<Node> species_a_genes = {a1_a, a2_a};
  std::sort(species_a_genes.begin(), species_a_genes.end());

  std::vector<Node> species_c_genes = {c1_c};
  std::sort(species_c_genes.begin(), species_c_genes.end());

  auto species_a_mapping = map0.getGenes(species_a);
  std::sort(species_a_mapping.begin(), species_a_mapping.end());

  species_b_mapping = map0.getGenes(species_b);
  std::sort(species_b_mapping.begin(), species_b_mapping.end());

  auto species_c_mapping = map0.getGenes(species_c);
  std::sort(species_c_mapping.begin(), species_c_mapping.end());

  ASSERT_TRUE(utils::comp_all(species_a_genes.begin(), species_a_genes.end(),
                              species_a_mapping.begin(), species_a_mapping.end()
  ));
  ASSERT_TRUE(utils::comp_all(species_b_genes.begin(), species_b_genes.end(),
                              species_b_mapping.begin(), species_b_mapping.end()
  ));
  ASSERT_TRUE(utils::comp_all(species_c_genes.begin(), species_c_genes.end(),
                              species_c_mapping.begin(), species_c_mapping.end()
  ));
  ASSERT_EQ(species_a, map0.getAssociatedSpecies(a2_a));

  // Then merge the two maps (map1 <- map0), with data duplications. We are
  // expected to avoid duplicates in maps
  map1.addMap(map0);

  ASSERT_EQ(3, map1.nspecies());
  ASSERT_EQ(2, map1.ngenes(species_a));
  ASSERT_EQ(3, map1.ngenes(species_b));
  ASSERT_EQ(1, map1.ngenes(species_c));

  species_a_mapping = map1.getGenes(species_a);
  std::sort(species_a_mapping.begin(), species_a_mapping.end());

  species_b_mapping = map1.getGenes(species_b);
  std::sort(species_b_mapping.begin(), species_b_mapping.end());

  species_c_mapping = map1.getGenes(species_c);
  std::sort(species_c_mapping.begin(), species_c_mapping.end());

  ASSERT_TRUE(utils::comp_all(species_a_genes.begin(), species_a_genes.end(),
                              species_a_mapping.begin(), species_a_mapping.end()
  ));
  ASSERT_TRUE(utils::comp_all(species_b_genes.begin(), species_b_genes.end(),
                              species_b_mapping.begin(), species_b_mapping.end()
  ));
  ASSERT_TRUE(utils::comp_all(species_c_genes.begin(), species_c_genes.end(),
                              species_c_mapping.begin(), species_c_mapping.end()
  ));
  ASSERT_TRUE(map0.comp(map1)); //map0 and map1 are similar
}

TEST_F(treeMapping_test, testMapping) {
  // We are going to use species named "a" "b" and "c" to test mapping.
  // But first, we count their number in gene tree.
  auto gene_leaf_names = genetree->getAllLeavesNames();
  auto n_a_species = 0;
  auto n_b_species = 0;
  auto n_c_species = 0;
  for(auto& gene_leaf_name: gene_leaf_names){
    if(gene_leaf_name.find("_a") !=std::string::npos) n_a_species++;
    else if(gene_leaf_name.find("_b") !=std::string::npos) n_b_species++;
    else if(gene_leaf_name.find("_c") !=std::string::npos) n_c_species++;
  }

  // Get node associated to a species name
  auto a_species = PhyloTreeToolBox::getNodeFromName(*speciestree, "a", true);
  auto b_species = PhyloTreeToolBox::getNodeFromName(*speciestree, "b", true);
  auto c_species = PhyloTreeToolBox::getNodeFromName(*speciestree, "c", true);

  //Mapping according to gene names with a separator and a position.
  auto genemap_with_gene_names =
      SpeciesGeneMapper::mapWithGenesNames(*genetree, *speciestree, "_", false);
  ASSERT_TRUE(genemap_with_gene_names.nspecies() > 0);
  ASSERT_EQ(n_a_species, genemap_with_gene_names.getGenes(a_species).size());
  ASSERT_EQ(n_b_species, genemap_with_gene_names.getGenes(b_species).size());
  ASSERT_EQ(n_c_species, genemap_with_gene_names.getGenes(c_species).size());
  ASSERT_EQ(speciestree->getNumberOfLeaves(),
            genemap_with_gene_names.nspecies());

  auto genemapwith_gene_names_str =
      SpeciesGeneMapper::mapWithGenesNames(
          genetree_str, speciestree_str, "_", false);
  ASSERT_EQ(0, genemapwith_gene_names_str.getGenes("A").size());
  ASSERT_EQ(4, genemapwith_gene_names_str.getGenes("a").size());
  ASSERT_EQ(1, genemapwith_gene_names_str.getGenes("b").size());
  ASSERT_EQ(2, genemapwith_gene_names_str.getGenes("c").size());
  ASSERT_EQ(3, genemapwith_gene_names_str.nspecies());

  //Mapping according to trees.
  auto genemap_with_trees
      = SpeciesGeneMapper::mapWithTrees(*genetree, *speciestree);
  ASSERT_TRUE(genemap_with_trees.nspecies() > 0);
  ASSERT_EQ(n_a_species, genemap_with_trees.getGenes(a_species).size());
  ASSERT_EQ(n_b_species, genemap_with_trees.getGenes(b_species).size());
  ASSERT_EQ(n_c_species, genemap_with_trees.getGenes(c_species).size());
  ASSERT_EQ(speciestree->getNumberOfLeaves(), genemap_with_trees.nspecies());

  auto genemap_with_trees_str
      = SpeciesGeneMapper::mapWithTrees(genetree_str, speciestree_str);
  ASSERT_EQ(0, genemap_with_trees_str.getGenes("A").size());
  ASSERT_EQ(4, genemap_with_trees_str.getGenes("a").size());
  ASSERT_EQ(1, genemap_with_trees_str.getGenes("b").size());
  ASSERT_EQ(2, genemap_with_trees_str.getGenes("c").size());
  ASSERT_EQ(3, genemap_with_trees_str.nspecies());

  //Mapping according to NHX tags.
  auto genemap_with_nhx
      = SpeciesGeneMapper::mapWithNhxTags(*genetree_nhx, *speciestree);
  ASSERT_TRUE(genemap_with_nhx.nspecies() > 0);
  ASSERT_EQ(n_a_species, genemap_with_nhx.getGenes(a_species).size());
  ASSERT_EQ(n_b_species, genemap_with_nhx.getGenes(b_species).size());
  ASSERT_EQ(n_c_species, genemap_with_nhx.getGenes(c_species).size());
  ASSERT_EQ(speciestree->getNumberOfLeaves(), genemap_with_nhx.nspecies());

  auto genemap_with_nhx_str
      = SpeciesGeneMapper::mapWithNhxTags(genetree_nhx_str, speciestree_str);
  ASSERT_EQ(0, genemap_with_nhx_str.getGenes("A").size());
  ASSERT_EQ(4, genemap_with_nhx_str.getGenes("a").size());
  ASSERT_EQ(1, genemap_with_nhx_str.getGenes("b").size());
  ASSERT_EQ(2, genemap_with_nhx_str.getGenes("c").size());
  ASSERT_EQ(3, genemap_with_nhx_str.nspecies());

  //Mapping according to a file
  auto genemap_with_file
      = SpeciesGeneMapper::mapFromFile(mapping_file, *speciestree, *genetree);
  ASSERT_TRUE(genemap_with_file.nspecies() > 0);
  ASSERT_EQ(n_a_species, genemap_with_file.getGenes(a_species).size());
  ASSERT_EQ(n_b_species, genemap_with_file.getGenes(b_species).size());
  ASSERT_EQ(n_c_species, genemap_with_file.getGenes(c_species).size());
  ASSERT_EQ(speciestree->getNumberOfLeaves(), genemap_with_file.nspecies());

  auto genemap_with_file_str
      = SpeciesGeneMapper::mapFromFile(genetree_str, speciestree_str, mapping_file);
  ASSERT_EQ(0, genemap_with_file_str.getGenes("A").size());
  ASSERT_EQ(4, genemap_with_file_str.getGenes("a").size());
  ASSERT_EQ(1, genemap_with_file_str.getGenes("b").size());
  ASSERT_EQ(2, genemap_with_file_str.getGenes("c").size());
  ASSERT_EQ(3, genemap_with_file_str.nspecies());

  //Actually, the three genemaps are the same
  ASSERT_TRUE(genemap_with_gene_names == genemap_with_trees);
  ASSERT_TRUE(genemap_with_trees == genemap_with_file);
  ASSERT_TRUE(genemap_with_gene_names == genemap_with_file);

  ASSERT_TRUE(genemapwith_gene_names_str == genemap_with_trees_str);
  ASSERT_TRUE(genemap_with_trees_str == genemap_with_file_str);
  ASSERT_TRUE(genemapwith_gene_names_str == genemap_with_file_str);
  ASSERT_TRUE(genemap_with_nhx_str == genemap_with_trees_str);
}

TEST_F(treeMapping_test, testSpeciesGeneMapper) {

  // First of all generate a map according to a file.
  auto map = SpeciesGeneMapper::loadCompara(
      ensembl_examples_path + "Compara.73.protein.nhx.emf",
      big_speciestree->getAllLeaves(),
      big_genetree->getAllLeaves(), false);

  // All leaves are mapped
  ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map, *big_genetree, *big_speciestree, true));

  // But only leaves
  ASSERT_FALSE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map, *big_genetree, *big_speciestree, false));

  // Then update their mapping to the inner nodes
  SpeciesGeneMapper::updateInnerNodeMapping(
      map, *big_genetree, *big_speciestree);

  // All leaves + internal nodes are mapped
  ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map, *big_genetree, *big_speciestree, false));

  // Create a copy of our map
  auto map_backup(map);

  // Which is the same as the original one
  ASSERT_TRUE(map.comp(map_backup));

  // If we change the root of the tree
  auto actual_root = big_genetree->getRoot();
  auto inner_nodes = big_genetree->getAllInnerNodes();
  big_genetree->rootAt( inner_nodes.at(( inner_nodes.size() + 1 )/2) );

  // Check if nothing changes during rerooting
  ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map, *big_genetree, *big_speciestree, false));

  // But update the map according to the rerooting
  SpeciesGeneMapper::updateInnerNodeMappingAfterRerooting(
      map, *big_genetree, *big_speciestree, actual_root);

  // Nothing changes in gene composition in map
  ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map, *big_genetree, *big_speciestree, false));

  // Now copy and backup are different
  ASSERT_FALSE(map.comp(map_backup));

  // Update backup according to the rerooting but with a "naive" function
  // (visit and reset all mapping)
  SpeciesGeneMapper::updateInnerNodeMapping(
      map_backup, *big_genetree, *big_speciestree);

  // After this update, all genes are mapped.
  ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(
      map_backup, *big_genetree, *big_speciestree, false));

  // And now, these two maps are equivalents
  ASSERT_TRUE(map.comp(map_backup));
}
