%% prototype 2
% Goal:
% * Rearrange receiver stages according to latest block diagram.
% * Be more rigorous with variable names, to mitigate confusion.
% * Test each stage in new sequence.

close all; clear;
params; % Load global parameters.

%%% TRANSMITTER
%% Constellation (normalized to unit energy E=d^2)
% Implemented constellations:
% 'BPSK'
% '4-QAM'
% '8-PSK'
% '16-QAM'
constellation = generateConstellation('4-QAM');

%% Preamble
delay = zeros(1, randi([1,100], 1)); % Random length delay. Simulates period when nothing transmitted.

B2 =  [1 -1];       % Barker-2
B3 =  [1 1 -1];     % Barker-3
B4 =  [1 1 -1 1];   % Barker-4
B5 =  [1 1 1 -1 1]; % Barker-5%
B7 =  [1 1 1 -1 -1 1 -1];               % Barker-7
B11 = [1 1 1 -1 -1 -1 1 -1 -1 1 -1];   % Barker 11
B13 = [1 1 1 1 1 -1 -1 1 1 -1 1 -1 1];  % Barker-13
% Nested barker is constructed from Kronecker product of any two Barker
% sequences.
preamble = kron(1, B13);   % Select preamble here. if 1-dimensional barker, set first argument to 1.
preamble_symbs = preamble; % OPTIONAL. Map preamble to symbols.

%% Data Source
% b = randsrc(1,N,[0 1]); % Random message bits.
string = 'If the universe of so big, then why wont it fight me?';
b = string2bitStream(string);

b = [b zeros(1, N-length(b))]; % Pad string with zeros to fill length 432 packet.
disp(['Sent message: ' string])

%% Bits to Messages
M = length(constellation);      % M unique messages/symbols.
bpsymb = log2(M);       % Bits per symbol
Rsymb  = Rb/bpsymb;
Tsymb  = 1/Rsymb;
sps    = fs/Rsymb; % Samples per symbol. [Sample/s / Symbol/s = Sample/Symbol] 
m      = buffer(b, bpsymb)'; % Group information bits into bits per symbol
m_idx  = bi2de(m, 'left-msb')'+1; % Bits to symbol index, msb: the Most Significant Bit

%% Message to symbol
message_symbs = constellation(m_idx);    % Map messages to constellation symbols using indices
symbs = [delay preamble_symbs message_symbs];

%% Pulse modulation
[pulse ~]  = rtrcpuls(rolloff, Tsymb, fs, span); % Generate basis pulse.
symbs_upsample = upsample(symbs, sps); % Upsample to fs.
s          = conv(pulse, symbs_upsample); % Pulse shaping the symbol-train by convolving with basis pulse.
t_s        = (0:length(s)-1).*Ts; % Signal time-axis ie. x[n]*n*Ts.

%% Normalize baseband amplitude to sqrt(2) (because pulse-convolution attenuates signal).
s = s.*(sqrt(2)/max(abs(s)));

%% Split into I/Q streams
I = real(s); 
Q = imag(s);

%% Upmix to carrier frequency (MAKE SURE THIS IS REAL AND NOT COMPLEX!)
tx_sig = I.*cos(2*pi*fc*(t_s)) - Q.*sin(2*pi*fc*(t_s));

SIMULATION = 0; % DEBUG: Will simulate channel and oversampling.
if(SIMULATION)
    %%% CHANNEL
    %% Simulate oversampled signal by interpolating
    tx_sig_interp = interp(tx_sig, OS);

    figure
    t_s_upsample = (0:length(tx_sig_interp)-1).*(1/rx_fs);
    plot(t_s_upsample, tx_sig_interp)
    
    %% Channel (AWGN)
    SNRdB = 100;
    s_noisy = awgn(tx_sig_interp, SNRdB, 'measured');
else
    my_sound = audiorecorder_example()
    audio_player = audioplayer(tx_sig, fs);
    playblocking(audio_player); 
    
    s_noisy = my_sound.getaudiodata'
end
    
%%% Receiver
%% Down-mix
phzOffset = 0; % DEBUG: Phase-offset.

t_s_OS = (0:length(s_noisy)-1).*(1/rx_fs); % Redefine time axis for our received signal % Signal time-axis ie. x[n]*n*Ts.
rx_I =  s_noisy.*cos(2*pi*fc*t_s_OS + phzOffset);
rx_Q = -s_noisy.*sin(2*pi*fc*t_s_OS + phzOffset);

figure
subplot(121)
plot(t_s_OS, rx_I)
title('RX-I, down-mixed')
subplot(122)
plot(t_s_OS, rx_Q)
title('RX-Q, down-mixed')

%% Recombine I/Q streams into Complex
rx_sig = rx_I + 1i*rx_Q;

%% Match-filter
[pulse_OS ~]  = rtrcpuls(rolloff, Tsymb, rx_fs, span); % Generate basis pulse.

MF = fliplr(conj(pulse_OS));        %create matched filter impulse response
MF_out = filter(MF, 1, rx_sig);
MF_out = MF_out(length(MF):end);   %remove transient

figure; plot(real(MF_out))

%% Preamble-detection
% Correlates the baseband (MF_out) with the pulse modulated preamble (preamble_pulse_mod).
[preamble_pulse_OS ~]  = rtrcpuls(rolloff, Tsymb, rx_fs, span); % Generate basis pulse.
preamble_upsample = upsample(preamble_symbs, OS*sps);
preamble_pulse_mod = conv(preamble_pulse_OS, preamble_upsample);

% Length of transient before and after pulse modulated preamble, due to convolution.
transient_length = (length(preamble_pulse_mod)-length(preamble_upsample))/2;
preamble_pulse_mod = preamble_pulse_mod(transient_length+1:end-(transient_length+sps*OS));    % Truncate transient ends caused by convolution

figure 
subplot(211)
plot(preamble)
subplot(212)
plot(preamble_pulse_mod)



corr = conv(MF_out, fliplr(preamble_pulse_mod));
[~,peak_idx] = max(corr); % Find index of correlation peak.


figure 
plot(abs(MF_out))
hold on
plot(peak_idx, abs(MF_out(peak_idx)), 'or')
title('Received Baseband')


figure
subplot(311)
plot(abs(corr))
hold on; 
plot(peak_idx, abs(corr(peak_idx)), 'or')
grid on
subplot(312)
plot(real(corr))
grid on
subplot(313)
plot(imag(corr))
grid on




%% Normalize signal using preamble energy
MF_out = MF_out.* (1./abs(MF_out(peak_idx)))

%% Downsample by oversample rate (OS)
rx_downsampled = downsample(MF_out, OS);

%% Synchronizer
%samples = rx_downsampled(1:sps:end);

avg = zeros(sps, 1);
figure
for delta = 0:sps-1
    shifted_samples = circshift(rx_downsampled, delta);
    shifted_sync_points = shifted_samples(1:sps:end);
    avg(delta+1) = sum(abs(shifted_sync_points))./length(shifted_sync_points);

    hold off
    plot(shifted_sync_points, 'ko')
    xlabel('In Phase'); ylabel('Quadrature');
    title(['Sync offset: ',  num2str(delta)])
    xlim([-1.5, 1.5]);
    ylim([-1.5, 1.5]);
    grid on
    pause(0.05)
end

[max_avg, max_idx] = max(avg);
sync_offset = max_idx-1;

%% Shift downsampled signal by number of samples to sync, then sample at that point.
rx_downsampled_sync = circshift(rx_downsampled, sync_offset);
samples = rx_downsampled_sync(1:sps:end);

scatterplot(samples)

figure
plot(0:sps-1,avg)
xlabel('sync offset [samples]'); ylabel('Average energy');

% Plot sample points on baseband signal.
figure
subplot(121)
plot(real(rx_downsampled))
hold on
stem(1:sps:length(rx_downsampled), real(samples))
grid on; title('Real samples')

subplot(122)
plot(imag(rx_downsampled))
hold on
stem(1:sps:length(rx_downsampled), imag(samples))
grid on; title('Imag samples')
hold off

%% Phase-Correct
phzError = angle(corr(peak_idx))
samples_phzCorrected = samples.*exp(1i*-phzError); % Compensate for phase error by rotating.

% Minimum euclidean distance. Unmap symbols to messages
d = abs(repmat(samples_phzCorrected', 1, length(constellation)) - repmat(constellation, length(samples_phzCorrected), 1)).^2; % Compute distance of every sample symbol to every constellation point.
[~, nearest_idx] = min(d, [], 2);  % Get index of value with minimum distance to a constellation point. Traverse along column.

rx_mapped = constellation(nearest_idx); % Choose corresponding symbol.


%% Symbols to messages
m_hat = rx_mapped; 
if (length(m_hat) < length(message_symbs)) % Check if we've received at least enough symbols for 1 message.
    disp('Not enough symbols read! Check preamble detection.');
    return
end
m_hat = m_hat(1:N/bpsymb);

SER = symerr(message_symbs, m_hat) 

%% Messages to bits
b_hat_buffer = de2bi(nearest_idx-1, bpsymb, 'left-msb')';
b_hat = b_hat_buffer(:)'; %write as a vector
b_hat = b_hat(1:N);
BER = biterr(b, b_hat) %count of bit errors

b_hat_trunc = b_hat(1:end-rem(length(b_hat), 7)); % Truncate received bitstream so that it's divisible by 7 (ASCII length).
disp(['Received message (quote marks added): "' bitStream2string(b_hat_trunc) '"'])



%% Diagnostic info
%fprintf("TX bits: %d\n", N) %etc







