import cirq
import numpy as np
from collections import defaultdict

from lib 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 from_angles_2(theta, phi):
    # phase difference
    theta, phi = State._normalize_angles(theta, phi)
    m0 = np.cos(theta / 2)
    m1 = np.sin(theta / 2) * np.power(np.e, (1j * (-phi)))
    m = m0 * _0 + m1 * _1
    return State(m.m)


def print_all_samples(name, all_samples):
    print("------------- ALL SAMPLES for {}".format(name))
    for k, v in sorted(all_samples.items(), key=lambda x: x[0]):
        print("{}: {}".format(k, v))
    print("==============================================")


def math_sim(q2func=State.from_bloch_angles, iterations=1000, sample_count=1):
    all_samples = defaultdict(int)
    for i in range(iterations):
        # print("Running iteration {}".format(i))
        # Generating uniform random points on a sphere is not trivial
        # https://www.bogotobogo.com/Algorithms/uniform_distribution_sphere.php
        # theta = np.arccos(2 * np.random.uniform(0, 1) - 1.0)
        # phi = 2 * np.pi * np.random.uniform(0, 1.0)

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

        # print("theta: {:.2f} ({:.2f} deg) | phi: {:.2f} ({:.2f} deg)".format(
        #     theta, np.rad2deg(theta),
        #     phi, np.rad2deg(phi)))
        q1 = State.from_bloch_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("math_sim", all_samples)


class MemoizedExp(object):
    def __init__(self, *funcs):
        self.memoized = dict()
        self.funcs = {func.__name__: func for func in funcs}
        self.reset()

    def call(self, func_name, *args, **kwargs):
        if func_name not in self.funcs:
            raise Exception("Func {} doesn't exist".format(func_name))
        if self.memoized[func_name]:
            return self.memoized[func_name]
        values = self.funcs[func_name](*args, **kwargs)
        self.memoized[func_name] = values
        return values

    def reset(self):
        for func_name in self.funcs.keys():
            self.memoized[func_name] = None


def gen_exp_for_cirq_0():
    # TODO: FIX THIS AS AN EXPONENT UNIFORM DISTRIBUTION
    #       Generating uniform random points on a sphere is not trivial
    # https://www.bogotobogo.com/Algorithms/uniform_distribution_sphere.php
    theta = round(np.random.uniform(0, 1), 10)
    phi = round(np.random.uniform(0, 1), 10)
    return theta, phi


def gen_exp_for_cirq_1():
    """TODO: How to generate the exponents for the second case"""
    theta = round(np.random.uniform(0, 1), 10)
    phi = round(np.random.uniform(0, 2), 10)
    return theta, phi


def cirq_sim(q1func, q2func, iterations=1000, sample_count=1):
    all_samples = defaultdict(int)
    memoized_exp = MemoizedExp(q1func, q2func)
    for i in range(iterations):
        theta1, phi1 = memoized_exp.call(q1func.__name__)
        theta2, phi2 = memoized_exp.call(q2func.__name__)
        memoized_exp.reset()

        q1 = cirq.GridQubit(0, 0)
        q2 = cirq.GridQubit(1, 0)

        # Create a circuit
        circuit = cirq.Circuit(
            cirq.XPowGate(exponent=theta1).on(q1),
            cirq.ZPowGate(exponent=phi1).on(q1),
            cirq.XPowGate(exponent=theta2).on(q2),
            cirq.ZPowGate(exponent=phi2).on(q2),
            cirq.CNOT(q1, q2),
            cirq.H(q1),

            cirq.measure(q1, key='q1'),  # Measurement.
            cirq.measure(q2, key='q2')  # Measurement.
        )
        # print(circuit)

        # Simulate the circuit several times.
        simulator = cirq.Simulator()
        result = simulator.run(circuit, repetitions=sample_count)
        for rep in range(len(result.measurements['q1'])):
            rq1 = result.measurements['q1'][rep][0]
            rq2 = result.measurements['q2'][rep][0]
            k = "{}{}".format(rq1, rq2)
            all_samples[k] += 1
    print_all_samples("cirq", all_samples)


def math_sim_0(*args, **kwargs):
    math_sim(q2func=State.from_bloch_angles, *args, **kwargs)


def math_sim_1(*args, **kwargs):
    math_sim(q2func=from_angles_1, *args, **kwargs)


def math_sim_2(*args, **kwargs):
    math_sim(q2func=from_angles_2, *args, **kwargs)


def cirq_sim_0(*args, **kwargs):
    cirq_sim(q1func=gen_exp_for_cirq_0,
             q2func=gen_exp_for_cirq_0,
             *args, **kwargs)


def cirq_sim_1(*args, **kwargs):
    cirq_sim(q1func=gen_exp_for_cirq_0,
             q2func=gen_exp_for_cirq_1,
             *args, **kwargs)


if __name__ == "__main__":
    # math_sim_0(iterations=5000)
    math_sim_2(iterations=500)
    # cirq_sim_0()