from collections import deque
import matplotlib.pyplot as plt
import numpy as np

from pyqcs import State, H, X, S, CZ
from pyqcs.graph.state import GraphState
from pyqcs.util.random_circuits import random_circuit

from measure_circuit import execution_statistics

def S_with_extra_arg(act, i):
    return S(act)

def test_scaling_qbits(state_factory
                    , nstart
                    , nstop
                    , ngates_per_qbit
                    , ncircuits
                    , **kwargs):
    results = deque()
    for qbits in range(nstart, nstop):
        circuits = [random_circuit(qbits, ngates_per_qbit * qbits, X, H, S_with_extra_arg, CZ)
                        for _ in range(ncircuits)]
        state = state_factory(qbits)

        print("running test with", qbits, "qbits")
        N, avg, std_dev = execution_statistics(circuits, state, scale=qbits, **kwargs)

        results.append([qbits, N, avg, std_dev])
    return np.array(results, dtype=np.double)

if __name__ == "__main__":
    nstart = 4
    nstop = 19
    ncircuits = 250
    ngates_per_qbit = 100

    np.random.seed(0xdeadbeef)
    results_naive = test_scaling_qbits(State.new_zero_state
                                    , nstart
                                    , nstop
                                    , ngates_per_qbit
                                    , ncircuits
                                    , repeat=10)
    np.random.seed(0xdeadbeef)
    results_graph = test_scaling_qbits(GraphState.new_zero_state
                                    , nstart
                                    , nstop
                                    , ngates_per_qbit
                                    , ncircuits
                                    , repeat=10)

    h0 = plt.errorbar(results_naive[:, 0], results_naive[:, 2], results_naive[:, 3]
                    , label=f"Dense Vector Simulator $N_c={int(results_naive[:, 1][0])}$ circuits"
                    , marker="o"
                    , color="black")
    h1 = plt.errorbar(results_graph[:, 0], results_graph[:, 2], results_graph[:, 3]
                    , label=f"Graphical Simulator $N_c={int(results_graph[:, 1][0])}$ circuits"
                    , marker="^"
                    , color="black")

    plt.legend(handles=[h0, h1])
    plt.xlabel("Number of Qbits $N_q$")
    plt.ylabel("Execution time per circuit [s]")
    plt.title(f"Execution Time for ${ngates_per_qbit}\\times N_q$ Gates with random Circuits (rescaled)")

    #plt.show()
    plt.savefig("Figure_1.png", dpi=400)