function waterwave
% WATER WAVE
% 2D Shallow Water Model
%
% Lax-Wendroff finite difference method.
% Reflective boundary conditions.
% Random water drops initiate gravity waves.
% Surface plot displays height colored by momentum.
% Plot title shows t = simulated time and tv = a measure of total variation.
% An exact solution to the conservation law would have constant tv.
% Lax-Wendroff produces nonphysical oscillations and increasing tv.
%
% See:
%    http://en.wikipedia.org/wiki/Shallow_water_equations
%    http://www.amath.washington.edu/~rjl/research/tsunamis
%    http://www.amath.washington.edu/~dgeorge/tsunamimodeling.html
%    http://www.amath.washington.edu/~claw/applications/shallow/www

% Parameters

n = 64;                  % grid size
g = 9.8;                 % gravitational constant
dt = 0.02;               % hardwired timestep
dx = 1.0;
dy = 1.0;
nplotstep = 8;           % plot interval
ndrops = 1;              % maximum number of drops
dropstep = 500;          % drop interval
D = droplet(1.5,21);     % simulate a water drop
gif_counter = 1;
gifname = 'testmovie.gif' % TODO: The file name the gif is saved under


% Initialize graphics

[surfplot,top,restart,stop] = initgraphics(n);

% Outer loop, restarts.

while get(stop,'value') == 0
   set(restart,'value',0)
   
   H = ones(n+2,n+2);   U = zeros(n+2,n+2);  V = zeros(n+2,n+2);
   Hx = zeros(n+1,n+1); Ux = zeros(n+1,n+1); Vx = zeros(n+1,n+1);
   Hy = zeros(n+1,n+1); Uy = zeros(n+1,n+1); Vy = zeros(n+1,n+1);
   ndrop = ceil(rand*ndrops);
   nstep = 0;

   % Set inital conditions
   w = size(D,1);

    % Droplet always starts in the corner
   i = ceil(1+(1:w));
   j = ceil(1+(1:w));
   H(i,j) = H(i,j) + D;
   
   %`Creating the terrain of the island
   island = zeros(n+2, n+2);
   island(1/3*(n+2):2/3*(n+2), 1/3*(n+2):2/3*(n+2)) = .4;
   
   % Creating the terrain of the slope 
   slope = zeros(n+2, n+2);
   for i = ceil((n+2)/2):n+2
    slope(i,:) = (i-ceil((n+2)/2))* 2/((n+2/2));
   end
   
   % TODO: Set to either island or terrain
   terrain = island;
   if terrain == slope
    is_slope= 1;
    is_island = 0;
    tall_island = 0;
   else
       is_slope = 0;
       is_island = 1;
       % TODO: Toggle to set tall or short island
       tall_island = 1;
   end
   
   slope_friction_constant = 50;
   
   combined = H + terrain;
   
   % Inner loop, time steps.
   while get(restart,'value')==0 && get(stop,'value')==0
       nstep = nstep + 1
     
       % Reflective boundary conditions
       H(:,1) = H(:,2);      U(:,1) = U(:,2);       V(:,1) = -V(:,2);
       H(:,n+2) = H(:,n+1);  U(:,n+2) = U(:,n+1);   V(:,n+2) = -V(:,n+1);
       H(1,:) = H(2,:);      U(1,:) = -U(2,:);      V(1,:) = V(2,:);
       H(n+2,:) = H(n+1,:);  U(n+2,:) = -U(n+1,:);  V(n+2,:) = V(n+1,:);
       
       % Reflective object conditions
       if tall_island
        H(1/3*(n+2):2/3*(n+2), ceil(1/3*(n+2))) = H(1/3*(n+2):2/3*(n+2), ceil(1/3*(n+2))-1);
        H(1/3*(n+2):2/3*(n+2), ceil(1/3*(n+2))) = H(1/3*(n+2):2/3*(n+2), ceil(2/3*(n+2))+1);
        H(ceil(1/3*(n+2)), 1/3*(n+2):2/3*(n+2)) = H(ceil(1/3*(n+2))-1, 1/3*(n+2):2/3*(n+2));
        H(ceil(1/3*(n+2)), 1/3*(n+2):2/3*(n+2)) = H(ceil(2/3*(n+2))+1, 1/3*(n+2):2/3*(n+2));
       end
       
%        H(ceil((n+2)/2), :) = .5*
       
       % First half step
   
       % x direction
       i = 1:n+1;
       j = 1:n;
   
       % height
       Hx(i,j) = (H(i+1,j+1)+H(i,j+1))/2 - dt/(2*dx)*(U(i+1,j+1)-U(i,j+1));
   
       % x momentum
       for f = 1:n+1
           for s = 1:n
               if terrain(f,s)
                   Ux(f,s) = (U(f+1,s+1)+U(f,s+1))/2 -  ...
                 dt/(2*dx)*((U(f+1,s+1).^2./H(f+1,s+1) + slope_friction_constant*g/2*H(f+1,s+1).^2) - ...
                            (U(f,s+1).^2./H(f,s+1) + slope_friction_constant*g/2*H(f,s+1).^2));
               else
                   Ux(f,s) = (U(f+1,s+1)+U(f,s+1))/2 -  ...
                 dt/(2*dx)*((U(f+1,s+1).^2./H(f+1,s+1) + g/2*H(f+1,s+1).^2) - ...
                            (U(f,s+1).^2./H(f,s+1) + g/2*H(f,s+1).^2));
               end
           end
       end
%        Ux(i,j) = (U(i+1,j+1)+U(i,j+1))/2 -  ...
%                  dt/(2*dx)*((U(i+1,j+1).^2./H(i+1,j+1) + g/2*H(i+1,j+1).^2) - ...
%                             (U(i,j+1).^2./H(i,j+1) + g/2*H(i,j+1).^2));
                        
       
   
       % y momentum
       Vx(i,j) = (V(i+1,j+1)+V(i,j+1))/2 - ...
                 dt/(2*dx)*((U(i+1,j+1).*V(i+1,j+1)./H(i+1,j+1)) - ...
                            (U(i,j+1).*V(i,j+1)./H(i,j+1)));
                        
       
       % y direction
       i = 1:n;
       j = 1:n+1;
   
       % height
       Hy(i,j) = (H(i+1,j+1)+H(i+1,j))/2 - dt/(2*dy)*(V(i+1,j+1)-V(i+1,j));
   
       % x momentum
       Uy(i,j) = (U(i+1,j+1)+U(i+1,j))/2 - ...
                 dt/(2*dy)*((V(i+1,j+1).*U(i+1,j+1)./H(i+1,j+1)) - ...
                            (V(i+1,j).*U(i+1,j)./H(i+1,j))- ...
                            100*g * (terrain(f+1,s) - terrain(f,s)));
                        
       for f = 1:n
           for s = 1:n+1
               if terrain(f,s)
                   Vy(f,s) = (V(f+1,s+1)+V(f+1,s))/2 - ...
                 dt/(2*dy)*((V(f+1,s+1).^2./H(f+1,s+1) + slope_friction_constant*g/2*H(f+1,s+1).^2) - ...
                            (V(f+1,s).^2./H(f+1,s) + slope_friction_constant*g/2*H(f+1,s).^2));
               else
                   Vy(f,s) = (V(f+1,s+1)+V(f+1,s))/2 - ...
                 dt/(2*dy)*((V(f+1,s+1).^2./H(f+1,s+1) + g/2*H(f+1,s+1).^2) - ...
                            (V(f+1,s).^2./H(f+1,s) + g/2*H(f+1,s).^2));
               end
           end
       end
       
%        Vy(i,j) = (V(i+1,j+1)+V(i+1,j))/2 - ...
%                  dt/(2*dy)*((V(i+1,j+1).^2./H(i+1,j+1) + g/2*H(i+1,j+1).^2) - ...
%                             (V(i+1,j).^2./H(i+1,j) + g/2*H(i+1,j).^2));
       
                        
       % Second half step
       i = 2:n+1;
       j = 2:n+1;
   
       % height
       H(i,j) = H(i,j) - (dt/dx)*(Ux(i,j-1)-Ux(i-1,j-1)) - ...
                         (dt/dy)*(Vy(i-1,j)-Vy(i-1,j-1));
       % x momentum
       U(i,j) = U(i,j) - (dt/dx)*((Ux(i,j-1).^2./Hx(i,j-1) + g/2*Hx(i,j-1).^2) - ...
                         (Ux(i-1,j-1).^2./Hx(i-1,j-1) + g/2*Hx(i-1,j-1).^2)) ...
                       - (dt/dy)*((Vy(i-1,j).*Uy(i-1,j)./Hy(i-1,j)) - ...
                         (Vy(i-1,j-1).*Uy(i-1,j-1)./Hy(i-1,j-1)));
       % y momentum
       V(i,j) = V(i,j) - (dt/dx)*((Ux(i,j-1).*Vx(i,j-1)./Hx(i,j-1)) - ...
                         (Ux(i-1,j-1).*Vx(i-1,j-1)./Hx(i-1,j-1))) ...
                       - (dt/dy)*((Vy(i-1,j).^2./Hy(i-1,j) + g/2*Hy(i-1,j).^2) - ...
                         (Vy(i-1,j-1).^2./Hy(i-1,j-1) + g/2*Hy(i-1,j-1).^2));
        
                     
       combined = H + terrain;
       if tall_island
            combined(1/3*(n+2):2/3*(n+2), 1/3*(n+2):2/3*(n+2)) = 2;
       end
       
       % Update plot
       if mod(nstep,nplotstep) == 0
          C = abs(U(i,j)) + abs(V(i,j));  % Color shows momemtum
          t = nstep*dt;
          tv = norm(C,'fro');
          set(surfplot,'zdata',combined(i,j),'cdata',C);
          set(top,'string',sprintf('t = %6.2f,  tv = %6.2f',t,tv))
          drawnow
          frame = getframe(1);
          im = frame2im(frame);
          [imind,cm] = rgb2ind(im,256);
          if gif_counter == 1
            imwrite(imind,cm,gifname,'gif', 'Loopcount',inf);
          else
            imwrite(imind,cm,gifname,'gif','WriteMode','append');
          end
          
          gif_counter = gif_counter+1;
       end
      
       if all(all(isnan(H))), break, end  % Unstable, restart
   end
end
close(gcf)

% ------------------------------------

function D = droplet(height,width)
% DROPLET  2D Gaussian
% D = droplet(height,width)
   [x,y] = ndgrid(-1:(2/(width-1)):1);
   D = height*exp(-5*(x.^2+y.^2));

% ------------------------------------

function [surfplot,top,restart,stop] = initgraphics(n);
% INITGRAPHICS  Initialize graphics for waterwave.
% [surfplot,top,start,stop] = initgraphics(n)
% returns handles to a surface plot, its title, and two uicontrol toggles.

   clf
   shg
   set(gcf,'numbertitle','off','name','Shallow_water')
   x = (0:n-1)/(n-1);
   surfplot = surf(x,x,ones(n,n),zeros(n,n));
   grid off
   axis([0 1 0 1 -1 3])
   caxis([-1 1])
   shading faceted
   c = (1:64)'/64;
   cyan = [0*c c c];
   colormap(cyan)
   top = title('Click restart');
   restart = uicontrol('position',[20 20 80 20],'style','toggle','string','restart');
   stop = uicontrol('position',[120 20 80 20],'style','toggle','string','stop');