def aggregate(alpha, measurements):

    if not 0 <= alpha <= 1:
        raise ValueError(f"alpha must be in [0, 1] but was {alpha}")

    # sort by values
    sorted_measurements = sorted(measurements, key=lambda x: x[1])
    accumulated_percent = 0.0  # once alpha is reached, stop
    cvar = 0.0
    for probability, value in sorted_measurements:
        cvar += value * min(probability, alpha - accumulated_percent)
        accumulated_percent += probability
        if accumulated_percent >= alpha:
            break

    return cvar / alpha

from qopt_best_practices.cost_function import evaluate_sparse_pauli


def cost_func_cvar_sampler(params, ansatz, hamiltonian, sampler, aggregation):

    job = sampler.run(ansatz, params)
    sampler_result = job.result()
    sampled = sampler_result.quasi_dists[0]

    # a dictionary containing: {state: (measurement probability, value)}
    evaluated = {
        state: (probability, evaluate_sparse_pauli(state, hamiltonian))
        for state, probability in sampled.items()
    }

    result = aggregate(aggregation, evaluated.values())

    return result

# BASIC STATEVECTOR SIMULATOR BACKEND --> primitives in qiskit.primtives

from qiskit.primitives import Sampler

sampler = Sampler()

# Import pre-computed Hamiltonian
import json
from qiskit.quantum_info import SparsePauliOp

graph_file = "data/graph_2layers_0seed.json"
data = json.load(open(graph_file, "r"))
hamiltonian = SparsePauliOp.from_list(data["paulis"])
print(hamiltonian)

# Build basic ansatz using the circuit library utility
from qiskit.circuit.library import QAOAAnsatz

ansatz = QAOAAnsatz(hamiltonian, reps=2)
ansatz.draw("mpl")
ansatz.measure_all()

# Define random initial point
import numpy as np

init_params = np.random.rand(ansatz.num_parameters)
print(init_params)

from scipy.optimize import minimize

result = minimize(
    cost_func_cvar_sampler,
    init_params,
    args=(ansatz, hamiltonian, sampler, 0.5),
    method="COBYLA",
)
print(result)

# auxiliary functions to sample most likely bitstring
def to_bitstring(integer, num_bits):
    result = np.binary_repr(integer, width=num_bits)
    return [int(digit) for digit in result]


def sample_most_likely_bitstring(state_vector, num_bits):
    values = list(state_vector.values())
    most_likely = np.argmax(np.abs(values))
    most_likely_bitstring = to_bitstring(most_likely, num_bits)
    most_likely_bitstring.reverse()
    return np.asarray(most_likely_bitstring)

import matplotlib.pyplot as plt
import networkx as nx

# auxiliary function to plot graphs
def plot_result(G, x):
    colors = ["r" if i == 0 else "b" for i in x]
    pos, default_axes = nx.spring_layout(G), plt.axes(frameon=True)
    nx.draw_networkx(G, node_color=colors, node_size=600, alpha=0.8, pos=pos)

from qopt_best_practices.utils import build_max_cut_graph

graph = build_max_cut_graph(data["paulis"])

qc = ansatz.assign_parameters(result.x)
samp_dist = sampler.run(qc, shots=int(1e4)).result().quasi_dists[0]

best_result = sample_most_likely_bitstring(samp_dist, len(graph))

plot_result(graph, best_result)