from __future__ import annotations

import time

import torch

from pyqtorch.core.batched_operation import batchedRX
from pyqtorch.modules import RX, QuantumCircuit, zero_state
from pyqtorch.modules.parametric import RotationGate
from pyqtorch.modules.primitive import ControlledOperationGate


def timeit(f, *args, niters=100):
    t = 0
    for _ in range(niters):
        t0 = time.time()
        out = f(*args)
        t1 = time.time()
        t += t1 - t0
    return t / niters, out


dtype = torch.cdouble
device = "cuda"
batch_size = 1000
qubits = [0]
n_qubits = 10

state = zero_state(n_qubits, batch_size=batch_size, device=device, dtype=dtype)
phi = torch.rand(batch_size, device=device, dtype=dtype)
thetas = {"phi": phi}

func_time, func_out = timeit(batchedRX, phi, state, qubits, n_qubits)

gate = RX(qubits, n_qubits, "phi").to(device=device, dtype=dtype)
mod_time, mod_out = timeit(gate, thetas, state)

print(f"Functional pyq: {func_time}")
print(f"Module pyq:     {mod_time}")
print(f"Same results:   {torch.allclose(func_out, mod_out)}")


circ = QuantumCircuit(
    n_qubits,
    [
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        RotationGate("Y", (0,), 10, "phi"),
        RotationGate("Y", (1,), 10, "phi"),
        RotationGate("Y", (2,), 10, "phi"),
        RotationGate("Y", (3,), 10, "phi"),
        RotationGate("Y", (4,), 10, "phi"),
        RotationGate("Y", (5,), 10, "phi"),
        RotationGate("Y", (6,), 10, "phi"),
        RotationGate("Y", (7,), 10, "phi"),
        RotationGate("Y", (8,), 10, "phi"),
        RotationGate("Y", (9,), 10, "phi"),
        RotationGate("X", (0,), 10, "phi"),
        RotationGate("X", (1,), 10, "phi"),
        RotationGate("X", (2,), 10, "phi"),
        RotationGate("X", (3,), 10, "phi"),
        RotationGate("X", (4,), 10, "phi"),
        RotationGate("X", (5,), 10, "phi"),
        RotationGate("X", (6,), 10, "phi"),
        RotationGate("X", (7,), 10, "phi"),
        RotationGate("X", (8,), 10, "phi"),
        RotationGate("X", (9,), 10, "phi"),
        ControlledOperationGate("X", (0, 1), 10),
        ControlledOperationGate("X", (2, 3), 10),
        ControlledOperationGate("X", (4, 5), 10),
        ControlledOperationGate("X", (6, 7), 10),
        ControlledOperationGate("X", (8, 9), 10),
        ControlledOperationGate("X", (1, 2), 10),
        ControlledOperationGate("X", (3, 4), 10),
        ControlledOperationGate("X", (5, 6), 10),
        ControlledOperationGate("X", (7, 8), 10),
    ],
).to(device=device, dtype=dtype)

circ(thetas, state)

print(timeit(circ, thetas, state))
