clear variables;
close all;
clc;

%% Defining params

% Bottom parameters

cw = 1500;      % sound speed in water, m/s
cb = 2000;      % sound speed in bottom, m/s
rhow = 1;       % density in water, g/cm^3
rhob = 2;       % density in bottom, g/cm^3
attb = 0.5;     % bottom attenuation, dB/wavelength
h0 = 90;        % h_0, m
alpha = 0.5;   % slope angle, degrees

%alpha = 2.86;   % slope angle, degrees


% Source

zs = 10;        % source depth, m
freq = 50;      % frequency, Hz

% Wavenumbers (Pekeris disp. relation is satisfied up to 10^(-10) )
% if cw, cb, rhow, rhob, h0 or freq are modified, kj must be recalculated!


kj = [  0.207245856833476;
        0.200306356759023;
        0.187717636741757;
        0.168196050824593  ];

    
% Receiver and grid:

zr = 10;        % receiver depth, m

zm = 500;       % domain size, m       
xm = 8000;
ym = 4000;

nz = 500;       % number of grid points
ny = 1000;
nx = 4000;

%% Computing the field

% some auxilliary variables...

gammaw = 1/rhow;
gammab = 1/rhob;

eta = 1/(40*pi*log10(exp(1)));
nub = 2*1i*eta*attb*(2*pi*freq/cb)^2;
omeg = 2*pi*freq;
tanA = tan(alpha*pi/180);

x = linspace(0,xm,nx);
y = linspace(-ym,ym,ny);
z = linspace(0,zm,nz);

nmod = length(kj);

Pxz(1:nx,1:nz)=0;
Pxy(1:nx,1:ny)=0;
Px(1:nx)=0;

phi(1:nz) = 0;

aPtrack(1:nx) = 0;      % acoustic field along the track -- for y=0, z=zr
aPxy(1:nx,1:ny) = 0;    % acoustic field P(x,y,z_r) at z = zr -- 2D array

Xx = repmat(x',1,ny);
Yy = repmat(y,nx,1);


%% This loop computes the sum of Eq. (12)

for jj = 1:nmod
    
    %---------------------------------
    % PART 1: MODAL FUNCTIONS \phi(z)
    %---------------------------------
    
    % vertical wavenumbers in water and bottom 
    
    k = kj(jj);
    kzw = sqrt( (omeg/cw)^2 - k^2 );
    kzb = sqrt( k^2 - (omeg/cb)^2 );
    
    % computing the norm to normalize mode functions
    
    mNorm = sqrt(-0.5*( cos(kzw*h0)*sin(kzw*h0) - kzw*h0 )/(rhow*kzw) + 0.5* (sin(kzw*h0))^2 /(rhob*kzb));
    izBot = find(z<=h0,1,'last');
    phi(1:izBot) = sin( kzw*z(1:izBot) )/mNorm;
    phi(izBot+1:nz) = sin( kzw*h0 )*exp( kzb*h0 )*exp( -kzb*z(izBot+1:nz) )/mNorm;
    
    
    phih0 = sin( kzw*h0 )/mNorm;        % \psi_j(h_0)
    phih0z = kzw*cos( kzw*h0 )/mNorm;   % derivative \psi_{jz}(h_0)
    phizs = sin( kzw*zs )/mNorm;        % \psi_j(z_s)
    phizr = sin( kzw*zr )/mNorm;        % \psi_j(z_r)
    
    %---------------------------------
    % PART 2: MODE AMPLITUDES A_j(x,y)
    %---------------------------------
    
    % the values a_j and b_j in Eq. (14)
    b =  ( (phih0^2)*( gammaw*(omeg/cw)^2 - gammab*(omeg/cb)^2 ) +...
            (phih0^2)*(gammab-gammaw)*k^2   -  (rhob - rhow)*(gammaw*phih0z)^2 );
    a = 0.5*nub*(sin(kzw*h0))^2 /(rhob*kzb*mNorm^2);
      
    % the Cauchy data parameters, 
    % sigma is w from Eq. (11), 
    % A0 is \bar A from Eq. (11),   
    sigma = 1/k;
    A0 = 1i*phizs/(2*sqrt(pi)*rhow);
    
    % A(x,0) -- i.e. on the track x=0
    % computed using Eq. (16)
    
    Atrack(1:nx) = A0*exp(1i*(a*x /(2*k)  - x.^3*(b*tanA)^2 /(24*k^3) )) .*...
                                    sqrt(sigma^2./ (sigma ^ 2 + (2*1i)*x/k )).* ...
                                exp(-(x.^4*(b*tanA)^2/(16*k^4) )./ (sigma^2 + 2*1i*x/k   ));
    
    % A(x,y), via Eq. (16)
    
    Axy(1:nx,1:ny) = A0*exp(1i*(a*Xx /(2*k)  + (b*tanA)*Yy.*Xx/(2*k) - Xx.^3*(b*tanA)^2 /(24*k^3) )) .*...
                                    sqrt(sigma^2./ (sigma ^ 2 + (2*1i)*Xx/k )).* ...
                                exp(- ( (k*Yy - Xx.^2*(b*tanA)/(4*k) ).^2 )./ (1 + 2*1i*k*Xx   ));
    
    %---------------------------------                        
    % PART 3: phase exponentials \exp(i k_j x) for the sum in Eq. (12)
    %---------------------------------
    
    theta(1:nx) = exp(1i*k*x);
    
    
    % assembling the field, summing up by Eq. (12)
    
    aPtrack(1:nx) = aPtrack(1:nx) + theta.*Atrack*phizr ;
    aPxy(1:nx,1:ny) = aPxy(1:nx,1:ny) + repmat( theta.', 1,ny ).*Axy*phizr ;
    
end;

% usual formula for transmission loss

TLtrack(1:nx)  = 20*log10( abs(4*pi*aPtrack) );
TLxy(1:nx,1:ny)  = 20*log10( abs(4*pi*aPxy) );

%% Computing the TL(\theta) = TL(R_0,\theta)

% opening angle

thetaM = pi/14;
thetaN = 1000;

theta = linspace(-thetaM,thetaM,thetaN);
r0 = 8000;

xr0 = r0*cos(theta);
yr0 = r0*sin(theta);

[Xx1,Yy1] = meshgrid(x,y);


% interpolation to the arc z = z_s, r = r0 is accomplished
TLr0 = interp2(Xx1,Yy1,TLxy.',xr0,yr0);


%% figures and comparison

% default fonts, screen size, etc...

set(0, 'DefaultAxesFontSize', 18, 'DefaultAxesFontName', 'Arial');
set(0, 'DefaultTextFontSize', 18, 'DefaultTextFontName', 'Arial');
set(0,'Units','pixels'); 
scnsize = get(0,'ScreenSize');
scnsize(2) = scnsize(2) + 300;
scnsize(4) = scnsize(4) - 300;

% output of results to textfiles


x = x/1000;
y = y/1000;

dlmwrite('TLy.tl', [y.' TLxy(end,:).'] ,'delimiter','\t','precision', 6);
dlmwrite('TLx.tl', [x.' TLtrack.'] ,'delimiter','\t','precision', 6);

%% COMPARISON WITH 3D PE, (c) Frederic Sturm, Ecole centrale de Lyon
% cf F. Sturm, “On the need for a leading-order cross term correction in
% three-dimensional parabolic equation models", JASA, 2015

if alpha==0
    TLref = dlmread('2D_PE_of_F_Sturm.dat');                    % if alpha = 0 then we a have a plane-paralllel waveguide, and compare with 2D PE
else
    TLref = dlmread('3D_PE_of_F_Sturm_with_cross_term.dat');    % if alpha <> 0 then compare with the 3D PE solution
end;

% plotting figures

hF1 = figure();
set(hF1,'OuterPosition',scnsize);

plot(x, TLtrack,'linewidth',2);
hold on;
plot(TLref(:,1)/1000,TLref(:,2),'linewidth',2,'color','red');
ylim([-80 -40]);
xlim([3 8]);
xlabel('r, km');
ylabel('TL, dB');
title(sprintf('TL(x), y=0, z=%.1f m (along the track), f=%d Hz',zr,freq));
grid on;

pAx = get(gca,'Position');
pAx(1) = 0.075;
pAx(3) = 0.85;
pAx(2) = 0.15;
pAx(4) = 0.75;
set(gca,'Position',pAx);
legend('MPE','3D WAPE');



hF2 = figure();
set(hF2,'OuterPosition',scnsize);

imagesc(x,y,TLxy');
caxis([-80 -40]);
colorbar;
pAx = get(gca,'Position');
pAx(1) = 0.075;
pAx(3) = 0.85;
pAx(2) = 0.15;
pAx(4) = 0.75;
set(gca,'Position',pAx);
xlim([x(1) x(end)]);
ylim([y(1) y(end)]);
xlabel('x, km');
ylabel('y, km');
title(sprintf('TL(x,y), z=%.1f m, f=%d Hz',zr,freq));


hF3 = figure();

plot(y,TLxy(end,:),'linewidth',2,'color','red');
pAx = get(gca,'Position');
pAx(1) = 0.075;
pAx(3) = 0.85;
pAx(2) = 0.15;
pAx(4) = 0.75;
set(gca,'Position',pAx);
grid on;
xlim([y(1) y(end)]);
ylim([-80 -55] );
xlabel('y, km');
ylabel('TL, dB');
title(sprintf('TL(y), z=%.1f m (across the track), x=%.1f km, f=%d Hz',zr,x(end),freq));
set(hF3,'OuterPosition',scnsize);


hF4 = figure();
thetaDeg = theta*180/pi;
plot(thetaDeg,TLr0,'linewidth',2,'color','red');
pAx = get(gca,'Position');
pAx(1) = 0.075;
pAx(3) = 0.85;
pAx(2) = 0.15;
pAx(4) = 0.75;
set(gca,'Position',pAx);
grid on;
xlim([thetaDeg(1) thetaDeg(end)]);
ylim([-80 -50] );
xlabel('\theta, deg');
ylabel('TL, dB');
title(sprintf('TL(\\theta), z=%.1f m, r=%.1f km, f=%d Hz',zr,r0/1000,freq));
set(hF4,'OuterPosition',scnsize);

dlmwrite('TLtheta.tl', [thetaDeg.' TLr0.'] ,'delimiter','\t','precision', 6);

