from qiskit import *
from qiskit.result import *
from qiskit.extensions import *
from qiskit.tools.visualization import *
import numpy as np

pi = np.pi

def BoxIsVain(n):
    qreg_q = QuantumRegister(2, 'q')
    creg_c = ClassicalRegister(2, 'c')
    qc = QuantumCircuit(qreg_q, creg_c)

    qc.reset(0)
    
    for i in range(n):
        qc.reset(1)
        qc.rx(pi/(2*n), 0)

        # Bomb is a CNot gate, will send the result to secundary qubit
        #qc.cx(0, 1)

        #qc.measure(0, 0)
        qc.measure(1, 0) # output from qubit 2 is measured

    # We'll run the program on a simulator
    backend = BasicAer.get_backend('qasm_simulator')
    job = execute(qc, backend,memory=True)
    output = job.result().get_memory()[0]

    display(plot_histogram(job.result().get_counts(qc)))
    return qc, output

def BoxIsBomb(n):
    qr = QuantumRegister(4, 'q')
    cr = ClassicalRegister(4, 'c')
    qc = QuantumCircuit(qr, cr)

    qc.reset(0)
    
    for i in range(n):
        qc.reset(1)
        qc.rx(pi/(2*n), 0)

        # Bomb is a CNot gate, will send the result to secundary qubit
        qc.cx(0, 1)
        # Or gate
        qc.ccx(1,2,3)
        qc.cx(1,3)
        qc.cx(2,3)

        #qc.measure(0, 0)
        qc.measure(3, 0) # output from qubit 4 is measured

    # We'll run the program on a simulator
    backend = BasicAer.get_backend('qasm_simulator')
    job = execute(qc, backend,memory=True)
    output = job.result().get_memory()[0]

    display(plot_histogram(job.result().get_counts(qc)))
    return qc, output

## Test the functions

n=1
print("Test with 1 iteration => [ n=1 ] :")
print('\n')

qc, output = BoxIsVain(n)
print('Box is Vain: Number of iterations n=',n,' gives output=Succeed Vain')
display(qc.draw())
print('\n')

qc, output = BoxIsBomb(n)
if output[3] == '0':
    print('Box is Bomb: Number of iterations n=',n,' gives output=Succeed Bomb')
elif output[3] == '1':
    print('Box is Bomb: Number of iterations n=',n,' gives output=blown up')
display(qc.draw())
print('\n')

n=2
print("Test with 2 iterations => [ n=2 ] :")
print('\n')

qc, output = BoxIsVain(n)
print('Box is Vain: Number of iterations n=',n,' gives output=Succeed Vain')
display(qc.draw())
print('\n')

qc, output = BoxIsBomb(n)
if output[3] == '0':
    print('Box is Bomb: Number of iterations n=',n,' gives output=Succeed Bomb')
elif output[3] == '1':
    print('Box is Bomb: Number of iterations n=',n,' gives output=blown up')
display(qc.draw())
print('\n')

n=500
print("Test with 500 iterations => [ n=500 ] :")
print('\n')

qc, output = BoxIsVain(n)
print('Box is Vain: Number of iterations n=',n,' gives output=Succeed Vain')
display(qc.draw())
print('\n')

qc, output = BoxIsBomb(n)
if output[3] == '0':
    print('Box is Bomb: Number of iterations n=',n,' gives output=Succeed Bomb')
elif output[3] == '1':
    print('Box is Bomb: Number of iterations n=',n,' gives output=blown up')
display(qc.draw())
print('\n')

