import cirq
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)


class MemoizedExp(object):
    def __init__(self):
        self.theta = None
        self.phi = None

    def gen_exp_for_cirq_0(self):
        if not (self.theta and self.phi):
            self.theta = round(np.random.uniform(0, 1), 10)
            self.phi = round(np.random.uniform(0, 1), 10)
        return self.theta, self.phi

    def reset(self):
        self.theta = None
        self.phi = None


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, 1), 10)
    return theta, phi


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


def math_sim(q2func=from_angles_1, iterations=1000, 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(all_samples)


def cirq_sim(q1func, q2func, memoized_exp, iterations=1000, sample_count=1):
    all_samples = defaultdict(int)
    for i in range(iterations):
        theta1, phi1 = q1func()
        theta2, phi2 = q2func()
        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(all_samples)


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


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


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


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


if __name__ == "__main__":
    cirq_sim_0()