#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan 31 13:37:47 2020

@author: kidiki
"""

import numpy as np
from qiskit import (
    QuantumCircuit,
    execute,
    Aer)
from qiskit.visualization import plot_histogram, plot_bloch_vector
from matplotlib import pyplot as plt

from scipy.stats import unitary_group
from collections import defaultdict

from qiskit.quantum_info.random.utils import random_state

# Use Aer's qasm_simulator
simulator = Aer.get_backend('qasm_simulator')

def create_bloch_vector(theta, phi):
    x =np.sin(theta)*np.cos(phi)
    y =np.sin(theta)*np.sin(phi)
    z =np.cos(theta)
    return [x,y,z]


def perform_exp(iterations=100):
    simulator = Aer.get_backend('qasm_simulator')
    all_counts = defaultdict(int)

    for i in range(iterations):
        qc = QuantumCircuit(2, 2)

        # uniform distribution on the sphere
        t = round(np.random.uniform(0, 1), 10)
        theta = np.arccos(1 - 2 * t)
        phi = round(np.random.uniform(0, 2 * np.pi), 10)

        # 0_case, identical states
        #  qc.r(theta, phi, 0)
        #  qc.r(theta, phi, 1)

        # 1_case, orthogonal states on the same meridian 
        #  qc.r(theta, phi, 0)
        #  qc.r(theta+np.pi, phi, 1)

        # my_case, phase difference DOESNT WORK
        phi1 = 2 * np.pi - phi
        qc.r(theta, phi, 0)
        qc.r(theta, phi1, 1)

       # print(theta, phi, phi1)

        # q0 = create_bloch_vector(theta, phi)
        # plot_bloch_vector(q0)
        # plt.show()
        #
        # q1 = create_bloch_vector(theta, phi1)
        # plot_bloch_vector(q1)
        # plt.show()


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

        # print(qc)

        job = execute(qc, simulator, shots=1)
        result = job.result()
        # print(result)
        counts = result.get_counts(qc)

        # print(counts.keys())

        for k, v in counts.items():
            all_counts[k] += v

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


if __name__ == "__main__":
    perform_exp()