% Sonar imaging : Exercise for Module 4
%
%   See the sonar_exercise.pdf in the current folder for the exercise text.
%
% Author: Ole Marius Hoel Rindal <21.09.2021> adapted from code and data
%         from Roy Edgar Hansen and Fabrice Prieur

clear all;
close all;

%% Read the data channel data
channel_data = uff.channel_data();
channel_data.read('./sonar_ping.uff','/channel_data');
channel_data.plot()

%% Read a beamformed image of the raw uncompressed signal
% This is an example of how the image beamformed fomr the raw uncompressed
% signal might look like in part c) of the exercise
b_data_raw= uff.beamformed_data();
b_data_raw.read('./sonar_ping.uff','/b_data_raw');
b_data_raw.plot([],['SONAR uncompressed'],[],[],[],[],'m','dark')
colormap default
caxis([-65 -10])

%% Parameters for the transmit signal a Linear Frequency Modulated (LFM) upchirp pulse
bw = 30000;     % Bandwidth in Hertz
t_p = 0.0080;   % Pulse Length in seconds


%% Define the theoretical transmit pulse, a Linear Frequency Modulated (LFM) pulse
t = 0:1/channel_data.sampling_frequency:t_p;
f0 = channel_data.modulation_frequency - bw/2; % Start frequency of the chirp
f1 = channel_data.modulation_frequency + bw/2; % End frequency of the chirp
tx_pulse = chirp(t,f0,t_p,f1,'linear');

%% Create a copy of the channel data object to hold the pulse compressed data
channel_data_compressed = uff.channel_data(channel_data);
% Change the initial time to half of the pulse length to "center" the compressed data
channel_data_compressed.initial_time = t_p/2;

%% Do Pulse Compression using FFT
% Create an empty buffer to hold the resulting data 
match_filtered_data = zeros(channel_data.N_samples, channel_data.N_elements);

% Prepare the matched filter (complex conjugate of tx_pulse)
% Remember to complex conjugate when using FFT-based method
matched_filter = conj(tx_pulse);

% Zero-pad the matched filter to the same length as the data
matched_filter_padded = zeros(channel_data.N_samples, 1);
matched_filter_padded(1:length(matched_filter)) = matched_filter;

% Compute FFT of the matched filter once (can reuse for all channels)
H = fft(matched_filter_padded);

% Loop over all channels and do the matched filtering (pulse compression)
for ch = 1:channel_data.N_elements
    % Get the raw data for this channel
    raw_data_ch = channel_data.data(:,ch);
    
    % Compute FFT of the raw data
    X = fft(raw_data_ch);
    
    % Multiply in frequency domain (equivalent to convolution in time domain)
    Y = X .* H;
    
    % Transform back to time domain
    compressed_data_ch = ifft(Y);
    
    % Store the compressed data in the buffer
    match_filtered_data(:,ch) = compressed_data_ch;
end

% replace the copied data in the uff.channel_data object with the
% pulse compressed data
channel_data_compressed.data = match_filtered_data;

%% Create figures directory if it doesn't exist
if ~exist('figs', 'dir')
    mkdir('figs');
end

%% Plot and compare channel data compressed and not compressed (Exercise b)
% Create Figure 2: Plot magnitude of a single channel before and after pulse compression
% Showing both full range and zoomed view (samples 4000-4400)
fig1 = figure('Name', 'Raw vs Pulse Compressed Channel Data');

% Select a single channel to display
ch = 1;

% Time vectors in milliseconds
time_raw = (channel_data.initial_time + (0:channel_data.N_samples-1)/channel_data.sampling_frequency) * 1000; % ms
time_compressed = (channel_data_compressed.initial_time + (0:channel_data_compressed.N_samples-1)/channel_data_compressed.sampling_frequency) * 1000; % ms

% Find indices where time is positive (only use positive delays)
idx_raw_positive = time_raw >= 0;
idx_compressed_positive = time_compressed >= 0;

% Zoom range (samples 4000 to 4400)
zoom_start = 4000;
zoom_end = 4400;

% Find the overlapping time range for both raw and compressed data in the zoom window
time_raw_zoom = time_raw(zoom_start:zoom_end);
time_compressed_zoom = time_compressed(zoom_start:zoom_end);
xlim_min = max(min(time_raw_zoom), min(time_compressed_zoom));
xlim_max = min(max(time_raw_zoom), max(time_compressed_zoom));

% Top left: Raw data - Full range (positive delays only)
subplot(2,2,1);
plot(time_raw(idx_raw_positive), abs(channel_data.data(idx_raw_positive,ch)));
xlabel('Time [ms]');
ylabel('Magnitude');
title(sprintf('Raw Data - Channel %d (Full Range)', ch));
grid on;
axis tight;

% Top right: Raw data - Zoomed (samples 4000-4400)
subplot(2,2,2);
plot(time_raw(zoom_start:zoom_end), abs(channel_data.data(zoom_start:zoom_end,ch)));
xlabel('Time [ms]');
ylabel('Magnitude');
title(sprintf('Raw Data - Channel %d (Zoomed)', ch));
grid on;
xlim([xlim_min, xlim_max]);

% Bottom left: Pulse Compressed data - Full range (positive delays only)
subplot(2,2,3);
plot(time_compressed(idx_compressed_positive), abs(channel_data_compressed.data(idx_compressed_positive,ch)));
xlabel('Time [ms]');
ylabel('Magnitude');
title(sprintf('Pulse Compressed - Channel %d (Full Range)', ch));
grid on;
axis tight;

% Bottom right: Pulse Compressed data - Zoomed (samples 4000-4400)
subplot(2,2,4);
plot(time_compressed(zoom_start:zoom_end), abs(channel_data_compressed.data(zoom_start:zoom_end,ch)));
xlabel('Time [ms]');
ylabel('Magnitude');
title(sprintf('Pulse Compressed - Channel %d (Zoomed)', ch));
grid on;
xlim([xlim_min, xlim_max]);

% Save this figure for the report
print(fig1, 'figs/figure1.pdf', '-dpdf', '-fillpage');
fprintf('Saved figure to figs/figure1.pdf\n');

%% Task c) Beamforming with USTB

%% Define a sector scan based on the theoretical angular resolution
% Calculate theoretical angular resolution
c = channel_data.sound_speed;  % Sound speed in water (m/s)
fc = channel_data.modulation_frequency;  % Center frequency (Hz)
lambda = c / fc;  % Wavelength (m)
D = channel_data.probe.element_width * channel_data.N_elements;  % Aperture size (m)
angular_resolution = lambda / D;  % Theoretical angular resolution (radians)

fprintf('Theoretical angular resolution: %.4f radians = %.2f degrees\n', ...
    angular_resolution, rad2deg(angular_resolution));

% Define sector scan parameters
% Maximum range from the data (use the maximum time delay)
max_range = c * max(time_raw) / 1000 / 2;  % Divide by 2 for round-trip, time is in ms
fprintf('Maximum range: %.1f m\n', max_range);

% Create sector scan
scan = uff.sector_scan();
scan.azimuth_axis = linspace(-30*pi/180, 30*pi/180, 256);  % Azimuth angles (radians)
scan.depth_axis = linspace(1, max_range, 512);  % Depth range (meters), resolution < 20 cm

fprintf('Azimuth resolution: %.2f degrees\n', rad2deg(diff(scan.azimuth_axis(1:2))));
fprintf('Depth resolution: %.2f cm\n', diff(scan.depth_axis(1:2))*100);

%% Set up the DAS beamformer for raw (uncompressed) data
% Create midprocess object
mid = midprocess.das();
mid.channel_data = channel_data;
mid.scan = scan;

% Use Hamming window apodization to suppress grating lobes
% This will reduce sidelobe levels significantly
mid.transmit_apodization.window = uff.window.hamming;
mid.receive_apodization.window = uff.window.hamming;

% Run beamforming on raw data
fprintf('Beamforming raw data with Hamming apodization to suppress grating lobes...\n');
b_data = mid.go();

%% Set up the DAS beamformer for pulse compressed data
mid_compressed = midprocess.das();
mid_compressed.channel_data = channel_data_compressed;
mid_compressed.scan = scan;

% Use Hamming window apodization to suppress grating lobes
mid_compressed.transmit_apodization.window = uff.window.hamming;
mid_compressed.receive_apodization.window = uff.window.hamming;

% Run beamforming on pulse compressed data
fprintf('Beamforming pulse compressed data with Hamming apodization to suppress grating lobes...\n');
b_data_compressed = mid_compressed.go();

%% Display the results
% Figure showing channel data (before beamforming) for both raw and compressed
fig2 = figure('Name', 'Channel Data Comparison (Before Beamforming)');

% Raw channel data
subplot(1,2,1);
time_plot = time_raw(idx_raw_positive);
imagesc(1:channel_data.N_elements, time_plot, 20*log10(abs(channel_data.data(idx_raw_positive,:))));
xlabel('Channel');
ylabel('Time [ms]');
title('Raw Channel Data');
colorbar;
caxis([-80 -20]);
axis xy;

% Pulse compressed channel data
subplot(1,2,2);
time_plot_compressed = time_compressed(idx_compressed_positive);
imagesc(1:channel_data_compressed.N_elements, time_plot_compressed, 20*log10(abs(channel_data_compressed.data(idx_compressed_positive,:))));
xlabel('Channel');
ylabel('Time [ms]');
title('Pulse Compressed Channel Data');
colorbar;
caxis([-80 -20]);
axis xy;

% Save figure
print(fig2, 'figs/channel_data_comparison.pdf', '-dpdf', '-fillpage');
fprintf('Saved figure to figs/channel_data_comparison.pdf\n');

%% Display beamformed images
% Raw data beamformed
fig3 = figure('Name', 'Beamformed Image - Raw Data');
b_data.plot([],['SONAR - Raw (Uncompressed)'],[],[],[],[],'m','dark');
colormap default;
caxis([-65 -10]);
print(fig3, 'figs/beamformed_raw.pdf', '-dpdf', '-fillpage');
fprintf('Saved figure to figs/beamformed_raw.pdf\n');

% Pulse compressed data beamformed
fig4 = figure('Name', 'Beamformed Image - Pulse Compressed');
b_data_compressed.plot([],['SONAR - Pulse Compressed'],[],[],[],[],'m','dark');
colormap default;
caxis([-65 -10]);
print(fig4, 'figs/beamformed_compressed.pdf', '-dpdf', '-fillpage');
fprintf('Saved figure to figs/beamformed_compressed.pdf\n');

%% Interactive comparison of beamformed images
fig5 = figure('Name', 'SONAR Comparison');
b_data_compare = uff.beamformed_data(b_data);
b_data_compare.data(:,1) = b_data.data./mean(b_data.data(:));
b_data_compare.data(:,2) = b_data_compressed.data./mean(b_data_compressed.data(:));
b_data_compare.plot([],['SONAR Comparison: 1 = raw, 2 = pulse compressed'],[],[],[],[],'m','dark');
caxis([-65 -10]);
colormap default;
print(fig5, 'figs/beamformed_comparison.pdf', '-dpdf', '-fillpage');
fprintf('Saved figure to figs/beamformed_comparison.pdf\n');

fprintf('\n=== Analysis Questions ===\n');
fprintf('1. Difference before/after beamforming:\n');
fprintf('   - Before: Channel data shows time-domain signals per element\n');
fprintf('   - After: Beamformed image shows spatial reflectivity map\n\n');

fprintf('2. Difference before/after pulse compression:\n');
fprintf('   - Before: Long pulse duration, poor range resolution\n');
fprintf('   - After: Compressed pulse, improved range resolution\n\n');

fprintf('3. Target at ~130m range:\n');
fprintf('   - Examine the pulse compressed beamformed image\n');
fprintf('   - Measure the target size in the azimuth and range dimensions\n');
