%%% solving the wave equation using leap-frog, x in [0,1], t in [0,T]
% MU''=-KU+F, K=A^T C A

%%% initialization
n=31; %user, number of points in x (or number of masses)
h=1/(n+1); %spacial step
T=5; %final time
nt=T*(n+1)*2; %user, number of steps taken in time
dt=T/nt; %time step
Minv=speye(n,n); % inverse of matrix M
F=zeros(n,nt+1); %user, forcing F(i,j)=f(i*h,j*dt)
%F(1,:)=0;
F(1,:)=.25*sin(2*pi*(0:dt:T)*4);
U=zeros(n,nt+1); % this will contain the solution: U(i,j)=u(i*h,(j-1)*dt)
bc=0; %user, bc=0 means fixed-fixed, bc=1 means free-fixed, bc==2 free-free


%%% boundary conditions and second difference
if bc==0,
    xhalf=h/2:h:1-h/2; %this is where the "springs" are
    C=spdiags(.5+2*(xhalf'>1/3),0,length(xhalf),length(xhalf)); %constitutive law
%    C=spdiags(xhalf'.^0,0,length(xhalf),length(xhalf)); %constitutive law
    A=spdiags([ones(length(xhalf),1) -ones(length(xhalf),1)],[0 -1],length(xhalf),n); % forward difference
elseif bc==1
    xhalf=3*h/2:h:1-h/2; %user, this is where the "springs" are
    C=spdiags(xhalf'.^0,0,length(xhalf),length(xhalf)); %user, constitutive law
    A=spdiags([ones(length(xhalf),1) -ones(length(xhalf),1)],[1 0],length(xhalf),n); % user, forward difference
elseif bc==2
    xhalf=3*h/2:h:1-3*h/2; %user, this is where the "springs" are
    C=spdiags(xhalf'.^0,0,length(xhalf),length(xhalf)); %user, constitutive law
    A=spdiags([ones(length(xhalf),1) -ones(length(xhalf),1)],[1 0],length(xhalf),n); % user, forward difference
end
K=A'*C*A; % approx -d/dx ( c(x) du/dx) * h^2
K=K/h^2;  % approx -d/dx ( c(x) du/dx)

%%% initial conditions
x=h:h:1-h;
%U0=.5*exp(-100*(x'-.25).^2); % user, u(x,0)
%U0=.25*sin(6*2*pi*x'); % user, u(x,0)
%U0=.5-abs(x'-.5); % user, u(x,0)
U0=zeros(length(x),1); % user, u(x,0)
U(:,1)=U0;
Up0=zeros(length(x),1); % user, u'(x,0)
U(:,2)=U0+dt*Up0; %forward difference to evaluate u(x,dt) using u'(x,0)
%%%

%%% leap-frog algorithm
for i=3:nt+1
    U(:,i)=2*U(:,i-1)-U(:,i-2)+dt^2*Minv*(-K*U(:,i-1)+(1/h^2)*C(1,1)*F(:,i-1)); % user, leap-frog step
%    U(:,i)=2*U(:,i-1)-U(:,i-2)+dt^2*Minv*(-K*U(:,i-1)+[1;zeros(n-1,1)].*(K*F(:,i-1))); % user, leap-frog step
%    U(:,i)=2*U(:,i-1)-U(:,i-2)+dt^2*Minv*(-K*U(:,i-1)+(1/h^2)*A'*C*[F(:,i-1);0]); % user, leap-frog step
end

%%% boundary conditions again
% we add a first row and last row to U to reflect boundary conditions
if bc==0
    U=[F(1,:);U]; %U(x=0,t)=forcing
%    U=[zeros(1,nt+1);U]; %U(x=0,t)=0
    U=[U;zeros(1,nt+1)]; %U(x=1,t)=0
elseif bc==1
    U=[U(1,:);U]; %user, U(x=0,t)=U(x=h,t), free end
    U=[U;zeros(1,nt+1)]; %U(x=1,t)=0
elseif bc==2
    U=[U(1,:);U]; %user, U(x=0,t)=U(x=h,t), free end
    U=[U;U(end,:)]; %U(x=1,t)=0
end

%%% plot the solution
xall=0:h:1;
for i=1:nt+1
    plot(xall,U(:,i),'-o'); grid on;axis([0 1 -1 1]);
    xlabel('x');ylabel('u');title(['solution u at time t=' num2str((i-1)*dt)]);
    pause(.1);
end
