import numpy as np
from qiskit import Aer
from qiskit_aqua import run_algorithm
from qiskit_aqua.input import EnergyInput
from qiskit_aqua.translators.ising import graphpartition
from qiskit_aqua.algorithms import ExactEigensolver

np.random.seed(100)
num_nodes = 4
w = graphpartition.random_graph(num_nodes, edge_prob=0.8, weight_range=10)
print(w)

def brute_force():
    # use the brute-force way to generate the oracle
    def bitfield(n, L):
        result = np.binary_repr(n, L)
        return [int(digit) for digit in result]  # [2:] to chop off the "0b" part

    L = num_nodes
    max = 2**L
    minimal_v = np.inf
    for i in range(max):
        cur = bitfield(i, L)

        how_many_nonzero = np.count_nonzero(cur)
        if how_many_nonzero * 2 != L:  # not balanced
            continue

        cur_v = graphpartition.objective_value(np.array(cur), w)
        if cur_v < minimal_v:
            minimal_v = cur_v
    return minimal_v

sol = brute_force()
print("objective value computed by the brute-force method is", sol)

qubit_op, offset = graphpartition.get_graphpartition_qubitops(w)
algo_input = EnergyInput(qubit_op)

params = {
    'problem': {'name': 'ising'},
    'algorithm': {'name': 'ExactEigensolver'}
}
result = run_algorithm(params, algo_input)
x = graphpartition.sample_most_likely(result['eigvecs'][0])
# check against the oracle
ising_sol = graphpartition.get_graph_solution(x)
np.testing.assert_array_equal(ising_sol, [0, 1, 0, 1])
print("objective value computed by ExactEigensolver is", graphpartition.objective_value(x, w))

algo = ExactEigensolver(algo_input.qubit_op, k=1, aux_operators=[])
result = algo.run()
x = graphpartition.sample_most_likely(result['eigvecs'][0])
# check against the oracle
ising_sol = graphpartition.get_graph_solution(x)
np.testing.assert_array_equal(ising_sol, [0, 1, 0, 1])
print("objective value computed by the ExactEigensolver is", graphpartition.objective_value(x, w))

algorithm_cfg = {
    'name': 'VQE',
    'operator_mode': 'matrix'
}

optimizer_cfg = {
    'name': 'SPSA',
    'max_trials': 300
}

var_form_cfg = {
    'name': 'RY',
    'depth': 5,
    'entanglement': 'linear'
}

params = {
    'problem': {'name': 'ising', 'random_seed': 10598},
    'algorithm': algorithm_cfg,
    'optimizer': optimizer_cfg,
    'variational_form': var_form_cfg
}
backend = Aer.get_backend('statevector_simulator')
result = run_algorithm(params, algo_input, backend=backend)
x = graphpartition.sample_most_likely(result['eigvecs'][0])
# check against the oracle
ising_sol = graphpartition.get_graph_solution(x)
np.testing.assert_array_equal(ising_sol, [1, 0, 0, 1])
print("objective value computed by VQE is", graphpartition.objective_value(x, w))