#!/usr/bin/env python3
# limit: 5000
'''
'''

import networkx as nx
import matplotlib.pyplot as plt
import copy
from networkx.algorithms.isomorphism import is_isomorphic


class Net(nx.DiGraph):
    NTypes = {  "P" : (1, 2, "P"),
                "S" : (2, 2, "S"),
                "B" : (1, 1, "B"),
                "V" : (1, 1, "V"),
                "R" : (1, 1, "R"),
                "C" : (2, 1, "R"),
                "I" : (1, 0, ""),
                "H" : (0, 1, "H"),
             }

    Once = { "C", "S", "P" }

    def __init__(self, *ap, **an):
        super().__init__(*ap, **an)
        if len(self.nodes) == 0:
            self += "HI"

    def __iadd__(self, E):
        A, B = self@E[0], self@E[1]
        self.add_edge(A, B, ntype=f"{self%A}{self%B}")
        return self

    def __matmul__(self, ntype):
        if len(ntype) > 1:
            return ntype
        if not ntype in self.NTypes:
            raise TypeError(f"No {ntype} node type")
        idxes = { self.nodes[h]['idx'] for h in self.nodes }
        idx = (min(set(range(max(idxes)+2)) - idxes)) if idxes else 0
        self.add_node(N := f"{ntype}{idx}", ntype=ntype, idx=idx)
        return N

    def replace(self, node, ntype):
        if not ntype in self.NTypes:
            raise TypeError(f"No {ntype} node type")
        if not node in self:
            raise KeyError(f"Node {node} is not in the graph")
        if self.NTypes[ntype][0] < len(self.in_edges(node)):
            raise IndexError(f"Minimum {len(self.in_edges(node))} inbound networks is allowed when replacing {node}, but {self.NTypes[ntype][0]} is given")
        if self.NTypes[ntype][1] < len(self.out_edges(node)):
            raise IndexError(f"Minimum {len(self.out_edges(node))} outbound networks is allowed when replacing {node}, but {ntype}={self.NTypes[ntype][1]} is given")
        nx.relabel_nodes(self, {node: (N:=f"{ntype}{self.nodes[node]['idx']}")}, copy=False)
        self.nodes[N]['ntype'] = ntype
        for i in range(len(self.in_edges(N)), self.NTypes[ntype][0]):
            self += "H", N
        for i in range(len(self.out_edges(N)), self.NTypes[ntype][1]):
            self += N, "I"
        for a, b in set(self.in_edges(N)) | set(self.out_edges(N)):
            self[a][b]['ntype'] = f"{self%a}{self%b}"
        return self

    def show(self, ax=None, Num=1):

        nodelist = [h for h in self if self%h != "I"]
        #labels = {h: (self.NTypes[self%h][2] if self%h != "I" else "") for h in self}
        if ax is None:
            fig, ax = plt.subplots()

        title = f"{Num}: " + " ".join(f"{a}→{b}" for a, b in sorted(self.edges) if self%b !="I")
        print(f" * {title}")
        ax.set_title(title)
        nx.draw_shell(self, ax=ax, 
                nodelist=nodelist,
                #labels=labels,
                with_labels=True, 
                node_size=600, 
                node_color="white", 
                edgecolors = "green", 
                node_shape="s")

    def __str__(self):
        return " ".join(f"{a}→{b}" for a, b in self.edges)+"\n"+\
                " ".join(f"{n}:{self.nodes[n]}" for n in self.nodes)

    def __mod__(self, node):
        return self.nodes[node]['ntype']

    def __invert__(self):
        return [ self % n for n in self ]

    def __pos__(self):
        return len([n for n in self if self%n != "I"])

    def __neg__(self):
        return [ f"{self%a}{self%b}" for a, b in self.edges ]

def gen(N, mx=5):
    if +N < mx:
        for H in N:
            if N%H in "IH":
                for T in list(N.NTypes)[:-2]:
                    if T not in N.Once & set(~N):
                        yield from gen(copy.deepcopy(N).replace(H, T), mx)
    elif +N == mx:
        yield N

def valid(N):
    C = ~N
    return      ("V" in C or "B" in C) \
            and ("S" in C or "P" in C) \
            and C.count("P")<2 \
            and "VV" not in -N \
            and "VI" not in -N

def like(N, NN):
    chk = lambda a, b: a['ntype'] == b['ntype']
    return set(~N) == set(~NN) or is_isomorphic(N, NN, chk, chk)

def makepool(size):
    Pool = []
    for N in gen(Net(), 5):
        if valid(N):
            if all(not like(N, NN) for NN in Pool):
                Pool.append(N)
    return Pool

def main():
    Pool = makepool(5)
    fig, ax = plt.subplots(3, 5)
    idx = 0, 12, 3, 8, 1, 5, 13, 2, 10, 9, 4, 7, 6, 11
    for n, (i, a) in enumerate(zip(idx, ax.reshape(15))):
        Pool[i].show(a, n+1)
    plt.show()
    return Pool

if __name__ == "__main__":
    main()