import numpy as np
import scipy.io as sio
import matplotlib.pyplot as plt
from scipy.signal import correlate as corr

# Problem 2

# Load the received signal and the training symbols
rx = sio.loadmat('rx_signal.mat')
rx_signal = rx['rx_signal'].flatten()

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

# Define parameters
N = 256  # number of carriers
L = 25  # length of the cyclic prefix


# Pre-process in an appropriate manner the training symbols, such that we
# can use them for finding the start of the cyclic prefix of the training
# OFDM block in the received samples.

ts_IFFT = np.fft.ifft(training_symbols, N)  # take the IFFT
ts_IFFT_with_CP = np.concatenate((ts_IFFT[-L:], ts_IFFT))  # add the cyclic prefix

# Find the start of the cyclic prefix of the training OFDM block in
# the received samples. Check point: you should find the beginning of the
# cyclic prefix at index 159 in the received samples.

R = corr(rx_signal, ts_IFFT_with_CP, 'valid', 'fft')
index = np.argmax(abs(R))
print('Sample index = ', index)

# plot for checking
fig, ax = plt.subplots()
ax.plot(np.arange(R.size), abs(R), marker='*')
ax.set_title('Correlation result')
ax.set_xlabel('Index')
ax.set_ylabel('|R|')
ax.grid()
fig.show()

# Implement the OFDM receiver for the portion of the received data starting
# with the training OFDM block. Note that you might need to trim the
# received data such that you obtain an integer number of OFDM symbols.

rx = rx_signal[index:]  # start at the training symbols
rx = rx[:-(rx.size % (N+L))]  # trim to obtain an integer number of OFDM symbols
rx = np.reshape(rx, (N+L, -1), 'F')
rx = rx[L:, :]  # remove the cyclic prefix
rxFFT = np.fft.fft(rx, N, axis=0)

# Using the LS approximation, estimate the vector lambda_LS of channel
# coefficients

lambda_LS = rxFFT[:, 0] / training_symbols  # LS is the component-wise division

# Plot the absolute value of the channel coefficients lambda_LS obtained above.
# The symbol-level impulse response of the channel is h = [1.621, 0.8448, 0.0283].
# Plot the corresponding lambdas (absolute value) on the same figure.
# Check point: lambda_LS should be a noisy version of the lambdas obtained from h.

h = np.array([1.621, 0.8448, 0.0283])
hFFT = np.fft.fft(h, N)  # lamdbas = DFT(h)

fig, ax = plt.subplots()
ax.plot(abs(lambda_LS), label='Lambdas with LS')
ax.plot(abs(hFFT), label='Lambdas with DFT(h)')
ax.set_xlabel('Frequency index')
ax.set_ylabel('|Lambdas|')
ax.set_title('Lambda estimates')
ax.grid()
ax.legend()
fig.show()

# Using lambda_LS, equalize the OFDM symbols which follow the training
# symbols. Plot (scatterplot) the result. Check point: you should obtain a
# noisy 4-QAM constellation (4 distinct clouds, no rotation).

rxFFT = rxFFT[:, 1:]  # remove the training symbols
rxEq = rxFFT / lambda_LS.reshape((-1, 1))
rxEq = rxEq.flatten('F')

fig, ax = plt.subplots()
ax.scatter(rxEq.real, rxEq.imag, marker='*')
ax.set_xlabel('Re')
ax.set_ylabel('Im')
ax.set_title('Received constellation after OFDM demodulation and equalization')
ax.grid()
fig.show()

plt.show()
