import networkx as nx
import pennylane as qml
from pennylane import numpy as np
from pennylane import qaoa
import sys
import time

def benchmark_qaoa_minvertex_heavy(n_wires):

    n_layers = 5
    graph = nx.complete_graph(n_wires)
    params = [[0.5] * n_layers, [0.5] * n_layers]
    device = "default.qubit"
    interface = "autograd"
    diff_method = "best"

    # if device name is given, create device
    if isinstance(device, str):
        device = qml.device(device, wires=len(graph.nodes), shots=1)

    options_dict = {"interface": interface, "diff_method": diff_method}

    H_cost, H_mixer = qaoa.min_vertex_cover(graph, constrained=False)

    n_wires = len(graph.nodes)

    def qaoa_layer(gamma, alpha):
        qaoa.cost_layer(gamma, H_cost)
        qaoa.mixer_layer(alpha, H_mixer)

    @qml.qnode(device)
    def circuit(params):
        for w in range(n_wires):
            qml.Hadamard(wires=w)
        qml.layer(qaoa_layer, n_layers, params[0], params[1])
        return [qml.sample(qml.PauliZ(i)) for i in range(n_wires)]

    circuit(params)
  

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

n_wires = int(sys.argv[1])
      
start = time.time()
#PROFILE_BEGIN
benchmark_qaoa_minvertex_heavy(n_wires)
#PROFILE_END
end = time.time()
print("time =",(end-start))