// 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 "examples_paths.h"
#include <treerecs/tools/treeReconciliation/TreeReconciliationConductor.h>
#include <treerecs/tools/IO/IO.h>
#include <treerecs/tools/SpeciesGeneMapper.h>

using namespace treerecs;

/*!
 * @class treeReconciliation_test
 * @brief Test TreeReconciliationConductor.
 */
class treeReconciliation_test  : public testing::Test {
protected:
  TreeReconciliationConductor trconductor;

  std::string map_file = multipolytomy_examples_path + "map.txt";
  std::string speciestree_file = multipolytomy_examples_path + "species_tree.txt";
  std::string genetree_file = multipolytomy_examples_path + "gene_tree.txt";
};

TEST_F(treeReconciliation_test, testTreeReconciliation_rearrangment) {
    // Tests gene tree rearrangment.
    auto genetrees = IO::readTreesFile(genetree_file, TextFormat::newick, 2,
                                       true, true, 0, 0);
    std::shared_ptr<bpp::PhyloTree> genetree = genetrees.front();
    auto speciestree = IO::readTreeFile(speciestree_file, TextFormat::newick,
                                        false, false);

    auto map = SpeciesGeneMapper::mapFromFile(map_file, *speciestree,
                                              *genetree);

    std::size_t number_of_leaves_before_rec = genetree->getNumberOfLeaves();

    double dupcost = 2;
    double losscost = 1;
    double threshold = 0.5;
    bool reroot = true;
    std::size_t sample_size = 10;
    bool keepArtificialGenes = false;
    bool computeDistances = true;
    bool printProgression = false;

    auto solutions = trconductor(*genetree, *speciestree, map, 0, false, dupcost,
                                 losscost, threshold, true, reroot, sample_size,
                                 false, false, keepArtificialGenes,
                                 computeDistances, printProgression, false);

    std::list<double> rearrangment_costs;

    for (auto &solution: solutions) {
        std::shared_ptr<const bpp::PhyloTree> generated_genetree
            = solution.genetree();

        rearrangment_costs.push_back(solution.cost());

        // Check if there is no monofurcation.
        ASSERT_FALSE(PhyloTreeToolBox::hasMonofurcations(*generated_genetree));

        // Check if there is no artificial gene.
        ASSERT_FALSE(PhyloTreeToolBox::hasArtificialGenes(*generated_genetree));

        // Check if there is no multifurcation.
        ASSERT_EQ(0,
                  PhyloTreeToolBox::findPolytomies(*generated_genetree).size());

        // Check if the number of leaves is the same as before
        ASSERT_EQ(number_of_leaves_before_rec,
                  generated_genetree->getNumberOfLeaves());

        // Check if the number of solutions is correct, as asked
        ASSERT_EQ(sample_size, solutions.size());

        // Check if edges have length equal to or inferior than 0.0.
        auto edges = generated_genetree->getAllEdges();
        ASSERT_FALSE(utils::contains(edges.begin(), edges.end(),
                                     [](const std::shared_ptr<bpp::PhyloBranch> &edge) {
                                       return utils::double_equal_or_inferior(
                                           edge->getLength(), 0.0);
                                     }
        ));
    }

    // Then, check if all solutions has the same minimal cost :
    double minimal_cost_reference = rearrangment_costs.front();
    for(auto cost: rearrangment_costs) {
        ASSERT_TRUE(cost > 0);
        ASSERT_TRUE(utils::double_equivalence(cost, minimal_cost_reference));
    }
}

TEST_F(treeReconciliation_test, testTreeReconciliation_full_reconciliation) {
    // Test with addition of gene losses in the generated gene trees.
    auto genetrees = IO::readTreesFile(genetree_file, TextFormat::newick, 2,
                                       true, true, 0, 0);
    std::shared_ptr<bpp::PhyloTree> genetree = genetrees.front();
    auto speciestree = IO::readTreeFile(speciestree_file, TextFormat::newick,
                                        false, false);

    auto map = SpeciesGeneMapper::mapFromFile(map_file, *speciestree,
                                              *genetree);

    std::size_t number_of_leaves_before_rec = genetree->getNumberOfLeaves();

    double dupcost = 2;
    double losscost = 1;
    double threshold = 0.5;
    bool reroot = true;
    std::size_t sample_size = 10;
    bool keepArtificialGenes = true; //
    bool computeDistances = false; //
    bool printProgression = false;

    auto solutions = trconductor(*genetree, *speciestree, map,
                            nullptr, false,
                            dupcost, losscost, threshold, true, reroot,
                            sample_size, false, false,
                            keepArtificialGenes, computeDistances,
                            printProgression, false);

    std::list<double> reconciliation_costs;

    for (auto &solution: solutions) {
        std::shared_ptr<const bpp::PhyloTree> generated_genetree
            = solution.genetree();

        reconciliation_costs.push_back(solution.cost());

        // Check if there is no monofurcation.
        ASSERT_FALSE(PhyloTreeToolBox::hasMonofurcations(*generated_genetree));

        // Check if there is artificial genes.
        ASSERT_TRUE(PhyloTreeToolBox::hasArtificialGenes(*generated_genetree));

        // Check if there is no multifurcation.
        ASSERT_EQ(0,
                  PhyloTreeToolBox::findPolytomies(*generated_genetree).size());

        // Check if the number of leaves is the same as before
        ASSERT_NE(number_of_leaves_before_rec,
                  generated_genetree->getNumberOfLeaves());

        // Check if the number of solutions is correct, as asked
        ASSERT_EQ(solutions.size(), sample_size);
    }

    // Test if each solution share the same reconciliation cost.
    double cost_reference = reconciliation_costs.front();
    for(auto rec_cost: reconciliation_costs) {
      ASSERT_TRUE(utils::double_equivalence(cost_reference, rec_cost));
    }

}

TEST_F(treeReconciliation_test, testTreeReconciliation_cost_estimation) {
// Test with addition of gene losses in the generated gene trees.
    auto genetrees = IO::readTreesFile(ensembl_examples_path + "phyml_tree.txt", TextFormat::newick, 0,
                                       true, true, 0, 0);
    std::shared_ptr<bpp::PhyloTree> genetree = genetrees.front();
    auto speciestree = IO::readTreeFile(ensembl_examples_path + "species_tree.txt", TextFormat::newick,
                                        false, false);

    auto map = SpeciesGeneMapper::mapFromFile(ensembl_examples_path + "phyml_tree.smap", *speciestree,
                                              *genetree);

    ASSERT_TRUE(SpeciesGeneMapper::checkIfAllGenesAreMapped(map, *genetree, *speciestree, true));

    double dupcost = 1.0;
    double losscost = 1.0;
    double threshold = 0.5;
    bool reroot = false;
    std::size_t n = 3;

    auto solutions = trconductor(*genetree, *speciestree, map,
                                 nullptr, true, dupcost, losscost,
                                 threshold, true, reroot, n,
                                 false, false, false,
                                 false, false, false);

    ASSERT_FALSE(utils::double_equivalence(trconductor.estimated_costs().first, 0.0));
    ASSERT_FALSE(utils::double_equivalence(trconductor.estimated_costs().second, 0.0));
    ASSERT_EQ(solutions.size(), n);

    // Check if total score is equal to ndup * dupcost + nloss * losscost
    ASSERT_EQ(trconductor.estimated_costs().first * solutions.front().ndup() +
              trconductor.estimated_costs().second * solutions.front().nloss(),
              solutions.front().cost());

    //ASSERT_EQ(0.824383, trconductor.estimated_costs().first);
    ASSERT_TRUE(utils::double_equivalence(0.824383,
                                          trconductor.estimated_costs().first,
                                          10e-6));
    ASSERT_EQ(121, solutions.front().ndup());

    //ASSERT_EQ(0.175617, trconductor.estimated_costs().second);
    ASSERT_TRUE(utils::double_equivalence(0.175617,
                                          trconductor.estimated_costs().second,
                                          10e-6));
    ASSERT_EQ(568, solutions.front().nloss());
}
