from collections import deque
import matplotlib.pyplot as plt
import numpy as np
import json
from random import shuffle

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

from measure_circuit import execution_statistics, measure_all

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

def test_scaling_circuits(state_factory
                        , nstart
                        , nstop
                        , step
                        , nqbits
                        , ncircuits
                        , **kwargs):
    trials = deque()

    N = (nstop - nstart) / step
    print()
    for n, ngates in enumerate(range(nstart, nstop, step)):
        print(f"generating test data... {int(n/N * 100)} %", end="\r", flush=True)
        measurement_circuit = list_to_circuit([M(i) for i in range(nqbits)])
        circuits = [random_circuit(nqbits, ngates, M, H, S_with_extra_arg, CZ)
                            | measurement_circuit
                        for _ in range(ncircuits)]
        state = state_factory(nqbits)
        for circuit in circuits:
            trials.append((ngates, circuit, state))
    print("generating test data... done ")

    print("randomizing tests...", end="", flush=True)
    shuffle(trials)
    print(" done")

    results = measure_all(trials, **kwargs)
    N, avg, std_dev = execution_statistics(results, scale=None, **kwargs)
    ngates = [[i] for i in sorted(results.keys())]
    N = [[i] for i in N]
    avg = [[i] for i in avg]
    std_dev = [[i] for i in std_dev]

    return np.hstack([ngates, N, avg, std_dev])


if __name__ == "__main__":
    nstart = 400
    nstop = 2800
    step = 50
    ncircuits = 100
    nqbits0 = 100
    nqbits1 = 50
    seed = 0xdeadbeef

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

    np.savetxt("circuit_scaling_graph0_measurements.csv", results_graph0)
    print("saved results0 to circuit_scaling_graph0_measurements.csv")
    np.savetxt("circuit_scaling_graph1_measurements.csv", results_graph1)
    print("saved results1 to circuit_scaling_graph1_measurements.csv")

    meta = {
            "nstart": nstart
            , "nstop": nstop
            , "step": step
            , "ncircuits": ncircuits
            , "nqbits0": nqbits0
            , "nqbits1": nqbits1
            , "seed": seed}

    with open("circuit_scaling_measurements_meta.json", "w") as fout:
        json.dump(meta, fout)
    print("saved meta to circuit_scaling_measurements_meta.json")

