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

# Problem 1

# Load the files
rx = sio.loadmat('rx_signal.mat')
rx_signal = rx['rx_signal'].flatten()

pn = sio.loadmat('pn_sequence.mat')
pn_sequence = pn['pn_sequence'].flatten()

# Define parameters
num_carriers = int(256)    # total number of carriers
prefix_length = int(20)    # length of the cyclic prefix

# Correlate the received signal with the P/N sequence in order to find this
# sequence within the received signal
R = corr(rx_signal, pn_sequence, 'valid', 'fft')

# Plot to check if we find a clear peak
plt.figure()
plt.plot(abs(R),'-*')
plt.grid()
plt.xlabel('Index')
plt.ylabel('abs(correlation)')
plt.title('Correlation at Rx')
plt.show()

index_start_pn_sequence = np.argmax(abs(R))
print('index_start_pn_sequence = ', index_start_pn_sequence)

# Trim the received signal and keep the "full" OFDM data blocks
rx_symbols = rx_signal[index_start_pn_sequence + pn_sequence.size:]
rx_symbols = rx_symbols[:-(rx_symbols.size%(num_carriers + prefix_length))]

# Process the OFDM blocks
rx_withCP = np.reshape(rx_symbols, (prefix_length+num_carriers, -1), 'F')
rx_noCP = rx_withCP[prefix_length + np.arange(num_carriers), :]  # remove the rows corresponding to cyclic prefix

# Go to the frequency domain
Rf = fft.fft(rx_noCP, num_carriers, axis=0)

# Plot the received constellation
plt.scatter(Rf.flatten('F').real, Rf.flatten('F').imag, marker='*', color='b')
plt.grid()
plt.title('Constellation at Rx')
plt.show()

# Decode the first block
first_block = Rf[:,1]
first_block_decoded = 2*(first_block.real > 0) - 1 + 1j*(2*(first_block.imag > 0) - 1)

# Estimate the lambdas with LS
lambdas = first_block/first_block_decoded

# Plot the lambdas
plt.figure()
plt.plot(abs(lambdas),'-*')
plt.grid()
plt.xlabel('Index')
plt.ylabel('Magnitude')
plt.title('Estimated Lambdas')
plt.show()
