"""
Tests for cassiopeia/tools/parameter_estimators.py
"""

import unittest

import networkx as nx
import numpy as np
import pandas as pd

import cassiopeia as cas
from cassiopeia.tools import parameter_estimators
from cassiopeia.mixins import ParameterEstimateError, ParameterEstimateWarning


class TestCassiopeiaTree(unittest.TestCase):
    def setUp(self):
        # A small network to test likelihood calculation
        small_net = nx.DiGraph()
        small_net.add_edges_from(
            [
                ("node5", "node0"),
                ("node5", "node1"),
                ("node6", "node2"),
                ("node6", "node3"),
                ("node6", "node4"),
                ("node7", "node5"),
                ("node7", "node6"),
            ]
        )
        self.small_net = small_net

        cm1 = pd.DataFrame.from_dict(
            {
                "node0": [0, -1, -1],
                "node1": [1, 1, -1],
                "node2": [1, -1, -1],
                "node3": [1, -1, -1],
                "node4": [1, -1, -1],
            },
            orient="index",
        )
        priors1 = {0: {1: 1}, 1: {1: 1}, 2: {1: 1}}
        self.discrete_tree = cas.data.CassiopeiaTree(
            tree=small_net, character_matrix=cm1, priors=priors1
        )

        cm2 = pd.DataFrame.from_dict(
            {
                "node0": [1, 0],
                "node1": [1, 1],
                "node2": [2, 3],
                "node3": [-1, 2],
                "node4": [-1, 1],
            },
            orient="index",
        )
        priors2 = {
            0: {1: 0.2, 2: 0.7, 3: 0.1},
            1: {1: 0.2, 2: 0.7, 3: 0.1},
            2: {1: 0.2, 2: 0.7, 3: 0.1},
        }
        self.continuous_tree = cas.data.CassiopeiaTree(
            tree=small_net, character_matrix=cm2, priors=priors2
        )
        self.continuous_tree.set_branch_length("node5", "node0", 1.5)
        self.continuous_tree.set_branch_length("node6", "node3", 2)

    def test_proportions(self):
        prop_mut = parameter_estimators.get_proportion_of_mutation(
            self.discrete_tree
        )
        prop_missing = parameter_estimators.get_proportion_of_missing_data(
            self.discrete_tree
        )

        self.assertEqual(prop_mut, 5 / 6)
        self.assertEqual(prop_missing, 0.6)

        prop_mut = parameter_estimators.get_proportion_of_mutation(
            self.continuous_tree
        )
        prop_missing = parameter_estimators.get_proportion_of_missing_data(
            self.continuous_tree
        )

        self.assertEqual(prop_mut, 7 / 8)
        self.assertEqual(prop_missing, 0.2)

    def test_estimate_mutation_rate(self):
        mut_rate = parameter_estimators.estimate_mutation_rate(
            self.discrete_tree, continuous=False
        )
        self.assertTrue(np.isclose(mut_rate, 0.44967879185089554))

        mut_rate = parameter_estimators.estimate_mutation_rate(
            self.discrete_tree,
            continuous=False,
            assume_root_implicit_branch=False,
        )
        self.assertTrue(np.isclose(mut_rate, 0.5917517095361371))

        mut_rate = parameter_estimators.estimate_mutation_rate(
            self.continuous_tree, continuous=True
        )
        self.assertTrue(np.isclose(mut_rate, 0.5917110077950752))

        mut_rate = parameter_estimators.estimate_mutation_rate(
            self.continuous_tree,
            continuous=True,
            assume_root_implicit_branch=False,
        )
        self.assertTrue(np.isclose(mut_rate, 0.9041050181216678))

    def test_estimate_missing_data_bad_cases(self):
        with self.assertRaises(ParameterEstimateError):
            temp = parameter_estimators.estimate_missing_data_rates(
                self.discrete_tree, continuous=False
            )

        with self.assertRaises(ParameterEstimateError):
            temp = parameter_estimators.estimate_missing_data_rates(
                self.discrete_tree,
                continuous=False,
                heritable_missing_rate=0.25,
                stochastic_missing_probability=0.2,
            )

        with self.assertRaises(ParameterEstimateError):
            self.discrete_tree.parameters["heritable_missing_rate"] = 0.25
            self.discrete_tree.parameters[
                "stochastic_missing_probability"
            ] = 0.2
            temp = parameter_estimators.estimate_missing_data_rates(
                self.discrete_tree, continuous=False
            )

        with self.assertRaises(ParameterEstimateWarning):
            self.discrete_tree.reset_parameters()
            self.discrete_tree.parameters["heritable_missing_rate"] = 0.5
            temp = parameter_estimators.estimate_missing_data_rates(
                self.discrete_tree, continuous=False
            )

        with self.assertRaises(ParameterEstimateWarning):
            self.continuous_tree.parameters[
                "stochastic_missing_probability"
            ] = 0.9
            temp = parameter_estimators.estimate_missing_data_rates(
                self.continuous_tree, continuous=True
            )

    def test_estimate_stochastic_missing_data_probability(self):
        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree, continuous=False, heritable_missing_rate=0.25
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 0.0518518518518518))

        self.discrete_tree.parameters["heritable_missing_rate"] = 0.25
        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree, continuous=False
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 0.0518518518518518))

        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree,
            continuous=False,
            assume_root_implicit_branch=False,
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 13 / 45))

        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree, continuous=True, heritable_missing_rate=0.05
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 0.046322071416968195))

        self.continuous_tree.parameters["heritable_missing_rate"] = 0.05
        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree, continuous=True
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 0.046322071416968195))

        s_missing_prob = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree,
            continuous=True,
            assume_root_implicit_branch=False,
        )[0]
        self.assertTrue(np.isclose(s_missing_prob, 0.10250124994244929))

    def test_estimate_heritable_missing_data_rate(self):
        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree,
            continuous=False,
            stochastic_missing_probability=0.12,
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.23111904017137075))

        self.discrete_tree.parameters["stochastic_missing_probability"] = 0.2
        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree, continuous=False
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.2062994740159002))

        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.discrete_tree,
            continuous=False,
            assume_root_implicit_branch=False,
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.2928932188134524))

        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree,
            continuous=True,
            stochastic_missing_probability=0.04,
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.05188011778689765))

        self.continuous_tree.parameters["stochastic_missing_probability"] = 0.1
        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree, continuous=True
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.0335154979510034))

        h_missing_rate = parameter_estimators.estimate_missing_data_rates(
            self.continuous_tree,
            continuous=True,
            assume_root_implicit_branch=False,
        )[1]
        self.assertTrue(np.isclose(h_missing_rate, 0.05121001550277538))


if __name__ == "__main__":
    unittest.main()
