# -*- coding: utf-8 -*-

"""Tests for simplifying latent variable DAGs."""

import unittest

import networkx as nx

from y0.algorithm.simplify_latent import (
    DEFAULT_SUFFIX,
    iter_latents,
    remove_redundant_latents,
    remove_widow_latents,
    transform_latents_with_parents,
)
from y0.graph import set_latent


class TestSimplify(unittest.TestCase):
    """Tests for the Robin Evans simplification algorithms."""

    def assert_latent_variable_dag_equal(self, expected, actual) -> None:
        """Check two latent variable DAGs are the same."""
        self.assertEqual(
            sorted(expected),
            sorted(actual),
            msg="Nodes are incorrect",
        )
        self.assertEqual(
            dict(expected.nodes.items()),
            dict(actual.nodes.items()),
            msg="Tags are incorrect",
        )
        self.assertEqual(
            sorted(expected.edges()),
            sorted(actual.edges()),
        )

    def test_remove_widows(self):
        """Test simplification 1 - removing widows."""
        graph = nx.DiGraph()
        graph.add_edges_from(
            [
                ("X1", "X2"),
                ("X2", "W"),
                ("U", "X2"),
                ("U", "X3"),
                ("U", "W"),
            ]
        )
        latents = {"U", "W"}
        set_latent(graph, latents)
        self.assertEqual(latents, set(iter_latents(graph)))

        # Apply the simplification
        _, removed = remove_widow_latents(graph)
        self.assertEqual({"W"}, removed)

        expected = nx.DiGraph()
        expected.add_edges_from(
            [
                ("X1", "X2"),
                ("U", "X2"),
                ("U", "X3"),
            ]
        )
        set_latent(expected, "U")

        self.assert_latent_variable_dag_equal(expected, graph)

    def test_transform_latents_with_parents(self):
        """Test simplification 2: latents with parents can be transformed."""
        graph = nx.DiGraph()
        graph.add_edges_from(
            [
                ("X1", "U"),
                ("X2", "U"),
                ("U", "Y1"),
                ("U", "Y2"),
                ("U", "Y3"),
            ]
        )
        set_latent(graph, "U")
        # Apply the simplification
        transform_latents_with_parents(graph)

        expected = nx.DiGraph()
        expected.add_edges_from(
            [
                ("X1", "Y1"),
                ("X1", "Y2"),
                ("X1", "Y3"),
                ("X2", "Y1"),
                ("X2", "Y2"),
                ("X2", "Y3"),
                (f"U{DEFAULT_SUFFIX}", "Y1"),
                (f"U{DEFAULT_SUFFIX}", "Y2"),
                (f"U{DEFAULT_SUFFIX}", "Y3"),
            ]
        )
        set_latent(expected, f"U{DEFAULT_SUFFIX}")

        self.assert_latent_variable_dag_equal(expected, graph)

    def test_remove_redundant_latents(self):
        """Test simplification 3 - remove redundant latents."""
        graph = nx.DiGraph()
        graph.add_edges_from(
            [
                ("U", "X1"),
                ("U", "X2"),
                ("U", "X3"),
                ("W", "X1"),
                ("W", "X2"),
            ]
        )
        set_latent(graph, ["U", "W"])
        # Apply the simplification
        remove_redundant_latents(graph)

        expected = nx.DiGraph()
        expected.add_edges_from(
            [
                ("U", "X1"),
                ("U", "X2"),
                ("U", "X3"),
            ]
        )
        set_latent(expected, ["U"])

        self.assert_latent_variable_dag_equal(expected, graph)
