import numpy as np
import scipy as sp
import matplotlib.pyplot as plt
from matplotlib import rc
from scipy.fftpack import fft, ifft
import time

# graphical parameters
###############################################################################
plt.style.use("ggplot")
rc("font", **{"family": "serif", "serif": ["Computer Modern Roman"], "size": "12"})
rc("text", usetex=True)
rc("lines", linewidth=2)
plt.rcParams["axes.facecolor"] = "w"
###############################################################################


np.random.seed(42)

# defines some functions first
def mu(t):
    return np.sin(2 * np.pi * t)


def Cov(t, s, rho=1 / 20):
    return np.exp(-np.abs(t - s) / rho)


# defines some parameters
N = 50
dt = 1 / N
t = np.linspace(0, 1, N)
muvec = mu(t)  # gets mean
sigma = np.zeros((N, N))
rho = 0.005
# notice that if we increase rho, the covariance becomes singular,
# try changing rho to 1/10

# Populates covatiance matrix
for i in range(N):
    for j in range(N):
        sigma[i, j] = Cov(t[i], t[j], rho)

print(np.linalg.det(sigma))  # checks if its singular
plt.contourf(sigma, 50, cmap="Spectral")
plt.title(r"Covariance, $\rho=" + str(rho) + "$")
plt.show()

U, D, V = np.linalg.svd(sigma)  # we do this since depending on the value of rho,
# \Sigma can become numerically singular
A = U @ np.diag(np.sqrt(D))

# samples random process with direct generation
proc = muvec + A @ np.random.standard_normal(N)
plt.plot(t, proc, label="Direct method")

# Samples random process with circulant embedding
# Construct FFT
c = np.hstack([sigma[0, :], sigma[0, -2:0:-1]])
labda = fft(c)

# Generate iid normal vectors and do IFFT
YR = np.random.standard_normal(2 * (N - 1))
YI = np.random.standard_normal(2 * (N - 1))
Y = np.array([complex(yr, yi) for yr, yi in zip(YR, YI)])
eta = np.sqrt(2 * (N - 1)) * np.sqrt(labda)
print(eta * Y)
Xtilde = ifft(eta * Y)
proc_circ = np.real(Xtilde[1 : N + 1])

plt.plot(t, proc_circ, label="Circulant embedding")
plt.title(r" $\mu(t)=\sin(2\pi t)$")
plt.legend()
plt.show()

# adds new data by adding one midpoint
tall = np.linspace(0, 1, N * 2)
Nall = len(tall)
M = Nall - N

# transforms to layout X=(Y,Z) where Z is the previous observation
zobs = proc
# gets the indices
idZ = np.arange(0, Nall, 2)
idY = np.arange(1, Nall, 2)
idP = np.concatenate((idY, idZ))
P = np.eye(Nall)
P = P[idP, :]
# mu for the fine process
mu_all = mu(tall)
sigma_all = np.zeros((Nall, Nall))

for i in range(Nall):
    for j in range(Nall):
        sigma_all[i, j] = Cov(tall[i], tall[j], rho)

sigmaX = P @ sigma_all @ np.transpose(P)
muX = P @ mu_all
# extracts the components
muY = muX[0:M]
muZ = muvec

# constructs sub covariance matrices
sigmaYY = sigmaX[:M, :M]
sigmaYZ = sigmaX[:M, M:]
sigmaZZ = sigmaX[M:, M:]

# computes conditional mean and variance

muYgZ = muY + sigmaYZ @ (np.linalg.solve(sigmaZZ, zobs - muZ))
sigmaYgZ = sigmaYY - sigmaYZ @ (np.linalg.solve(sigmaZZ, sigmaYZ.T))

# samples again
U, D, V = np.linalg.svd(sigmaYgZ)
A = U @ np.diag(np.sqrt(D))
Y = muYgZ + A @ np.random.standard_normal(M)
X = np.concatenate((Y, zobs))
proc_all = np.linalg.solve(P, X)

tg = np.linspace(0, 1, 10 * N)

plt.plot(t, proc, "-o")
plt.plot(tall, proc_all, "-*")
plt.plot(tg, mu(tg), "--", color="k")
plt.legend(
    ["Coarse", "Fine", r"$\mu(t)=\sin(2\pi t)$"],
    fancybox=True,
    framealpha=0.0,
    loc="upper center",
    bbox_to_anchor=(0.5, -0.05),
    shadow=False,
    ncol=3,
)
plt.show()

rho = 0.005
# Compare direct method and circulant embedding costs
N_grid = np.arange(2000, 5000, 500)
num_rep = 10
times_dir = []
times_circ = []

for N in N_grid:
    dt = 1 / N
    t = np.linspace(0, 1, N)
    muvec = mu(t)  # gets mean
    sigma = Cov(
        t.reshape(-1, 1), t.reshape(1, -1)
    )  # faster way to populate covariance matrix

    time_dir = 0
    time_circ = 0
    for i in range(num_rep):
        # Samples random process with direct generation
        start = time.time()
        U, D, V = sp.linalg.svd(sigma)
        A = U @ np.diag(np.sqrt(D))
        proc = muvec + A @ np.random.standard_normal(N)
        end = time.time()
        time_dir += end - start

        # Samples random process with circulant embedding
        c = np.hstack([sigma[0, :], sigma[0, -2:0:-1]])
        start = time.time()
        labda = fft(c)
        YR = np.random.standard_normal(2 * (N - 1))
        YI = np.random.standard_normal(2 * (N - 1))
        Y = np.array([complex(yr, yi) for yr, yi in zip(YR, YI)])
        eta = np.sqrt(2 * (N - 1)) * np.sqrt(labda)
        Xtilde = ifft(eta * Y)
        proc_circ = np.real(Xtilde[1 : N + 1])
        end = time.time()
        time_circ += end - start
    times_dir.append(time_dir / num_rep)
    times_circ.append(time_circ / num_rep)

plt.plot(N_grid, times_dir, label="Direct method")
plt.plot(N_grid, times_circ, label="Circulant embedding")
xx = np.log(N_grid)
yyd = np.log(times_dir)
yyc = np.log(times_circ)
aad, bbd = np.polyfit(xx, yyd, deg=1)
aac, bbc = np.polyfit(xx, yyc, deg=1)
plt.plot(
    np.exp(xx),
    np.exp(aad * xx + bbd),
    "-.",
    label=r"$\mathcal{O}(N^{" + str(round(aad, 2)) + "})$",
    color="gray",
)
plt.plot(
    np.exp(xx),
    np.exp(aac * xx + bbc),
    "--",
    label=r"$\mathcal{O}(N^{" + str(round(aac, 2)) + "})$",
    color="gray",
)

plt.yscale("log")
plt.xscale("log")
plt.title("Cost")
plt.legend()
plt.show()