from typing import List, Dict

from generator_relaxed.customtests import CustomTests

from generator_relaxed.tree.infoset import Infoset

from misc.game_structures import PlayerId, PID_TO_CARDINAL


class Treeplex:
    """
    A treeplex is a tree of infosets
    """
    def __init__(self, player: PlayerId, infosets: List[Infoset]):
        self.num_sequences = Treeplex.validate(infosets)
        self.player = player
        self.infosets = infosets

    def num_infosets(self) -> int:
        return len(self.infosets)

    def has_sequence(self, sequence: int) -> bool:
        return sequence < self.num_sequences

    def empty_sequence_id(self) -> int:
        return self.num_sequences - 1

    def __str__(self):
        s = f"Treeplex of {PID_TO_CARDINAL[self.player]}, number of sequences = {self.num_sequences}\n"
        for infoset in self.infosets:
            s += "\t" + str(infoset) + "\n"
        return s

    def short_str(self) -> str:
        s = f"=== {PID_TO_CARDINAL[self.player]}\n"
        for infoset in self.infosets:
            s += "{}\n".format(infoset.short_str())
        return s

    @staticmethod
    def validate(infosets: List[Infoset]) -> int:
        """
        1. The DFS visit of the tree induced by the infosets coincides with the list of infosets
        2. The sequence numbers are laid out in order
        """
        # maps each sequence id with the infosets generated by it -> sibilings
        adjacency_map: Dict[int, List[Infoset]] = {}
        empty_sequence_id = 0

        # Point 1
        for infoset in infosets:
            ps_id = infoset.parent_sequence_id
            if ps_id not in adjacency_map.keys():
                adjacency_map[ps_id] = []
            adjacency_map[ps_id].append(infoset)

            if ps_id > empty_sequence_id:
                empty_sequence_id = ps_id

        expected_infoset_order: List[Infoset] = []

        def visit_tree(root: int, adj_map: Dict[int, List[Infoset]], expected_infoset_order: List[Infoset]):
            if root in adj_map.keys():
                next_infosets = adj_map[root]
                for next_infoset in next_infosets:
                    for next_sequence_id in range(next_infoset.start_sequence_id, next_infoset.end_sequence_id + 1):
                        visit_tree(next_sequence_id, adj_map, expected_infoset_order)
                    expected_infoset_order.append(next_infoset)

        visit_tree(empty_sequence_id, adjacency_map, expected_infoset_order)

        # asserting that infosets and expected_infoset_order are the same list (order counts)
        assert len(infosets) == len(expected_infoset_order)
        test_instance = CustomTests()
        test_instance.assert_list_equal(infosets, expected_infoset_order)

        # Point 2
        expected_next_sequence_id = 0
        for infoset in infosets:
            test_instance.assert_equal(infoset.start_sequence_id, expected_next_sequence_id)
            assert infoset.start_sequence_id <= infoset.end_sequence_id
            assert infoset.parent_sequence_id > infoset.end_sequence_id

            expected_next_sequence_id = infoset.end_sequence_id + 1
        test_instance.assert_equal(empty_sequence_id, expected_next_sequence_id)

        return empty_sequence_id + 1
