function [sig,xi,w,b]=baro(z,Uz,tz,f,k,l,imode)
%
% USAGE: [sig,xi,w,b]=baro(z,Uz,tz,f,k,l,imode)
% Stability analysis for a flow in thermal wind balance
%
% INPUTS:
% z = vertical coordinate vector (evenly spaced)
% Uz = shear (constant); code assumes that U=Uz*z;
% tz = vertical buoyancy gradient, can be a single value or a profile
% k,l = wave vector
% imode = mode choice (default imode=1)
%         imode=0: output all modes, sorted by growth rate
%
% OUTPUTS:
% sig = growth rate of FGM
% eigenfunctions (optional):
% xi = vertical vorticity
% w = vertical velocity 
% b = buoyancy
%
% CALLS:
% ddz,ddz2
%
% CITATION:
%   Instability in Geophysical Flows
%   W.D. Smyth & J.R. Carpenter
%   Cambridge University Press, 2019
%
% W. Smyth, March 2014, updated January 2019

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Stage 1: Preliminaries
%

% check for equal spacing
if abs(std(diff(z))/mean(diff(z)))>.000001
    disp(['baro: values not evenly spaced!'])
    sig=NaN;
    return
end

% defaults
if ~exist('imode');imode=1;end

% define constants
ii=complex(0.,1.);
del=mean(diff(z));N=length(z);
kt=sqrt(k^2+l^2);

U=Uz*z; % compute U from uniform shear;
By=-f*Uz; % compute By from thermal wind balance

% expand tz into a vector if necessary
if length(tz)==1;
    tz=tz*ones(size(z));
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Stage 2: Derivative matrices and BCs

% D2: 2nd derivative for w (impermeable boundary)
D2=ddz2(z);
D2(1,:)=0;D2(1,1:2)=[-2 1]/del^2;
D2(N,:)=0;D2(N,N-1:N)=[1 -2]/del^2;

% D1w: 1st derivative for w
D1w=ddz(z);
D1w(1,:)=0; D1w(1,2)=1/(2*del); % impermeable
D1w(N,:)=0; D1w(N,N-1)=-1/(2*del);

% D1z: 1st derivative for xi
D1z=ddz(z);
D1z(1,:)=0;
D1z(N,:)=0;
D1z(1,[1 2])=[-1 1]*2/(3*del);
D1z(N,[N-1 N]) =[-1 1]*2/(3*del);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Stage 3: Assemble stability matrices A and B
%
Id=eye(N);
L=D2-kt^2*Id;

A=[Id Id*0 Id*0 ; Id*0 L Id*0; Id*0 Id*0 Id];

b11=-ii*k*diag(U);
b12=ii*l*Uz*Id+f*D1w;
b13=Id*0;

b21=-f*D1z;
b22=-ii*k*diag(U)*L;
b23=-Id;

b31=ii*k*By*Id;
b32=-ii*l*By*D1w-kt^2*diag(tz);
b33=-ii*k*diag(U);

B=[b11 b12 b13; b21 b22 b23; b31 b32 b33];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Stage 4: Solve eigenvalue problem and extract results
%
% Solve generalized eigenvalue problem

if nargout==1 % compute eigvals only (faster)
    sigma=eig(B,A);
    [sr,ind]=sort(real(sigma),1,'descend');
    sigma=sigma(ind);
    if imode==0
        sig=sigma;
    elseif imode>0
        sig=sigma(imode);
    end
elseif nargout>1 % compute eigvals and eigvecs
    [v,e]=eig(B,A);sigma=diag(e);
    [sr,ind]=sort(real(sigma),1,'descend');
    sigma=sigma(ind);
    v=v(:,ind);
    if imode==0
        sig=sigma;
        xi=v(1:N,:);
        w=v(N+1:2*N,:);
        b=v(2*N+1:3*N,:)/kt^2;
    elseif imode>0
        sig=sigma(imode);
        xi=v(1:N,imode);
        w=v(N+1:2*N,imode);
        b=v(2*N+1:3*N,imode)/kt^2;
    end
end
return