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_circuits(state_factory
                        , nstart
                        , nstop
                        , step
                        , nqbits
                        , ncircuits
                        , **kwargs):
    results = deque()

    for ngates in range(nstart, nstop, step):
        circuits = [random_circuit(nqbits, ngates, X, H, S_with_extra_arg, CZ)
                        for _ in range(ncircuits)]
        state = state_factory(nqbits)

        print("running test with", ngates, "gates on", nqbits, "qbits")

        N, avg, std_dev = execution_statistics(circuits, state, scale=1, **kwargs)
        results.append([ngates, N, avg, std_dev])

    return np.array(results, dtype=np.double)


if __name__ == "__main__":
    nstart = 400
    nstop = 1600
    step = 50
    ncircuits = 250
    nqbits0 = 100
    nqbits1 = 50

    np.random.seed(0xdeadbeef)
    results_graph0 = test_scaling_circuits(GraphState.new_zero_state
                                    , nstart
                                    , nstop
                                    , step
                                    , nqbits0
                                    , ncircuits
                                    , repeat=10)
    np.random.seed(0xdeadbeef)
    results_graph1 = test_scaling_circuits(GraphState.new_zero_state
                                    , nstart
                                    , nstop
                                    , step
                                    , nqbits1
                                    , ncircuits
                                    , repeat=10)

    h0 = plt.errorbar(results_graph0[:, 0], results_graph0[:, 2], results_graph0[:, 3]
                    , label=f"Graphical Simulator $N_q={nqbits0}$ Qbits"
                    , marker="^"
                    , color="black")
    h1 = plt.errorbar(results_graph1[:, 0], results_graph1[:, 2], results_graph1[:, 3]
                    , label=f"Graphical Simulator $N_q={nqbits1}$ Qbits"
                    , marker="o"
                    , color="black")

    plt.legend(handles=[h0, h1])
    plt.xlabel("Number of gates in circuit")
    plt.ylabel("Execution time per circuit [s]")
    plt.title(f"Execution Time for random Circuits")

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