from itertools import product, combinations, cycle
from subprocess import run
import networkx as nx
import matplotlib.pyplot as plt


class AcesEightsPuzzle:
    def __init__(self, names=None, draw=None, out=True):
        if names is None:
            names = ["Alice", "Bob", "Carl"]
        elif len(names) != 3 and any(not isinstance(name, str) for name in names):
            raise ValueError("Wrong agent naming!\n")

        self.out = out
        self.agents = names
        self.agent_colors = {names[0]: "red", names[1]: "green", names[2]: "blue"}
        self.G = nx.Graph()
        self.G.add_nodes_from(["".join(var)
                               for var in list(product(["A", "E", "M"], repeat=3))
                               if ((var.count('A') == 2 or var.count('E') == 2) and
                                   var.count('M') == 0) or (var.count('A') < 2 and
                                                            var.count('E') < 2)])
        if draw is None:
            self.draw = "AEA"
        elif not isinstance(draw, str) or draw.upper() not in self.G.nodes:
            raise ValueError("Incorrect draw given!\n")
        else:
            self.draw = draw.upper()

        self.G.add_edges_from([(edge[0], edge[1], {'color': 'red'})
                               for edge in combinations(self.G.nodes, 2)
                               if edge[0][1:] == edge[1][1:]])

        self.G.add_edges_from([(edge[0], edge[1], {'color': 'green'})
                               for edge in combinations(self.G.nodes, 2)
                               if edge[0][0] == edge[1][0] and edge[0][-1] == edge[1][-1]])

        self.G.add_edges_from([(edge[0], edge[1], {'color': 'blue'})
                               for edge in combinations(self.G.nodes, 2)
                               if edge[0][:-1] == edge[1][:-1]])
        if out:
            print(f'{10*"="}Puzzle start{10*"="}\nAgents: {names} (ordered)\nDraw: {self.draw}')
            print(32*"=")

    def agent_2_color(self, agent):
        # returns assigned color for agent's alias
        if agent not in self.agents:
            raise ValueError("Invalid agent entered!\n")
        return self.agent_colors.get(agent)

    def remove_colored(self, agent):
        # removes states where the agent could be sure of his draw
        agent_color = self.agent_2_color(agent.capitalize())
        if agent_color != "red" and agent_color != "green" and agent_color != "blue":
            raise ValueError("Weird color!\n")
        removed_nodes = []
        for node in self.G.nodes():
            if not any(neighbour["color"] == agent_color for neighbour in self.G[node].values()):
                removed_nodes.append(node)
        for node in removed_nodes:
            self.G.remove_node(node)
        return removed_nodes

    def can_agent_tell(self, agent):
        # checks whether the agent could know his draw
        agent_color = self.agent_2_color(agent.capitalize())
        return not any([p['color'] == agent_color for p in self.G[self.draw].values()])

    def play(self):
        # simulates the game by browsing the agents
        pool = cycle(self.agents)
        agent = next(pool)
        steps = 1
        while not self.can_agent_tell(agent):
            removed = self.remove_colored(agent)
            if self.out:
                print(f'{agent} cannot tell!')
                print(f'Removed {removed}.')
            agent = next(pool)
            steps += 1

        if self.out:
            print(f'\nNow {agent} can tell that the draw is {self.draw}.\n{agent} wins!')
            print(f'{10*"="}Puzzle ends{11*"="}')
        return steps, agent

    def visualize(self, export=False):
        # visualises the current puzzle state
        # every agent's knowledge is shown with different color
        colors = nx.get_edge_attributes(self.G, 'color').values()
        nx.draw(self.G, nx.shell_layout(self.G), edge_color=colors, with_labels=True)
        plt.axis("off")
        if export:
            filename = input("filename :")
            plt.savefig(filename + ".svg", format="SVG", transparent=True)
            command = r""""C:\Program Files\Inkscape\bin\inkscape.exe" --export-type="pdf" """
            command += filename + """.svg"""
            run(command)
            plt.close()
        else:
            plt.show()


if __name__ == "__main__":
    Draw = input("Enter draw: ")
    t = AcesEightsPuzzle(draw=Draw)
    t.play()
    t.visualize()