%
%  finite-difference wave code
%
%  code9a.m -- wave equation with speed discontinuity
%           -- and interface conditions:
%                           u & u_x continuous 
%

disp(' ')
disp('wave equation for snell''s law ')
disp(' ')
disp('work = 0/1 (ok/better) & pchoice = 0/1 (contour/pcolor)')
work = 0
pchoice = 1

L = 1.0;     N = 64*(2^work);
ds = L/N;

M  = 168*(2^work);  Msave = 8*(2^work);

CN  = 1.0;  CP = 0.5;
dt = 0.6*ds/max([CN CP]);
A2N = (CN*dt/ds)^2;   A2P = (CP*dt/ds)^2;

% coordinates (F denotes Full vectors, otherwise Interior)

xF =  0:ds:L;    xG =  0:ds:2*L-ds;    
x = xF(1:end-1);

yN = -L+ds:ds: -ds;
yP =    ds:ds:L-ds;

yF = [-L yN 0 yP L];

[xg ygN] = meshgrid(x,yN);
[xg ygP] = meshgrid(x,yP);

% initial conditions (ang = angle of incidence)

ang = pi/4;

gauss = exp(-200*((xg-0.5).^4 + (ygN+0.5).^4));
gauss = gauss .* (gauss>10^(-4));
uNi  = cos(30*(cos(ang)*xg + sin(ang)*ygN)) .* gauss;
duNi = 30*CN*sin(30*(cos(ang)*xg + sin(ang)*ygN)) .* gauss;

uPi  = zeros(size(uNi));
duPi = zeros(size(uNi));

uZi  = zeros(1,N);
duZi = zeros(1,N);

% initial plot

zz = zeros(1,N  );
zF = zeros(1,2*N);
con = -1:0.2:1;

uNiF = [uNi, uNi];
uZiF = [uZi, uZi];
uPiF = [uPi, uPi];

uplot = [zF ; uNiF ; uZiF ; uPiF ; zF];

clf
if(pchoice==0)
  subplot(1,2,1)
  contour(xG,yF,uplot);axis('equal');axis([0 L -L L])
     caxis([-1 1]);colorbar
  title('Initial Wave');xlabel('x-axis');ylabel('y-axis')
  hold on;  plot([0 2],[0 0],'k--');  hold off
  subplot(1,2,2)
  contour(xG,yF,uplot);axis('equal');axis([0 L -L L])
  caxis([-1 1]);colorbar
else
  pcolor(xG,yF,uplot);axis('equal');axis([0 2*L -L L]);
  shading flat
  caxis([-1 1]);colorbar
end
hold on;  plot([0 2],[0 0],'k--');  hold off
title('Snell''s Law at an Interface');xlabel('x-axis');ylabel('y-axis')

% first timestep, zero dirichlet boundary conditions

u1Ni = uNi + dt*duNi ...
     + (A2N/2) * ( diff([uNi(:,end) , uNi , uNi(:,1)],2,2) ...
                 + diff([zz ; uNi ; uZi]  ,2,1) );

u1Pi = uPi + dt*duPi ...
     + (A2P/2) * ( diff([uPi(:,end) , uPi , uPi(:,1)],2,2) ...
                 + diff([uZi ; uPi ; zz]  ,2,1) );

u1Zi = 0.5*(u1Ni(end,:) + u1Pi(1,:));

% set-up for time loops

uNold = uNi;   uNnew = u1Ni;
uPold = uPi;   uPnew = u1Pi;
uZold = uZi;   uZnew = u1Zi;

% timestep loop

for j = 2:M
   uN = 2*uNnew - uNold ...
      + (A2N) * ( diff([uNnew(:,end) , uNnew , uNnew(:,1)],2,2) ...
                + diff([zz ; uNnew ; uZnew]               ,2,1) );

   uP = 2*uPnew - uPold ...
      + (A2P) * ( diff([uPnew(:,end) , uPnew , uPnew(:,1)],2,2) ...
                + diff([uZnew ; uPnew ; zz]               ,2,1) );

   uNold = uNnew;
   uNnew = uN;

   uPold = uPnew;
   uPnew = uP;

   uZold = uZnew;
   uZnew = 0.5*(uN(end,:) + uP(1,:));

   if (mod(j,Msave)==0)
      j 

      uNiF = [uNnew, uNnew];
      uZiF = [uZnew, uZnew];
      uPiF = [uPnew, uPnew];

      uplot = [zF ; uNiF ; uZiF ; uPiF ; zF];

   figure(1)
   if(pchoice==0)
     contour(xG,yF,uplot,con);axis('equal'),axis([0 L -L L])
     caxis([-1 1]);colorbar
   else
     pcolor(xG,yF,uplot);axis('equal');axis([0 2*L -L L])
     shading flat
     caxis([-1 1]);colorbar
   end
   hold on;  plot([0 2],[0 0],'k--');  hold off
   title('Snell''s Law at an Interface');xlabel('x-axis');ylabel('y-axis')

   end
end