import numpy as np
import json

from qiskit import Aer

from qiskit_aqua import run_algorithm
from qiskit_aqua.input import EnergyInput
from qiskit_aqua.translators.ising import setpacking
from qiskit_aqua.algorithms import ExactEigensolver


input_file = 'sample.setpacking'
with open(input_file) as f:
    list_of_subsets = json.load(f)
    print(list_of_subsets)
    qubitOp, offset = setpacking.get_setpacking_qubitops(list_of_subsets)
    algo_input = EnergyInput(qubitOp)

def brute_force():
    # brute-force way: try every possible assignment!
    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 = len(list_of_subsets)
    max = 2**L
    max_v = -np.inf
    for i in range(max):
        cur = bitfield(i, L)
        cur_v = setpacking.check_disjoint(cur, list_of_subsets)
        if cur_v:
            if np.count_nonzero(cur) > max_v:
                max_v = np.count_nonzero(cur)
    return max_v

size = brute_force()
print("size of set packing", size)

params = {
    'problem': {'name': 'ising'},
    'algorithm': {'name': 'ExactEigensolver'}
}
result = run_algorithm(params, algo_input)
x = setpacking.sample_most_likely(len(list_of_subsets), result['eigvecs'][0])
ising_sol = setpacking.get_solution(x)
np.testing.assert_array_equal(ising_sol, [0, 1, 1])
print("size of set packing", np.count_nonzero(ising_sol))

algo = ExactEigensolver(algo_input.qubit_op, k=1, aux_operators=[])
result = algo.run()
x = setpacking.sample_most_likely(len(list_of_subsets), result['eigvecs'][0])
ising_sol = setpacking.get_solution(x)
np.testing.assert_array_equal(ising_sol, [0, 1, 1])
oracle = brute_force()
print("size of set packing", np.count_nonzero(ising_sol))

algorithm_cfg = {
    'name': 'VQE',
    'operator_mode': 'paulis'

}

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

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

params = {
    'problem': {'name': 'ising', 'random_seed': 100},
    'algorithm': algorithm_cfg,
    'optimizer': optimizer_cfg,
    'variational_form': var_form_cfg
}
backend = Aer.get_backend('qasm_simulator')
result = run_algorithm(params, algo_input, backend=backend)
x = setpacking.sample_most_likely(len(list_of_subsets), result['eigvecs'][0])
ising_sol = setpacking.get_solution(x)
print("size of set packing", np.count_nonzero(ising_sol))