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[~bit_mask] = qm_state[bit_mask];"
                "result_py[bit_mask] = qm_state[~bit_mask];"
                "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;"
                    "bit_mask = np.array([1 if (i & (1 << 0)) else 0 for i in range(2**10)])"
            , repeat=5
            , number=1_000_000
            )
print(" done")

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

bit_mask = np.array([1 if (i & (1 << 0)) else 0 for i in range(2**10)]).astype(np.bool)
result_py = np.zeros(2**10, dtype=np.cdouble)
cl_py = np.zeros(10, dtype=np.int8)
result_py[~bit_mask] = qm_state[bit_mask]
result_py[bit_mask] = qm_state[~bit_mask]

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)