import numpy as np
import matplotlib.pyplot as plt

N = 100
P_prior = np.zeros((N,N))

x_j = np.arange(N)
x_i = x_j

W = np.array([[200, 150], [150, 400]])*2
inv_W = np.linalg.inv(W)
b = np.array([65, 55])

for i in range(N):
  for j in range(N):
    x = [x_j[j], x_i[i]]
    P_prior[i,j] = np.exp(-(x-b).T@inv_W@(x-b))

P_prior[:50,:45] = 0
P_prior[80:,60:] = 0
P_prior /= np.sum(P_prior)

plt.figure(figsize=(8,6), dpi=100)
plt.imshow(P_prior, cmap='gray')
plt.colorbar()
plt.xlabel(r'$j$-coordinate')
plt.ylabel(r'$i$-coordinate')
plt.title('Prior belief of location')
plt.show()

# your code here


# your code here


P_observe = np.zeros((N,N))

for i in range(N):
  for j in range(N):
    dist1 = np.linalg.norm([x_j[j]-80, x_i[i]-0])
    prob1 = np.exp(-(dist1-50)**2/200)*0.9
    dist2 = (x_j[j]+x_i[i])
    prob2 = np.exp(-(dist2-50)**2/1000)*0.5
    P_observe[i,j] = np.max([prob1, prob2])

plt.figure(figsize=(8,6), dpi=100)
plt.imshow(P_observe, cmap='gray')
plt.colorbar()
plt.xlabel(r'$j$-coordinate')
plt.ylabel(r'$i$-coordinate')
plt.title('Probability of observing two bars')
plt.show()

# your code here


# your code here