import numpy as np
from qiskit import QuantumCircuit, ClassicalRegister, QuantumRegister
from qiskit import execute
from qiskit import BasicAer
π = np.pi

q = QuantumRegister(6)
c = ClassicalRegister(2)
hhl = QuantumCircuit(q, c)

# Superposition
hhl.h(q[1])
hhl.h(q[2])
# Controlled-U0
hhl.cu3(-π / 2, -π / 2, π / 2, q[2], q[3])
hhl.cu1(3 * π / 4, q[2], q[3])
hhl.cx(q[2], q[3])
hhl.cu1(3 * π / 4, q[2], q[3])
hhl.cx(q[2], q[3])
# Controlled-U1
hhl.cx(q[1], q[3]);

hhl.swap(q[1], q[2])
hhl.h(q[2])
hhl.cu1(-π / 2, q[1], q[2])
hhl.h(q[1]);

hhl.swap(q[1], q[2]);

hhl.cu3(0.392699, 0, 0, q[1], q[0])  # Controlled-RY0
hhl.cu3(0.19634955, 0, 0, q[2], q[0]);  # Controlled-RY1

hhl.swap(q[1], q[2])
hhl.h(q[1])
hhl.cu1(π / 2, q[1], q[2])  # Inverse(Dagger(Controlled-S))
hhl.h(q[2])
hhl.swap(q[2], q[1])
# Inverse(Controlled-U1)
hhl.cx(q[1], q[3])
# Inverse(Controlled-U0)
hhl.cx(q[2], q[3])
hhl.cu1(-3 * π / 4, q[2], q[3])
hhl.cx(q[2], q[3])
hhl.cu1(-3 * π / 4, q[2], q[3])
hhl.cu3(-π / 2, π / 2, -π / 2, q[2], q[3])
# End of Inverse(Controlled-U0)
hhl.h(q[2])
hhl.h(q[1]);

# Target state preparation
hhl.rz(-π, q[4])
hhl.u1(π, q[4])
hhl.h(q[4])
hhl.ry(-0.9311623288419387, q[4])
hhl.rz(π, q[4])
# Swap test
hhl.h(q[5])
hhl.cx(q[4], q[3])
hhl.ccx(q[5], q[3], q[4])
hhl.cx(q[4], q[3])
hhl.h(q[5])

hhl.barrier(q)
hhl.measure(q[0], c[0])
hhl.measure(q[5], c[1]);

def get_psuccess(counts):
    '''Compute the success probability of the HHL protocol from the statistics

    :return: (float) The success probability.
    '''
    try:
        succ_rotation_fail_swap = counts['11']
    except KeyError:
        succ_rotation_fail_swap = 0
    try:
        succ_rotation_succ_swap = counts['01']
    except KeyError:
        succ_rotation_succ_swap = 0
    succ_rotation = succ_rotation_succ_swap + succ_rotation_fail_swap
    try:
        prob_swap_test_success = succ_rotation_succ_swap / succ_rotation
    except ZeroDivisionError:
        prob_swap_test_success = 0
    return prob_swap_test_success

backend = BasicAer.get_backend('qasm_simulator')
job = execute(hhl, backend, shots=100)
result = job.result()
counts = result.get_counts(hhl)
print(get_psuccess(counts))