function [u p h A B c k uA uB pA pB]=CELT_Q(H,X,A0,B0,N2,omega,f,nu,x,z,Q)
% USAGE:  [u p h A B c k]=CELT(H,X,A0,B0,N2,omega,f,nu,x,z)
% Solve the Coupling Equation for Linear Tides (CELT)
%
% The solution assumes linear, Boussinesq, nonhydrostatic, f-plane mechanics.
% The barotropic mode propagates as a shallow water wave.
%
% INPUTS:
% H    [Nx+1 x 1] Height of flats (positive)
% X    [Nx  x 1]  Location of steps
% A0   [Nm  x 1]  Intenral-tide forcing from left
% B0   [Nm  x 1]  Internal-tide forcing from right
% N2   [Nx+1 x Nz] Stratification (from shallow to deep)
% omega [1  x 1]  Frequency of waves
% f    [1   x 1]  Intertial frequency
% nu   [1   x 1]  Vertical viscosity
% x    [nx  x 1]  Horizontal coordinates of output
% z    [nz  x 1]  Vertical coordinates of output (positive)
% Q    [1  x 1]   Barotropic volume flux (conserved over each step)
%
% OUTPUTS:
% u    [nx x nz]  Complex amplitude of internal-tide velocity
% p    [nx x nz]  Complex amplitude of internal-tide pressure
% h    [nx x  1]  Topography mapped to output coordinates
% A    [Nx x Nm]  Amplitudes of right-going waves
% B    [Nx x Nm]  Amplitudes of left-going waves
% c    [Nx x Nm]  group speed
% k    [Nx x Nm]  wavenumber
% uA   [nx x nz]  Complex amplitude of right-going internal-tide velocity
% uB   [nx x nz]  Complex amplitude of left-going internal-tide velocity
% pA   [nx x nz]  Complex amplitude of right-going internal-tide pressure
% pB   [nx x nz]  Complex amplitude of left-going internal-tide pressure
%
%
% This model is described in detail by:
%
% S. M. Kelly, N. L. Jones and J. D. Nash (2013), A coupled model for 
% Laplace's tidal equations in a fluid with one horizontal dimension and 
% variable depth. J. Phys. Oceangr., 43, 1780-1797.
%
% Common problems:
%
% (1) With a fine vertical grid and a lot of steps the program spends a lot
% of time computeing mode shapes. Try using a coaser vertical grid, but see
% (3).
%
% (2) A typical laptop will start to choke on inversions much bigger than
% a 10,000 x 10,000 matrix.
%
% (3) The vertical grid is uniform and the same over all steps. Threfore
% you must choose a fine-enough grid so that all of the modes can be
% resolved on the shallowest step. E.g. You wont be able to match 5 modes 
% if the shallowest step is H=100 m and the grid spacing is dz=25 m. A step
% of H=0 will produce a reflection boundary condition.
%
%
% Sam Kelly, 21 AUG 2015 (smkelly@d.umn.edu)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
Nm=length(A0);
Nz=size(N2,2);
Nx=length(X);
nx=length(x);
nz=length(z);

ii=complex(0,1);
g=9.81;
Z=linspace(-max(H),0,Nz+1)';
Z=(Z(2:end)+Z(1:end-1))/2;
dz=nanmean(diff(Z));

% Check if there's enough vertical resolution 
Nm0=Nz-dsearchn(Z,-min(H(H~=0)));
if Nm>Nm0-1
    disp(['ERROR: Not enough vertical resolution to match ',num2str(Nm),' modes'])
    u=NaN; p=[]; h=[]; A=[]; B=[]; c=[]; k=[]; uA=[]; uB=[]; pA=[]; pB=[];    
    return
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Generate sructure functions
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Computing vertical modes');

N2=N2.';
N2=flipud(N2);
for i=1:Nx+1;
    if H(i)~=0
        % Get bottom index
        ind.b(i)=dsearchn(Z,-H(i));
        
        % Calculate modes: solve eigenvalue problem with depth-varying stratification
        [phi{i} c{i}]=MODES(dz,N2(ind.b(i):end,i),omega);
        
        % Keep necessary modes
        phi{i}=phi{i}(:,2:Nm+1)';
        c{i}=c{i}(2:Nm+1);
        
        % Pad with zeros below bottom
        phi{i}=[zeros([Nm ind.b(i)-1]) phi{i}];
                
        % Convert from egienspeed to group speed and wavenumber
        c{i}=sqrt(1-f^2/omega^2)*c{i};
        k{i}=(1-f^2/omega^2)*omega./c{i};
    else
        phi{i}=zeros([Nm Nz]);
        c{i}=zeros([Nm 1]);
        k{i}=zeros([Nm 1]);
    end
    PROGRESS_BAR(i,1:Nx+1);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Set up systems of equations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Calculating matrix coefficients');
G=zeros([Nm*Nx*2 Nm*(Nx+1)*2]);
F=zeros([Nm*Nx*2 1]);

% Cycle through steps 
for i=1:Nx  
    N20=nanmean(N2(:,i));

    % Viscous attenuation of incoming wave from left
    if i==1; rL=ones([Nm 1]); else
        dx=X(i)-X(i-1);
        rL=exp(-(dx./c{i}).*nu.*(1-f^2/omega^2).*(N20-omega^2)./c{i}.^2);
    end
    
    % Viscous attenuation of incoming wave from right
    if i==Nx; rR=ones([Nm 1]); else
        dx=X(i+1)-X(i);
        rR=exp(-(dx./c{i+1}).*nu.*(1-f^2/omega^2).*(N20-omega^2)./c{i+1}.^2);
    end    
    
    % Velocity and Pressure coefficients
    if H(i)>H(i+1) % PART I: Getting shallower        
        % left velocity constraints
        u1=diag(rL.*exp(ii*X(i)*k{i})); % left incoming wave (A)
        u2=diag(exp(-ii*X(i)*k{i})); % left outgoing wave (B)      
        
        % right velocity constraints
        if H(i+1)==0 % If vertical wall                   
            u3=zeros(Nm); % right outgoing wave (A)            
            u4=zeros(Nm); % right incoming wave (B)
        else 
            for n=1:Nm
                for m=1:Nm
                    u3(n,m)=nansum(phi{i+1}(m,:).*phi{i}(n,:)/H(i)*dz,2).*exp(ii*X(i)*k{i+1}(m));  % right outgoing wave (A)  
                    u4(n,m)=nansum(phi{i+1}(m,:).*phi{i}(n,:)/H(i)*dz,2).*rR(m).*exp(-ii*X(i)*k{i+1}(m)); % right incoming wave (B)
                end
            end
        end
        
        % left pressure constraints
        if H(i+1)==0 % If vertical wall (treat as if flat)
            p1=diag(rL.*c{i}.*exp(ii*X(i)*k{i})); % left incoming wave (A)
            p2=diag(c{i}.*exp(-ii*X(i)*k{i})); % left outgoing wave (B)
        else
            for n=1:Nm
                for m=1:Nm
                    p1(n,m)=nansum(c{i}(m)*phi{i}(m,:).*phi{i+1}(n,:)/H(i+1)*dz,2).*rL(m).*exp(ii*X(i)*k{i}(m)); % left incoming wave (A)
                    p2(n,m)=nansum(c{i}(m)*phi{i}(m,:).*phi{i+1}(n,:)/H(i+1)*dz,2).*exp(-ii*X(i)*k{i}(m)); % left outgoing wave (B)
                end
            end
        end
        
        % right pressure constraints
        if H(i+1)==0 % If vertical wall (treat as if a mirror)
            p3=-diag(c{i}.*exp(ii*X(i)*k{i})); % right outgoing wave (A)
            p4=-diag(rL.*c{i}.*exp(-ii*X(i)*k{i})); % right incoming wave (B)
        else
            p3=diag(c{i+1}.*exp(ii*X(i)*k{i+1})); % right outgoing wave (A)
            p4=diag(rR.*c{i+1}.*exp(-ii*X(i)*k{i+1})); % right incoming wave (B)
        end       
        
    else % PART II: Getting deeper        
        % left velocity constraints
        if H(i)==0 % If vertical wall 
            u1=zeros(Nm); % left outgoing wave (A)
            u2=zeros(Nm); % left incoming wave (B)
        else
            for n=1:Nm
                for m=1:Nm
                    u1(n,m)=nansum(phi{i}(m,:).*phi{i+1}(n,:)/H(i+1)*dz,2).*rL(m).*exp(ii*X(i)*k{i}(m)); % left incoming wave (A)
                    u2(n,m)=nansum(phi{i}(m,:).*phi{i+1}(n,:)/H(i+1)*dz,2).*exp(-ii*X(i)*k{i}(m)); % left outgoing wave (B)
                end
            end
        end
        
        % right velocity constraints
        u3=diag(exp(ii*X(i)*k{i+1})); % right outgoing wave (A)
        u4=diag(rR.*exp(-ii*X(i)*k{i+1})); % right incoming wave (B)
        
        % left pressure contraints
        if H(i)==0 % If vertical wall (treat as if a mirror)
            p1=-diag(rR.*c{i+1}.*exp(ii*X(i)*k{i+1})); % left incoming wave (A)
            p2=-diag(c{i+1}.*exp(-ii*X(i)*k{i+1})); % left outgoing wave (B)
        else            
            p1=diag(rL.*c{i}.*exp(ii*X(i)*k{i}));% left incoming wave (A)
            p2=diag(c{i}.*exp(-ii*X(i)*k{i}));% left outgoing wave (B)
        end
                
        % right pressure contraints
        if H(i)==0 % If vertical wall (treat as if flat)
            p3=diag(c{i+1}.*exp(ii*X(i)*k{i+1})); % right outgoing wave (A)
            p4=diag(rR.*c{i+1}.*exp(-ii*X(i)*k{i+1})); % right incoming wave (B)
        else
            for n=1:Nm
                for m=1:Nm
                    p3(n,m)=nansum(c{i+1}(m)*phi{i+1}(m,:).*phi{i}(n,:)/H(i)*dz,2).*exp(ii*X(i)*k{i+1}(m)); % right outgoing wave (B)
                    p4(n,m)=nansum(c{i+1}(m)*phi{i+1}(m,:).*phi{i}(n,:)/H(i)*dz,2).*rR(m).*exp(-ii*X(i)*k{i+1}(m)); % right incoming wave (B)
                end
            end
        end
    end
    
    % Fill in the matrix of coefficient matrices
    ind.row=(i-1)*2*Nm+1:i*2*Nm;
    ind.col=(i-1)*2*Nm+[1:4*Nm];
    G(ind.row,ind.col)=[[u1  u2 -u3 -u4];...
                        [p1 -p2 -p3  p4]];
    
    % Fill in forcing vector
    if i==1
        F(1:2*Nm)=F(1:2*Nm)+[[-u1*A0];...
                   [-p1*A0]];
    end
    if i==Nx
        F(end+1-2*Nm:end)= F(end+1-2*Nm:end)+[[ u4*B0];...
                           [-p4*B0]];
    end
    
    % Barotropic Forcing
    for n=1:Nm
        if H(i)==0;
            F((i-1)*2*Nm+n)=F((i-1)*2*Nm+n);
        elseif H(i)>H(i+1)
            F((i-1)*2*Nm+n)=F((i-1)*2*Nm+n)+nansum(Q/H(i)*phi{i}(n,ind.b(i+1):end)*1/H(i+1)*dz,2); 
        else
            F((i-1)*2*Nm+n)=F((i-1)*2*Nm+n)-nansum(Q/H(i+1)*phi{i+1}(n,ind.b(i):end)*1/H(i)*dz,2); 
        end
    end
    
    % Track progress
    PROGRESS_BAR(i,1:Nx);
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Solve system
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
G=G(:,Nm+1:end-Nm); % Trim coefficient matrix

disp(['Inverting ',num2str(size(G,1)),' x ',num2str(size(G,2)),' matrix']);pause(.01);
%tic
AB=sparse(G)\F; % Solve matrix problem
%AB=G\F; % Solve matrix problem
%toc
AB=[A0; AB; B0]; % Add forcing to output

% Parse solution
AB=reshape(AB.',[2*Nm Nx+1]);
A=AB(1:Nm,:);
B=AB(Nm+1:end,:);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Construct fields
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
disp('Constructing the fields');pause(.01);
h=zeros([nx  1]);
u=zeros([nx nz]);
uA=zeros([nx nz]);
uB=zeros([nx nz]);
p=zeros([nx nz]);
pA=zeros([nx nz]);
pB=zeros([nx nz]);

% New structure functions for output
for i=1:Nx+1;
    phi0{i}=zeros([Nm nz]);
    for n=1:Nm;
        phi0{i}(n,:)=interp1(Z,phi{i}(n,:),-z);
    end
end

% Assemble fields
for i=1:Nx+1
    N20=nanmean(N2(:,i));
    
    % Find locations between steps
    if i==1 
        ind.x=x<X(1);         
    elseif i==Nx+1
        ind.x=X(Nx)<=x; 
    else
        ind.x=X(i-1)<=x & x<X(i);
    end
    h(ind.x)=H(i);
   
    for n=1:Nm
        % Calculate viscous attenuation
        if i==1
            Aeps=1;            
            Beps=exp(-(X(i)-x(ind.x))./c{i}(n)*nu.*(1-f^2/omega^2).*(N20-omega^2)./c{i}(n).^2);
        elseif i==Nx+1
            Aeps=exp(-(x(ind.x)-X(i-1))./c{i}(n)*nu.*(1-f^2/omega^2).*(N20-omega^2)./c{i}(n).^2);
            Beps=1;
        else
            Aeps=exp(-(x(ind.x)-X(i-1))./c{i}(n)*(n*pi/H(i))^2*nu);
            Beps=exp(-(X(i)-x(ind.x))./c{i}(n)*(n*pi/H(i))^2*nu);
        end
        
        % Assemble fields        
        uA(ind.x,:)=uA(ind.x,:)+A(n,i)*Aeps.*exp( ii*x(ind.x)*k{i}(n))*phi0{i}(n,:);
        uB(ind.x,:)=uB(ind.x,:)+B(n,i)*Beps.*exp(-ii*x(ind.x)*k{i}(n))*phi0{i}(n,:);
        u(ind.x,:)=uA(ind.x,:)+uB(ind.x,:);
        
        pA(ind.x,:)=pA(ind.x,:)+c{i}(n)*A(n,i)*Aeps.*exp( ii*x(ind.x)*k{i}(n))*phi0{i}(n,:);
        pB(ind.x,:)=pB(ind.x,:)-c{i}(n)*B(n,i)*Beps.*exp(-ii*x(ind.x)*k{i}(n))*phi0{i}(n,:);
        p(ind.x,:)=pA(ind.x,:)+pB(ind.x,:);        
    end
end
k=cell2mat(k);
c=cell2mat(c);

return


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% [phi C]=MODES(dz,N2)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [phi C]=MODES(dz,N2,omega)
% USAGE: [phi C]=MODES(dz,N2,omega)
% Obtain vertical modes for arbitrary stratification
%
% INPUTS:
% dz  [1 x 1]   vertical spacing (positive)
% N2  [Nz x 1]  stratification 
% omega [Nz x 1] wave frequency
%
% OUTPUTS:
% phi [Nz x Nz]  orthonormal pressure and velocity structure eigenfunctions 
% C   [Nz x  1]  eigenspeeds (i.e., group/phase speed if f=0)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
N=length(N2);
H=dz*N;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Second derivative matrix (2nd order accurate)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
D2=zeros(N-1,N-1);
for i=2:N-2
    D2(i,i-1)=1/dz^2;
    D2(i,i)=-2/dz^2;
    D2(i,i+1)=1/dz^2;
end

% Upper boundary condition: rigid lid
D2(1,1)=-2/dz^2; 
D2(1,2)=1/dz^2;

% Bottom boundary condition: flat bottom
D2(N-1,N-1)=-2/dz^2;
D2(N-1,N-2)=1/dz^2;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% System of equations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% "A" Matrix
A=diag(-(N2(1:N-1)-omega^2));

% Solve generalized eigenvalue problem
[phi k2]=eig(D2,A);

% Sort modes by eigenspeed
k2=diag(k2);
k2(k2<0)=Inf;
C=1./sqrt(k2);
C(C>1000)=0; % Remove modes with c faster 1000 m/s
[C,ind]=sort(C,1,'descend');
phi=phi(:,ind);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Derive U and P structures
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Add in boundary values
phi=[zeros([1 N-1]); phi; zeros([1 N-1])]; 

% Take derivative to get U and P structure
phi=-diff(phi,1)./dz;

% Add in surface mode
phi=[ones([N 1]) phi];
C=[1; C]; % Leave mode-0 eigenspeed for later

% Normalize 
A=repmat(nansum(phi.^2.*dz,1)./H,[N 1]).^(1/2);
A(A==0)=Inf;
phi=phi./A;
phi(:,phi(N,:)<0)=-phi(:,phi(N,:)<0);

return


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% PROGRESS_BAR(ind,all_inds)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function PROGRESS_BAR(ind,all_inds)
% vital function to show progress, needs update to hour glass or spinning wheel
global counters

% Initiate
if ind==all_inds(1)    
    disp(['0%       10%       20%       30%       40%       50%       60%       70%       80%       90%      100%'])
    counters.percent=0; 
    counters.N=1;
end

% Print progress
temp=round(counters.N/length(all_inds)*100);
while counters.percent<=temp
    fprintf('|');
    counters.percent=counters.percent+1;    
end
counters.N=counters.N+1;

% Closeout
if ind==all_inds(end)
    clear global counters
    fprintf('\n'); 
end

return

