% Script to compare different ways of calculating the intensity.

% author: Bradley Treeby
% date: 24th November 2020
% last update: 24th November 2020

clearvars;

%% ========================================================================
% SIMULATION PROPERTIES
% =========================================================================

% grid size
Nx = 256;
Ny = 256;
x_sz = 25e-3;
pml_size = 20;

% source properties 
source_freq = 1e6;
source_mag = 10e6;

% properties of the propagation medium
medium.sound_speed  = 1500;
medium.density = 1000;
medium.alpha_power = 1.5;
medium.alpha_coeff = 5;  

% time array
points_per_period = 100;
t_end = 30e-6;
periods_to_record = 3;

%% ========================================================================
% SETUP SIMULATION
% =========================================================================

% create the computational grid
dx = x_sz/(Nx - 2*pml_size);    	
kgrid = makeGrid(Nx, dx, Ny, dx);

% define transducer
radius = 90;
source.p_mask = zeros(Nx, Ny);
source.p_mask(pml_size + 1, Ny/2 - radius + 1:Ny/2 + radius) = 1;

% calculate the time length of one acoustic period
T = 1/source_freq;

% create the time array using a nice number of points
dt = T/points_per_period;
kgrid.t_array = 0:dt:t_end;

% double check the CFL
disp(['CFL = ' num2str(dt * medium.sound_speed / dx)]);

% define a time varying sinusoidal source
source_signal = source_mag * sin(2 * pi * source_freq * kgrid.t_array);

% smooth the source
source_signal = filterTimeSeries(kgrid, medium, source_signal);

% create the time delays to focus in the middle of the grid
source.p = focus(kgrid, source_signal, source.p_mask, [0, 0], medium.sound_speed);

% create a display mask to display the transducer
display_mask = source.p_mask;

% create a sensor mask covering the entire computational domain
sensor.mask = ones(Nx, Ny);

% set start point to record the final few periods
sensor.record_start_index = kgrid.Nt - periods_to_record * points_per_period + 1;

% set the recorded parameters
sensor.record = {'p', 'u', 'p_rms', 'u_rms', 'u_non_staggered', 'I_avg'};

%% ========================================================================
% RUN SIMULATION
% =========================================================================

% assign the input options 
input_args = {...
    'DataCast', 'single', ...
    'DataRecast', true, ...
    'DisplayMask', display_mask, ...
    'PlotPML', false, ...
    'PlotScale', [-source_mag, source_mag], ...
    };

% run the simulation
sensor_data = kspaceFirstOrder2D(kgrid, medium, source, sensor, input_args{:});

%% ========================================================================
% PROCESS OUTPUT DATA
% =========================================================================

% extract the amplitudes
[p_amp, p_ph] = extractAmpPhase(sensor_data.p, 1/kgrid.dt, source_freq, 'FFTPadding', 1, 'Window', 'Rectangular');
[ux_amp, ux_ph] = extractAmpPhase(sensor_data.ux_non_staggered, 1/kgrid.dt, source_freq, 'FFTPadding', 1, 'Window', 'Rectangular');
[uy_amp, uy_ph] = extractAmpPhase(sensor_data.uy_non_staggered, 1/kgrid.dt, source_freq, 'FFTPadding', 1, 'Window', 'Rectangular');

% compute intensity using spectral components
Ix_spect = p_amp .* ux_amp .* cos(p_ph - ux_ph)/2;
Iy_spect = p_amp .* uy_amp .* cos(p_ph - uy_ph)/2;
I_spect = sqrt(Ix_spect.^2 + Iy_spect.^2);
I_spect = reshape(I_spect, Nx, Ny);

% compute using a plane wave assumption
I_pw = p_amp.^2 ./ (2 * medium.sound_speed * medium.density);
I_pw = reshape(I_pw, Nx, Ny);

% temporally shift the particle velocity
ux_non_staggered = fourierShift(sensor_data.ux_non_staggered, 1/2);
uy_non_staggered = fourierShift(sensor_data.uy_non_staggered, 1/2);

% calculate the instantaneous intensity using p and u
Ix = sensor_data.p .* ux_non_staggered;
Iy = sensor_data.p .* uy_non_staggered;

% calculate the discrete time average
Ix_disc = mean(Ix, 2);
Iy_disc = mean(Iy, 2);
I_disc = sqrt(Ix_disc.^2 + Iy_disc.^2);
I_disc = reshape(I_disc, Nx, Ny);

% calculate the time average of p and u separately and multiply
p_rms = reshape(sensor_data.p_rms, Nx, Ny);

% ux_rms = sqrt(mean(sensor_data.ux_non_staggered.^2, 2));
% uy_rms = sqrt(mean(sensor_data.uy_non_staggered.^2, 2));
% ux_rms = reshape(ux_rms, Nx, Ny);
% uy_rms = reshape(uy_rms, Nx, Ny);

ux_rms = reshape(sensor_data.ux_rms, Nx, Ny);
uy_rms = reshape(sensor_data.uy_rms, Nx, Ny);
ux_rms = fourierShift(ux_rms, -1/2, 1);
uy_rms = fourierShift(uy_rms, -1/2, 2);

I_sep = p_rms .* sqrt( ux_rms.^2 + uy_rms.^2 );

% reshape the pressure and phases for plotting
p_amp = reshape(p_amp, Nx, Ny);
p_ph = reshape(p_ph, Nx, Ny);
ux_ph = reshape(ux_ph, Nx, Ny);
uy_ph = reshape(uy_ph, Nx, Ny);

% trim the source from everything
trim_points = 25;
end_trim = 5;
p_amp = p_amp(trim_points:end - end_trim, :);
p_ph = p_ph(trim_points:end - end_trim, :);
ux_ph = ux_ph(trim_points:end - end_trim, :);
uy_ph = uy_ph(trim_points:end - end_trim, :);
I_spect = I_spect(trim_points:end - end_trim, :);
I_disc = I_disc(trim_points:end - end_trim, :);
I_pw = I_pw(trim_points:end - end_trim, :);
I_sep = I_sep(trim_points:end - end_trim, :);

% compute some differences
I_ref = I_spect;
I_spect_diff = 100 * abs(I_spect - I_ref) / max(I_ref(:));
I_disc_diff = 100 * abs(I_disc - I_ref) / max(I_ref(:));
I_pw_diff = 100 * abs(I_pw - I_ref) / max(I_ref(:));
I_sep_diff = 100 * abs(I_sep - I_ref) / max(I_ref(:));

% plot axis
x_ax = kgrid.x_vec(trim_points:end - end_trim) * 1e3;
x_ax = x_ax - x_ax(1);
y_ax = kgrid.y_vec * 1e3;

%% ========================================================================
% PLOTTING
% =========================================================================

% plot the fields
figure;

subplot(4, 2, 1);
imagesc(y_ax, x_ax, I_spect);
colorbar;
axis image;
title('I spectral');

subplot(4, 2, 2);
imagesc(y_ax, x_ax, I_spect_diff);
colorbar;
axis image;
title('I spectral diff');

subplot(4, 2, 3);
imagesc(y_ax, x_ax, I_disc);
colorbar;
axis image;
title('I discrete');

subplot(4, 2, 4);
imagesc(y_ax, x_ax, I_disc_diff);
colorbar;
axis image;
title('I discrete diff');

subplot(4, 2, 5);
imagesc(y_ax, x_ax, I_pw);
colorbar;
axis image;
title('I plane wave');

subplot(4, 2, 6);
imagesc(y_ax, x_ax, I_pw_diff);
colorbar;
axis image;
title('I plane wave diff');

subplot(4, 2, 7);
imagesc(y_ax, x_ax, I_sep);
colorbar;
axis image;
title('I sep');

subplot(4, 2, 8);
imagesc(y_ax, x_ax, I_sep_diff);
colorbar;
axis image;
title('I sep diff');

% plot the phase components for p and ux
figure;

subplot(2, 3, 1);
imagesc(y_ax, x_ax, p_ph);
colorbar;
axis image;
title('p phase');

subplot(2, 3, 2);
imagesc(y_ax, x_ax, ux_ph);
colorbar;
axis image;
title('ux phase');

subplot(2, 3, 3);
imagesc(y_ax, x_ax, uy_ph);
colorbar;
axis image;
title('uy phase');

subplot(2, 3, 5);
imagesc(y_ax, x_ax, cos(p_ph - ux_ph));
colorbar;
axis image;
title('cos(p phase - ux phase)');

subplot(2, 3, 6);
imagesc(y_ax, x_ax, cos(p_ph - uy_ph));
colorbar;
axis image;
title('cos(p phase - uy phase)');