import random

import numpy as np
from qiskit import (
    QuantumCircuit,
    execute,
    Aer)
from collections import defaultdict

CASE_IDENTICAL = "identical"
CASE_ORTHOGONAL = "orthogonal"


def perform_exp(case):
    # Use Aer's qasm_simulator
    simulator = Aer.get_backend('qasm_simulator')
    all_counts = defaultdict(int)

    qc = QuantumCircuit(2, 2)

    # produce angles with uniform distribution on the sphere
    t = round(np.random.uniform(0, 1), 10)
    theta0 = np.arccos(1 - 2 * t)
    phi0 = round(np.random.uniform(0, 2 * np.pi), 10)

    # rotate the 0th qubit in (theta, phi)
    qc.r(theta0, phi0, 0)

    # rotate the 1st qubit depending on the case we are exploring...
    if case == CASE_IDENTICAL:
        # ... for identical we are rotating the 1st qubit the same angles as
        # the 0th
        theta1, phi1 = theta0, phi0
    elif case == CASE_ORTHOGONAL:
        # for orthogonal we rotate the 1st qubit in 90 degrees
        # orthogonal to the first
        theta1, phi1 = theta0 + np.pi, phi0

    # perform the rotation of the 1st qubit
    qc.r(theta1, phi1, 1)

    # Measure in the Bell's basis
    qc.cx(0, 1)
    qc.h(0)

    # measuring maps the 0/1 qubit register to the 0/1 classical register
    qc.measure([0, 1], [0, 1])

    # execute the job
    job = execute(qc, simulator, shots=1)
    result = job.result()
    counts = result.get_counts(qc)
    return [i for i in counts.keys()][0]


if __name__ == "__main__":
    correct = 0
    results = defaultdict(int)
    for i in range(100):
        case = random.choice([CASE_IDENTICAL, CASE_ORTHOGONAL])
        result = perform_exp(case=case)
        results[result] += 1
        if result == '11':
            guess = CASE_ORTHOGONAL
        else:
            guess = CASE_IDENTICAL
        if guess == case:
            correct += 1
    print("Correct: {}".format(correct))
    # print("Results: {}".format(results))