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

from pyqcs import State, H, X, S, CZ
from pyqcs.graph.state import GraphState
from pyqcs.util.random_circuits import random_circuit

def S_with_extra_arg(act, i):
    return S(act)

np.random.seed(0xdeadbeef)

circuit = random_circuit(50, 700, X, H, S_with_extra_arg, CZ)

state = circuit * GraphState.new_plus_state(50)

vops, edges = state._g_state.to_lists()

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_edges_str + "\n}"

print(dot_str)