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

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

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

// Include containers
#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 algorithms as Polytomysolver and ProfileNJ.
#include <treerecs/algorithms/Polytomysolver.h>
#include <treerecs/algorithms/ProfileNJ.h>

// Examples files path
#include "examples_paths.h"

using namespace treerecs;

class profilenj_test : public testing::Test {
protected:
  virtual void SetUp() {
  }

  virtual void TearDown() {
  }

  ProfileNJ pnj;
};

TEST_F(profilenj_test, profileNJTest) {
    bool verbose = false;

    // Load trees
    auto genetree = IO::readTreeFile(polytomy_examples_path + "gene_polytomy.txt", TextFormat::newick, true, verbose);
    auto speciestree = IO::readTreeFile(polytomy_examples_path + "species_polytomy.txt", TextFormat::newick, false,
                                        verbose);

    // Get species nodes
    auto species_a = PhyloTreeToolBox::getNodeFromName(*speciestree, "a");
    auto species_b = PhyloTreeToolBox::getNodeFromName(*speciestree, "b");
    auto species_c = PhyloTreeToolBox::getNodeFromName(*speciestree, "c");
    auto species_d = speciestree->getFather(species_a);
    auto species_e = speciestree->getFather(species_d);

    // Get gene nodes
    auto a1_a = PhyloTreeToolBox::getNodeFromName(*genetree, "a1_a");
    auto a2_a = PhyloTreeToolBox::getNodeFromName(*genetree, "a2_a");
    auto b1_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b1_b");
    auto b2_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b2_b");
    auto b3_b = PhyloTreeToolBox::getNodeFromName(*genetree, "b3_b");
    auto c1_a = PhyloTreeToolBox::getNodeFromName(*genetree, "c1_c");

    // Build SpeciesGeneMap
    SpeciesGeneMap genemap = SpeciesGeneMapper::mapWithTrees(*genetree, *speciestree);
    ASSERT_EQ(genemap.getGenes().size(), genetree->getNumberOfLeaves());
    ASSERT_EQ(genemap.getSpecies().size(), speciestree->getNumberOfLeaves());

    // Build Distances Matrix
  DistanceMatrix D = IO::readDistanceMatrix(
        polytomy_examples_path + "dmatrix.txt", *genetree, verbose);
    D.addRow(std::vector<double>(D.ncol(), 0.0), genetree->getRoot());
    D.addCol(std::vector<double>(D.ncol()+1, 0.0), genetree->getRoot());

    // Prune tree and create a new map accordingly
  PhyloTreeToolBox::pruneTree_v0(*speciestree, genemap, false);
    genemap = SpeciesGeneMapper::mapWithTrees(*genetree, *speciestree);
    genemap.deleteSpecies(speciestree->getRoot());

    // Use Polytomysolver's algorithm
    Polytomysolver solver;
    auto cost_table = solver.computeCostTable(*speciestree,
                                              genemap, 1, 1, verbose);
    ASSERT_EQ(speciestree->getNumberOfNodes(), cost_table.nrow());
    ASSERT_EQ(5, cost_table.ncol());
    ASSERT_EQ(1, cost_table(species_a, 1).value);
    ASSERT_EQ(0, cost_table(species_a, 2).value);
    ASSERT_EQ(1, cost_table(species_a, 3).value);
    ASSERT_EQ(2, cost_table(species_a, 4).value);
    ASSERT_EQ(2, cost_table(species_b, 1).value);
    ASSERT_EQ(1, cost_table(species_b, 2).value);
    ASSERT_EQ(0, cost_table(species_b, 3).value);
    ASSERT_EQ(1, cost_table(species_b, 4).value);
    ASSERT_EQ(0, cost_table(species_c, 1).value);
    ASSERT_EQ(1, cost_table(species_c, 2).value);
    ASSERT_EQ(2, cost_table(species_c, 3).value);
    ASSERT_EQ(3, cost_table(species_c, 4).value);
    ASSERT_EQ(2, cost_table(species_d, 1).value);
    ASSERT_EQ(1, cost_table(species_d, 2).value);
    ASSERT_EQ(1, cost_table(species_d, 3).value);
    ASSERT_EQ(2, cost_table(species_d, 4).value);
    ASSERT_EQ(2, cost_table(species_e, 1).value);
    ASSERT_EQ(2, cost_table(species_e, 2).value);
    ASSERT_EQ(3, cost_table(species_e, 3).value);
    ASSERT_EQ(4, cost_table(species_e, 4).value);

    // Test the count vector V
    auto V = solver.getCountVector(*speciestree, cost_table, genemap);
    ASSERT_EQ(2, V.at(species_a));
    ASSERT_EQ(2, V.at(species_b));
    ASSERT_EQ(1, V.at(species_c));
    ASSERT_EQ(1, V.at(species_d));
    ASSERT_EQ(1, V.at(species_e));

    // Apply the ProfileNJ algorithm
    pnj(*genetree, genetree->getRoot(), *speciestree, D, V, genemap, true, verbose);

    std::shared_ptr<bpp::PhyloTree>solution = IO::readTreeFile(
        polytomy_examples_path + "solution_polytomy.txt");

    ASSERT_TRUE(PhyloTreeToolBox::compare_trees_with_days_algorithm(*solution, *genetree));
}
