import sys, os
import numpy as np
import scipy
from scipy.stats import norm
sys.path.append("../../../quantum-mc") # see os.getcwd()
import matplotlib.pyplot as plt
import quantum_mc.calibration.fitting as ft
import quantum_mc.calibration.time_series as ts
import matplotlib.dates as mdates
from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,
                               AutoMinorLocator)
# AAPL, MSFT, SPX
ticker = "MSFT"
data = ts.get_data(ticker)
((cdf_x, cdf_y), sigma) = ft.get_cdf_data(ticker)
(x, y) = ft.get_fit_data(ticker, norm_to_rel = False)
pl = ft.fit_piecewise_linear(x, y)
pc = ft.fit_piecewise_cubic(x, y)
xf = np.linspace(-3, 3, 100)


cm = 1 / 2.54
years = mdates.YearLocator()   # every year
months = mdates.MonthLocator()  # every month
years_fmt = mdates.DateFormatter('%Y')

fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])

#fig, ax = plt.subplots(nrows=1, figsize=(8*cm, 6*cm))
fig.set_facecolor('w')
ax.plot(data.index, data["Close"], linewidth=0.5)
ax.set_xlabel('Time (years)')
ax.set_ylabel('Closing price')

# format the ticks
ax.xaxis.set_major_locator(years)
ax.xaxis.set_major_formatter(years_fmt)
ax.xaxis.set_minor_locator(months)
fig.savefig('hist_series.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True) 


fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(cdf_x, cdf_y, label="Empirical")
ax.plot(cdf_x, norm.cdf(cdf_x), '--', label="Normal")
ax.set_xlim(-5, 4)
ax.xaxis.set_major_locator(MultipleLocator(1))
ax.yaxis.set_major_locator(MultipleLocator(0.2))
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Cumulative probability')
ax.legend()
fig.savefig('hist_cdf_linear.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)


fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(cdf_x, cdf_y, label="Empirical")
ax.plot(cdf_x, norm.cdf(cdf_x), '--', label="Normal")
ax.set_xlim(-5, 4)
ax.set_ylim(0.005, 1.0)
ax.set_yscale("log")
ax.xaxis.set_major_locator(MultipleLocator(1))
#ax.yaxis.set_major_locator(MultipleLocator(0.2))
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Cumulative probability')
ax.legend()
fig.savefig('hist_cdf_log.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)

cm = 1 / 2.54
fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(x, y, "o", markersize = 2, label="Observed")
ax.set_xlabel('Normal return ($\sigma$)')
ax.set_ylabel('Empirical return ($\sigma$)')
ax.xaxis.set_major_locator(MultipleLocator(1))
ax.yaxis.set_major_locator(MultipleLocator(2))
fig.savefig('hist_scatter.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)

cm = 1 / 2.54
fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(x, y, "o", markersize = 2, label="Observed")
ax.plot(xf, pl(xf), label="Fitted", color = "green")
ax.set_xlabel('Normal return ($\sigma$)')
ax.set_ylabel('Empirical return ($\sigma$)')
ax.xaxis.set_major_locator(MultipleLocator(1))
ax.yaxis.set_major_locator(MultipleLocator(2))
ax.legend()
fig.savefig('hist_scatter_fitted.pdf', formt='pdf', facecolor=fig.get_facecolor(), transparent=True)

cm = 1 / 2.54
fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(x, y, "o", markersize = 2, label="Observed")
ax.plot(xf, pc(xf), label="Fitted", color = "green")
ax.set_xlabel('Normal return ($\sigma$)')
ax.set_ylabel('Empirical return ($\sigma$)')
ax.xaxis.set_major_locator(MultipleLocator(1))
ax.yaxis.set_major_locator(MultipleLocator(2))
ax.legend()
fig.savefig('hist_scatter_spline_fitted.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)

fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
ax.plot(cdf_x, cdf_y, label="Empirical")
ax.plot(cdf_x, norm.cdf(cdf_x), '--', label="Normal")
ax.plot(pc(xf), norm.cdf(xf), ':', label="Fitted") # plot the fitted 
ax.set_xlim(-5, 4)
ax.set_ylim(0.005, 1.0)
ax.set_yscale("log")
ax.xaxis.set_major_locator(MultipleLocator(1))
#ax.yaxis.set_major_locator(MultipleLocator(0.2))
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Cumulative probability')
ax.legend()
fig.savefig('hist_cdf_fitted_log.pdf', formt='pdf', facecolor=fig.get_facecolor(), transparent=True)

lookup = scipy.interpolate.interp1d(norm.cdf(cdf_x), cdf_x)
print(lookup(0.01))
lookup = scipy.interpolate.interp1d(cdf_y, cdf_x)
print(lookup(0.01))
lookup = scipy.interpolate.interp1d(norm.cdf(xf), pc(xf))
print(lookup(0.01))

rets = ts.returns(ticker) # 10 day log returns
rets = rets - np.mean(rets)      
# normalize returns into units of (maximum-likelihood-estimated) standard deviations
sig = np.std(rets)
rets = rets / sig

# create discretised version
import scipy
cdf_fn = scipy.interpolate.interp1d(pc(xf), norm.cdf(xf))
def pdf_fn(x):
    return (cdf_fn(x + 1e-9) - cdf_fn(x - 1e-9)) / 2e-9
xf2 = np.linspace(-4.8, 3.4, 100)
#plt.plot(xf2, cdf_fn(xf2))

fig = plt.figure(figsize=(8*cm,6*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
plt.plot(xf2, pdf_fn(xf2), label = "Fitted")
ax.hist(rets, density = True, bins = 100, label = "Empirical")
ax.set_xlim(-6, 4)
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Probability density')
ax.legend()
fig.savefig('hist_pdf_spline_fitted.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)

from qiskit import execute, Aer, QuantumCircuit, QuantumRegister, ClassicalRegister, AncillaRegister
from qiskit.aqua.algorithms import IterativeAmplitudeEstimation
from qiskit.circuit.library import NormalDistribution, LogNormalDistribution, LinearAmplitudeFunction, IntegerComparator, WeightedAdder
from qiskit.visualization import plot_histogram

from quantum_mc.probability_distributions.gaussian_copula import GaussianCopula
rho = ts.correl(ts.returns("AAPL"), ts.returns("MSFT"))

mu = [0, 0]
sigma = [[1, rho], [rho, 1]]

num_qubits = [4, 4]
bounds = [(-3.4, 3.4), (-3.4, 3.4)] 

def F(x):
    return cdf_fn(x)
    #return norm.cdf(x)

def f(x):
    return pdf_fn(x)
    #return norm.pdf(x)

cdfs = [F, F]
pdfs = [f, f]

pdf_circ = GaussianCopula(num_qubits, cdfs, sigma=sigma, bounds=bounds, pdfs = pdfs)
#pdf_circ = NormalDistribution(num_qubits, mu=mu, sigma=sigma, bounds=bounds)

qr_state1 = QuantumRegister(4, 'state1')
qr_state2 = QuantumRegister(4, 'state2')
state_out = ClassicalRegister(4, 'state_out')
circ = QuantumCircuit(qr_state1, qr_state2, state_out)
circ.append(pdf_circ, qr_state1[:] + qr_state2[:])
circ.measure(qr_state1, state_out)

circ.draw()

pdf_circ._probabilities
pdf_circ._values
pdf_circ._values[0]

comb = list(zip(pdf_circ._values, pdf_circ._probabilities))
print(comb[0])

values = sorted(set(map(lambda x:x[0], pdf_circ._values)))
#newlist = [sum([y[0] for y in values if y[1]==x]) for x in values]
print(values)
probs = [sum([y[1] for y in comb if y[0][0]==x]) for x in values]
print(probs)

#plt.plot(values, probs)
fig = plt.figure(figsize=(2*8*cm,8*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')
rects = ax.bar(values, probs, width = 0.3)
ax.set_ylim(0, 0.3)
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Probability')
labels = [r"$ | " + "{:04b}".format(i) + r" \rangle $" for i in range(len(values))]
ax.bar_label(rects, padding=4, labels = labels, rotation=60)
fig.savefig('hist_state.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)
#plt.plot(xf2, pdf_fn(xf2), label = "Fitted")

#fig = plt.figure(figsize=(8*cm,6*cm))  
counts = execute(circ, Aer.get_backend('qasm_simulator', shots = 10000)).result().get_counts()
#plot_histogram(counts, figsize=(9,5))

fig = plt.figure(figsize=(2*8*cm,8*cm))  
ax = fig.add_axes([0.2, 0.19, 0.77, 0.77])
fig.set_facecolor('w')

cnts = [counts.get(('{0:0%sb}' % 4).format(i), 0) for i, v in enumerate(values)] 
cnts = np.array(cnts)
cnts = cnts / sum(cnts)
rects = ax.bar(values, cnts, width = 0.3)
ax.set_ylim(0, 0.3)
ax.set_xlabel('Return ($\sigma$)')
ax.set_ylabel('Probability')
labels = [r"$ | " + "{:04b}".format(i) + r" \rangle $" for i in range(len(values))]
ax.bar_label(rects, padding=4, labels = labels, rotation=60)
fig.savefig('hist_state_simulated.pdf', format='pdf', facecolor=fig.get_facecolor(), transparent=True)

shots=1000
#sim_statevector = Aer.get_backend('statevector_simulator')
job = execute(circ, Aer.get_backend('statevector_simulator'))
result = job.result()

outputstate = result.get_statevector(circ, decimals=3)

from qiskit.visualization import plot_state_city
plot_state_city(outputstate)



plt.plot(np.linspace(-4, 4, 100), np.exp(np.linspace(-4, 4, 100) * sigma) - 1)

Aer.backends()

