function [lambda,u,p,alpha]=CTW_MODES(omega,f,dy,dz,H,N2,Nm)
% Function to determine coastal trapped wave modes. The waves propagate in
% the positive x direction, so u(y,z,t) is along-shore velocity
%
% Notes: 
% (1) walls are placed at y=0 and y=y_max
% (2) u and p are staggered, with u at the boundaries. The orthogonality
% condition is to average u to the center of each cell before computing
% flux
% (3) There is a linear free surface (this depends on one line of code for 
% the z derivative)
%
% INPUTS:
% omega [1 x 1 ]  wave frequency (should be less than f)
% f     [1 x 1]   inertial frequency (assumed positive)  
% dy    [1 x 1]   spacing for offshore grid 
% dz    [1 x 1]   spacing for the depth grid (positive)
% H     [Ny x 1]  topographic depth (positive)
% N2    [Nz x 1]  stratification (should be greater than f and omega)
% Nm    [1 x 1]   number of modes to find (includes BT mode and offshore
%                 wall modes). Must be even.
%
% OUTPUTS:
% lambda    [Nm+1 x 1]  modal wavelengths  
% u         [Ny x Nz x Nm]  along shore velocity 
% p         [Ny x Nz x Nm]  reduced pressure 
%
% Sam Kelly, 28 JUL 2021 (smkelly@d.umn.edu)

Ny=length(H);
Nz=length(N2);

% Create depth grid
z=(1:Nz)'*dz-dz/2';

% Velocity depth is maximum of adjacent cell depths
Hu=max([H(1:end-1) H(2:end)],[],2);
Hu=[H(1); Hu; H(end)];

% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Build K matrix
% The equations are written starting in the upper left corner, then moving right. When the row is done, the location moves down. 
N=Ny*Nz;
Nu=(Ny+1)*Nz;

% Initiate matricies
K11=(omega^2/f^2-1)*sparse(diag(ones([Nu 1])));
K12=sparse(Nu,N);
K21=sparse(N,Nu);
K22=sparse(N,N);

M11=sparse(Nu,Nu);
M12=sparse(Nu,N);
M21=sparse(N,Nu);
M22=sparse(N,N);

% negative y-derivative of pressure (only defined where u is unknown)
i=1;
j=1;
indp=1;
for ind=1:Nu
    
    % Couple pressure if it's above ground
    if i==1
        if z(j)<Hu(i) 
            M12(ind,indp)=1/2;
        end
    elseif i==Ny+1
        if z(j)<Hu(i)
            M12(ind,indp-1)=1/2;
        end
    else
        if z(j)<H(i-1)
            M12(ind,indp-1)=1/2;
        end
        if z(j)<H(i)
            M12(ind,indp)=1/2;
        end
    end
    
    % Y derivative of pressure
    if z(j)<Hu(i) % check if u is above the bottom
        if i==1 || i==Ny+1 || z(j)>H(i-1) || z(j)>H(i) % check if there is a wall, then -dp/dy=fu
            K11(ind,ind)=(K11(ind,ind)+1)/2; % Division by two comes from 2-point average in the M matrix
        else % open water
            K12(ind,indp-1)=K12(ind,indp-1)+1/dy;
            K12(ind,indp)=K12(ind,indp)-1/dy;
        end
    end
   
    % update indicies
    i=i+1;
    indp=indp+1;
    if i>Ny+1
        i=1; % Return to coast
        j=j+1; % Move down
        indp=indp-1; % go back one since p only has Ny nodes at each depth
    end 
end
    
% y-derivative of u velocity 
i=1;
j=1;
indu=1;
for ind=1:N
    
    % Average velocity to p nodes
    if z(j)<H(i) % check if p is above the bottom
        M21(ind,indu)=1/2;
        M21(ind,indu+1)=1/2;
    end
    
    if z(j)<H(i) % check if p is above the bottom
        if i~=Ny 
            if z(j)<H(i+1) % check if front is open
                K21(ind,indu+1)=K21(ind,indu+1)+1/dy;            
            end
        end
        if i~=1
            if z(j)<H(i-1) % check if back is open
                K21(ind,indu)=K21(ind,indu)-1/dy;
            end
        end            
    end 

    % update indicies
    i=i+1;
    indu=indu+1;
    if i>Ny
        i=1; % Return to coast
        j=j+1; % Move down
        indu=indu+1; % go forward since u has Ny+1 nodes at each depth
    end 
end

% Pressure y and z derivatives
i=1;
j=1;
for ind=1:N  
   
    % Pressure second y-derivative
    if z(j)<H(i) 
        if (i==1 && z(j)>H(i+1)) || (i==Ny && z(j)>H(i-1)) || ((i~=1 && i~=Ny) && (z(j)>H(i-1) && z(j)>H(i+1)))
            % You're in a trench! do nothing.
        elseif i==1 || z(j)>H(i-1) % If there is a wall, there are some cancellations with the 2nd derivative of pressure                      
            K22(ind,ind)=K22(ind,ind)-1/(dy^2);
            K22(ind,ind+1)=K22(ind,ind+1)+1/(dy^2);            
        elseif i==Ny || z(j)>H(i+1) % Outer boundary is a wall   
            K22(ind,ind-1)=K22(ind,ind-1)+1/(dy^2);
            K22(ind,ind)=K22(ind,ind)-1/(dy^2);        
        else % Open water
            K22(ind,ind-1)=K22(ind,ind-1)+1/(dy^2);
            K22(ind,ind)=K22(ind,ind)-2/(dy^2);
            K22(ind,ind+1)=K22(ind,ind+1)+1/(dy^2);
        end
    end 
    
    % Pressure z derivatives
    if z(j)<H(i)
        if j==1 && z(j+1)>H(i) % Point is at the top and bottom
            % Do nothing!
        elseif j==1 % Point is at the top
            K22(ind,ind)=   K22(ind,ind)    +1/(dz^2*(N2(j)+N2(j+1))/2/omega^2) + 1/(dz*9.81/omega^2);
            K22(ind,ind+Ny)=K22(ind,ind+Ny) -1/(dz^2*(N2(j)+N2(j+1))/2/omega^2);
        elseif j==Nz || z(j+1)>H(i) % Point is at the bottom
            K22(ind,ind-Ny)=K22(ind,ind-Ny) -1/(dz^2*(N2(j-1)+N2(j))/2/omega^2);
            K22(ind,ind)=   K22(ind,ind)    +1/(dz^2*(N2(j-1)+N2(j))/2/omega^2);
        else % Point is in the interior water column
            K22(ind,ind-Ny)=K22(ind,ind-Ny) -1/(dz^2*(N2(j-1)+N2(j))/2/omega^2);
            K22(ind,ind)=   K22(ind,ind)    +1/(dz^2*(N2(j-1)+N2(j))/2/omega^2) + 1/(dz^2*(N2(j)+N2(j+1))/2/omega^2);
            K22(ind,ind+Ny)=K22(ind,ind+Ny) -1/(dz^2*(N2(j)+N2(j+1))/2/omega^2);
        end
    end
           
    % update indicies
    i=i+1;
    if i>Ny
        i=1;
        j=j+1;
    end 
end

K=[K11 K12;...
   K21 K22];

K=f/omega^2*K;

M=[M11 M12;...
   M21 M22];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Solve the eigenvalue problem

% Shrink the problem to above ground nodes
good=(sum(M)~=0);
M2=M(good,good);
K2=K(good,good);

% Guess Kelvin wave
c_theory=sqrt(9.81*max(H));

% % Solve eigenvalue problem for positive speeds
opts.maxit=1000;
[PHI c]=eigs(M2,K2,Nm/2,c_theory,opts);
if sum(imag(diag(c))~=0)>0 %Try doubling Nm if there are complex speeds
    [PHI c]=eigs(M2,K2,2*(Nm-1),c_theory,opts);
    bad=imag(diag(c))~=0;
    c(bad,bad)=0;
end
lambda=2*pi*diag(c)/omega*1e-3;
[~,I]=sort(real(lambda),'descend');
lambda=lambda(I);
PHI=PHI(:,I);
lambda1=lambda(1:Nm/2);
PHI1=PHI(:,1:Nm/2);

% Solve eigenvalue problem for negative speeds
opts.maxit=1000;
[PHI c]=eigs(M2,K2,Nm/2,-c_theory,opts);
if sum(imag(diag(c))~=0)>0 %Try doubling Nm if there are complex speeds
    [PHI c]=eigs(M2,K2,2,-c_theory,opts);
    bad=imag(diag(c))~=0;
    c(bad,bad)=0;
end
lambda=2*pi*diag(c)/omega*1e-3;
[~,I]=sort(real(lambda),'descend');
lambda=lambda(I);
PHI=PHI(:,I);
lambda2=lambda(1:Nm/2);
PHI2=PHI(:,1:Nm/2);

% Combine
PHI=zeros([N+Nu Nm]);
PHI(good,:)=cat(2,PHI1,PHI2);
lambda=cat(1,lambda1,lambda2);

% Normalize
alpha=ones([Nm 1]);
for i=1:Nm  
    F=PHI(:,i)'*M*PHI(:,i)*f*dy*dz;    
    PHI(:,i)=PHI(:,i)/sqrt(F);
    
    if F<0
       alpha(i)=-1;
    end
    
    ind=find(abs(PHI(:,i))>1e-10,1,'first');
    if PHI(ind,i)<0
        PHI(:,i)=-PHI(:,i);
    end
end

% Reshape
mask=z'<H;

u=PHI(1:Nu,:); 
u=reshape(u,[Ny+1 Nz Nm]);
u=(u(1:end-1,:,:)+u(2:end,:,:))/2;

p=PHI(Nu+1:end,:)*f; 
p=reshape(p,[Ny Nz Nm]);

% Nan below bottom 
u(repmat(~mask,[1 1 Nm]))=0+0i; 
p(repmat(~mask,[1 1 Nm]))=NaN+NaN*1i; 

return
