import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np

from qiskit import *

from qiskit.tools.visualization import plot_histogram
from IPython.display import display, Math, Latex
from qiskit.visualization import plot_state_qsphere

def BuildBell(x, y):
    U = QuantumCircuit(2)
    
    U.h(0)
    U.cx(0, 1)
    
    if x == 1:
        U.cz(0, 1)
    if y == 1:
        U.x(1)
        
    U = U.to_gate()
    U.name = 'Build Bell'
    #ctl_U = U.control() make it a controlled gate
    
    return U

backend = BasicAer.get_backend('statevector_simulator')

n = 2

#bell state 00
[x, y] = [0, 0]
b00 = QuantumCircuit(n, n)
b00.append(BuildBell(x, y), range(n))

#bell state 01
[x, y] = [0, 1]
b01 = QuantumCircuit(n, n)
b01.append(BuildBell(x, y), range(n))

#bell state 10
[x, y] = [1, 0]
b10 = QuantumCircuit(n, n)
b10.append(BuildBell(x, y), range(n))

#bell state 11
[x, y] = [1, 1]
b11 = QuantumCircuit(n, n)
b11.append(BuildBell(x, y), range(n))

bqs00 = execute(b00, backend).result()
bqs01 = execute(b01, backend).result()
bqs10 = execute(b10, backend).result()
bqs11 = execute(b11, backend).result()

#GENERAL CIRCUIT b_00
bxy = QuantumCircuit(n, n)
bxy.h(0)
bxy.cx(0, 1)
bxy.measure(range(n), range(n))

backend = BasicAer.get_backend('qasm_simulator')
atp = 1024
res = execute(bxy, backend=backend, shots=atp).result()
ans = res.get_counts()

bxy.draw('mpl')

plot_histogram(ans)

plot_state_qsphere(bqs00.get_statevector(b00)) #bell state 00 qsphere

plot_state_qsphere(bqs01.get_statevector(b01)) #bell state 01 qsphere

plot_state_qsphere(bqs10.get_statevector(b10)) #bell state 10 qsphere

plot_state_qsphere(bqs11.get_statevector(b11)) #bell state 11 qsphere