import numpy as np
from collections import defaultdict

from lib_q_computer_math import State, QuantumCircuit, QuantumProcessor, C, H, x, _, _0, _1


def from_angles_1(theta, phi):
    theta, phi = State._normalize_angles(theta, phi)
    m0 = -np.sin(theta / 2)
    m1 = np.cos(theta / 2) * np.power(np.e, (1j * phi))
    m = m0 * _0 + m1 * _1
    return State(m.m)


def krisi(q2func, iterations=100, sample_count=1):
    all_samples = defaultdict(int)
    for i in range(iterations):
        # print("Running iteration {}".format(i))
        theta = round(np.random.uniform(0, np.pi), 3)
        phi = round(np.random.uniform(0, 2 * np.pi), 3)
        # print("theta: {:.2f} ({:.2f} deg) | phi: {:.2f} ({:.2f} deg)".format(
        #     theta, np.rad2deg(theta),
        #     phi, np.rad2deg(phi)))
        q1 = State.from_angles(theta, phi)
        q2 = q2func(theta, phi)

        qc = QuantumCircuit(2, [[q1, q2], ])
        qc.add_row([C, H])
        qc.add_row([x, _])
        qp = QuantumProcessor(qc)
        this_samples = qp.get_sample(sample_count)
        for k, v in this_samples.items():
            all_samples[k] += v
    print("------------- ALL SAMPLES for {}".format(q2func.__name__))
    for k, v in sorted(all_samples.items(), key=lambda x: x[0]):
        print("{}: {}".format(k, v))
    print("==============================================")


def krisi_0():
    krisi(q2func=State.from_angles)


def krisi_1():
    krisi(q2func=from_angles_1)


if __name__ == "__main__":
    krisi_0()