import numpy as np
from qiskit_textbook.tools import random_state
# Importing standard Qiskit libraries
from qiskit import QuantumCircuit,QuantumRegister, ClassicalRegister, transpile, Aer, IBMQ, assemble, execute
from qiskit.tools.jupyter import *
from qiskit.visualization import *
from ibm_quantum_widgets import *
from qiskit_textbook.tools import random_state
from qiskit.extensions import Initialize
from qiskit.providers.aer.noise import NoiseModel
import qiskit.quantum_info as qi

# Loading your IBM Quantum account(s)
provider = IBMQ.load_account()

qr = QuantumRegister(3, name="q")    # Protocol uses 3 qubits
crz = ClassicalRegister(1, name="crz") # and 2 classical bits
crx = ClassicalRegister(1, name="crx") # in 2 different register
teleportation_circuit = QuantumCircuit(qr, crz, crx)


def psi_00(circuit, q1, q2):
    circuit.h(q1)
    circuit.cx(q1,q2)
    return circuit




def Bell_measurement(circuit, q1, q2,crx, crz):
    # FILL YOUR CODE HERE
    circuit.cx(q2,q1)
    circuit.h(q2)
    circuit.measure(q1,crx)
    circuit.measure(q2,crz)
    return circuit
    

    
    


def conditional_gates(circuit, q0, crx, crz):
    circuit.x(q0).c_if(crx,1)
    circuit.z(q0).c_if(crz,1)
    return circuit
    


qr = QuantumRegister(3, name="q")    # Protocol uses 3 qubits
crz, crx = ClassicalRegister(1, name="crz"), ClassicalRegister(1, name="crx") #Classical Registers
teleportation_circuit = QuantumCircuit(qr, crz, crx)


#STEP 1:  Create the Bell state between q0 and q1
psi_00(teleportation_circuit, 0, 1 )

#STEP 2: Bell Measurement
teleportation_circuit.barrier()
Bell_measurement(teleportation_circuit, 1, 2, crx, crz)

#STEP 3: Conditional Gates
teleportation_circuit.barrier()
conditional_gates(teleportation_circuit,0, crx, crz)
teleportation_circuit.draw()


def new_conditional_gates(circuit, q0):
    # FILL YOUR CODE HERE
    

def new_Bell_measurement(circuit, q1, q2):
    # FILL YOUR CODE HERE
    

    
inverse_init_gate = init_gate.gates_to_uncompute()

psi = random_state(1)
init_gate = Initialize(psi)
init_gate.label = "init"

qr = QuantumRegister(3)    # Protocol uses 3 qubits
cr=ClassicalRegister(1)
teleportation_circuit = QuantumCircuit(qr, cr)


# STEP 0: let's initialize Bob's q2
teleportation_circuit.append(init_gate, [2])
teleportation_circuit.barrier()

#STEP 1:  Create the Bell state between q0 and q1
psi_00(teleportation_circuit, 0, 1 )

#STEP 2: Bell Measurement
teleportation_circuit.barrier()
new_Bell_measurement(teleportation_circuit, 1, 2)

#STEP 3: Conditional Gates
teleportation_circuit.barrier()
new_conditional_gates(teleportation_circuit,0)

# teleportation_circuit.append(inverse_init_gate, [0])
teleportation_circuit.measure(0,0)
teleportation_circuit.draw()

# Build noise model from backend properties
# provider = IBMQ.load_account()
backend = provider.get_backend('ibmq_belem')
noise_model = NoiseModel.from_backend(backend)

# Get coupling map from backend
coupling_map = backend.configuration().coupling_map

# Get basis gates from noise model
basis_gates = noise_model.basis_gates

# Perform a noise simulation
result = execute(teleportation_circuit, Aer.get_backend('qasm_simulator'),
                 coupling_map=coupling_map,
                 basis_gates=basis_gates,
                 noise_model=noise_model).result()
counts = result.get_counts()
plot_histogram(counts)

#Without Noise Model
# backend = provider.get_backend('aer_simulator')
result = execute(teleportation_circuit, Aer.get_backend('qasm_simulator'),
                 coupling_map=coupling_map).result()
counts2 = result.get_counts()
plot_histogram(counts2)

# FILL YOUR CODE HERE to calculate fidelity





plot_histogram(counts)

print(state1)

