% Cartoon of ptychography experiment. Scattering object is part of a
% tropical swallowtail butterfly wing.

clear
close all
figure('units','pixels','position',[0 0 1920 1080],'ToolBar','none');
set(0,'defaultfigurecolor',[1 1 1]);

vid = VideoWriter('ptychographyCartoon.mp4','MPEG-4');
vid.Quality = 100;
vid.FrameRate = 30;
open(vid);

myGold = [0.8 0.64 0];
lp = [-2 -2 1];
lp2 = [1 -2 -1]; 
pos1 = [-0.1 0 1 1];

% Dimensions in matlab units
incRadLength = 7; % Distance upstream to which incident radiation extends
FZPpos = -1; % Position upstream of FZP relative to focus @ (0,0,0)
samplePos = 0.35; % Position of sample downstream of focus 
DetPos = 10; % Position of detector downstream
halfDetWidth = 2; % Half width of detector
sampleHt = 1.5; % Height of butterfly sample

% Set up cones representing x-rays
% Incident x-rays, slightly divergent
t = 1:-0.001:0.8;
[xC1,yC1,zC1] = cylinder(t,800);
% Cone of focussed x-rays and divergent scattered x-rays
t = 0:0.001:1;
[xC2,yC2,zC2] = cylinder(t,800);

% Image with transparent background - for displaying in cartoon
[butterflyIm1,~,alphaMap] = imread('butterflyWing.png');
alphaMap = 0.8*(alphaMap); % alpha map of sample image
samplePixWidth = size(butterflyIm1,2); % Image width in pixels
samplePixHt = size(butterflyIm1,1); % Image height in pixels 
xPixCentre = floor(samplePixWidth/2); % Image central x-pixel position 
yPixCentre = floor(samplePixHt/2); % Image central y-pixel position 
imgAspectRatio = samplePixWidth/samplePixHt;

% Image with black background - for cropping and weighting with 2D Gaussian
% for performing an FT
butterflyIm2 = imread('butterflyWing2.png');
% imshow(butterflyIm1)

illumStartx = 525; % x-position of central pixel of first illuminated region
illumStarty = 650; % y-position of central pixel of first illuminated region
xScan = 1536; % Range of movement of sample in x-direction
yScan = 1536; % Range of movement of sample in y-direction
Delx = 64; % Shift in x between neighbouring positions
Dely = Delx; % Shift in y between neighbouring rows
xSteps = xScan/Delx; % Number of steps in x-direction = number of images - 1 per row
ySteps = yScan/Dely; % Number of steps in y-direction = number of rows - 1

% Parameters for Gaussian weighting function in units of pixels
cropHalfWidth = 128; % Half width of Gaussian weighting function for cropped region of original image
fwhm = 0.83*cropHalfWidth; % FWHM of Gaussian weighting function to avoid edge ringing in FTs
sigma = fwhm/(8*log(2))^0.5; % SD of Gaussian

% Set up grid for weighting Gaussian
[xGauss, yGauss] = meshgrid(-cropHalfWidth:cropHalfWidth, ...
    -cropHalfWidth:cropHalfWidth);
zGauss = zeros(size(xGauss)); % Initialize all Gaussian values to zero
% Calculate the weighting function zGauss for cropped image
for ii = 1:size((xGauss),1)
    for jj = 1:size((yGauss),1)
        rGauss = (xGauss(ii,jj)^2 + yGauss(ii,jj)^2)^0.5;
        zGauss(ii,jj) = exp(-rGauss^2/(2*sigma^2));
    end
end

% _________________________________________________________________________
%
% Now nested for-loops to scan sample and record FTs

for yy = 0:ySteps % Move down rows 
    for xx = 0:xSteps % Move along a row 
        hold off 
        subplot('position',pos1);
        newplot 
        if (round(yy/2) == yy/2) % Scan developing to the left 
            illumNowx = illumStartx + xx*Delx; % Central pixel position of illuminated area 
        else % Scan developing to the right 
            illumNowx = illumStartx + xScan - xx*Delx; % Central pixel position of illuminated area 
        end
        illumNowy = illumStarty + yy*Dely; 
        % Crop to rectangular region to be sampled
        butterflyImCropped = ...
            butterflyIm2(illumNowx-cropHalfWidth:illumNowx+cropHalfWidth, ...
            illumNowy-cropHalfWidth:illumNowy+cropHalfWidth,:);
        butterflyImCropped = im2double(butterflyImCropped); % Convert from uint8 to double
        buttImCropBW = rgb2gray(butterflyImCropped); % Convert from color to B/W
        % Weight the cropped image with zGauss
        buttCropWeight = buttImCropBW.*zGauss;
        % FT of weighted, cropped image
        FoTr = fft2(buttCropWeight);
        speckle = log(abs(fftshift(FoTr)));
        speckleMax = max(speckle(:));
        speckleMin = min(speckle(:));
        
        % FZP
        FZPth = 0:pi/100:2*pi;
        imax = 16;
        for jj = imax:-2:1
            hold on
            r1 = 0.125*(jj)^0.5;
            r2 = 0.125*(jj-1)^0.5;
            z_circle1 = r1*cos(FZPth);
            y_circle1 = r1*sin(FZPth);
            x_circle1 = 0.0*y_circle1;
            z_circle2 = r2*cos(FZPth);
            y_circle2 = r2*sin(FZPth);
            x_circle2 = 0.0*y_circle2;
            z_circle = [z_circle1 z_circle2];
            y_circle = [y_circle1 y_circle2];
            x_circle = [x_circle1 x_circle2];
            zone = fill3(x_circle+FZPpos,y_circle,z_circle,'k','LineStyle','none');
        end
        
        hold on
        
        % Incident, slightly divergent, radiation cone
        maxFZPrad = imax^0.5*0.125;
        incRad = surf(maxFZPrad*xC1,maxFZPrad*yC1,-FZPpos+incRadLength*zC1,...
            'FaceColor',myGold,'LineStyle',...
            'none','FaceLighting','flat','DiffuseStrength',1);
        rotate(incRad,[0 1 0],-90,[0,0,0]);
        % Vary alpha of divergent radiation in x-direction with offset so it
        % remains semitransparent for all elements of surf = divRad
        alpha(incRad,0.25*incRad.XData/halfDetWidth-2);
        
        % Focussed radiation cone after FZP
        focRad = surf(maxFZPrad*xC2,maxFZPrad*yC2,-FZPpos*zC2,...
            'FaceAlpha',0.25,'FaceColor',myGold,'LineStyle',...
            'none','FaceLighting','flat','DiffuseStrength',1);
        rotate(focRad,[0 1 0],-90,[0,0,0]);
        
        % Divergent radiation cone after focus
        maxFZPrad = imax^0.5*0.125;
        divRad = surf((1+halfDetWidth)*xC2,(1+halfDetWidth)*yC2,(DetPos)*zC2,...
            'FaceColor',myGold,'FaceAlpha',0.28,'LineStyle',...
            'none','FaceLighting','flat','DiffuseStrength',1);
        rotate(divRad,[0 1 0],90,[0,0,0]);
        % Vary alpha of divergent radiation in x-direction with offset so it
        % remains semitransparent for all elements of surf = divRad
        alpha(divRad,-0.25*divRad.XData/halfDetWidth-1.6);
        
        % Butterfly wing sample
        XX = samplePos*[1 1; 1 1]; % Position of sample in front of focus 
        % YY in matlab coords = xx in sample coords 
        pixShiftXnow = xPixCentre - illumNowx; 
        sampleShiftXnow = pixShiftXnow*sampleHt/samplePixWidth;
        YY = sampleHt*imgAspectRatio*[1 -1; 1 -1] - sampleShiftXnow; 
        % ZZ in matlab coords = yy in sample coords 
        pixShiftYnow = yPixCentre - illumNowy; 
        sampleShiftYnow = pixShiftYnow*sampleHt/samplePixWidth;
        ZZ = sampleHt*[1 1;-1 -1] - sampleShiftYnow;
        surf(XX,YY,ZZ,butterflyIm1,'FaceColor','texturemap','FaceAlpha','texturemap','AlphaData',alphaMap,'EdgeColor','none','AlphaDataMapping','none');
        
        % FT on detector
        XX = [DetPos DetPos; DetPos DetPos];
        YY = [halfDetWidth -halfDetWidth; halfDetWidth -halfDetWidth];
        ZZ = [halfDetWidth halfDetWidth;-halfDetWidth -halfDetWidth];
        surf(XX,YY,ZZ,speckle,'FaceColor','texturemap','EdgeColor','none');
        colormap jet
        caxis([speckleMin speckleMax])
        
        % Detector
        [V1,F1] = platonic_solid(2,0.88); % Cube
        V1(:,1) = 0.1*V1(:,1)+10.06; % Thin detector 1 r.l.u. downstream of (0 0 0)
        V1(:,2) = 2.*halfDetWidth*V1(:,2); % Thin detector 1 r.l.u. downstream of (0 0 0)
        V1(:,3) = 2.*halfDetWidth*V1(:,3); % Thin detector 1 r.l.u. downstream of (0 0 0)
        ps2 = patch('Faces',F1,'Vertices',V1,'FaceColor',[0.1 0.1 0.1],...
            'FaceAlpha',1,'EdgeColor','none','FaceLighting','flat',...
            'DiffuseStrength',1,'AmbientStrength',1,'SpecularStrength',0);
        [V1,F1] = platonic_solid(2,0.88); % Cube
        V1(:,1) = V1(:,1)+10.61; % Thin detector 1 r.l.u. downstream of (0 0 0)
        V1(:,2) = 2.1*halfDetWidth*V1(:,2); % Thin detector 1 r.l.u. downstream of (0 0 0)
        V1(:,3) = 2.1*halfDetWidth*V1(:,3); % Thin detector 1 r.l.u. downstream of (0 0 0)
        ps2 = patch('Faces',F1,'Vertices',V1,'FaceColor',[0.25 0.25 0.25],...
            'FaceAlpha',1,'EdgeColor','none','FaceLighting','flat',...
            'DiffuseStrength',1,'AmbientStrength',1,'SpecularStrength',0);
        
        view(-44,10);
        set(gca,'Projection','perspective');
        light('Position',lp,'Style','infinite');
        axis equal
        axis off 
        frame = getframe(gcf);
        writeVideo(vid,frame);        
    end
end

close(vid);