% Iterative scheme for the TAT problem

clc; clear; close all;

N = 150;
h = 2/N;

grid.N = N;
grid.h = h;

x = -1:h:1;

[xxsq,yysq] = meshgrid(x,x);
xxsq = reshape(xxsq, (N+1)^2, 1);
yysq = reshape(yysq, (N+1)^2, 1);

bsq = find(abs(xxsq)==1 | abs(yysq)==1);

xxsq = reshape(xxsq, N+1, N+1);
yysq = reshape(yysq, N+1, N+1);

grid.xxsq = xxsq;
grid.yysq = yysq;
grid.bsq = bsq;

%%%%%%%%%%%%%%%%%%% PML parameters

delta = 0.3;
b = 500; 

PML.delta = delta;
PML.b = b;

%%%%%%%%%%%%%%%%%% sound speeds

% the sound speed for the forward problem need not be the same as that of
% the backprojection.
% sound speed for forward waves

ctype = 'c1paper';

[c,c_caracs] = soundSpeedMap(ctype,N,xxsq,yysq);

c_switch = c_caracs.c_switch;

if c_switch 
    fc = c_caracs.fc;
    fdxc = c_caracs.fdxc;
    fdyc = c_caracs.fdyc;
end

% sound speed for backward waves

c2type = 'c1paper';

[c2,c2_caracs] = soundSpeedMap(c2type,N,xxsq,yysq);

c2_switch = c2_caracs.c_switch;

if c2_switch 
    fc2 = c2_caracs.fc;
    fdxc2 = c2_caracs.fdxc;
    fdyc2 = c2_caracs.fdyc;
end


%%%%%%%%%%%%%%%%% initial condition
% can be 'disk', 'gaussian', 'gaussianspiral'

u0type = 'gaussianspiral';
[u0,u0_sing_switch,points] = initCondMap(u0type,N,xxsq,yysq);

u0max = max(max(u0));

%%%%%%%%%%%%% display sound speed and initial condition

figure; 
subplot(1,3,1); myImagesc(c,'forward sound speed');
subplot(1,3,2); myImagesc(c2,'backward sound speed');
subplot(1,3,3); myImagesc(u0,'initial condition');

pause

%%%%%%%%%%%%%%%%% differential operators of interest

[DX,DY,Lap] = diffOps(h,N,bsq);

%%%%%%%%%%%%%%%%% computation of T_0

disp('eikonal equation')

NGS = 2; % number of Gauss-Seidel iterations

TT = eikonal(c,grid,NGS,0);

figure; 
subplot(1,2,1); myImagesc(c,'forward sound speed');
subplot(1,2,2); myImagesc(TT,'distance map'); 

pause

%%%%%%% time variables

T = 1.3;
dtdisp = 0.03;

dt = .3*h/max(max(c));
nt = floor(T/dt);
ntdisp = floor(dtdisp/dt);

time.T = T;
time.dtdisp = dtdisp;
time.dt = dt;
time.nt = nt;
time.ntdisp = ntdisp;

%%%%%%%%%%%%%%%%%% forward solving

disp('forward solver');

[u,data] = wavePML(u0,c,grid,PML,time);

figure; subplot(1,2,1);
myImagesc(u,'u(T) cropped over [-1,1]^2')

%%%%%%%%%%%%%%% processing of the data at time T

postproc = 'HE';

switch postproc
    case 'HE'
        disp('harmonic extension of trace of u(T)');
        
        uT = u;
        
        rhs = zeros((N+1)^2,1);
        rhs(bsq) = data(end,:);
        
        sol = Lap\rhs;
        sol = reshape(sol, N+1, N+1);
        
        subplot(1,2,2); 
        myImagesc(sol, 'harmonic ext. of trace of u(T)');
        
        uend = sol;
        

    otherwise
        disp('time-reversal post-processing: do nothing');
        uend = zeros(N+1,N+1);
end
        
%----------- backward problem 

disp('backward wave problem')

u = backwave(data, uend, c2, DX, DY, grid, time, u0max);

figure;
subplot(2,2,1); myImagesc(u0, 'u_0');
subplot(2,2,2); myImagesc(u, 'u');
subplot(2,2,3); myImagesc(u-u0, 'u-u0')
subplot(2,2,4); cut = floor(N/2);
plot(x,u(cut,:),'b',x,u0(cut,:),'r'); legend('u','u_0');  title('cut at y=0'); 

% figure(10);
% subplot(1,3,1); surf(xxsq,yysq,u0); 
% shading flat; title('u_0')
% subplot(1,3,2); surf(xxsq,yysq,u); 
% shading flat; title('u')
% subplot(1,3,3); surf(xxsq,yysq,u0-u); 
% shading flat; title('u_0-u')

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%------------ further iterations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

ufinal = u;

recons{1} = ufinal;

Niter = 10;
time.dtdisp = 0; % no more display, only the iterations
time.ntdisp = 0;

relErr = zeros(1, 1+Niter);
relErr(1) = norm(u-u0,'fro')/norm(u0,'fro');

figure(12); 
subplot(2,5,1); surf(xxsq,yysq,u); shading flat; title('u, k=1')
subplot(2,5,6); surf(xxsq,yysq,u-u0); shading flat; title('u-u0, k=1')

cut = floor(N/2);
figure(13); subplot(1,2,1);
plot(x,u0(cut,:),'k',x,u(cut,:),'b-'); hold on;

for k = 1:Niter,
    
    disp(['k = ',num2str(k+1),' / ',num2str(Niter+1)]);
    
    %----- forward wave, "Lambda" operator
%     figure(10); subplot(1,2,1);
    disp('forward solver')
    [Lambda_u,data] = wavePML(u,c,grid,PML,time);
    
    %----- postprocessing of data at time T    
    switch postproc
        case 'HE'
            disp('harmonic extension')
            rhs = zeros((N+1)^2,1);
            rhs(bsq) = data(end,:);
            
            PLambda_u = Lap\rhs;

        otherwise
            disp('time-reversal post-processing: do nothing');
            PLambda_u = zeros(N+1,N+1);
    end
    
    %----- backward wave
    disp('backward wave')
%     subplot(1,2,2);
    ALambda_u = backwave(data, PLambda_u, c2, DX, DY, grid, time, u0max);
    
    %----- plotting results
    u = u - ALambda_u;
    ufinal = ufinal + u; % partial sum
    
    recons{k+1} = ufinal;
    
    relErr(k+1) = norm(ufinal-u0,'fro')/norm(u0,'fro');
%     
%     figure(12); 
%     subplot(2,5,mod(k,5)+1); surf(xxsq,yysq,ufinal); 
%     shading flat; title(['u_k, k = ', num2str(k+1)]);
%     subplot(2,5,mod(k,5)+6); surf(xxsq,yysq,ufinal-u0); 
%     shading flat; title(['u_k-u0, k = ', num2str(k+1)])
%     
%     figure(13); subplot(1,2,1); 
%     plot(x,ufinal(cut,:),'b'); hold on;
%     title('cut at y=0')
%     subplot(1,2,2); 
%     plot(1:(Niter+1),relErr,'b'); title('relative L_2 error')
%     
end

figure;
slider(recons,u0)

