import numpy as np
import matplotlib.pyplot as plt
import scipy.linalg as linalg
import time
from matplotlib import rc

# 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)

N = int(1e6)
mu = np.array([2, 1])
Sigma = np.array([[1, 2], [2, 5]])
p = 2

# ### Part 1
#
# Generates samples based on the Cholesky factorization

mu = np.ones((p, N))
mu[0, :] = 2.0 * np.ones(N)
mu[1, :] = 1.0 * np.ones(N)

tstart = time.time()
A = np.linalg.cholesky(Sigma)
xi = np.zeros([2, N])
xi[0, :] = np.random.standard_normal(N)
xi[1, :] = np.random.standard_normal(N)
X = mu + A @ xi
telapsed = time.time() - tstart
print("time needed  " + str(telapsed) + " " + "seconds")


# Computes a 2D histogram

plt.figure(1)
plt.hist2d(X[0, :], X[1, :], bins=70, cmap="Spectral")
plt.colorbar()
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
plt.title("2d Histogram Cholesky")
plt.show()
# ### part 2
#
#
# Now we try to the same with the second problem. Note that the Cholesky
# decomposition for $\Sigma=\begin{bmatrix}1 & 2\\2 & 4\end{bmatrix}$ fails.
# Thus, we do an eigenvalue decomposition (e.g. via svd):

Sigma = np.array([[1, 2], [2, 4]])
tstart = time.time()
# Computes the SVD
U, s, Vh = linalg.svd(Sigma)
A = U * np.sqrt(s)
xi = np.zeros([2, N])
xi[0, :] = np.random.standard_normal(N)
xi[1, :] = np.random.standard_normal(N)
X = mu + A @ xi
telapsed = time.time() - tstart
print("time needed  " + str(telapsed) + " " + "seconds")

plt.figure(2)
plt.hist2d(X[0, :], X[1, :], bins=70, cmap="Spectral")
plt.colorbar()
plt.xlabel("$x_1$")
plt.ylabel("$x_2$")
plt.title("2d Histogram, SVD")
plt.show()