close all
clear


%% Select 'role' of this script.
% 'Both' transmits and records/receives.
% TX and RX transmit only and recieve only, resp.
BOTH_MODE = 0;
TX_MODE   = 1;
RX_MODE   = 2;
MODE = RX_MODE;


rx_fs = 48000;

%% Input parameters

Rb = rx_fs/6%/12.5;   % bit rate [bit/sec]
fs = rx_fs; % Sampling frequency [Sample/s]
Ts = 1/fs; % Sample time [s/Sample]
fc = 21000;  % Carrier frequency [Hz]
span = 6;   % Pulse-width in symbol times.


decimation_factor = rx_fs/fs;

rolloff = 0.4; % Roll-off factor (ie. alpha) for Raised Cosine pulse.

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

if length(constellation) > 2 % Voronoi need at least 3 constellation points.
    voronoi(real(constellation), imag(constellation))
    title('Constellation plot')
    xlabel("Phase"); ylabel("Quadrature");
end



% Preamble (including initial delay where nothing is tx'd)
delay_time = randi([1,200], 1);       % Generate a random delay between 1 and 20 samples long.
delay = zeros(1, delay_time);


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
preamble = kron(B7,B7);   % Select preamble here. if 1-dimensional barker, set first argument to 1.


%preamble_symbs = preamble.*(1+1i)./sqrt(2); % Map preamble to symbols, so that they are 180deg apart! MUST BE MODIFIED IF USING DIFFERENT CONSTELLATION.
preamble_symbs = preamble; % Do NOT need to map to valid symbols


%% Data source


b = genAllTenBit();
%b = genAllASCII();

N = length(b);    % Number of bits to transmit
% N = 432
%b = randsrc(1,N,[0 1]); % Random message bits.

string = 'To be, or not to be: that is the question: Whether  "tis nobler in the mind to suffer The slings and arrows of outrageous fortune, Or to take arms against a sea of troubles, And by opposing end them? To die: to sleep; No more; and by a sleep to say we end The heart-ache and the thousand natural shocks That flesh is heir to,  "tis a consummation Devoutly to be wish "d. To die, to sleep; To sleep: perchance to dream: ay, there "s the rub; For in that sleep of death what dreams may come When we have shuffled off this mortal coil, Must give us pause: there "s the respect That makes calamity of so long life; For who would bear the whips and scorns of time, The oppressor "s wrong, the proud man "s contumely, The pangs of despised love, the law "s delay, The insolence of office and the spurns That patient merit of the unworthy takes, When he himself might his quietus make With a bare bodkin? who would fardels bear, To grunt and sweat under a weary life, But that the dread of something after death, The undiscover "d country from whose bourn No traveller returns, puzzles the will And makes us rather bear those ills we have Than fly to others that we know not of? Thus conscience does make cowards of us all; And thus the native hue of resolution Is sicklied o "er with the pale cast of thought, And enterprises of great pith and moment With this regard their currents turn awry, And lose the name of action.--Soft you now! The fair Ophelia! Nymph, in thy orisons Be all my sins remember "d!';
b = string2bitStream(string);
N = length(b);
if length(b) > N
    disp('Source too long! Truncating...')
    b = b(1:N);
end
b = [b zeros(1, N-length(b))]; % Pad string with zeros to fill length 432 packet.


%% 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 ]; % Prepend variable delay and preamble to message.
symbs = [preamble_symbs message_symbs];




% Generate basic pulse
[pulse t] = rtrcpuls(rolloff, Tsymb, fs, span); % Generate basis pulse.



if MODE == RX_MODE || MODE == BOTH_MODE 
    %% Record!
    my_sound = audiorecorder_example()
    if MODE == RX_MODE
        pause(5)
    end
end

% Play carrier wave in speakers!
% Create an audioplayer object to play the noise at a given fs.
if MODE == TX_MODE || MODE == BOTH_MODE 
end

if MODE == RX_MODE || MODE == BOTH_MODE 
    figure
    plot(my_sound.getaudiodata); title('Recorded audio');


    s_noisy = my_sound.getaudiodata.';



    %% It's all good, man

    %% Receiver
    %phzOffs = 2*pi/7/2;% Phase offset
    phzOffs = 0;

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

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

    % Lowpass filter to give baseband signal.
    rx_I_filtered = lowpass(rx_I, Rsymb, rx_fs, 'ImpulseResponse','iir','Steepness',0.75);
    rx_Q_filtered = lowpass(rx_Q, Rsymb, rx_fs, 'ImpulseResponse','iir','Steepness',0.75);

    rx_I_filtered = rx_I_filtered;
    rx_Q_filtered = rx_Q_filtered;

    figure
    subplot(121)
    plot(t_s, rx_I_filtered)
    title('RX-I, filtered')
    subplot(122)
    plot(t_s, rx_Q_filtered)
    title('RX-Q, filtered')


    scatterplot(rx_I_filtered+i*rx_Q_filtered)
    
    % IMPORTANT: lowpass doesn't seem to introduce any latency!
    % No need to remove any transient etc...
    rx_I_filtered = downsample(circshift(rx_I_filtered, 0), decimation_factor);
    rx_Q_filtered = downsample(circshift(rx_Q_filtered, 0), decimation_factor);

    % %% Preamble detection before downsampling %%%%%%%%%%%%%%%%%%%%%%%
    % sig_tmp = rx_I_filtered + i*rx_Q_filtered;
    % 
    % % Zero-stuff
    % %preamble_upsampled = upsample(preamble, decimation_factor*sps)
    % % Stepwise upsample (ie. same value maintained) !!BEST!!
    % preamble_upsampled = reshape(repmat(preamble, decimation_factor*sps, 1), 1, []);
    % % Interpolate !!DOES NOT WORK!!!
    % %preamble_upsampled = interp(preamble, decimation_factor*sps)
    % 
    % corr = conv(sig_tmp, fliplr(preamble_upsampled))
    % [~,peak_idx] = max(corr);
    % 
    % figure
    % plot( abs(corr))
    % hold on
    % plot(peak_idx, abs(corr(peak_idx)), 'or')
    % 
    % % Cut off preamble, advance by one symbol time.
    % sig_tmp = sig_tmp(peak_idx:end);
    % 
    % sig_tmp = downsample(sig_tmp, decimation_factor)
    % 
    % plot(real(sig_tmp))
    % plot(imag(sig_tmp))
    % 
    % rx_I_filtered = real(sig_tmp);
    % rx_Q_filtered = imag(sig_tmp);
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    % Downsample here???



    % Matched filter
    MF = fliplr(conj(pulse));        %create matched filter impulse response
    MF_I_output = filter(MF, 1, rx_I_filtered); 
    MF_I_output = MF_I_output(length(MF):end); %remove transient
    MF_Q_output = filter(MF, 1, rx_Q_filtered); 
    MF_Q_output = MF_Q_output(length(MF):end); %remove transient

    figure
    subplot(121)
    plot(MF_I_output)
    hold on
    grid on; title('RX-I, matched-filtered')
    subplot(122)
    plot(MF_Q_output)
    grid on; title('RX-Q, matched-filtered')

    % Synchroniser (TODO)
    rx_symbs = [MF_I_output(1:sps:end)', MF_Q_output(1:sps:end)']; % This assumes PERFECT synchronization.

    % Synchroniser LOL %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

    avg = zeros(sps, 1);

    MF_output_complex = MF_I_output + i.*MF_Q_output;

    scatterplot(MF_output_complex)    

    %MF_output_complex = circshift(MF_output_complex, 12) % DEBUG: intentionally "unsynchronise" symbols.
        figure
        hold off
    for delta = 1:sps
        shifted_MF_output = circshift(MF_output_complex, delta-1);
        shifted_samples = shifted_MF_output(1:sps:end);
        avg(delta) = sum(abs(shifted_samples))./length(shifted_samples);

        % Visual animation of constellation points converging (if
        % synchronisation possible)

%         plot(shifted_samples, 'ko')
%         xlabel('In Phase'); ylabel('Quadrature');
%         title(['Sync offset: ',  num2str(delta-1)])
%         pause(0.01)
    end

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

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

    MF_I_output_shifted = circshift(MF_I_output, sync_offset);
    MF_Q_output_shifted = circshift(MF_Q_output, sync_offset);
    rx_symbs = [MF_I_output_shifted(1:sps:end)', MF_Q_output_shifted(1:sps:end)'];
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


    % Plot sample points on baseband signal.
    figure
    subplot(121)
    plot(MF_I_output_shifted)
    hold on
    stem(1:sps:length(MF_I_output_shifted), MF_I_output_shifted(1:sps:end))
    grid on; title('RX-I, sampled')
    subplot(122)
    plot(MF_Q_output_shifted)
    hold on
    stem(1:sps:length(MF_Q_output_shifted), MF_Q_output_shifted(1:sps:end))
    grid on; title('RX-Q, sampled')
    hold off

    % Unmapping to symbols
    samples = rx_symbs(:,1) + 1i.*rx_symbs(:,2);
    figure
    hold on
    if length(constellation)>2
        voronoi(real(constellation), imag(constellation))
    end
    plot(samples, 'o')
    hold off
    axis square; grid on
     xlim([-1 1]);
     ylim([-1 1]);
    title('Received samples');

    % Preamble correlation
    sample_corr = conv(samples, fliplr(preamble));
    [~,peak_idx] = max(sample_corr);
    figure
    subplot(311)
    plot(abs(sample_corr))
    hold on; plot(peak_idx, abs(sample_corr(peak_idx)), 'or')
    grid on
    subplot(312)
    plot( real(sample_corr) )
    grid on
    subplot(313)
    plot(imag(sample_corr))
    grid on
    
    % Get received preamble symbols (based on preamble detection)
    rx_preamble_symbs =  samples(peak_idx-(length(preamble)-1):peak_idx);
    % Average angle
    % rx_preamble_avg_angle = angle(mean(rx_preamble_symbs))
    angles = angle(rx_preamble_symbs);
    neg_angles = angles < 0 ;
    angles(neg_angles) =angles(neg_angles)+pi;
    rx_preamble_avg_angle = mean(angles);
    
    if angle(samples(peak_idx)) < 0
        rx_preamble_avg_angle = rx_preamble_avg_angle+pi;
    end
    
    figure
    plot([exp(i*rx_preamble_avg_angle), 1*exp(i*(rx_preamble_avg_angle+pi))])
    title('Phase offset')
    
    % Attempt to use preamble to find phase errors/ambiguities.
    % eg. if there is a phase error of 90deg, there will be a correlation, but
    % negative in one of the I/Q streams.
    phzError = rx_preamble_avg_angle%angle(sample_corr(peak_idx))%-pi/4 % Offset by nominal phase of preamble.
    samples_phzCorrected = samples.*exp(1i*-phzError); 

    % Normalize signal amplitude using preamble
    % samples_phzCorrected = samples_phzCorrected./abs(samples(peak_idx));
    samples_phzCorrected = samples_phzCorrected./( sum(abs(samples))./length(samples)); % Normalize to average signal energy.

   
    % Compute average energy of all preamble symbols
    avg_preamble_energy = mean(abs(samples_phzCorrected(peak_idx-(length(preamble)-1):peak_idx)))
    
    %% Normalize received sample energy
    samples_phzCorrected = samples_phzCorrected.* 1./avg_preamble_energy;
    
    % Plot phase-corrected samples.
    figure(100)
    hold on
    plot(samples_phzCorrected(peak_idx+1:peak_idx+ceil(N/bpsymb)), 'ob'); title('Phase-corrected rx. Decision regions overlaid.')
    if length(constellation)>2
        voronoi(real(constellation), imag(constellation))
    end
    %plot(real(constellation), imag(constellation), 'or')
    grid on; axis square;
    xlim([-2 2])
    ylim([-2 2])
    hold off
    
    figure
    stem(abs(samples_phzCorrected(peak_idx+1:peak_idx+ceil(N/bpsymb)))); title('Received samples')

    % Minimum euclidean distance. Unmap symbols to messages
    d = abs(repmat(samples_phzCorrected, 1, length(constellation)) - repmat(constellation, length(samples), 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.

    % Preamble detection
    %corr = conv(rx_mapped, fliplr(preamble)); % Compare preamble to received symbol vector. Peak at end of preamble.
    %[~,peak_idx] = max(corr);





    % Symbols to messages
    m_hat = rx_mapped(peak_idx+1:end); % Get symbols after preamble -> actual message.
    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:ceil(N/bpsymb));

    [SER, SER_ratio, SER_loc] = symerr(message_symbs, m_hat);
    

    if SER > 0
        figure
        if length(constellation)>2
            voronoi(real(constellation), imag(constellation))
        end
        hold on
        plot(samples_phzCorrected(peak_idx+1:peak_idx+ceil(N/bpsymb))' .* SER_loc, '*r')
        title('Incorrect Symbols')
        axis square
    end

    
    
    % Messages to bits
    b_hat_buffer = de2bi(nearest_idx(peak_idx+1:end)-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

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

    figure
      [pxx, f] = pwelch(rx_I + rx_Q*1i ,1024,768,1024, rx_fs); % note that pwr_spect.f will be normalized frequencies
    f = fftshift(f); %shift to be centered around fs
    f(1:length(f)/2) = f(1:length(f)/2) - rx_fs; % center to be around zero
    p = fftshift(10*log10(pxx/max(pxx))); % shift, normalize and convert PSD to dB
    plot(f,p)
    

    %% Eye diagram
    % WARNING! Heavy Computation required.
%     eye_vec = MF_I_output_shifted+i*MF_Q_output_shifted;
%     eye_vec = eye_vec.*exp(i*-phzError);
%     eyediagram(eye_vec, sps, Tsymb)

    binscatter( rx_I_filtered, rx_Q_filtered , [50 50])
    
end
