from qiskit import BasicAer
from qiskit_aqua import QuantumInstance
from qiskit_aqua import Operator, run_algorithm
from qiskit_aqua.input import EnergyInput
from qiskit_aqua.translators.ising import portfolio
from qiskit_aqua.algorithms import VQE, QAOA, ExactEigensolver
from qiskit_aqua.components.optimizers import COBYLA
from qiskit_aqua.components.variational_forms import RY
import numpy as np

from qiskit import IBMQ
IBMQ.load_accounts()

# set number of assets (= number of qubits)
num_assets = 4
# get random expected return vector (mu) and covariance matrix (sigma)
mu, sigma = portfolio.random_model(num_assets, seed=42)

q = 0.5 # set risk factor
budget = int(num_assets / 2) # set budget
penalty = num_assets # set parameter to scale the budget penalty term

qubitOp, offset = portfolio.get_portfolio_qubitops(mu, sigma, q, budget, penalty)
algo_input = EnergyInput(qubitOp)

def index_to_selection(i, num_assets):
    s = "{0:b}".format(i).rjust(num_assets)
    x = np.array([1 if s[i]=='1' else 0 for i in reversed(range(num_assets))])
    return x

def print_result(result):
    selection = portfolio.sample_most_likely(result['eigvecs'][0])
    value = portfolio.portfolio_value(selection, mu, sigma, q, budget, penalty)
    print('Optimal: selection {}, value {:.4f}'.format(selection, value))

    probabilities = np.abs(result['eigvecs'][0])**2
    i_sorted = reversed(np.argsort(probabilities))
    print('\n----------------- Full result ---------------------')
    print('selection\tvalue\t\tprobability')
    print('---------------------------------------------------')
    for i in i_sorted:
        x = index_to_selection(i, num_assets)
        value = portfolio.portfolio_value(x, mu, sigma, q, budget, penalty)    
        probability = probabilities[i]
        print('%10s\t%.4f\t\t%.4f' %(x, value, probability))

exact_eigensolver = ExactEigensolver(qubitOp, k=1)
result = exact_eigensolver.run()

""" the equivalent if using declarative approach
algorithm_cfg = {
    'name': 'ExactEigensolver'
}

params = {
    'problem': {'name': 'ising'},
    'algorithm': algorithm_cfg
}
result = run_algorithm(params, algo_input)
"""

print_result(result)

backend = BasicAer.get_backend('statevector_simulator')
seed = 50

cobyla = COBYLA()
cobyla.set_options(maxiter=250)
ry = RY(qubitOp.num_qubits, depth=3, entanglement='full')
vqe = VQE(qubitOp, ry, cobyla, 'matrix')
vqe.random_seed = seed

quantum_instance = QuantumInstance(backend=backend, seed=seed, seed_mapper=seed)

result = vqe.run(quantum_instance)

"""declarative approach
algorithm_cfg = {
    'name': 'VQE',
    'operator_mode': 'matrix'
}

optimizer_cfg = {
    'name': 'COBYLA',
    'maxiter': 250
}

var_form_cfg = {
    'name': 'RY',
    'depth': 3,
    'entanglement': 'full'
}

params = {
    'problem': {'name': 'ising', 'random_seed': seed},
    'algorithm': algorithm_cfg,
    'optimizer': optimizer_cfg,
    'variational_form': var_form_cfg
}
result = run_algorithm(params, algo_input, backend=backend)
"""
print_result(result)

backend = BasicAer.get_backend('statevector_simulator')
seed = 50

cobyla = COBYLA()
cobyla.set_options(maxiter=250)
qaoa = QAOA(qubitOp, cobyla, 3, 'matrix')
qaoa.random_seed = seed

quantum_instance = QuantumInstance(backend=backend, seed=seed, seed_mapper=seed)

result = qaoa.run(quantum_instance)

"""declarative approach
algorithm_cfg = {
    'name': 'QAOA.Variational',
    'p': 3,
    'operator_mode': 'matrix'
}

optimizer_cfg = {
    'name': 'COBYLA',
    'maxiter': 250
}

params = {
    'problem': {'name': 'ising', 'random_seed': seed},
    'algorithm': algorithm_cfg,
    'optimizer': optimizer_cfg
}
result = run_algorithm(params, algo_input, backend=backend)
"""
print_result(result)

