%++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
% Convolution 
%                                                -- wr, 21 Sep 2017 
% ========================================================
% Code is written to work in Octave. 
% You need to have teh geometry package from octave-forge 
% for the arrows to show. 
% To get the geometry package, from the octave prompt do 
% the following: 
%       >pkg install -forge geometry 
%       >pkg load geometry 
% 
% You also need GhostScript for saving the figures. 
% ========================================================
%++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

clc 
clear all
%close all

N = 64;   % Number of time-samples 

% Define the range: 
t = linspace(-pi,pi,N); 
tspan = t(N) - t(1); 

% Choose your first function f1(t): 
% I choose Gaussian: f1(t) = exp[-((t - t0)^2)/(2 x sigma^2)]
A1 = 0.5; 
sigma1 = 0.5; 
t0 = 0; 
f1 = A1*exp(-((t-t0).^2)/(2*sigma1^2)); %/(sigma*sqrt(2*pi));

% Define the second function (I choose a second Gaussian function): 
A2 = 1.0 
sigma2 = 1.0; 
t0 = 0; 
f2 = A2*exp(-((t-t0).^2)/(2*sigma2^2)); %/(sigma*sqrt(2*pi)); 
plot(t,f1,t,f2)

% Generate the coordinates/axes for the shifted version of f2: 
N_full = length(f1) + length(f2) - 1;
tspan_full = tspan*N_full/N;
t_full = linspace(-tspan_full/2,tspan_full/2,N_full);

% Iitialise the result array: 
summation = zeros(1,N_full);
figure(1) 

for i = 1:N_full 
	ishift = t(N/2+1) - t_full(i);
        %f2_shifted = exp(-((t_full - t_full(i)).^2)/(2*sigma^2)); %/(sigma*sqrt(2*pi)); 
        f2_shifted = A2*exp(-((t_full(i) - t_full).^2)/(2*sigma2^2)); %/(sigma*sqrt(2*pi)); 
	% Define function 1 on the same grid: 
        f1_here = A1*exp(-((t_full-t0).^2)/(2*sigma1^2)); %/(sigma*sqrt(2*pi)); 

	% At a given t, compute the product
	product = f2_shifted.*f1_here; 
	summation(i) = sum(product); 
	subplot(2,1,1),plot(t_full,f1_here,t_full,f2_shifted) 
        title('Convolution of two functions')
        ylabel('f1(x) & f2(x)')
	axis([-10,10,0,1])
	legend('f1(x)','f2(x)')
	grid on 
	d = input('Press a key to continue... ','%s\n')
	subplot(2,1,2),plot(t_full,summation,'g+') 
	axis([-10,10,0,6])
	grid on 
	xlabel('X') 
	ylabel('Convolution') 
	outfile = strcat('PNG/convolution_',num2str(i),'.png');
	print(outfile,'-dpng')
	pause(0.1)
end
figure(2) 
%c = conv(f1,f2); 
%plot(t_full,summation,'r-o',t_full,c,'b-+')
c_norm = summation/max(summation); 
plot(t,f1,'b',t,f2,'r',t_full,c_norm,'g')
grid on 
xlabel('X')
ylabel('f1(x) & f2(x)')
legend('f1(x)','f2(x)','Conv Function') 
outfile = 'PNG/convolution.png'; 
figure(2) 
print(outfile,'-dpng')

figure(3) 
% Let us look at a 2D convolution: (we will use Octave function conv2) 

N = 64;
Nx = N; 
Ny = N; 
x = linspace(1,Nx,Nx); 
y = linspace(1,Ny,Ny); 

% Lets convolve a point source with a Gaussian beam: 

% Generate the point Point source + Noise 
sky = 0.00*randn(Nx,Ny); 
xref = floor(Nx/2);
yref = ceil(Ny/2); 
sky(xref,yref) = sky(xref,yref) + 20; 

subplot(2,3,1),surface(x,y,sky')
shading('interp')
xlabel(' ')
ylabel('Y-axis ( Sky )')
axis([1,Nx,1,Ny])
axis('equal')

% Generate the function you want to convolve with: (the beam)
% We will generate a gaussian with sigma = 5 pixels
beam = abs(besselj(0,2*t).*besselj(0,2*t)');
subplot(2,3,2),surface(x,y,beam')
shading('flat')
xlabel(' ')
ylabel('( Dty Beam ) ')
axis([1,Nx,1,Ny])
axis('equal')

% Perform the convolution using octave function conv2: 
DtyIm = conv2(sky,beam,'same'); % Keep the output conv array of same size as input
% Now plot the convolution: 
subplot(2,3,3),surface(x,y,DtyIm')
shading('flat')
xlabel(' ')
ylabel('( Dty Image ) ')
axis([1,Nx,1,Ny])
axis('equal')


% Restored beam: 
sigma = 3; 
restoredBeam = exp(-((x - xref).^2 + (y - yref)'.^2)/(2*sigma^2));
% Put one more point source: 
xref2 = floor(Nx/3);
yref2 = ceil(3*Ny/4); 
sky(xref2,yref2) = sky(xref2,yref2) + 20; 
restoredIm = conv2(sky,restoredBeam,'same'); % Keep the output conv array of same size as input
pause(0.1) 

subplot(2,3,4),surface(x,y,sky')
shading('interp')
xlabel('X-axis')
ylabel('Y-axis ( Clean Components )')
axis([1,Nx,1,Ny])
axis('equal')
pause(0.1) 

subplot(2,3,5),surface(x,y,restoredBeam')
shading('flat')
xlabel('X-axis')
ylabel('( Restoring Beam ) ')
axis([1,Nx,1,Ny])
axis('equal')
pause(0.1) 

subplot(2,3,6),surface(x,y,restoredIm')
shading('flat')
xlabel('X-axis')
ylabel(' (Restored Image )')
%axis([1,Nx,1,Ny])
axis('equal')
pause(0.1) 

outfile = 'PNG/convolution_examples_imaging.png'; 
figure(3) 
print(outfile,'-dpng')
