from graphviz import Digraph
import os


class layer:
    def __init__(self, graph, name, size, color):
        self.name = name
        self.size = size
        self.color = color
        self.graph = graph
    
    def draw():
        pass

class input_layer(layer):
    def __init__(self, graph, size):
        super().__init__(graph, f"Input Layer({size})", size)
        self.graph.node(self.name, shape='circle', style='filled', fillcolor=self.color, label=" ")
        self.graph.attr(label=f'{self.name} Layer({self.size})', fontname='Times New Roman', fontweight='bold', fontsize='36')


def draw_neural_net(input_size, hidden_sizes, num_classes, show_hidden=3):
    g = Digraph('G', filename='neural_network', format='png')
    g.attr(rankdir='LR', size='10,8', nodesep='1', ranksep='2', bgcolor='transparent', dpi='300')

    # Input layer
    with g.subgraph(name='cluster_input') as c:
        c.attr(color='white')
        for i in range(input_size):
            c.node(f'input_{i}', shape='circle', style='filled', fillcolor='darkorange:orange', label=" ")
        c.attr(label=f'Input Layer({input_size})', fontname='Times New Roman', fontweight='bold', fontsize='36')

    # Hidden layers
    previous_layer = 'input'
    previous_layer_size = input_size
    for layer_idx, hidden_size in enumerate(hidden_sizes):
        with g.subgraph(name=f'cluster_hidden_{layer_idx}') as c:
            c.attr(color='white')
            for i in range(show_hidden):
                c.node(f'hidden_{layer_idx}_{i}', shape='circle', style='filled', fillcolor='darkgreen:lightgreen', label=" ")
            if hidden_size > show_hidden * 2:
                c.node(f'ellipsis_{layer_idx}', shape='plaintext', label='...')
            for i in range(hidden_size - show_hidden, hidden_size):
                c.node(f'hidden_{layer_idx}_{i}', shape='circle', style='filled', fillcolor='darkgreen:lightgreen', label=" ")
            c.attr(label=f'Hidden Layer {layer_idx + 1}({hidden_size})', fontname='Times New Roman', fontweight='bold', fontsize='36')

          # Add edges from previous layer to current hidden layer
        if layer_idx == 0:  # Only connect input layer to first hidden layer
            for i in range(previous_layer_size):
                for j in range(show_hidden):
                    g.edge(f'{previous_layer}_{i}', f'hidden_{layer_idx}_{j}')
                for j in range(hidden_size - show_hidden, hidden_size):
                    g.edge(f'{previous_layer}_{i}', f'hidden_{layer_idx}_{j}')
        else:
            for i in range(show_hidden):
                for j in range(show_hidden):
                    g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
                for j in range(hidden_size - show_hidden, hidden_size):
                    g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
            for i in range(hidden_size - show_hidden, hidden_size):
                for j in range(show_hidden):
                    g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')
                for j in range(hidden_size - show_hidden, hidden_size):
                    g.edge(f'hidden_{layer_idx - 1}_{i}', f'hidden_{layer_idx}_{j}')

        previous_layer = f'hidden_{layer_idx}'
        previous_layer_size = hidden_size

    # Output layer
    with g.subgraph(name='cluster_output') as c:
        c.attr(color='white')
        for i in range(num_classes):
            c.node(f'output_{i}', shape='circle', style='filled', fillcolor='darkorange:orange', label=" ")
        c.attr(label=f'Output Layer({num_classes})', fontname='Times New Roman', fontweight='bold', fontsize='36')

    # Add edges from last hidden layer to output layer
    # for i in range(previous_layer_size):
    #     for j in range(num_classes):
    #         g.edge(f'{previous_layer}_{i}', f'output_{j}')
    # # Add edges
    # # Add edges from input to visible hidden nodes
    # for i in range(input_size):
    #     for j in range(show_hidden):
    #         g.edge(f'input_{i}', f'hidden_{j}')
    # for i in range(input_size):
    #     for j in range(hidden_size - show_hidden, hidden_size):
    #         g.edge(f'input_{i}', f'hidden_{j}')

    # # Add edges from visible hidden nodes to output layer
    # for i in range(show_hidden):
    #     for j in range(num_classes):
    #         g.edge(f'hidden_{i}', f'output_{j}')
    # for i in range(hidden_size - show_hidden, hidden_size):
    #     for j in range(num_classes):
    #         g.edge(f'hidden_{i}', f'output_{j}')
    # Add edges from last hidden layer to output layer
    for i in range(show_hidden):
        for j in range(num_classes):
            g.edge(f'{previous_layer}_{i}', f'output_{j}')
    for i in range(previous_layer_size - show_hidden, previous_layer_size):
        for j in range(num_classes):
            g.edge(f'{previous_layer}_{i}', f'output_{j}')

    return g

if __name__ == '__main__':
    g = draw_neural_net(7, [60, 60], 7)
    output_path = g.render(view=False)
    print(output_path)
    os.system(f'explorer.exe neural_network.png')