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

CASE_IDENTICAL = "identical"
CASE_ORTHOGONAL = "orthogonal"


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

    for i in range(iterations):
        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)

        # update the counts of all experiments with the current run
        for k, v in counts.items():
            all_counts[k] += v

    # gather all counts
    print("\nTotal counts for experiment '{}' are:\n".format(case))
    for k, v in sorted(all_counts.items(), key=lambda x: x[0]):
        print("{}: {}".format(k, v))


if __name__ == "__main__":
    perform_exp(iterations=100, case=CASE_IDENTICAL)
    perform_exp(iterations=100, case=CASE_ORTHOGONAL)