# %pip install --upgrade pip
# %pip install -r requirements.txt

import seaborn as sn
import matplotlib.pyplot as plt
from qiskit.visualization import plot_histogram
from datetime import datetime
import numpy as np
from itertools import permutations

from qiskit_finance.data_providers import YahooDataProvider
import json # Necessary for YahooDataProvider
import http.client # Necessary for YahooDataProvider

import pandas as pd
import pandas_datareader as web

import os

from docplex.mp.model import Model

from qiskit_optimization.translators import from_docplex_mp
from qiskit_optimization.algorithms import MinimumEigenOptimizer
from qiskit.algorithms import NumPyMinimumEigensolver
from qiskit_optimization.converters import LinearEqualityToPenalty

from qiskit.circuit.library import TwoLocal # For Ansatz

from qiskit.algorithms.optimizers import COBYLA, ADAM, NELDER_MEAD, POWELL, SPSA, QNSPSA
from qiskit.opflow import CVaRExpectation
from qiskit.opflow import PauliExpectation

from qiskit.algorithms import VQE
from qiskit.algorithms import QAOA

from qiskit import Aer

dfMonthlyReturns = pd.read_csv('data/Equity-Benchmarks-Monthly-Return.csv',parse_dates=["Date"],index_col=0)
dfMonthlyReturnsRiskScaled = pd.read_csv('data/Equity-Benchmarks-Monthly-Risk-Scaled.csv',parse_dates=["Date"],index_col=0)

assets = dfMonthlyReturns.columns.to_list()
numAssets = len(assets)
print("We have " + str(numAssets) + " Assets:\n")
[print(asset + "  ") for asset in assets];

q = 0.5 # risk apetite
B = 5 # budget

#Whole Period
startDate=datetime(2001, 1, 31)
endDate=datetime(2021, 12, 31)

# #Pre 2008
# startDate=datetime(2001, 1, 31)
# endDate=datetime(2008, 6, 30)

# #QE Era
# startDate=datetime(2010, 1, 31)
# endDate=datetime(2019, 12, 31)

# #Pandemic
# startDate=datetime(2020, 6, 30)
# endDate=datetime(2021, 12, 31)

startStr = str(startDate.date())
endStr = str(endDate.date())

backend = Aer.get_backend('statevector_simulator') # statevector_simulator, qasm_simulator, aer_simulator_statevector

maxIter = 10000
# fidelity = QNSPSA.get_fidelity(ansatz) # only for QNSPSA
optimizer = SPSA(maxiter=500)  

# Classical optimizer – COBYLA(maxiter=maxIter), ADAM(maxiter=maxIter), NELDER_MEAD(maxiter=maxIter), 
# POWELL(maxiter=maxIter), SPSA(maxiter=100), QNSPSA(fidelity,maxiter=100)

penalty = numAssets # scaling of penalty term. This can be changed at your discretion.

mu = dfMonthlyReturns[startDate:endDate].mean().to_numpy()
sigma = dfMonthlyReturns[startDate:endDate].cov().to_numpy()

muRisk = dfMonthlyReturnsRiskScaled[startDate:endDate].mean().to_numpy()
sigmaRisk = dfMonthlyReturnsRiskScaled[startDate:endDate].cov().to_numpy()

covHeatmap = sn.heatmap(sigma,cmap="mako")
covHeatmap.set(xticklabels=assets,yticklabels=assets)
plt.yticks(rotation='horizontal')
plt.xticks(rotation='vertical')
plt.title("Covariance of Monthly Returns\n"+startStr+" to " +endStr)
plt.show()

covHeatmapRisk = sn.heatmap(sigmaRisk,cmap="mako")
covHeatmapRisk.set(xticklabels=assets,yticklabels=assets)
plt.yticks(rotation='horizontal')
plt.xticks(rotation='vertical')
plt.title("Covariance of Monthly Returns (Risk Scaled)\n"+startStr+" to " +endStr)
plt.show()

covHeatmap = sn.heatmap(sigma,cmap="mako",vmin=0,vmax=0.004)
covHeatmap.set(xticklabels=assets,yticklabels=assets)
plt.yticks(rotation='horizontal')
plt.xticks(rotation='vertical')
plt.title("Covariance of Monthly Returns\n"+startStr+" to " +endStr)
plt.show()

fig = covHeatmap.get_figure()
fig.savefig("figures/CovarianceMonthlyReturns-" + startStr + "-to-" + endStr + ".pdf",bbox_inches='tight')

covHeatmapRisk = sn.heatmap(sigmaRisk,cmap="mako",vmin=0,vmax=0.00011)
covHeatmapRisk.set(xticklabels=assets,yticklabels=assets)
plt.yticks(rotation='horizontal')
plt.xticks(rotation='vertical')
plt.title("Covariance of Monthly Returns (Risk Scaled)\n"+startStr+" to " +endStr)
plt.show()

figRisk = covHeatmapRisk.get_figure()
figRisk.savefig("figures/CovarianceMonthlyReturnsRisk-" + startStr + "-to-" + endStr + ".pdf",bbox_inches='tight')

mdl = Model("Non-Scaled") # You can input a string into the argument if you want to name the model.

x = mdl.binary_var_list('x{}'.format(i) for i in range(numAssets)) # binary decision variables

# Objective Function to be maximized
objective = mdl.sum([mu[i]*x[i] for i in range(numAssets)])
objective -= q * mdl.sum([sigma[i,j]*x[i]*x[j] for i in range(numAssets) for j in range(numAssets)])

mdl.maximize(objective)
mdl.add_constraint(mdl.sum(x[i] for i in range(numAssets)) == B) # Adding the budget constraint

qp = from_docplex_mp(mdl) # Conversion to a quadratic program

qp

mdlRisk = Model("Risk-Scaled") # You can input a string into the argument if you want to name the model.

y = mdlRisk.binary_var_list('y{}'.format(i) for i in range(numAssets)) # binary decision variables

# Objective Function to be maximized
objectiveRisk = mdlRisk.sum([muRisk[i]*y[i] for i in range(numAssets)])
objectiveRisk -= q * mdlRisk.sum([sigmaRisk[i,j]*y[i]*y[j] for i in range(numAssets) for j in range(numAssets)])

mdlRisk.maximize(objectiveRisk)
mdlRisk.add_constraint(mdlRisk.sum(y[i] for i in range(numAssets)) == B) # Adding the budget constraint

qpRisk = from_docplex_mp(mdlRisk) # Conversion to a quadratic program

qpRisk

classicalResult = MinimumEigenOptimizer(NumPyMinimumEigensolver()).solve(qp)
# classicalResultEigenstate = (classicalResult.min_eigen_solver_result.eigenstate._primitive.data.real).tolist()
classicalResultEigenstate = classicalResult.min_eigen_solver_result.eigenstate.to_dict_fn().primitive
classicalResult

# making the labels for plotting
# labels = []
# for i in range(2**numAssets):
#     labels.append(format(i,'0'+str(numAssets)+'b'))
# fig = plot_histogram(dict(zip(labels, classicalResultEigenstate)),figsize=(60,10))
fig = plot_histogram(classicalResultEigenstate,figsize=(60,10))
fig.savefig("figures/ClassicalEigensolver-"+startStr+"-"+endStr+".pdf")
# fig

classicalResultRisk = MinimumEigenOptimizer(NumPyMinimumEigensolver()).solve(qpRisk)
#classicalResultRiskEigenstate = (classicalResultRisk.min_eigen_solver_result.eigenstate._primitive.data.real).tolist()
classicalResultRiskEigenstate = classicalResultRisk.min_eigen_solver_result.eigenstate.to_dict_fn().primitive
classicalResultRisk


# making the labels for plotting
# labels = []
# for i in range(2**numAssets):
#     labels.append(format(i,'0'+str(numAssets)+'b'))
# fig = plot_histogram(dict(zip(labels, classicalResultRiskEigenstate)),figsize=(60,10))
fig = plot_histogram(classicalResultRiskEigenstate,figsize=(60,10))
fig.savefig("figures/ClassicalEigensolverRisk-"+startStr+"-"+endStr+".pdf")
# fig

linearToPenalty = LinearEqualityToPenalty(penalty=penalty)
qubo = linearToPenalty.convert(qp)
quboRisk = linearToPenalty.convert(qpRisk)

qubo

quboRisk

reps = 2

ansatz = TwoLocal(num_qubits=numAssets, rotation_blocks='rx', entanglement_blocks= 'cz', reps=reps, entanglement='full')
ansatz.decompose().draw()

maxIter = 10000
# fidelity = QNSPSA.get_fidelity(ansatz) # only for QNSPSA
optimizer = SPSA(maxiter=500)  

# Classical optimizer – COBYLA(maxiter=maxIter), ADAM(maxiter=maxIter), NELDER_MEAD(maxiter=maxIter), 
# POWELL(maxiter=maxIter), SPSA(maxiter=100), QNSPSA(fidelity,maxiter=100)

def pauliDict(H):
    tempKey = [] 
    tempVal = [] 
    tempDict = {}
    for i in range(len(H)):
        tempDict[H.primitive.to_list()[i][0]] = H.primitive.to_list()[i][1]
    return tempDict

H, offset = qubo.to_ising()
pauliDict(H)

HRisk, offsetRisk = quboRisk.to_ising()
pauliDict(HRisk)

def getCountDictStateVecSim(vqeResult):
    tempState = []
    tempProb = []
    for sample in vqeResult.samples:
        tempProb.append(sample.probability)
        tempStr = ''
        for stateElement in reversed(sample.x):
            tempStr += str(int(stateElement))
        tempState.append(tempStr)
    return dict(zip(tempState, tempProb))

vqeInit = VQE(ansatz=ansatz,optimizer=optimizer, quantum_instance=backend)
vqe = MinimumEigenOptimizer(vqeInit)
vqeResult = vqe.solve(qubo)

if backend.name() == 'qasm_simulator':
    counts = vqeResult.min_eigen_solver_result.eigenstate
elif backend.name() == 'statevector_simulator':
    counts = getCountDictStateVecSim(vqeResult)
#endif

vqeResult

vqeFig = plot_histogram(counts,figsize=(60,10))
vqeFig.savefig("figures/VQE"+startStr+"-"+endStr+".pdf")

vqeFig

## ALTERNATIVE PLOT SCHEME

# samples_for_plot = {
#     " ".join(f"{vqeResult.variables[i].name}={int(v)}" for i, v in enumerate(sample.x)): sample.probability
#     for sample in vqeResult.samples
# }
# plot_histogram(samples_for_plot,figsize=(60,5))

vqeRiskInit = VQE(ansatz=ansatz,optimizer=optimizer, quantum_instance=backend)
vqeRisk = MinimumEigenOptimizer(vqeRiskInit)
vqeRiskResult = vqeRisk.solve(quboRisk)

if backend.name() == 'qasm_simulator':
    countsRisk = vqeRiskResult.min_eigen_solver_result.eigenstate
elif backend.name() == 'statevector_simulator':
    countsRisk = getCountDictStateVecSim(vqeRiskResult)
#endif

vqeRiskResult

vqeRiskFig = plot_histogram(countsRisk,figsize=(60,10))
vqeRiskFig.savefig("figures/VQERisk"+startStr+"-"+endStr+".pdf")

vqeRiskFig

qaoaInit = QAOA(optimizer=optimizer, reps=reps, quantum_instance=backend)
qaoa = MinimumEigenOptimizer(qaoaInit)
qaoaResult = qaoa.solve(qubo)

if backend.name() == 'qasm_simulator':
    counts = qaoaResult.min_eigen_solver_result.eigenstate
elif backend.name() == 'statevector_simulator':
    counts = getCountDictStateVecSim(qaoaResult)
#endif

qaoaResult

qaoaFig = plot_histogram(counts,figsize=(60,10))
qaoaFig.savefig("figures/QAOA-"+startStr+"-"+endStr+".pdf")

qaoaFig

qaoaRiskInit = QAOA(optimizer=optimizer, reps=reps, quantum_instance=backend)
qaoaRisk = MinimumEigenOptimizer(qaoaRiskInit)
qaoaRiskResult = qaoaRisk.solve(quboRisk)

if backend.name() == 'qasm_simulator':
    countsRisk = qaoaRiskResult.min_eigen_solver_result.eigenstate
elif backend.name() == 'statevector_simulator':
    countsRisk = getCountDictStateVecSim(qaoaRiskResult)
#endif

qaoaRiskResult

qaoaRiskFig = plot_histogram(countsRisk,figsize=(60,10))
qaoaRiskFig.savefig("figures/QAOARisk-"+startStr+"-"+endStr+".pdf")

qaoaRiskFig

classicalResult

vqeResult

qaoaResult

classicalResultRisk

vqeRiskResult

qaoaRiskResult

def getCombinations(inputVal):
    combinations = []
    if type(inputVal) == str:
        for perm in set(permutations(inputVal)):
            combinations.append(''.join(perm))
        #endFor
    elif type(inputVal) == list:
        for i in range(len(inputVal)):
            inputVal[i] = int(inputVal[i])
        #endFor
        for perm in set(permutations(inputVal)):
            combinations.append(list(perm))
        #endFor
    else:
        print("Input a string or list of ones and zeroes")
    return combinations

def countNonZeroElemInList(inputList):
    numNonZeroElem = 0.
    for x in inputList:
        if x != 0:
            numNonZeroElem += 1
        else:
            numNonZeroElem += 0
        #endIf
    #endFor
    return numNonZeroElem

def assetTotalPortfolioGrowth(initialPortfolioValue,selectedAssets,testDF):
    amountSplit = initialPortfolioValue/countNonZeroElemInList(selectedAssets)
    assetGrowth = []
    for i in range(len(selectedAssets)):
        if selectedAssets[i] == 1:
            tempDF = testDF.iloc[:,i]
            amountTemp = amountSplit
            for pctChg in tempDF:
                amountTemp += pctChg*amountTemp
            assetGrowth.append(amountTemp)
        else:
            assetGrowth.append(0.)
        #endIf
    #endFor
    return assetGrowth

def assetGrowthTimeSeries(initialPortfolioValue,selectedAssets,testDF):
    amountSplit = initialPortfolioValue/countNonZeroElemInList(selectedAssets)
    assetGrowthTS = []
    for i in range(len(selectedAssets)):
        if selectedAssets[i] == 1:
            tempVec = []
            tempDF = testDF.iloc[:,i]
            amountTemp = amountSplit
            for pctChg in tempDF:
                amountTemp += pctChg*amountTemp
                tempVec.append(amountTemp)
            assetGrowthTS.append(tempVec)
        else:
            tempVec = []
            tempDF = testDF.iloc[:,i]
            for pctChg in tempDF:
                tempVec.append(0.0)
            assetGrowthTS.append(tempVec)
        #endIf
    #endFor
    return assetGrowthTS

def assetGrowthTimeSeriesWeighted(initialPortfolioValue,selectedAssets,testDF,weights):
    assetGrowthTS = []
    for i in range(len(selectedAssets)):
        if selectedAssets[i] == 1:
            tempVec = []
            tempDF = testDF.iloc[:,i]
            amountTemp = initialPortfolioValue*weights[i]
            for pctChg in tempDF:
                amountTemp += pctChg*amountTemp
                tempVec.append(amountTemp)
            assetGrowthTS.append(tempVec)
        else:
            tempVec = []
            tempDF = testDF.iloc[:,i]
            for pctChg in tempDF:
                tempVec.append(0.0)
            assetGrowthTS.append(tempVec)
        #endIf
    #endFor
    return assetGrowthTS


def portfolioGrowthTimeSeries(portfolioTS):
    portfolioValueTS = []
    for i in range(len(portfolioTS[0])):
        tempVal = 0.0
        for j in range(len(portfolioTS)):
            tempVal += portfolioTS[j][i]
        portfolioValueTS.append(tempVal)
    return portfolioValueTS

assetTest = qaoaRiskResult.x.tolist() 
amountInit = 1000 

allAssetCombinations = getCombinations(assetTest)

dateList = dfMonthlyReturns.index

# # Gives total number of chosen assets, AKA budget. More of a sanity check.
# numAssetsChosen = countNonZeroElemInList(assetTest)

# # Gives each selected asset's total final value in a list
# finalAssetValues = assetTotalPortfolioGrowth(amountInit,assetTest,dfMonthlyReturns)

# # Gives the time series of each portfolio element as a list of lists. In each list, 
# # the final element is the final value of the held asset.
# chosenAssetsTimeSeries = assetGrowthTimeSeries(amountInit,assetTest,dfMonthlyReturns)

# # Gives the time series for the total portfolio value. The final element is the final portfolio value.
# portfolioTimeSeries = portfolioGrowthTimeSeries(chosenAssetsTimeSeries)

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1) 
for i in range(len(allAssetCombinations)):
    chosenAssetsTimeSeries = assetGrowthTimeSeries(amountInit,allAssetCombinations[i],dfMonthlyReturns)
    portfolioTimeSeries = portfolioGrowthTimeSeries(chosenAssetsTimeSeries)
    label = ''
    for element in allAssetCombinations[i]:
        label+=str(element)
    #endFor
    ax.plot(dateList,portfolioTimeSeries,color='blue',linestyle='--',alpha=0.1,linewidth=0.5)
#endFor

winningAssetCombination = assetTest
winningAssetTimeSeries = assetGrowthTimeSeries(amountInit,winningAssetCombination,dfMonthlyReturns)
winningPortfolioTimeSeries = portfolioGrowthTimeSeries(winningAssetTimeSeries)

label = ''
for element in winningAssetCombination:
    label+=str(element)
#endFor
ax.plot(dateList,winningPortfolioTimeSeries,label=label,color='red',alpha=1,linewidth=2)

plt.xlim(startDate,endDate)
plt.ylabel('Total Portfolio Value')
plt.xlabel('Date')
plt.legend()#bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)
plt.show()

fig.savefig('figures/portfolio-'+label+'risk-'+str(q)+'-'+startStr+"-"+endStr+'.pdf',bbox_inches='tight')

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1) 
for i in range(len(allAssetCombinations)):
    chosenAssetsTimeSeries = assetGrowthTimeSeries(amountInit,allAssetCombinations[i],dfMonthlyReturns)
    portfolioTimeSeries = portfolioGrowthTimeSeries(chosenAssetsTimeSeries)
    label = ''
    for element in allAssetCombinations[i]:
        label+=str(element)
    #endFor
    ax.plot(dateList,portfolioTimeSeries,color='blue',linestyle='--',alpha=0.1,linewidth=0.5)
#endFor

winningAssetCombination = assetTest
winningAssetTimeSeries = assetGrowthTimeSeries(amountInit,winningAssetCombination,dfMonthlyReturns)
winningPortfolioTimeSeries = portfolioGrowthTimeSeries(winningAssetTimeSeries)

label = ''
for element in winningAssetCombination:
    label+=str(element)
#endFor

label += ' equally weighted'
ax.plot(dateList,winningPortfolioTimeSeries,label=label,color='red',alpha=1,linewidth=2)

portfolioWeights = [0.223214286,0.233800935,0.2042007,0.0,0.0,0.0,0.172583826,0.165250236]
weightedWinningAssetTimeSeries = assetGrowthTimeSeriesWeighted(amountInit,assetTest,dfMonthlyReturns,portfolioWeights)
weightedWinningPortfolioTimeSeries = portfolioGrowthTimeSeries(weightedWinningAssetTimeSeries)

label = ''
for element in winningAssetCombination:
    label+=str(element)
#endFor

label += ' risk-parity weighted'

ax.plot(dateList,weightedWinningPortfolioTimeSeries,label=label,color='purple',alpha=1,linewidth=2)

plt.xlim(startDate,endDate)
plt.ylabel('Total Portfolio Value')
plt.xlabel('Date')
plt.legend()#bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0)
plt.show()

fig.savefig('figures/weighted-comparison-portfolio-'+label+'risk-'+str(q)+'-'+startStr+"-"+endStr+'.pdf',bbox_inches='tight')

from qiskit import QuantumCircuit, assemble, Aer
from math import pi, sqrt
from qiskit.visualization import plot_bloch_multivector, plot_histogram, plot_bloch_vector
sim = Aer.get_backend('aer_simulator')

qc = QuantumCircuit(1)

qc.draw()

qc.save_statevector()
qobj = assemble(qc)
state = sim.run(qobj).result().get_statevector()
plot_bloch_multivector(state)

qc = QuantumCircuit(1)
qc.draw()
qc.save_statevector()
qobj = assemble(qc)
state = sim.run(qobj).result().get_statevector()
fig=plot_bloch_multivector(state)
fig.savefig('figures/bloch-no-rotate.pdf', bbox_inches='tight')

qc = QuantumCircuit(1)
qc.rx(theta=pi/2,qubit=0)
qc.draw()
qc.save_statevector()
qobj = assemble(qc)
state = sim.run(qobj).result().get_statevector()
fig=plot_bloch_multivector(state)
fig.savefig('figures/bloch-pi-over-two-x-rotate.pdf', bbox_inches='tight')

