import numpy as np
import matplotlib.pyplot as plt
import scipy.io as sio

# Problem 3

# Load the various files
rx = sio.loadmat('received_OFDM_symbols.mat')
received_OFDM_symbols = rx['received_OFDM_symbols']

ts = sio.loadmat('training_symbols.mat')
training_symbols = ts['training_symbols'].flatten()

lmda = sio.loadmat('lambdas.mat')
lambdas = lmda['lambdas'].flatten()


# Using the first column of the received symbols matrix and the training
# sequence, estimate the noise variance. The noise is assumed to have
# independent and identically distributed components.
tmp = received_OFDM_symbols[:, 0].flatten() - training_symbols * lambdas
noise_var = np.mean(np.square(np.abs(tmp)))
print('noise_var = ', noise_var)


# If you do not manage to compute the noise variance, you can use a value
# of 0.01 for it in the following.

# Compute the LS estimate of the transmitted data symbols
D = np.diag(lambdas)
# Since D is diagonal, any of the following will work, they are the same.
data_LS_estim1 = received_OFDM_symbols[:, 1:] / lambdas[:, np.newaxis]
data_LS_estim2 = np.matmul(np.linalg.pinv(D), received_OFDM_symbols[:, 1:])
data_LS_estim3 = np.linalg.lstsq(D, received_OFDM_symbols[:, 1:], rcond=None)[0]


# Plot the estimated data symbols. You should get a noisy 4-QAM
# constellation.
suf_data = data_LS_estim1.flatten('F')
plt.scatter(suf_data.real, suf_data.imag, marker='*', color='b')
plt.ylabel('Imaginary')
plt.xlabel('Real')
plt.title('The constellation of the estimated data symbols (LS)')
plt.grid()
plt.show()


# Compute the LMMSE estimate of the transmitted data symbols
Kay = D.conj().T
Ky = np.matmul(D, D.conj().T) + np.diag(noise_var * np.ones(lambdas.size))
tmp_matrix = np.linalg.lstsq(Ky, received_OFDM_symbols[:, 1:], rcond=None)[0]
data_LMMSE_estim = np.matmul(Kay, tmp_matrix)


# Plot the estimated data symbols. You should get a noisy 4-QAM
# constellation.
suf_data = data_LMMSE_estim.flatten('F')
plt.scatter(suf_data.real, suf_data.imag, marker='*', color='b')
plt.ylabel('Imaginary')
plt.xlabel('Real')
plt.title('The constellation of the estimated data symbols (LMMSE)')
plt.grid()
plt.show()


# Compare the values of the estimated symbols obtained with the two methods
# above. For example, compute something like
# np.mean(np.abs(data_LS_estim - data_LMMSE_estim)).
# If your estimates were correct, the value above should be around 0.0077.
avg_diff_abs = np.mean(np.abs(data_LS_estim1.flatten('F') - data_LMMSE_estim.flatten('F')))
print('avg_diff_abs = ', avg_diff_abs)
