import numpy as np
from scipy.linalg import toeplitz
import matplotlib.pyplot as plt
import scipy.io as sio

# Problem 3

# Load variables
var_tmp = sio.loadmat('variables.mat')
x = var_tmp['x'].flatten()
y = var_tmp['y'].flatten()
h = var_tmp['h'].flatten()
sigma = var_tmp['sigma'].flatten()
K_h = var_tmp['K_h']

# Define parameters
L = h.size

# Create the Toeplitz matrix needed for both estimators
first_column = np.concatenate((x, np.zeros(L-1, dtype=complex)))
first_row = np.zeros(L, dtype=complex)
first_row[0] = x[0]
C = toeplitz(first_column, first_row)

# LS estimator
# One can use @ as a shorthand for np.matmul
tmp_matrix = np.matmul(C.conj().T, C)
tmp_vector = np.matmul(C.conj().T, y.T)
h_LS = np.linalg.lstsq(tmp_matrix, tmp_vector, rcond=None)[0]
# Or this way, since np.linalg.lstsq returns the LS solution anyway
# h_LS = np.linalg.lstsq(C, y.T, rcond=None)[0]

# MMSE estimator
# Compute the covariance matrices
Khy = np.matmul(K_h, C.conj().T)
Ky = np.matmul(C, np.matmul(K_h, C.conj().T)) + (sigma**2)*np.eye(y.size)

# MMSE estimate
tmp_matrix = np.linalg.lstsq(Ky, y, rcond=None)[0]
h_MMSE = np.matmul(Khy, tmp_matrix)

# Plot to check the estimates
l1, = plt.plot(np.abs(h_LS), '-ob')
l2, = plt.plot(np.abs(h_MMSE), '-*r')
l3, = plt.plot(np.abs(h), '--dk')
plt.title('Channel impulse response')
plt.xlabel('Sample index')
plt.ylabel('Magnitude')
plt.grid()
plot_lines = [l1, l2, l3]
legend1 = plt.legend(plot_lines, ["LS estimation", "MMSE estimation", "Actual channel"], loc=1)
plt.gca().add_artist(legend1)
plt.show()
