import random
from collections import defaultdict

import numpy as np

from lib import *


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=bell_basis)
        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>", name="B0"),
        State(np.cos(beta) * b_0 + np.sin(beta) * s("|001>"), name="B1"),
        State(-np.sin(beta) * b_0 + np.cos(beta) * s("|001>"), name="B2"),
        State((1 / np.sqrt(2)) * (s("|100>") - s("|010>")), name="B3"),
        State((1 / np.sqrt(2)) * (s("|011>") - s("|101>")), name="B4"),
        State(np.cos(beta) * b_1 + np.sin(beta) * s("|001>"), name="B5"),
        State(-np.sin(beta) * b_1 + np.cos(beta) * s("|001>"), name="B6"),
        s("|111>", name="B7"),
    ]

    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)
        q1 = 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

        q2 = State.from_bloch_angles(theta1, phi1)

        q1q2_st = State(q1*q2)
        # meas1 = q1q2_st.measure(basis=bell_basis)
        # if meas1 == '11':
        #     return 'discarded'

        P_minus = b_psi_m.x(b_psi_m)
        P_s = b_psi_p.x(b_psi_p) + b_phi_m.x(b_phi_m) + b_phi_p.x(b_phi_p)

        # vec = MeasurementOperator(P_minus.m).on(q1q2_st)
        # if np.all(vec.m.transpose() == np.zeros(4)):
        #     return 'identical'
        vec2 = MeasurementOperator(P_s.m).on(q1q2_st)
        post_state = State(State.normalize(vec2))

        if case == CASE_ID_ID:
            assert post_state == q1q2_st
        elif case in [CASE_ORT_ID, CASE_ORT_ORT]:
            assert post_state != q1q2_st

        q3 = State.from_bloch_angles(theta2, phi2)
        meas = State(post_state * q3).measure(basis=basis)

        # # TODO: This part measures the first two qubits in Bell basis, then:
        # #       1. constructs a new pure state
        # #           (e.g. if result was 11, constructs |11>)
        # #       2. Combines this new pure state with the third qubit
        # MEASURE_FIRST_BELL = False
        # if MEASURE_FIRST_BELL:
        #     # Measure in the Bell's basis
        #     st = State(q1 * q2)
        #     meas = st.measure(basis=bell_basis)
        #     new_st = s("|" + meas + ">")
        #     st = State(new_st * q3)
        # else:
        #     # TODO: -OR- alternatively - measure all three without measuring in
        #     #                            Bell first
        #     st = State(q1 * q2 * q3)
        #
        # # Measure in the 8D basis
        # meas = st.measure(basis=basis)
        return meas

    def krisi_3_format_results(results):
        all_pos = basis
        print("Results:")
        for case in case_choices:
            rv = results.get(case)
            print("{}:".format(case))
            print("   raw  : {}".format(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)*100
                print("   {}  : {:3d} ({:.1f}%)".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()