from collections import deque
import matplotlib.pyplot as plt
import numpy as np
import json

from pyqcs import State, H, X, S, CZ, list_to_circuit
from pyqcs.graph.state import GraphState


circuit_CZ = list_to_circuit([CZ(0, i) for i in range(1, 5)])
circuit_H = list_to_circuit([H(i) for i in range(1, 5)])
circuit = circuit_CZ | circuit_H

state = circuit * GraphState.new_plus_state(5)

vops, edges = state._g_state.to_lists()

VOP_strs = {0: "H", 2: "I"}

dot_vertex_str = "\n".join((f"{i} [label=\"{i}, VOP = {VOP_strs[v]}\"]" for i,v in enumerate(vops)))

handled_edges = set()
dot_edges = deque()

for i, ngbhd in enumerate(edges):
    for j in ngbhd:
        if((i,j) not in handled_edges):
            dot_edges.append(f"{i} -- {j}")
            handled_edges |= {(i,j), (j,i)}

dot_edges_str = "\n".join(dot_edges)

dot_str = "graph graphical_state{\n" + dot_vertex_str + "\n" + dot_edges_str + "\n}"

print(dot_str)