% Script to illustrate the frequency shift between the source and recorded
% sensor data in different dimensions due to the time derivative that
% appears in the source term in the wave equation, and the convolution with
% the Green's function.
%
% author: Bradley Treeby
% date: 24 November 2020
% last update: 25 November 2020

clearvars;

% set number of dimensions
dim = 2;

% type of source 'p' or 'u'
source_type = 'p';

% create the computational grid
Nx = 64;
Ny = 32;
Nz = 32;
dx = 0.1e-3;
switch dim
    case 1
        kgrid = kWaveGrid(Nx, dx);        
    case 2
        kgrid = kWaveGrid(Nx, dx, Ny, dx);        
    case 3
        kgrid = kWaveGrid(Nx, dx, Ny, dx, Nz, dx);
end

% define the properties of the propagation medium
medium.sound_speed = 1500;
medium.density = 1000;

% define source and sensor mask
pml_size = 10;
pml_offset = 4;
switch dim
    case 1
        source_mask = zeros(Nx, 1);
        source_mask(pml_size + pml_offset) = 1;
        sensor.mask = zeros(Nx, 1);
        sensor.mask(Nx - pml_size - pml_offset) = 1;        
    case 2
        source_mask = zeros(Nx, Ny);
        source_mask(pml_size + pml_offset, ceil(Ny/2)) = 1;
        sensor.mask = zeros(Nx, Ny);
        sensor.mask(Nx - pml_size - pml_offset, ceil(Ny/2)) = 1;
    case 3
        source_mask = zeros(Nx, Ny, Nz);
        source_mask(pml_size + pml_offset, ceil(Ny/2), ceil(Nz/2)) = 1;
        sensor.mask = zeros(Nx, Ny, Nz);
        sensor.mask(Nx - pml_size - pml_offset, ceil(Ny/2), ceil(Nz/2)) = 1;        
end

% set time step
Fs = 60e6;
kgrid.setTime(350, 1/Fs);

% define source
tone_burst_freq = 1e6; % [Hz]
tone_burst_cycles = 1;
source_sig = toneBurst(Fs, tone_burst_freq, tone_burst_cycles);
source_sig = source_sig/max(abs(source_sig));

% filter source
source_sig = [source_sig, zeros(1, 3*length(source_sig))];
source_sig = filterTimeSeries(kgrid, medium, source_sig, 'PPW', 6);

% assign
if strcmp(source_type, 'p')
    source.p_mask = source_mask;
    source.p = source_sig;
else
    source.u_mask = source_mask;
    source.ux = source_sig / (medium.sound_speed * medium.density);
end

% run the simulation
input_args = {...
    'PMLSize', pml_size, ...
    'PlotPML', true, ...
    'DataCast', 'single', ...
    'PlotScale', [-1, 1] * 3 / 3^dim, ...
    };
switch dim
    case 1
        sensor_data = kspaceFirstOrder1D(kgrid, medium, source, sensor, input_args{:});
    case 2
        sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});
    case 3
        sensor_data = kspaceFirstOrder3D(kgrid, medium, source, sensor, input_args{:});
end
        
% calculate spectrum
[f1, as1] = spect(source_sig, Fs, 'FFTLength', 5 * length(sensor_data));
[f2, as2] = spect(sensor_data, Fs, 'FFTLength', 5 * length(sensor_data));

% plot normalised traces
figure; 
plot(source_sig ./ max(source_sig));
hold on;
plot(sensor_data ./ max(sensor_data), '--');
xlabel('Time index');
ylabel('Pressure');
legend('source', 'sensor', 'Location', 'best');
title([num2str(dim) 'D, ' source_type ' source']);

% plot normalised spectrum
figure;
plot(f1 * 1e-6, as1 ./ max(as1));
hold on;
plot(f2 * 1e-6, as2 ./ max(as2), '--');
set(gca, 'XLim', [0, 7.5]);
xlabel('Frequency [MHz]');
ylabel('Amplitude Spectrum');
legend('source', 'sensor', 'Location', 'best');
title([num2str(dim) 'D, ' source_type ' source']);