import random
from collections import defaultdict

import numpy as np

from lib import State, b_phi_p, b_phi_m, b_psi_p, b_psi_m, s, generate_bins


def test_krisi_measurement_2():
    CASE_IDENTICAL = "identical"
    CASE_ORTHOGONAL = "orthogonal"

    def perform_exp(case):
        # 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)
        q0 = State.from_bloch_angles(theta0, phi0)

        # 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
        else:
            # for orthogonal we rotate the 1st qubit in 90 degrees
            # orthogonal to the first
            theta1, phi1 = theta0 + np.pi, phi0

        q1 = State.from_bloch_angles(theta1, phi1)

        # Measure in the Bell's basis
        st = State(q0 * q1)
        meas = st.measure(basis=[b_phi_p, b_phi_m, b_psi_p, b_psi_m])
        return meas

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


def test_krisi_measurement_3():
    iterations = 400
    beta = np.pi * 0.62

    CASE_ID_ID = "id_id"
    CASE_ID_ORT = "id_ort"
    CASE_ORT_ID = "ort_id"
    CASE_ORT_ORT = "ort_ort"

    case_choices = [CASE_ID_ID, CASE_ID_ORT, CASE_ORT_ID, CASE_ORT_ORT]

    b_0 = State(1 / np.sqrt(2) * (s("|100>") - s("|010>")))
    b_1 = State(1 / np.sqrt(2) * (s("|011>") - s("|101>")))
    basis = [
        s("|000>"),
        State(np.cos(beta) * b_0 + np.sin(beta) * s("|001>")),
        State(-np.sin(beta) * b_0 + np.cos(beta) * s("|001>")),
        b_0,
        b_1,
        State(np.cos(beta) * b_1 + np.sin(beta) * s("|001>")),
        State(-np.sin(beta) * b_1 + np.cos(beta) * s("|001>")),
        s("|111>"),
    ]

    def perform_exp(case):
        # 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)
        q0 = State.from_bloch_angles(theta0, phi0)

        # rotate the 1st qubit depending on the case we are exploring...
        if case == CASE_ID_ID:
            theta1, phi1 = theta0, phi0
            theta2, phi2 = theta0, phi0
        elif case == CASE_ID_ORT:
            theta1, phi1 = theta0, phi0
            theta2, phi2 = theta0 + np.pi, phi0
        elif case == CASE_ORT_ID:
            theta1, phi1 = theta0 + np.pi, phi0
            theta2, phi2 = theta0, phi0
        else:
            # CASE_ORT_ORT
            theta1, phi1 = theta0 + np.pi, phi0
            theta2, phi2 = theta0 + np.pi, phi0

        q1 = State.from_bloch_angles(theta1, phi1)
        q2 = State.from_bloch_angles(theta2, phi2)

        # Measure in the arbitrary basis
        st = State(q0 * q1 * q2)
        meas = st.measure(basis=basis)

        return meas

    def krisi_3_format_results(results):
        all_pos = generate_bins(8)
        print("Results:")
        for case in case_choices:
            rv = results.get(case)
            print("{}:".format(case))
            print("   raw  : {}".format(sorted(rv.items())))
            case_total = sum(rv.values())
            print("   total: {}".format(case_total))
            for pos in all_pos:
                value = rv.get(pos, 0)
                percent = value / case_total
                print("   {}  : {:3d} ({:.2f}%)".format(pos, value, percent))

    results = defaultdict(lambda: defaultdict(int))
    for i in range(iterations):
        case = random.choice(case_choices)
        result = perform_exp(case=case)
        results[case][result] += 1

    krisi_3_format_results(results)


if __name__ == "__main__":
    # test_krisi_measurement_2()
    test_krisi_measurement_3()