# %%

x = np.linspace(2, 8, 1000)
y = np.linspace(25, 225, 1000)

X, Y = np.meshgrid(x, y)
Z = np.zeros((1000, 1000))

d1 = Normal(loc=5, scale=1)
d2 = Normal(loc=125, scale=20)

dist = JointIndependent(marginals=[d1, d2])

for i in range(len(x)):
    Z[i, :] = dist.pdf(
        np.append(np.atleast_2d(X[i, :]), np.atleast_2d(Y[i, :]), 0).T)

fig, ax = plt.subplots()
CS = ax.contour(X, Y, Z, 15)
plt.plot(m, k_hi, 'k')
plt.plot(m, k_lo, 'k')
# plt.fill_between(m,k_lo,k_hi)
plt.xlim([mu_m - 3 * sigma_m, mu_m + 3 * sigma_m])
plt.ylim([mu_k - 3 * sigma_k, mu_k + 3 * sigma_k])
plt.xlabel(r'Mass ($m$)')
plt.ylabel(r'Stiffness ($k$)')
plt.grid(True)
plt.tight_layout()
plt.show()

fig, ax = plt.subplots()
#%% md
#
# Plot the pdf of the distribution before and after the copula
# -------------------------------------------------------------
#

#%%

fig, ax = plt.subplots(ncols=2, figsize=(10, 4))

x = np.arange(-3, 3, 0.1)
y = np.arange(-3, 3, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist_1.pdf(
    x=np.concatenate([X.reshape((-1,
                                 1)), Y.reshape((-1, 1))], axis=1))
CS = ax[0].contour(X, Y, Z.reshape(X.shape))
ax[0].clabel(CS, inline=1, fontsize=10)
ax[0].set_title('Contour plot of pdf - independent normals')

x = np.arange(-3, 3, 0.1)
y = np.arange(-3, 3, 0.1)
X, Y = np.meshgrid(x, y)
Z = dist_2.pdf(
    x=np.concatenate([X.reshape((-1,
                                 1)), Y.reshape((-1, 1))], axis=1))
CS = ax[1].contour(X, Y, Z.reshape(X.shape))
ax[1].clabel(CS, inline=1, fontsize=10)
ax[1].set_title('Contour plot of pdf - normals with Gumbel copula')
plt.show()