function SWE_solver_IY(wave_type,ModelPrms,IntegPrms,HandAxes)

%Solve the shallow water equations
% u_t+u*u_x+v*u_y-f*v=-g*(h+b)_x
% v_t+u*v_x+v*v_y+f*u=-g*(h+b)_y
% h_t+(u*h)_x+(h*v)_y=0

%%
FlagStop = 0;
uicontrol('style','pushbutton','units','norm','position',[0.77 0.07 0.06 0.05],'string','Stop','callback',{@stopb});
uicontrol('style','pushbutton','units','norm','position',[0.77 0.14 0.06 0.05],'string','Pause','callback',{@pauseb});
delete(findall(gcf,'Tag','time_an')); % remove time annotation
switch wave_type
    case 'Ke'
        wave_type='kelvin';
    case 'Po'
        wave_type='poincare';
    case 'Ro'
        wave_type='rossby';
    case 'Ga'
        wave_type='gaussian';
    case 'Ha'
        wave_type='hat';
end

%% Physical constants

Omega = 2 * pi / 86400;
R_earth = 6371000;
g = 0.01;

%% Solver parameters

lat = ModelPrms.latitude;             % Mid-latitude (deg)
Lx = ModelPrms.Lx*1000;               % Channel length (m)
Ly = ModelPrms.Ly*1000;               % Channel width (m)
H0 = ModelPrms.H0;                          % Channel depth (m)
A = 1;                             % Wave amplitude
nl = 1;                            % Linear or non-linear
ah = 0.e-3;                         % Diffusion constant
bf = 0.e-10;                        % Bottom friction
dt = ModelPrms.time_step;           % Time step

T = ModelPrms.Tmax;

N = T*3600/dt;                        % Number of iterations

plot_rate = 2880/12;                 % sampling every plot_rate iterations

dx = Lx/50;                          % Spatial resolution x-axis (m)
dy = Ly/50;                          % Spatial resoluti

% Grid properties
% Zonal and meridional vectors
x = -Lx / 2 : dx : (Lx / 2 + dx);
y = 0 : dy : (Ly + dy);
% Grid size
nx = length(x);
ny = length(y);
% Relevant matrix indices.
i = 2 : (nx - 1);
im = i - 1;
ip = i + 1;
im(1) = nx - 1; ip(end) = 2; % Periodical boundaries in zonal direction
j = 2 : (ny - 1);
jm = j - 1;
jp = j + 1;

% Coriolis parameters
f0 = 2 * Omega * sind(lat);

if ModelPrms.IsBetaPlane
    beta = 2 * Omega * cosd(lat) / R_earth;
else
    beta=0;
end

fU = repmat(f0 + beta * (y - 0.5 * dy), nx, 1);
fV = repmat(f0 + beta * y, nx, 1);

% Rossby radius of deformation
R = abs(sqrt(g * H0) / f0);

% Other variables
d2x = 2 * dx;
d4x = 4 * dx;
dx2 = dx^2;
d2y = 2 * dy;
d4y = 4 * dy;
dy2 = dy^2;
tm = 1; tn = 2; tp = 3; % Time indices
delt = dt;
U = zeros(nx, ny, 3);
V = zeros(nx, ny, 3);
H = zeros(nx, ny, 3);
grav = zeros(nx, ny);
vel = zeros(nx, ny);
adv = zeros(nx, ny);
xadv1 = zeros(nx, ny);
xadv2 = zeros(nx, ny);
xadv = zeros(nx, ny);
diff_x = zeros(nx, ny);
diff_y = zeros(nx, ny);
DH1 = zeros(nx, ny);
DH2 = zeros(nx, ny);


%% Initial conditions 

n_x = ModelPrms.n_x;
n_y = ModelPrms.n_y;

kx = n_x * pi / Lx;
ky = n_y * pi / Ly;

switch wave_type
    
    case 'kelvin'
        n_x = 4;
        n_y = 2;

        kx = n_x * pi / Lx;
        ky = n_y * pi / Ly;

        h0_x = cos(kx * x).*(heaviside(x+3*Lx/(2*n_x))-heaviside(x-3*Lx/(2*n_x)));
        h0_y = exp(f0*(y-Ly)/sqrt(g*H0))+exp(-f0*y/sqrt(g*H0));
        u0_x = cos(kx * x).*(heaviside(x+3*Lx/(2*n_x))-heaviside(x-3*Lx/(2*n_x)));
        u0_y = -g/sqrt(g*H0)*(exp(f0*(y-Ly)/sqrt(g*H0))-exp(-f0*y/sqrt(g*H0)));
        v0_x = 0;
        v0_y = 0;
    
        k_x = [0:kx/100:kx*2];
        w = sqrt(g*H0).*k_x;
    
    case 'poincare'
        w_n = sqrt(f0^2+g*H0*(kx^2+ky^2));
        h0_x = cos(kx * x);
        h0_y = cos(ky * y)-kx/ky*f0/w_n*sin(ky * y);
        u0_x = cos(kx * x);
        u0_y = g*kx/w_n*cos(ky * y )-f0/(H0*ky)*sin(ky * y);
        v0_x = sin(kx * x);
        v0_y = -(f0^2+g*H0*ky^2)/(H0*w_n*ky)*sin(ky * y);
    
        k_x = [kx/2:kx/100:kx*2];
        w = sqrt(f0^2+g*H0*(k_x.^2+ky^2));
        
    case 'rossby'
        h0_x = cos(kx * x);
        h0_y = sin(ky * y);
        u0_x = cos(kx * x);
        u0_y = -g * ky * cos(ky * y) / f0;
        v0_x = sin(kx * x);
        v0_y = -g * kx * sin(ky * y ) / f0;
    
        k_x = [0:kx/100:kx*2];
        w = -k_x * beta ./ (k_x.^2 + ky^2+1/R^2);
        
    case 'gaussian'
        h0_x = exp(-(y-y(ny/2)).^2/256/R^2);
        h0_y = exp(-(x-x(nx/2)).^2/256/R^2);
        u0_x = h0_x;
        u0_y = -g/f0* (-2*(y-y(ny/2))/256/R^2 .*h0_y);
        v0_x = g/f0* (-2*(x-x(ny/2))/256/R^2 .*h0_x);
        v0_y = h0_y;
        
    case 'hat'
        D = Lx/4;
        h0_x = heaviside(x+D/2).*heaviside(D/2-x).*(cos(pi*x/D));
        h0_y = ones(size(y));
        u0_x = 0;
        u0_y = 0;
        v0_x = 0;
        v0_y = 0;
end

% Normalization of wave function according to model amplitude.
normf = max(max(h0_x' * h0_y));
H(:, :, tm) = A / normf * h0_x' * h0_y + H0;
U(:, :, tm) = A / normf * (u0_x' * u0_y) .* H(:, :, tm);
V(:, :, tm) = A / normf * (v0_x' * v0_y) .* H(:, :, tm);

if ModelPrms.IsBox
    % Rigid boundaries in zonal direction.
    U(1 , : , tm) = 0; U(nx-1 , : , tm) = 0;
else
    % Periodic boundaries in zonal direction.
    H(1, :, tm) = H(nx - 1, :, tm); H(nx, :, tm) = H(2, :, tm);
    U(1, :, tm) = U(nx - 1, :, tm); U(nx, :, tm) = U(2, :, tm);
    V(1, :, tm) = V(nx - 1, :, tm); V(nx, :, tm) = V(2, :, tm);
end

% Rigid boundaries in meridional direction.
V(:, 1, tm) = 0; V(:, ny - 1, tm) = 0;

% Setting current and past same to be the same for first iteration.
H(:, :, tn) = H(:, :, tm);
U(:, :, tn) = U(:, :, tm);
V(:, :, tn) = V(:, :, tm);



%% Time stepping loop

axes(HandAxes);
% time_an = annotation('textbox',[0.82 0.25 0.1 0.1],'String','time = 0 hours','EdgeColor','none','tag','time_an');
hold off;
qq = [-1, -0.7, -0.4, -0.11, 0.12, 0.4, 0.7, 1];
for t = 0 : N
    % Plotting
    if (mod(t, plot_rate) == 0) || t == N
        cla;
        contourf(x(i) / 1000, y(j) / 1000, (H(i, j, tn) - H0)',qq);
        X1 = x(i)/1000;
        Y1 = y(j)/1000;
        U1 = U(i,j,tn)'/H0;
        V1 = V(i,j,tn)'/H0;
        hold on;
        q = 2;
        quiver(X1(1:q:end,2:q:end) , Y1(1:q:end,2:q:end) , U1(2:q:end,2:q:end) , V1(2:q:end,2:q:end) ,0.5, 'k');        
        
        grid on;
        colormap(cmap_white_at_zero(-A,A,9));
        colorbar;

        caxis([-A A])
        xlabel('x [km]');
        ylabel('y [km]');
        title(sprintf('time  =  %g  hours',round(t*dt/3600)))   
       
        drawnow;    
       
    end
        
% Advancing H.
    H(i, j, tp) = H(i, j, tm) - delt * ((U(i, j, tn) - U(im, j, tn)) / dx + (V(i, j, tn) - V(i, jm, tn)) / dy);
    if ModelPrms.IsBox==0
        H(1, :, tp) = H(nx - 1, :, tp); H(nx, :, tp) = H(2, :, tp); % Periodic boundaries in zonal direction
    end                                        
    % Advancing U.
    grav(i, j) = (H(ip, j, tn) - H(i, j, tn)) .* (H(ip, j, tn) + H(i, j, tn)) / d2x; % Gravitation
    vel(i, j) = (V(i, j, tn) + V(i, jm, tn) + V(ip, j, tn) + V(ip, jm, tn)) / 4; % Average velocity
    adv(i, j) = ((U(ip, j, tn) + U(i, j, tn)).^2 ./ H(ip, j, tn) - (U(i, j, tn) + U(im, j, tn)).^2 ./ H(i, j, tn)) / d4x; % Advection
    % Cross advection
    DH1(i, j) = H(i, j, tn) + H(i, jp, tn) + H(ip, j, tn) + H(ip, jp, tn);
    DH2(i, j) = H(i, j, tn) + H(i, jm, tn) + H(ip, j, tn) + H(ip, jm, tn);
    xadv1(i, j) = (U(i, jp, tn) + U(i, j, tn)) .* (V(i, j, tn) + V(ip, j, tn)) ./ DH1(i, j);
    xadv1(:, ny - 1) = 0;
    xadv2(i, j) = (U(i, j, tn) + U(i, jm, tn)) .* (V(ip, jm, tn) + V(i, jm, tn)) ./ DH2(i, j);
    xadv2(:, 2) = 0;
    xadv(i, j) = (xadv1(i, j) - xadv2(i, j)) / dy;
    % Diffusion
    diff_x(i, j) = (U(im, j, tn) - 2 * U(i, j, tn) + U(ip, j, tn)) / dx2;
    diff_y(i, j) = (U(i, jm, tn) - 2 * U(i, j, tn) + U(i, jp, tn)) / dy2;
    diff_y(:, 2) = 0; diff_y(:, ny - 1) = 0;
    % Finite difference equation
    U(i, j, tp) = U(i, j, tm) - delt *(g * grav(i, j) - fU(i, j) .* vel(i, j) + bf * U(i, j, tn) - ah * (diff_x(i, j) + diff_y(i, j))) - ...
                  nl * delt * (adv(i, j) + xadv(i, j));
    if ModelPrms.IsBox
        U(1 , : , tp) = 0; U(nx-1 , : , tp) = 0;
    else
        U(1, :, tp) = U(nx - 1, :, tp); U(nx, :, tp) = U(2, :, tp); % Periodic boundaries in zonal direction
    end
    
    % Advancing V.
    grav(i, j) = (H(i, jp, tn) - H(i, j, tn)) .* (H(i, jp, tn) + H(i, j, tn)) / d2y; % Gravitation
    vel(i, j) = (U(i, j, tn) + U(im, j, tn) + U(i, jp, tn) + U(im, jp, tn)) / 4; % Average velocity
    adv(i, j) = ((V(i, jp, tn) + V(i, j, tn)).^2 ./ H(i, jp, tn) - (V(i, j, tn) + V(i, jm, tn)).^2 ./ H(i, j, tn)) / d4y; % Advection
    % Cross advection
    DH1(i, j) = H(i, j, tn) + H(i, jp, tn) + H(ip, j, tn) + H(ip, jp, tn);
    DH2(i, j) = H(i, j, tn) + H(im, j, tn) + H(i, jp, tn) + H(im, jp, tn);
    xadv1(i, j) = (U(i, jp, tn) + U(i, j, tn)) .* (V(i, j, tn) + V(ip, j, tn)) ./ DH1(i, j);
    xadv2(i, j) = (U(im, j, tn) + U(im, jp, tn)) .* (V(i, j, tn) + V(im, j, tn)) ./ DH2(i, j);
    xadv(i, j) = (xadv1(i, j) - xadv2(i, j)) / dx;
    % Diffusion
    diff_x(i, j) = (V(im, j, tn) - 2 * V(i, j, tn) + V(ip, j, tn)) / dx2;
    diff_y(i, j) = (V(i, jm, tn) - 2 * V(i, j, tn) + V(i, jp, tn)) / dy2;
    diff_y(:, 2) = 0; diff_y(:, ny - 2) = 0;
    % Finite difference equation
    V(i, j, tp) = V(i, j, tm) - delt *(g * grav(i, j) + fV(i, j) .* vel(i, j) + bf * V(i, j, tn) - ah * (diff_x(i, j) + diff_y(i, j))) - ...
                  nl * delt * (adv(i, j) + xadv(i, j));
    if ModelPrms.IsBox==0
        V(1, :, tp) = V(nx - 1, :, tp); V(nx, :, tp) = V(2, :, tp); % Periodic boundaries in zonal direction
    end
    V(:, 1, tp) = 0; V(:, ny - 1, tp) = 0; % Rigid boundaries in meridional direction
    
    % Robert-Asselin Filter.
    if (t > 5)
        U(:, :, tn) = U(:, :, tn) + 0.05 * (U(:, :, tp) - 2 * U(:, :, tn) + U(:, :, tm));
        V(:, :, tn) = V(:, :, tn) + 0.05 * (V(:, :, tp) - 2 * V(:, :, tn) + V(:, :, tm));
        H(:, :, tn) = H(:, :, tn) + 0.05 * (H(:, :, tp) - 2 * H(:, :, tn) + H(:, :, tm));
    end
    
    % Moving on in time.
    delt = 2 * dt;
    tm = mod(tm, 3) + 1;
    tn = mod(tn, 3) + 1;
    tp = mod(tp, 3) + 1;
    
    if FlagStop
        break;
    end
end
%delete(findall(gcf,'Tag','time_an')); % remove time annotation

%% stop and pause buttons 
   function stopb(~,~)
        FlagStop=1;
   end

   function pauseb(~,~)
        pause();
   end


end

