from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer, IBMQ
from qiskit.circuit.library import StatePreparation, RYGate
from qiskit.quantum_info import Statevector
import qiskit_tools as qt
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
from tqdm import tqdm
from datetime import datetime

# Shake things up by changing the random seed
rseed = int(datetime.now().timestamp())
np.random.seed(rseed) 

from sklearn.gaussian_process.kernels import RBF
from sklearn.gaussian_process import GaussianProcessRegressor

matplotlib.rcParams['mathtext.fontset'] = 'stix'
matplotlib.rcParams['font.family'] = 'STIXGeneral'
width=0.75
color='black'
fontsize=28
ticksize=22
figsize=(10,8)

N = 4

n = int(np.ceil(np.log2(N)))

xmin, xmax = 0, 8
X_train = np.linspace(start=xmin, stop=xmax, num=2**n).reshape(-1, 1)

def y_targ(x):
    return np.squeeze(x*np.sin(x))

noise_std = 0.75

y_train = y_targ(X_train) + np.random.normal(0,noise_std,N)

Nast = 16

X = np.linspace(start=xmin, stop=xmax, num=Nast).reshape(-1, 1)
y = y_targ(X)

def SEK(X1, X2, ls=1.0, sig=1.0):
    sqdist = np.sum(X1**2, 1).reshape(-1, 1) + np.sum(X2**2, 1) - 2 * np.dot(X1, X2.T)
    return sig * np.exp(-0.5 * sqdist / ls**2)

sigma = 2*noise_std
ls = 2.**(-(xmax-xmin)/Nast)

def posterior(X_s, X_train, Y_train, l=1.0, sigma_f=1.0, sigma_y=1e-8):
    K = SEK(X_train, X_train, l, sigma_f) + sigma_y**2 * np.eye(len(X_train))
    K_s = SEK(X_train, X_s, l, sigma_f)
    K_ss = SEK(X_s, X_s, l, sigma_f)
    K_inv = np.linalg.inv(K)
    
    mu_s = K_s.T.dot(K_inv).dot(Y_train)
    cov_s = K_ss - K_s.T.dot(K_inv).dot(K_s)
    
    return mu_s, cov_s

mean_prediction, cov_s = posterior(X, X_train, y_train, l=ls, sigma_f=sigma, sigma_y=noise_std)
std_prediction = np.diag(cov_s)

fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111)

ax.plot(X, y, label=r"$f(x) = x \sin(x)$", linestyle="dotted")
ax.errorbar(
    X_train,
    y_train,
    noise_std,
    linestyle="None",
    color="tab:blue",
    marker=".",
    markersize=10,
    label="Observations",
)
ax.plot(X, mean_prediction, label="Mean prediction", marker='.')
ax.fill_between(
    X.ravel(),
    mean_prediction - 1.96 * std_prediction,
    mean_prediction + 1.96 * std_prediction,
    color="tab:orange",
    alpha=0.5,
    label=r"95% confidence interval",
)
ax.legend()
ax.set_xlabel("$x$")
ax.set_ylabel("$f(x)$")

plt.show()

# Pick the middle x* to evaluate
u_ind = (len(X)//2) - 1

# Define v and u
v = y_train
u = SEK(X, X_train, ls, sig=sigma)[u_ind]

# Determine normalisation of v and u
v_norm = np.sqrt(np.sum(np.abs(v)**2))
u_norm = np.sqrt(np.sum(np.abs(u)**2))

# Normalised vectors v and u
v_ = v/v_norm
u_ = u/u_norm

# Determine A
A = SEK(X_train, X_train, ls, sig=sigma) + (np.eye(len(X_train))*noise_std**2)

kappa = np.linalg.cond(A)
lambda_min = np.min(np.abs(np.linalg.eigvals(A)))
lambda_max = np.max(np.abs(np.linalg.eigvals(A)))

# Choose a success probability of determining eigenvalue estimation
p_suc = 0.99
# Determine the number of ancillary qubits to ensure the given success probability
nanc = n + int(np.ceil(np.log2(2. + 1./(2.*(1-p_suc)))))

phase = True

# Calculate the new scale of the relative magnitudes of maximum and minimum eigenvalues
rescale = int((2**(nanc-phase)-1)/kappa)
# Make sure this scaling is not vanishingly small
if rescale < 1.e-7:
    rescale = 1
    
# determine the number of integer qubits required
nint = qt.get_nint(rescale)
# Assert that the number of integer qubits is equal to or greater than the number of ancillary qubits assigned
if nint < nanc - phase:
    nint = nanc - phase
    
# We can determine the scaling of the rotation of step 3 (variable c)
scaling = qt.bin_to_dec(qt.my_binary_repr(rescale, nint=nint, n=nint, phase=False), nint=0, phase=False)
t0 = scaling / (lambda_min * (2**phase))
t = 2 * np.pi * t0

qcont = QuantumRegister(1, 'cont')
qtarg = QuantumRegister(1, 'targ')
qreg = QuantumRegister(n, 'x')
qanc = QuantumRegister(nanc, 'anc')

circ = QuantumCircuit(qanc, qtarg, qreg, qcont)

circ.h(qcont);
circ.barrier();

prepv_gate = StatePreparation(Statevector(v_)).control(1)
circ.append(prepv_gate, [qcont, *qreg]);
circ.barrier();

circ.x(qcont);
circ.cx(qcont, qtarg);
circ.barrier();

prepu_gate = StatePreparation(Statevector(u_)).control(1)
circ.append(prepu_gate, [qcont, *qreg]);
circ.x(qcont);
circ.barrier();

hhl_gate = qt.HHL(circ, qreg, qanc, qtarg, A, t=t, scaling=scaling, wrap=True).control(1)
circ.append(hhl_gate, [qcont, *qreg, *qanc, qtarg]);

circ.draw('latex')

backend = Aer.get_backend('statevector_simulator')
job = execute(circ, backend)
result = job.result()
state_vector = np.asarray(result.get_statevector())

state_v = np.asarray(state_vector).reshape((2**1, 2**n, 2**1, 2**nanc)).T
state_v = state_v[0,1,:,1]

targ_v = np.matmul(np.linalg.inv(A),v_)*np.sqrt(2.)

fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111)

plt.scatter(np.arange(2**n), targ_v, marker='o', color='red')
plt.scatter(np.arange(2**n), state_v.real, marker='.', color='black')
ax.set_ylabel('Amplitude', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

state_u = np.asarray(state_vector).reshape((2**1, 2**n, 2**1, 2**nanc)).T
state_u = state_u[0,1,:,0]

targ_u = u_/np.sqrt(2.)

fig = plt.figure(figsize=np.array(figsize))
ax = fig.add_subplot(111)

plt.scatter(np.arange(2**n), targ_u, marker='o', color='red')
plt.scatter(np.arange(2**n), state_u.real, marker='.', color='black')
ax.set_ylabel('Amplitude', fontsize=fontsize);
ax.tick_params(axis='both', labelsize=ticksize);

state_v = np.asarray(state_vector).reshape((2**1, 2**n, 2**1, 2**nanc)).T
state_v = state_v[0,1].T

X_matrix = np.array([[0,1],[1,0]])

mu_out = np.matmul(state_v.flatten(),np.matmul(X_matrix, state_v.real).flatten()).real*(u_norm*v_norm*scaling/np.sqrt(2.))

print(mu_out, mean_prediction[u_ind])

mus = []
stds = []

for i in tqdm(np.arange(Nast)):
    
    v = y_train
    u = SEK(X, X_train, ls, sig=sigma)[i]

    v_norm = np.sqrt(np.sum(np.abs(v)**2))
    u_norm = np.sqrt(np.sum(np.abs(u)**2))

    v_ = v/v_norm
    u_ = u/u_norm
    
    v_gate = StatePreparation(Statevector(v_))
    u_gate = StatePreparation(Statevector(u_))

    qcont = QuantumRegister(1, 'cont')
    qtarg = QuantumRegister(1, 'targ')
    qreg = QuantumRegister(n, 'reg')
    qanc = QuantumRegister(nanc, 'anc')

    circ = QuantumCircuit(qanc, qtarg, qreg, qcont)
        
    QGPR_gate = qt.QGPR(circ, qreg, qanc, qtarg, qcont, v_gate, u_gate, A, t, scaling, wrap=True)
    circ.append(QGPR_gate, [*qreg, *qanc, qtarg, qcont]);
    
    backend = Aer.get_backend('statevector_simulator')
    job = execute(circ, backend)
    result = job.result()
    state_vector = np.asarray(result.get_statevector())
    
    state_v = np.asarray(state_vector).reshape((2**1, 2**n, 2**1, 2**nanc)).T[0,1].T
    mus.append(np.matmul(state_v.flatten(),np.matmul(X_matrix, state_v.real).flatten()).real*(u_norm*v_norm*scaling/np.sqrt(2)))
    
    qcont = QuantumRegister(1, 'cont')
    qtarg = QuantumRegister(1, 'targ')
    qreg = QuantumRegister(n, 'reg')
    qanc = QuantumRegister(nanc, 'anc')

    circ = QuantumCircuit(qanc, qtarg, qreg, qcont)
        
    QGPR_gate = qt.QGPR(circ, qreg, qanc, qtarg, qcont, u_gate, u_gate, A, t, scaling, wrap=True)
    circ.append(QGPR_gate, [*qreg, *qanc, qtarg, qcont]);
    
    backend = Aer.get_backend('statevector_simulator')
    job = execute(circ, backend)
    result = job.result()
    state_vector = np.asarray(result.get_statevector())
    
    state_v = np.asarray(state_vector).reshape((2**1, 2**n, 2**1, 2**nanc)).T[0,1].T
    M_exp = np.matmul(state_v.flatten(),np.matmul(X_matrix, state_v.real).flatten()).real*((u_norm**2)*scaling/np.sqrt(2))

    stds.append(SEK(X, X, ls, sig=sigma)[i,i]-M_exp)
    
mus = np.array(mus)
stds = np.array(stds)

fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111)

ax.plot(X, y, label=r"$f(x) = x \sin(x)$", linestyle="dotted")
ax.errorbar(
    X_train,
    y_train,
    noise_std,
    linestyle="None",
    color="tab:blue",
    marker=".",
    markersize=10,
    label="Observations",
)
ax.plot(X, mean_prediction, marker='.', label="Classical mean prediction")
ax.fill_between(
    X.ravel(),
    mean_prediction - 1.96 * std_prediction,
    mean_prediction + 1.96 * std_prediction,
    color="tab:orange",
    alpha=0.5,
    label=r"95% confidence interval",
)
ax.plot(X, mus, color='violet', marker='.', label="Quantum mean prediction")
ax.fill_between(
    X.ravel(),
    mus - 1.96 * stds,
    mus + 1.96 * stds,
    color="tab:purple",
    alpha=0.5,
    label=r"95% confidence interval",
)

ax.legend()
ax.set_xlabel("$x$")
ax.set_ylabel("$f(x)$")

plt.show()