% Task 2: First, create a MATLAB script, that simulates waves through the system. This is similar to the simulations we did in class, but now we have multiple elements, so you have to use successive image propagations through multiple planes. (Input image -> L1 -> filter -> L2 -> Output screen) Explain the exact methods/approximations you use. 

X_dim = 200; % problem size x
Y_dim = 400; % problem size y
d_x = 50.0;  % x discretization 
d_y = 50.0;  % y discretization 
[X,Y] = meshgrid((1:X_dim)*d_x, (1:Y_dim)*d_y); % mesh
lambda = 500;   % wavelength in vacuum
y_b = 100;      % lens 1 position
y_b2 = 200;     % slit position         
y_b3 = 300;     % lens 2 position

w_slit = lambda;   % slit width
phi = pi/180*0;    % incident angle 

f = 100*d_y; % focal length
R = ((1:X_dim)-X_dim/2)*d_x;

k1 = 2*pi/lambda*([sin(phi),cos(phi)]); % wavevector
M_1 = exp(1i*(k1(1)*X+k1(2)*Y)); % incident wave

Z = circ(R/f).*exp(-1i*2*pi/lambda/2/f*R.^2);      % lens 1
Z2 = (abs(R)<w_slit);                              % single slit
Z3 = circ(R/f).*exp(-1i*2*pi/(lambda)/2/f*R.^2);   % lens 2

M_2 = 0*M_1;
for x = 1:X_dim
    if Z(x)~=0
        M_2 = M_2 + huygens_source(Z(x)*M_1(y_b,x),Y,X,y_b,x,d_y,d_x,lambda);
    end
end

M_3 = 0*M_2;
for x = 1:X_dim
    if Z2(x)~=0
        M_3 = M_3 + huygens_source(Z2(x)*M_2(y_b2,x),Y,X,y_b2,x,d_y,d_x,lambda);
    end
end

M_4 = 0*M_3;
for x = 1:X_dim
    if Z3(x)~=0
        M_4 = M_4 + huygens_source(Z3(x)*M_3(y_b3,x),Y,X,y_b3,x,d_y,d_x,lambda);
    end
end

M_M = M_1;
M_M(y_b+1:Y_dim,:) = M_2(y_b+1:Y_dim,:);
M_M(y_b2+1:Y_dim,:) = M_3(y_b2+1:Y_dim,:);
M_M(y_b3+1:Y_dim,:) = M_4(y_b3+1:Y_dim,:);

figure(1)
imagesc(X(1,:)/1000, Y(:,1)/1000,real(M_M))
axis image
set(gca,'Ydir','reverse')
xlabel('x axis')
ylabel('y axis')
title('Plot the waves')
colormap gray
hold on
for i = 1:X_dim-1
    plot(X(1,i:i+1)/1000,[d_y*y_b/1000 d_y*y_b/1000],'-','Color',[0 0 0 1-angle(Z(i))/2/pi-0.5],'Linewidth',5) % use this for the lens
    plot(X(1,i:i+1)/1000,[d_y*y_b2/1000 d_y*y_b2/1000],'-','Color',[0 0 0 1-Z2(i)],'Linewidth',5)
    plot(X(1,i:i+1)/1000,[d_y*y_b3/1000 d_y*y_b3/1000],'-','Color',[0 0 0 1-angle(Z3(i))/2/pi-0.5],'Linewidth',5)
end
ylim([0 Y_dim*d_y/1000])
hold off

figure(2)
imagesc(X(1,:)/1000, Y(:,1)/1000,abs(M_M))
set(gca,'Ydir','reverse')
xlabel('x axis')
ylabel('y axis')
title('Plot the amplitude of the waves')
colormap gray
hold on
for i = 1:X_dim-1
    plot(X(1,i:i+1)/1000,[d_y*y_b/1000 d_y*y_b/1000],'-','Color',[0 0 0 1-angle(Z(i))/2/pi-0.5],'Linewidth',5)
    plot(X(1,i:i+1)/1000,[d_y*y_b2/1000 d_y*y_b2/1000],'-','Color',[0 0 0 1-Z2(i)],'Linewidth',5)
    plot(X(1,i:i+1)/1000,[d_y*y_b3/1000 d_y*y_b3/1000],'-','Color',[0 0 0 1-angle(Z3(i))/2/pi-0.5],'Linewidth',5)
end
ylim([0 Y_dim*d_y/1000])
hold off

figure(3)
imagesc(X(1,:)/1000,[0 1],angle(Z))
set(gca,'Ydir','normal')
hold on
plot(X(1,:)/1000,angle(Z),'Linewidth',2)
xlabel('x axis')
ylabel('y axis')
title('Plot the zone plates')
colormap gray
hold off