import gym
import gym_quantcircuit
import numpy as np

env = gym.make('quantcircuit-v0')

num_qubits = 3
test_goal_state = [0j] * (2**num_qubits - 1) + [1+0j]

env.var_init(num_qubits,
             unitary=False,
             gate_group='pauli',
             connectivity='fully_connected',
             goal_state=test_goal_state)

env.render()

env.plot_connectivity_graph()

action = env.sample()
env.gate_list[action]

env.step(action)

env.render()

for _ in range(5):
    env.step(env.sample())
env.render()

env.reset()
env.render()

# Environment must be created before this step

custom_gate_group = [
    env.qcircuit.x,
    env.qcircuit.h,
    env.qcircuit.t
]

env.set_gate_group('custom',custom_gate_group)

env.gate_list = env._create_gates()

env.gate_list

# Nearest neighbour
nnenv = gym.make('quantcircuit-v0')
nnenv.var_init(5,connectivity='nearest_neighbour')
nnenv.plot_connectivity_graph()

# Fully Connected
fc_env = gym.make('quantcircuit-v0')
fc_env.var_init(5,connectivity='fully_connected')
fc_env.plot_connectivity_graph()

# IBM - limited to 5, 14 or 20 qubits (these are the only current physical architectures)
ibm_env = gym.make('quantcircuit-v0')
ibm_env.var_init(5,connectivity='ibm')
ibm_env.plot_connectivity_graph()

# Custom
custom_conn = np.array([[1,1,1,0],[1,1,1,0],[1,1,1,1],[0,0,1,1]])
custom_env = gym.make('quantcircuit-v0')
custom_env.var_init(4,connectivity='custom',custom_connectivity=custom_conn)
custom_env.plot_connectivity_graph()

num_gates = 2
curriculum, tracker = env.make_curriculum(num_gates)

curriculum

tracker