import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.image import NonUniformImage
%matplotlib
A = 4
mu = np.array([1, 1])
cov = np.array([[2, 0],[0, 2]])
cov_inverse = np.linalg.inv(cov)
x = np.arange(mu[0]-2*cov.max(),mu[0]+2*cov.max(),step)
y = np.arange(mu[1]-2*cov.max(),mu[1]+2*cov.max(),step)
X,Y = np.meshgrid(x,y)
Z =A* 1/np.sqrt(2*math.pi*np.linalg.det(cov)) * np.exp(-1/2*(
(X-mu[0])**2*cov_inverse[0,0]
+(Y-mu[1])**2*cov_inverse[1,1]
+(X-mu[0])*(Y-mu[1])*cov_inverse[0,1]
+(Y-mu[1])*(X-mu[0])*cov_inverse[1,0]
))
step=0.05
alpha = 0.6
xylevel = 50
zlevel=int(len(x)/(2*step*100))
label_size=16
title_size=int(1.5*label_size)
fig = plt.figure(6,figsize=(10,10))
ax = fig.add_subplot(111,projection='3d')
surface = ax.plot_surface(X,Y,Z,cmap='jet',alpha=alpha)
fig.colorbar(surface,shrink=0.8,aspect=10)
ax.set_xlim(mu[0]-2*cov.max(),mu[0]+2*cov.max())
ax.set_ylim(mu[1]-2*cov.max(),mu[1]+2*cov.max())
ax.set_zlim(-np.around(Z.max(),decimals=2),np.around(Z.max(),decimals=2))
ax.contour(X,Y,Z,zdir='x',offset=X.min(),levels=xylevel)
ax.contour(X,Y,Z,zdir='y',offset=Y.max(),levels=xylevel)
ax.contour(X,Y,Z,zdir='z',offset=-Z.max(),cmap='jet',alpha=alpha,levels=zlevel)
ax.contourf(X,Y,Z,zdir='z',offset=-Z.max(),cmap='jet',alpha=alpha,levels=zlevel)
ax.set_xlabel('x',fontsize=label_size)
ax.set_ylabel('y',fontsize=label_size)
ax.set_zlabel('z',fontsize=label_size)
ax.set_title(r"$Z=\frac{1}{\sqrt{2\pi \det(\Sigma)}}\exp\left\{ -\frac{1}{2}\left[ (x-\mu)^\top \Sigma^{-1} (x-\mu)\right]\right\}$",fontsize=title_size)
plt.show()