# ! pip install qiskit qiskit-finance==0.2.1

import warnings
warnings.filterwarnings("ignore")
import qiskit.tools.jupyter
%qiskit_version_table

# from qiskit import IBMQ
# IBMQ.save_account('MY_API_TOKEN')

# ! pip install eikon

import eikon as ek
print("Eikon version: ", ek.__version__)

eikon_key = open("eikon.txt","r")
ek.set_app_key(str(eikon_key.read()))
eikon_key.close()

import numpy as np
import pandas as pd
import datetime
import matplotlib.pyplot as plt
import seaborn as sns 
from qiskit import Aer
from qiskit_finance.data_providers._base_data_provider import BaseDataProvider
from qiskit.finance.applications.ising import portfolio
from qiskit.circuit.library import TwoLocal
from qiskit.aqua import QuantumInstance
from qiskit.optimization.applications.ising.common import sample_most_likely
from qiskit.aqua.algorithms import VQE, QAOA, NumPyMinimumEigensolver
from qiskit.aqua.components.optimizers import COBYLA

class EikonDataProvider(BaseDataProvider):
    """
    The abstract base class for Eikon data_provider.
    """
    def __init__(self, stocks_list, start_date, end_date):
        '''
        stocks -> List of interested assets
        start -> start date to fetch historical data
        end -> 
        '''
        super().__init__()
        self._stocks = stocks_list
        self._start = start_date
        self._end = end_date
        self._data = []
        self.stock_data = pd.DataFrame()
        
    def run(self):
        self._data = []
        stocks_notfound = []
        stock_data = ek.get_timeseries(self._stocks,
                  start_date=self._start, 
                  end_date=self._end, 
                  interval='daily',
                  corax='adjusted')
        for ticker in self._stocks:
            stock_value = stock_data[ticker]['CLOSE']
            self.stock_data[ticker] = stock_data[ticker]['CLOSE']
            if stock_value.dropna().empty:
                stocks_notfound.append(ticker)
            self._data.append(stock_value)

# List of stocks 
stock_list = ['FB.O', 'AAPL.O', 'AMZN.O', 'NFLX.O', 'GOOGL.O']

# Start Date
start_date = datetime.datetime(2019,6,10)

# End Date
end_date = datetime.datetime(2021,10,20)

# Set number of equities to the number of stocks
num_assets = len(stock_list)

# Set the risk factor
risk_factor = 0.7

# Set budget
budget = 2

# Scaling of budget penalty term will be dependant on the number of assets
penalty = num_assets

data = EikonDataProvider(stocks_list = stock_list, start_date=start_date, end_date=end_date)
data.run()

# Top 5 rows of data 
df = data.stock_data
df.head()

df.describe()

# Closing Price History
fig, ax = plt.subplots(figsize=(15, 8))
ax.plot(df)
plt.title('Close Price History')
plt.xlabel('Date',fontsize =20)
plt.ylabel('Price in USD',fontsize = 20)
ax.legend(df.columns.values)
plt.show()

# # Mean value of each equities
# mean_vector = data.get_mean_vector()

# sns.barplot(y=mean_vector, x = stock_list)
# plt.show()

mu = data.get_period_return_mean_vector()

sns.barplot(y=mu, x = stock_list)
plt.show()

# Covariance Matrix
sigma = data.get_period_return_covariance_matrix()

f, ax = plt.subplots(figsize=(10, 8))
sns.heatmap(sigma, mask=np.zeros_like(sigma, dtype=np.bool), annot=True,
            square=True, ax=ax, xticklabels=stock_list, yticklabels=stock_list)
plt.title("Covariance between Equities")
plt.show()

# Correlation Matrix
corr = df.corr()

f, ax = plt.subplots(figsize=(10, 8))
sns.heatmap(corr, mask=np.zeros_like(corr, dtype=np.bool), annot=True,
            square=True, ax=ax, xticklabels=stock_list, yticklabels=stock_list)
plt.title("Correlation between Equities")
plt.show()

sns.pairplot(df)
plt.show()

# returns = np.log(df/df.shift(1)).dropna()
# returns.head()

# sns.pairplot(returns)
# plt.show()

# Retrieve Hamiltonian
qubitOp, offset = portfolio.get_operator(mu, sigma, risk_factor, budget, penalty)

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 selection_to_picks(num_assets, selection):
    purchase = []
    for i in range(num_assets):
        if selection[i] == 1:
            purchase.append(stock_list[i])
    return purchase

# def print_result(result):
#     selection = sample_most_likely(result.eigenstate)
#     value = portfolio.portfolio_value(selection, mu, sigma, q, budget, penalty)
#     print(f"!!! Optimal: portfolio holdings !!!")
#     return selection

def print_result(result):
    selection = sample_most_likely(result.eigenstate)
    value = portfolio.portfolio_value(selection, mu, sigma, risk_factor, budget, penalty)
    print('Optimal: selection {}, value {:.4f}'.format(selection, value))

    eigenvector = result.eigenstate if isinstance(result.eigenstate, np.ndarray) else result.eigenstate.to_matrix()
    probabilities = np.abs(eigenvector)**2
    i_sorted = reversed(np.argsort(probabilities))
    print('\n----------------- Full result ---------------------')
    print('selection\tvalue\t\tprobability')
    print('---------------------------------------------------')
    states, values, probs = [], [], []
    for i in i_sorted:
        x = index_to_selection(i, num_assets)
        value = portfolio.portfolio_value(x, mu, sigma, risk_factor, budget, penalty)
        probability = probabilities[i]
        print('%10s\t%.4f\t\t%.4f' %(x, value, probability))
        states.append(''.join(str(i) for i in x))
        values.append(value)
        probs.append(probability)
    return selection, states, values, probs

def numpyEigensolver(qubitOp):
    selections = []
    exact_eigensolver = NumPyMinimumEigensolver(qubitOp)
    result = exact_eigensolver.run()
    selection, state, values, probabilities = print_result(result)
    print(selection_to_picks(num_assets, selection))
    return state, values, probabilities

#Classical Benchmark
%time
state, values, probabilities = numpyEigensolver(qubitOp) # ['FB.O', 'AAPL.O', 'AMZN.O', 'NFLX.O', 'GOOGL.O']

# f, ax = plt.subplots(figsize=(20, 8))
# optimized_value = sns.barplot(state, values)
# for item in optimized_value.get_xticklabels():
#     item.set_rotation(45)
# plt.title("Optimized Values of Each Possible Combination of Assets", fontsize=25) #Lower the value is, best optimal set of asset it is
# plt.xlabel('Possible Combinations of Assets',fontsize =20)
# plt.ylabel('Optimized Values',fontsize = 20)
# plt.show()

f, ax = plt.subplots(figsize=(20, 8))
optimized_value = sns.barplot(state, probabilities)
for item in optimized_value.get_xticklabels():
    item.set_rotation(45)
plt.title("Probabilites of Each Possible Combination of Assets", fontsize=25) # Probability of Max Return at given risk level
plt.xlabel('Possible Combinations of Assets',fontsize =20)
plt.ylabel('Probability',fontsize = 20)
plt.show()

def vqe(qubitOp):
    backend = Aer.get_backend('statevector_simulator') # You can switch to different backends by providing the name of backend.
    seed = 50
    cobyla = COBYLA()
    cobyla.set_options(maxiter=1000)
    ry = TwoLocal(qubitOp.num_qubits, 'ry', 'cz', reps=3, entanglement='full')
    vqe = VQE(qubitOp, ry, cobyla)
    vqe.random_seed = seed
    quantum_instance = QuantumInstance(backend=backend, seed_simulator=seed, seed_transpiler=seed)
    result = vqe.run(quantum_instance)
    selection, state, values, probabilities = print_result(result)
    print(selection_to_picks(num_assets, selection))
    return state, values, probabilities

# Variational Quantum EigenSolver
%time
state, values, probabilities = vqe(qubitOp) # ['FB.O', 'AAPL.O', 'AMZN.O', 'NFLX.O', 'GOOGL.O']

# f, ax = plt.subplots(figsize=(20, 8))
# optimized_value = sns.barplot(state, values)
# for item in optimized_value.get_xticklabels():
#     item.set_rotation(45)
# plt.title("Optimized Values of Each Possible Combination of Assets", fontsize=25) #Lower the value is, best optimal set of asset it is
# plt.xlabel('Possible Combinations of Assets',fontsize =20)
# plt.ylabel('Optimized Values',fontsize = 20)
# plt.show()

f, ax = plt.subplots(figsize=(20, 8))
optimized_value = sns.barplot(state, probabilities)
for item in optimized_value.get_xticklabels():
    item.set_rotation(45)
plt.title("Probabilites of Each Possible Combination of Assets", fontsize=25) # Probability of Max Return at given risk level
plt.xlabel('Possible Combinations of Assets',fontsize =20)
plt.ylabel('Probability',fontsize = 20)
plt.show()

def qaoa(qubitOp):
    backend = Aer.get_backend('statevector_simulator') # You can switch to different backends by providing the name of backend.
    seed = 50
    cobyla = COBYLA()
    cobyla.set_options(maxiter=1000)
    qaoa = QAOA(qubitOp, cobyla, 3)
    qaoa.random_seed = seed
    quantum_instance = QuantumInstance(backend=backend, seed_simulator=seed, seed_transpiler=seed)
    results= qaoa.run(quantum_instance)
    selection, state, values, probabilities = print_result(results)
    print(selection_to_picks(num_assets, selection))
    return state, values, probabilities

# QAOA (Quantum Approximate Optimization Algorithm)
%time
state, values, probabilities = qaoa(qubitOp) # ['FB.O', 'AAPL.O', 'AMZN.O', 'NFLX.O', 'GOOGL.O']

# f, ax = plt.subplots(figsize=(20, 8))
# optimized_value = sns.barplot(state, values)
# for item in optimized_value.get_xticklabels():
#     item.set_rotation(45)
# plt.title("Optimized Values of Each Possible Combination of Assets", fontsize=25) #Lower the value is, best optimal set of asset it is
# plt.xlabel('Possible Combinations of Assets',fontsize =20)
# plt.ylabel('Optimized Values',fontsize = 20)
# plt.show()

f, ax = plt.subplots(figsize=(20, 8))
optimized_value = sns.barplot(state, probabilities)
for item in optimized_value.get_xticklabels():
    item.set_rotation(45)
plt.title("Probabilites of Each Possible Combination of Assets", fontsize=25) # Probability of Max Return at given risk level
plt.xlabel('Possible Combinations of Assets',fontsize =20)
plt.ylabel('Probability',fontsize = 20)
plt.show()