%pip install yfinance
%pip install cplex
%pip qiskit_optimization

%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

from docplex.mp.model import Model

colors = plt.cm.tab20(range(20))
label_size = 22
plt.rcParams['xtick.labelsize'] = label_size 
plt.rcParams['ytick.labelsize'] = label_size 
plt.rcParams['axes.labelsize'] = label_size 
plt.rcParams['legend.fontsize'] = label_size 
plt.rcParams['axes.titlesize'] = label_size 

stocks_info = {"Stocks":["MSFT", "AAPL", "GOOG", "PFE", "MRNA", "JPM", "C", "XOM", "HAL", "HON"],
          "Name":["Microsoft", "Apple Inc.", "Alphabet", "Pfizer", "Moderna","JPMorgan Chase", "Citigroup", "ExxonMobil","Halliburton", "Honeywell"],
          "Sector":3*["Information Technology"] + 2 * ["Health Care"] + 2* ["Financial"] + 2 * ["Energy"]+ ["Industrial"]}
pd.DataFrame(stocks_info)

import yfinance as yf
stocks = stocks_info["Stocks"]
tickers = yf.Tickers(stocks)

tickers.tickers["AAPL"].history(period="12mo")

fig, ax = plt.subplots(2,2, figsize=(24,17))
colors_case = {"Open":colors[0], "High":colors[5], "Low":colors[10], "Close":colors[19]}
linestyle = {"Open":"-", "High":"--", "Low":":", "Close":"-."}

for case in ["Open", "High","Low", "Close"]:
    ax[0,0].plot(tickers.tickers["AAPL"].history(period="12mo")[case], color = colors_case[case], linestyle=linestyle[case], label=case)
ax[0,0].legend()
ax[0,1].plot(tickers.tickers["AAPL"].history(period="12mo")["Volume"], label=case)
ax[1,0].plot(tickers.tickers["AAPL"].history(period="12mo")["Dividends"], label=case)
ax[1,1].plot(tickers.tickers["AAPL"].history(period="12mo")["Stock Splits"], label=case)
names = [["Stock Price", "Volume"], ["Dividends", "Stock Splits"]]
for i in range(2):
    for j in range(2):
        ax[i,j].set_xlabel("date")
        ax[i,j].set_ylabel(names[i][j])
        ax[i,j].set_title("Apple Inc.")

# Price of all the stocks of this example
Pt = np.array([tickers.tickers[stock].history(period="12mo")["Close"] for stock in stocks])

mu = (Pt[:,1:] - Pt[:,:-1]) / Pt[:,:-1]

plt.figure()
plt.plot(mu[0])
plt.xlabel("Days")
plt.ylabel(r"$\mu_t$")
plt.title("Apple Inc.")

N_stocks, N_days = mu.shape
sigma = np.zeros((N_stocks, N_stocks))
for i in range(N_stocks-1):
    for j in range(i, N_stocks):
        sigma[i,j] = ((mu[i] - mu[i].mean()) * (mu[j] - mu[j].mean())).sum()/(N_time - 1)
sigma += sigma.T
# sigma = np.cov(mu, varcol=True)

fig, ax = plt.subplots(1, 2, figsize=(22,9))
for i in range(N_stocks):
    ax[0].plot(range(N_days+1), Pt[i], color=colors[i], label=stocks[i])
ax[0].set_xlabel("day")
ax[0].set_ylabel("Stock Price [$]")
ax[0].legend(loc='upper center', bbox_to_anchor=(0.5, 1.2), ncol=5, fancybox=True, shadow=True)
im = ax[1].imshow(sigma, cmap="coolwarm", vmin=-0.0002, vmax=0.0006)
fig.colorbar(im)
ax[1].set_yticks(range(N_stocks))
ax[1].set_xticks(range(N_stocks))
ax[1].set_xticklabels(stocks, rotation=45)
ax[1].set_yticklabels(stocks, rotation=45)
ax[1].set_title(r"Covariance Matrix $\Sigma$ associated to the risk")

q = 0.1 # Risk aversion
Budget = 7 # I can invest in 7 stocks
c = N_stocks * [1] # Minimum we can invest in the different assets c_i

mdl = Model('Portfolio Optimization')

# Binary set of variables that represent the stocks
x = np.array(mdl.binary_var_list(N_stocks, name=stocks)) # x vector in numpy array for matrix multiplication

# Portfolio optimization function
objective_function = mu.mean(axis=1) @ x - q * x.T @ sigma @ x
mdl.maximize(objective_function)

# Budget constraint
mdl.add_constraint(c @ x == Budget, ctname='budget')

# Printing the docplex model
mdl.prettyprint()

sol_docplex = mdl.solve()
sol_docplex = {st: int(sol_docplex.get_value(st)) for st in stocks}
sol_docplex

from qiskit_optimization.translators import from_docplex_mp
from qiskit_optimization.runtime import QAOAClient, VQEClient
from qiskit.algorithms import VQE, QAOA
from qiskit_optimization.algorithms import MinimumEigenOptimizer
from qiskit import Aer  # local simulator
from qiskit.algorithms.optimizers import SPSA, COBYLA

backend = Aer.get_backend("qasm_simulator")
qubo = from_docplex_mp(mdl)

def optimization_QAOA(qubo, reps=1, optimizer=SPSA(maxiter=50), backend=None,
                      shots=1024, provider=None):
    intermediate_info = {'nfev': [],
                         'parameters': [],
                         'stddev': [],
                         'mean': []
                             }
    
    def callback(nfev, parameters, mean, stddev):
        intermediate_info['nfev'].append(nfev)
        intermediate_info['parameters'].append(parameters)
        intermediate_info['mean'].append(mean)
        intermediate_info['stddev'].append(stddev)
    
    qaoa_mes = QAOA(optimizer=optimizer, reps=reps, quantum_instance=backend,
                    callback=callback)
    qaoa = MinimumEigenOptimizer(qaoa_mes)
    result = qaoa.solve(qubo)
    return result, intermediate_info

sol_spsa, info_spsa = optimization_QAOA(qubo, backend=backend)

sol_cobyla, info_cobyla = optimization_QAOA(qubo, backend=backend, optimizer=COBYLA())

solutions = {"Stocks":stocks_info["Stocks"], "CPLEX":sol_docplex.values(), "COBYLA":sol_cobyla.x, "SPSA":sol_spsa.x}
pd.DataFrame(solutions)

plt.figure()
plt.plot(info_spsa["nfev"], info_spsa["mean"], label="SPSA", color=colors[5])
plt.plot(info["nfev"], info["mean"], label = "COBYLA", color=colors[10])
plt.legend()
plt.xlabel("num func evals")
plt.ylabel("Cost func.")

from qiskit.circuit.library import TwoLocal

ansatz = TwoLocal(N_stocks, rotation_blocks='ry', entanglement_blocks='cz')

def Optimization_VQE(qubo, ansatz, optimizer=SPSA(maxiter=50), backend=None, shots=1024):

    intermediate_info = {'nfev': [],
                         'parameters': [],
                         'stddev': [],
                         'mean': []
                         }

    def callback(nfev, parameters, mean, stddev):
        intermediate_info['nfev'].append(nfev)
        intermediate_info['parameters'].append(parameters)
        intermediate_info['mean'].append(mean)
        intermediate_info['stddev'].append(stddev)
        
    vqe_mes = VQE(ansatz=ansatz, quantum_instance=backend, callback=callback, optimizer=optimizer)
    vqe = MinimumEigenOptimizer(vqe_mes)
    result = vqe.solve(qubo)
    return result, intermediate_info

import qiskit.tools.jupyter
%qiskit_version_table


