import time
import pennylane as qml
from pennylane import numpy as pnp
from numpy.random import random
import tensorflow as tf
import sys


def benchmark_circuit(layers):

    wires = ['a', 'b', 'c']
    device = qml.device('default.qubit', wires=wires)
    n_wires = len(wires)
    n_layers = layers
    params = pnp.random.random(4**3-1)
    measurement = qml.probs(['a', 'c'])
    interface = "autograd"
    diff_method = "best"
    
    @qml.qnode(device, interface=interface, diff_method=diff_method)
    def circuit(params_):
        qml.templates.ArbitraryUnitary(params_, wires=wires)
        measurement.queue()
        return measurement

    # turn parameters into tensor from interface
    if interface == "autograd":
        params = pnp.array(params, requires_grad=True)
    elif interface == "tf":
        params = tf.Variable(params)
    elif interface == "torch":
        params = torch.tensor(params)
    circuit(params)

if len(sys.argv) != 2:
    print("python benchmark.py <layers>")
    sys.exit(1)

layers = int(sys.argv[1])

start = time.time()
#PROFILE_BEGIN
benchmark_circuit(layers)
#PROFILE_END
end = time.time()
print("time =",(end-start))