import os
from qiskit import *
import qiskit.tools.visualization as qt
from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
from qiskit.quantum_info import DensityMatrix
from qiskit.visualization import plot_state_city
from numpy import pi
import matplotlib.pyplot as plt
%matplotlib inline

qreg_q = QuantumRegister(2, 'q')
creg_c = ClassicalRegister(2, 'c')

circuits = []
for i in range(0, 4):
    circuits.append(QuantumCircuit(qreg_q, creg_c))

def entlanging(circuit, qreg_q):
    circuit.barrier()
    circuit.h(qreg_q[0])
    circuit.cnot(qreg_q[0], qreg_q[1])
    circuit.barrier()

# 00 state
circuits[0].reset(qreg_q[0])
circuits[0].reset(qreg_q[1])
entlanging(circuits[0], qreg_q)

# 01 state
circuits[1].reset(qreg_q[0])
circuits[1].x(qreg_q[1])
entlanging(circuits[1], qreg_q)

# 10 state
circuits[2].x(qreg_q[0])
circuits[2].reset(qreg_q[1])
entlanging(circuits[2], qreg_q)

# 11 state
circuits[3].x(qreg_q[0])
circuits[3].x(qreg_q[1])
entlanging(circuits[3], qreg_q)

simulator = Aer.get_backend('qasm_simulator')

circuits[0].measure(qreg_q[0], creg_c[0])
circuits[0].measure(qreg_q[1], creg_c[1])
result0 = execute(circuits[0], backend=simulator, shots=1024).result()
counts0 = result0.get_counts()

qt.plot_histogram(counts0, color="#5C9DFF", title="Entangled qubits values for input |00>")

circuits[0].draw(output='mpl')

circuits[1].measure(qreg_q[0], creg_c[0])
circuits[1].measure(qreg_q[1], creg_c[1])
result1 = execute(circuits[1], backend=simulator, shots=1024).result()
counts1 = result1.get_counts()

qt.plot_histogram(counts1, color="#5C9DFF", title="Entangled qubits values for input |01>")

circuits[1].draw(output='mpl')

circuits[2].measure(qreg_q[0], creg_c[0])
circuits[2].measure(qreg_q[1], creg_c[1])
result2 = execute(circuits[2], backend=simulator, shots=1024).result()
counts2 = result2.get_counts()

qt.plot_histogram(counts2, color="#5C9DFF", title="Entangled qubits values for input |10>")

circuits[2].draw(output='mpl')

circuits[3].measure(qreg_q[0], creg_c[0])
circuits[3].measure(qreg_q[1], creg_c[1])
result3 = execute(circuits[3], backend=simulator, shots=1024).result()
counts3 = result3.get_counts()

qt.plot_histogram(counts3, color="#5C9DFF", title="Entangled qubits values for input |11>")

circuits[3].draw(output='mpl')