import sys
sys.path.append('/Users/kirais/Documents/GitHub/Qiskit-for-GameDev/pygame/')

from qgame import CircuitGridModel, CircuitGridNode
from qgame import circuit_node_types as node_types

from qiskit import QuantumCircuit

circuit_grid_model = CircuitGridModel(qubit_count=3, circuit_depth=6)
print(circuit_grid_model)

qasm_str = circuit_grid_model.create_qasm_for_circuit()
print(qasm_str)

circuit = QuantumCircuit.from_qasm_str(qasm_str)
circuit.draw()

print(CircuitGridNode(node_types.X))

circuit_grid_model.set_node(qubit_index=1,depth_index=0,circuit_grid_node=CircuitGridNode(node_types.X))
print(circuit_grid_model)

qasm_str = circuit_grid_model.create_qasm_for_circuit()
print(qasm_str)

circuit = QuantumCircuit.from_qasm_str(qasm_str)
circuit.draw()

circuit_grid_model.set_node(0,1,CircuitGridNode(node_types.H))
circuit_grid_model.set_node(2,1,CircuitGridNode(node_types.H))
print(circuit_grid_model)

qasm_str = circuit_grid_model.create_qasm_for_circuit()
print(qasm_str)
circuit = QuantumCircuit.from_qasm_str(qasm_str)
circuit.draw()

from qiskit import BasicAer, execute, ClassicalRegister

def paddle_before_measurement(circuit, qubit_num, shot_num):
        backend_sv_sim = BasicAer.get_backend('statevector_simulator')
        job_sim = execute(circuit, backend_sv_sim, shots=shot_num)
        result_sim = job_sim.result()
        state = result_sim.get_statevector(circuit, decimals=3)
        probability = (state*state).real
        return probability

probability = paddle_before_measurement(circuit, 3, 1000)
print(probability)

from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
plt.figure(figsize=(20,1))
ax = plt.gca()
for i in range(8):
    ax.add_patch(Rectangle((i/8, 0), 1/8, 1, alpha=probability[i]+0.05, fc='k', ec='k'))

# add statvector labels
ax.text(0/8+1/40,-1,r"$\left|000\right\rangle$", fontsize=36)
ax.text(1/8+1/40,-1,r"$\left|001\right\rangle$", fontsize=36)
ax.text(2/8+1/40,-1,r"$\left|010\right\rangle$", fontsize=36)
ax.text(3/8+1/40,-1,r"$\left|011\right\rangle$", fontsize=36)
ax.text(4/8+1/40,-1,r"$\left|100\right\rangle$", fontsize=36)
ax.text(5/8+1/40,-1,r"$\left|101\right\rangle$", fontsize=36)
ax.text(6/8+1/40,-1,r"$\left|110\right\rangle$", fontsize=36)
ax.text(7/8+1/40,-1,r"$\left|110\right\rangle$", fontsize=36)
plt.axis('off')
plt.show()

from copy import deepcopy
def paddle_after_measurement(circuit, qubit_num, shot_num):
        backend_sv_sim = BasicAer.get_backend('qasm_simulator')
        circuit.measure(circuit.qregs[0], circuit.cregs[0])
        job_sim = execute(circuit, backend_sv_sim, shots=shot_num)
        result_sim = job_sim.result()
        counts = result_sim.get_counts(circuit)
        
        return int(list(counts.keys())[0], 2)

result = paddle_after_measurement(circuit, 3, 100)
print(result)

from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
plt.figure(figsize=(20,1))
ax = plt.gca()
for i in range(8):
    if i == result:
        ax.add_patch(Rectangle((i/8, 0), 1/8, 1, alpha=1, fc='k', ec='k'))
    else:
        ax.add_patch(Rectangle((i/8, 0), 1/8, 1, alpha=0, fc='k', ec='k'))

# add statvector labels
ax.text(0/8+1/40,-1,r"$\left|000\right\rangle$", fontsize=36)
ax.text(1/8+1/40,-1,r"$\left|001\right\rangle$", fontsize=36)
ax.text(2/8+1/40,-1,r"$\left|010\right\rangle$", fontsize=36)
ax.text(3/8+1/40,-1,r"$\left|011\right\rangle$", fontsize=36)
ax.text(4/8+1/40,-1,r"$\left|100\right\rangle$", fontsize=36)
ax.text(5/8+1/40,-1,r"$\left|101\right\rangle$", fontsize=36)
ax.text(6/8+1/40,-1,r"$\left|110\right\rangle$", fontsize=36)
ax.text(7/8+1/40,-1,r"$\left|110\right\rangle$", fontsize=36)
plt.axis('off')
plt.show()

