import timeit

import numpy as np

from pyqcs.gates.gate import BuiltinGate

cl_state = np.zeros(10, dtype=np.int8)
qm_state = np.zeros(2**10, dtype=np.cdouble)
qm_state[0] = 1

gate_uf = BuiltinGate('X', 0, 0, 0.0)

print("running benchmarks ...", end="", flush=True)

time_uf = timeit.repeat("result_uf = gate_uf(qm_state, cl_state)"
            , setup="import numpy as np;"
                    "from pyqcs.gates.gate import BuiltinGate;"
                    "cl_state = np.zeros(10, dtype=np.int8);"
                    "qm_state = np.zeros(2**10, dtype=np.cdouble);"
                    "qm_state[0] = 1;"
                    "gate_uf = BuiltinGate('X', 0, 0, 0.0);"
            , repeat=5
            , number=1_000_000
            )

time_py = timeit.repeat("result_py = np.zeros(2**10, dtype=np.cdouble);"
                    "result_py[0::2] = qm_state[1::2];"
                    "result_py[1::2] = qm_state[0::2];"
                    "cl_py = np.zeros(10, dtype=np.int8)"
            , setup="import numpy as np;"
                    "qm_state = np.zeros(2**10, dtype=np.cdouble);"
                    "qm_state[0] = 1;"
            , repeat=5
            , number=1_000_000
            )
print(" done")

print("running test ...", end="", flush=True)
result_uf, cl, m = gate_uf(qm_state, cl_state);

result_py = np.zeros(2**10, dtype=np.cdouble)
cl_py = np.zeros(10, dtype=np.int8)
result_py[0::2] = qm_state[1::2]
result_py[1::2] = qm_state[0::2]

assert np.allclose(result_py, result_uf)
print(" done")


best_uf = min(time_uf) / 1_000_000
best_py = min(time_py) / 1_000_000

print("time for ufunc gate (best out of 5, 1_000_000 runs):", best_uf)
print("time for python gate (best out of 5, 1_000_000 runs):", best_py)