# useful additional packages 
import matplotlib.pyplot as plt
import matplotlib.axes as axes
%matplotlib inline
import numpy as np
import networkx as nx

from qiskit import BasicAer
from qiskit.tools.visualization import plot_histogram
from qiskit.aqua import run_algorithm
from qiskit.aqua.input import EnergyInput
from qiskit.aqua.translators.ising import max_cut, tsp
from qiskit.aqua.algorithms import VQE, ExactEigensolver
from qiskit.aqua.components.optimizers import SPSA
from qiskit.aqua.components.variational_forms import RY
from qiskit.aqua import QuantumInstance

from qiskit.quantum_info import Pauli
from qiskit.aqua.operators import WeightedPauliOperator
from collections import OrderedDict
import math

# setup aqua logging
import logging
from qiskit.aqua import set_qiskit_aqua_logging

def get_knapsack_operator(values, weights, max_weight):
    if len(values) != len(weights):
        raise ValueError("The values and weights must have the same length")

    if any(v < 0 for v in values) or any(w < 0 for w in weights):
        raise ValueError("The values and weights cannot be negative")

    if all(v == 0 for v in values):
        raise ValueError("The values cannot all be 0")

    if max_weight < 0:
        raise ValueError("max_weight cannot be negative")

    y_size = int(math.log(max_weight, 2)) + 1 if max_weight > 0 else 1
    print(y_size)
    n = len(values)
    num_values = n + y_size
    pauli_list = []
    shift = 0

    M = 2000000  #10 * np.sum(values)

    # term for sum(x_i*w_i)**2
    for i in range(n):
        for j in range(n):
            coefficient = -1 * 0.25 * weights[i] * weights[j] * M

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[j] = not zp[j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[i] = not zp[i]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            coefficient = -1 * coefficient
            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[i] = not zp[i]
            zp[j] = not zp[j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

    # term for sum(2**j*y_j)**2
    for i in range(y_size):
        for j in range(y_size):
            coefficient = -1 * 0.25 * (2 ** i) * (2 ** j) * M

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[n + j] = not zp[n + j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[n + i] = not zp[n + i]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            coefficient = -1 * coefficient
            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[n + i] = not zp[n + i]
            zp[n + j] = not zp[n + j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

    # term for -2*W_max*sum(x_i*w_i)
    for i in range(n):
        coefficient = max_weight * weights[i] * M

        xp = np.zeros(num_values, dtype=np.bool)
        zp = np.zeros(num_values, dtype=np.bool)
        zp[i] = not zp[i]
        pauli_list.append([coefficient, Pauli(zp, xp)])
        shift -= coefficient

    # term for -2*W_max*sum(2**j*y_j)
    for j in range(y_size):
        coefficient = max_weight * (2 ** j) * M

        xp = np.zeros(num_values, dtype=np.bool)
        zp = np.zeros(num_values, dtype=np.bool)
        zp[n + j] = not zp[n + j]
        pauli_list.append([coefficient, Pauli(zp, xp)])
        shift -= coefficient

    for i in range(n):
        for j in range(y_size):
            coefficient = -1 * 0.5 * weights[i] * (2 ** j) * M

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[n + j] = not zp[n + j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[i] = not zp[i]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

            coefficient = -1 * coefficient
            xp = np.zeros(num_values, dtype=np.bool)
            zp = np.zeros(num_values, dtype=np.bool)
            zp[i] = not zp[i]
            zp[n + j] = not zp[n + j]
            pauli_list.append([coefficient, Pauli(zp, xp)])
            shift -= coefficient

    # term for sum(x_i*v_i)
    for i in range(n):
        coefficient = 0.5 * values[i]

        xp = np.zeros(num_values, dtype=np.bool)
        zp = np.zeros(num_values, dtype=np.bool)
        zp[i] = not zp[i]
        pauli_list.append([coefficient, Pauli(zp, xp)])
        shift -= coefficient

    return WeightedPauliOperator(paulis=pauli_list), shift


def get_solution(x, values):
    return x[:len(values)]


def knapsack_value_weight(solution, values, weights):
    value = np.sum(solution * values)
    weight = np.sum(solution * weights)
    return value, weight

def sample_most_likely(state_vector):
    if isinstance(state_vector, dict) or isinstance(state_vector, OrderedDict):
        # get the binary string with the largest count
        binary_string = sorted(state_vector.items(), key=lambda kv: kv[1])[-1][0]
        x = np.asarray([int(y) for y in reversed(list(binary_string))])
        return x
    else:
        n = int(np.log2(state_vector.shape[0]))
        k = np.argmax(np.abs(state_vector))
        x = np.zeros(n)
        for i in range(n):
            x[i] = k % 2
            k >>= 1
        return x


# values = [680, 120, 590, 178]
# weights = [13, 6, 15, 9]
# w_max = 32

values = [8, 2]
weights = [3, 1]
w_max = 3

qubitOp, offset = get_knapsack_operator(values, weights, w_max)

algo_input = EnergyInput(qubitOp)
ee = ExactEigensolver(qubitOp, k=1)
result = ee.run()

most_lightly = result['eigvecs'][0]
x = sample_most_likely(most_lightly)

print(x)
print('result=' + str(x[:len(values)]))

seed = 10598

spsa = SPSA(max_trials=10000)
ry = RY(qubitOp.num_qubits, depth=5, entanglement='linear')
vqe = VQE(qubitOp, ry, spsa)

backend = BasicAer.get_backend('statevector_simulator')
quantum_instance = QuantumInstance(backend, seed_simulator=seed, seed_transpiler=seed)

result_statevector = vqe.run(quantum_instance)

most_lightly_sv = result_statevector['eigvecs'][0]
x_statevector = sample_most_likely(most_lightly_sv)

print(x_statevector)
print('result usig statevector_simulator =' + str(x_statevector[:len(values)]))

# run quantum algorithm with shots
seed = 10598

spsa = SPSA(max_trials=1000)
ry = RY(qubitOp.num_qubits, depth=5, entanglement='linear')
vqe = VQE(qubitOp, ry, spsa)

backend = BasicAer.get_backend('qasm_simulator')
quantum_instance = QuantumInstance(backend, shots=1024, seed_simulator=seed, seed_transpiler=seed)

result_shots = vqe.run(quantum_instance)

most_lightly_shots = result_shots['eigvecs'][0]
x_shots = sample_most_likely(most_lightly_shots)

print(x_shots)
print('result usig qasm_simulator =' + str(x_shots[:len(values)]))

math.log(8, 2)