from qiskit import Aer
from qiskit import QuantumCircuit
from VehicleRouting.standard.CostCalculator import CostCalculator


def maxcut_obj(x, G):
    """
    Given a bitstring as a solution, this function returns
    the number of edges shared between the two partitions
    of the graph.

    Args:
        x: str
           solution bitstring

        G: networkx graph

    Returns:
        obj: float
             Objective
    """
    obj = 0
    for i, j in G.couplings():
        if x[i] != x[j]:
            obj -= 1

    return obj


def compute_expectation(counts, G):
    """
    Computes expectation value based on measurement results

    Args:
        counts: dict
                key as bitstring, val as count

        G: networkx graph

    Returns:
        avg: float
             expectation value
    """

    avg = 0
    sum_count = 0
    for bitstring, count in counts.items():
        obj = maxcut_obj(bitstring, G)
        avg += obj * count
        sum_count += count

    return avg / sum_count


# We will also bring the different circuit components that
# build the qaoa circuit under a single function
def create_qaoa_circuit(graph, theta):
    """
    Creates a parametrized qaoa circuit

    Args:
        graph: networkx graph
        theta: list
               unitary parameters

    Returns:
        qc: qiskit circuit
    """

    nqubits = len(graph.nodes())
    p = len(theta) // 2  # number of alternating unitaries
    qc = QuantumCircuit(nqubits)

    beta = theta[:p]
    gamma = theta[p:]

    # initial_state
    for i in range(0, nqubits):
        qc.h(i)

    for irep in range(0, p):

        # problem unitary
        for pair in list(graph.couplings()):
            qc.rzz(2 * gamma[irep], pair[0], pair[1])

        qc.barrier()

        # mixer unitary
        for i in range(0, nqubits):
            qc.rx(2 * beta[irep], i)

    qc.measure_all()

    return qc


# Finally we write a function that executes the circuit on the chosen backend
def get_expectation(graph, p, shots=512):
    """
    Runs parametrized circuit

    Args:
        graph: networkx graph
        p: int,
           Number of repetitions of unitaries
    """

    backend = Aer.get_backend('qasm_simulator')
    backend.shots = shots

    def execute_circ(theta):
        qc = create_qaoa_circuit(graph, theta)
        counts = backend.run(qc, seed_simulator=10,
                             nshots=512).result().get_counts()

        return compute_expectation(counts, graph)

    return execute_circ



# Finally we write a function that executes the circuit on the chosen backend
def get_execute_circuit(graph, shots=512):
    """
    Runs parametrized circuit

    Args:
        G: networkx graph
        p: int,
           Number of repetitions of unitaries
    """

    backend = Aer.get_backend('qasm_simulator')
    backend.shots = shots

    def execute_circ(theta):
        qc = create_qaoa_circuit(graph, theta)
        counts = backend.run(qc, seed_simulator=10,
                             nshots=512).result().get_counts()

        return compute_expectation(counts, graph)

    return execute_circ