% 2-D acoustic gradient
% Y. Pan, last update 22.10.2018

close all; 
clear all;
%----------------------------------------------------------
% define grid parameters
Lx=4000;     % Model size [meter]
Lz=5000;
dh=15.000;		% Grid spacing [meter]
nx=floor(Lx/dh);		% number of grid points in x-direction
nz=floor(Lz/dh);		% number of grid points in x-direction
T=1.000;			% propagation time [s]

dt=2.00e-3;		% time step [seconds]


fc=10.0;			% center frequency of Ricker wavelet
xsrc=1500;			% source position in grid points in 
zsrc=2500;
xrec=2500;			% receiver position in grid points in meter
zrec=[1200:50:3800];			% receiver position in grid points in meter
nsrc=length(zsrc);
ntr=length(zrec);

framerate=0.050;  % frames per second
framerate=round(framerate/dt);

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

% time and spatial vector
t=dt:dt:T;
x=(1:nx)*dh;
z=(1:nz)*dh;
nt=length(t);

xs=round(xsrc/dh);
zs=round(zsrc/dh);
xr=round(xrec/dh);
zr=round(zrec/dh);

% Initialize array for pressure
p=zeros(nx,nz);
vx=zeros(nx,nz);
vz=zeros(nx,nz);

% Model definition
c0=zeros(nx,nz);
c0(:,:)=3500.0;           % propagation velocity
xst=round(1800/dh);
xe=round(2200/dh);
ys1=round(1800/dh);
ys2=round(2800/dh);
ye1=round(2200/dh);
ye2=round(3200/dh);
c0(xst:xe,ys1:ye1)=2800.0;
c0(xst:xe,ys2:ye2)=4200.0;
c0_or=c0;

pho=zeros(nx,nz);
pho(:,:)=2200;
b=1./pho;
lambda=pho.*c0.*c0;
cfl=c0.*dt*sqrt(2)/dh;

% Check stability criterion
cflnum=max(max(cfl));
disp(['CFL number:',num2str(cflnum)]);
if (cflnum>1) disp(['CFL number is greater than one. Stop.']); 
    return; 
end

% check numerical dispersion according to minimum wavelength
lambda_min=min(min(c0)/(2*fc));
disp([' Minimum wavelength is ',num2str(lambda_min),' m']);
disp([' Number of gridpoints per minimum wavelength is ', num2str(lambda_min/dh)]);
if (num2str(lambda_min/dh)<3) disp(['Finer grid size is required. Stop.']); 
    return; 
end

% define ricker wavelet
delay=0.0/fc; 
shift=1.5/fc; 
ts=t-delay-shift; 
tau=pi*fc*ts;
ricker=(1.0-2.0*tau.*tau).*exp(-tau.*tau);
scale=1000000;

% sponge ABC
% reference: Cerjan et al., 1985, A nonreflecting boundary condition for
%             discrete acoustic and elastic wave equations
abc=ones(nx,nz);
abl=45;  % number of grids for ABC
%%% comment the part below to remove ABC %%%
for j=1:abl
    abc(1:nx,j)=exp(-(0.0053*(abl-j))^2);
end
for j=1:abl
    abc(j,abl:nz-abl)=exp(-(0.0053*(abl-j))^2);
end
abc(1:nx,nz:-1:nz-abl+1)=abc(1:nx,1:abl);
abc(nx:-1:nx-abl+1,abl:nz-abl)=abc(1:abl,abl:nz-abl);
%%% comment the part above to remove ABC %%%

dh2=1.0/(dh*dh);
dtx=dt/dh;
dtx2=dtx*dtx;
c02=c0.*c0;
b=b*dtx;
lambda=lambda*dtx;

rec=zeros(53,nt);
thred=scale/200;
drt2d=zeros(nx,nz);

% Loop over time steps and grid points
for n=1:nt  % First level of loop for time stepping
    for i=3:nx-2
        for j=3:nz-2
            %-------------2nd order eq. with 2nd order FD--------------
            %  --note: p=p^(n+1),vx=p^(n),vz=p^(n-1)
%             drt=vx(i+1,j)-2*vx(i,j)+vx(i-1,j)+vx(i,j+1)-2*vx(i,j)+vx(i,j-1);
%             p(i,j)=2*vx(i,j)-vz(i,j)+dtx2*c02(i,j)*drt;
             drt2d(i,j)=vx(i+1,j)-2*vx(i,j)+vx(i-1,j)+vx(i,j+1)-2*vx(i,j)+vx(i,j-1);
        end
    end
    drt2d=drt2d.*c02.*abc;  % attenuate it
    p=2*vx-vz+dtx2*drt2d;
    p=p.*abc;              % attenuate it
    for sr=1:nsrc
        p(xs,zs(sr))=p(xs,zs(sr))+scale*ricker(n);
    end
    vz=vx;
    vx=p;
    for tr=1:ntr
        rec(tr,n)=p(xr,zr(tr));
    end
end
obs=rec;

%initial model
c0(:,:)=3500.0; 
c02=c0.*c0;
forwardwave=zeros(nx,nz,nt);
p(:,:)=0.0;
vx(:,:)=0.0;
vz(:,:)=0.0;

for n=1:nt  % First level of loop for time stepping
    for i=3:nx-2
        for j=3:nz-2
            %-------------2nd order eq. with 2nd order FD--------------
            %  --note: p=p^(n+1),vx=p^(n),vz=p^(n-1)
            drt2d(i,j)=vx(i+1,j)-2*vx(i,j)+vx(i-1,j)+vx(i,j+1)-2*vx(i,j)+vx(i,j-1);
        end
    end
    drt2d=drt2d.*c02.*abc;  % attenuate it
    p=2*vx-vz+dtx2*drt2d;
    p=p.*abc;              % attenuate it
    forwardwave(:,:,n)=p-2*vx+vz;  %\partial_{tt}{p} without dt^2
    p(xs,zs)=p(xs,zs)+scale*ricker(n);
    vz=vx;
    vx=p;
    for tr=1:ntr
        rec(tr,n)=p(xr,zr(tr));
    end
end
syn=rec;
residual=syn-obs;
crtmisf=norm(residual)^2;

% backpropagate residual
p(:,:)=0.0;
vx(:,:)=0.0;
vz(:,:)=0.0;
backwardwave=zeros(nx,nz,nt);

for n=nt:-1:1  % First level of loop for time stepping 
    for i=3:nx-2
        for j=3:nz-2
            %-------------2nd order eq. with 2nd order FD--------------
            %  --note: p=p^(n+1),vx=p^(n),vz=p^(n-1)
            drt2d(i,j)=vx(i+1,j)-2*vx(i,j)+vx(i-1,j)+vx(i,j+1)-2*vx(i,j)+vx(i,j-1);
        end
    end
    drt2d=drt2d.*c02.*abc;  % attenuate it
    p=2*vx-vz+dtx2*drt2d;
    p=p.*abc;              % attenuate it
    backwardwave(:,:,n)=p;  %{p}^{\dagger} 
    for tr=1:ntr
        p(xr,zr(tr))=p(xr,zr(tr))+residual(tr,n); 
    end
    vz=vx;
    vx=p;
end


filename='snapmovie.gif';
% aviobj=avifile('test.avi','compression','None');
h=figure;
imagenum=1;

% crosscorrelation of forward and backward wavefields
grad=zeros(nx,nz);
for n=1:nt 
    grad(:,:)=grad(:,:)+forwardwave(:,:,n).*backwardwave(:,:,n);
%   if(mod(n,framerate)==0)  % output figure
%         imagesc([0,Lx],[0,Lz],forwardwave(:,:,n)');
%         title(['snapshot at T=',num2str(n*dt),'s']);
%         xlabel(' x [m]');
%         ylabel(' z [m]');
%         caxis([-thred thred]);
%         colormap(jet);
%         drawnow;
%         frame=getframe(h);
%         im=frame2im(frame);
%         [imind,cm]=rgb2ind(im,256);
%         pause(0.1);
%        if(imagenum==1)
%             imwrite(imind,cm,filename,'gif','Loopcount',Inf);
%             imagenum=imagenum+1;
%        else
%             imwrite(imind,cm,filename,'gif','Writemode','append');
%             imagenum=imagenum+1;
%        end
%   end 
end
grad(:,:)=grad(:,:)./c0(:,:);
grad(xs,zs)=0.0;

figure(2);
imagesc(grad');
hold on;

% Benchmark using FD
% gradbm=zeros(nx,nz);
% deltam=10.0;
% for xind=50:110
%     for yind=40:160
%         rec(:,:)=0.0;
%         c0(:,:)=3500.0; 
%         c0(xind,yind)=3500.0+deltam;
%         c02=c0.*c0;
%         forwardwave=zeros(nx,nz,nt);
%         p(:,:)=0.0;
%         vx(:,:)=0.0;
%         vz(:,:)=0.0;
%         % Loop over time steps and grid points
%         for n=1:nt  % First level of loop for time stepping
%             for i=3:nx-2
%                 for j=3:nz-2
%                     drt=vx(i+1,j)-2*vx(i,j)+vx(i-1,j)+vx(i,j+1)-2*vx(i,j)+vx(i,j-1);
%                     p(i,j)=2*vx(i,j)-vz(i,j)+dtx2*c02(i,j)*drt;
%                 end
%             end
%             for sr=1:nsrc
%                 p(xs,zs(sr))=p(xs,zs(sr))+scale*ricker(n);
%             end
%             vz=vx;
%             vx=p;
%             for tr=1:ntr
%                 rec(tr,n)=p(xr,zr(tr));
%             end
%         end
%         misfnew=norm(rec-obs)^2;
%         gradbm(xind,yind)=misfnew-crtmisf;
%     end
%     xind
% end
% 
% figure(3);
% gradbm(60,100)=0;
% gradbm(:,:)=gradbm(:,:)/deltam;
% imagesc(gradbm');