% Video cartoon of bunch compression in an XFEL 

clear; close all;

vid = VideoWriter('bunchCompression.mp4','MPEG-4');
vid.Quality = 100;
vid.FrameRate = 30;
open(vid);
figure('units','pixels','position',[0 0 1920 1080],'ToolBar','none');
set(0,'defaultfigurecolor',[1 1 1]);
set(gca,'linewidth',7);

myBlue = [0.4 0.44 0.73];
myCopper = [0.73 0.45 0.20];

LL = 10;

% Create a modulated cylinder that looks like a LINAC cavity
nCavs = 7; % Number of cavity units in LINAC 
t = 0:pi/10:nCavs*2*pi;
r = 0.5+(0.1 + (sin(t/2)).^2).^0.25;
[x,y,z] = cylinder(r,50);

% Delete half the cavity to see wot's goin' on, loik 
for ii = 141:-1:1
    for jj = 51:-1:1
        if (y(ii,jj)<0)
            x(ii,jj) = NaN;
            y(ii,jj) = NaN;
            z(ii,jj) = NaN;
        end
    end
end

viewAng = 20; 
cavLength = 7; % Length of cavity 
lChic = 7; % Length of chicane. 1 unit between cavity and chicane 
lp = [7 5 4]; % Lighting direction 
light('Position',lp,'Style','infinite');

nBunches = 84; % Number of bunches, each sort of representing a fuzzy electron 
ebunchxPosTemp = 0.22*randn(nBunches,1); % Distribution of bunches in x-direction
ebunchxPos = sort(ebunchxPosTemp); % Distribution of bunches in x-direction, sorted
ebunchyPos = 0.05*randn(nBunches,1); % Distribution of bunches in y-direction
ebunchzPos = 0.05*randn(nBunches,1); % Distribution of bunches in z-direction

ebunchxEndPos = 0.02*randn(nBunches,1); % Small variation in end positions
                                        % of bunches after bunch compression

M = 100; % Number of scatter points making up fuzzy cloud of bunch/electron 

% Initialize rgb numbers 
rcomp = zeros(nBunches,1);
gcomp = zeros(nBunches,1);
bcomp = zeros(nBunches,1);

nSteps = cavLength*20; % 20 steps per cavity unit 
cavStep = cavLength/nSteps;
viewStep = 20/((cavLength+lChic+2)/cavStep);

for jj = 0:cavStep:cavLength 
    viewAng = viewAng+viewStep;
    newplot
    hold on
    plot3([-1 cavLength+9.1], [-4 -4], [-4 -4],'LineWidth',0.1,'Color',[1 1 1]);
    plot3([-1 -1], [-2 -2], [-2 3],'LineWidth',0.1,'Color',[1 1 1]);
    cavity = surf(x,y,z*cavLength,'FaceAlpha',1,'FaceColor',myCopper,'LineStyle',...
        'none','FaceLighting','gouraud','DiffuseStrength',1);
    rotate(cavity,[0 1 0],90,[0,0,0]); % Transform from vertical to horizontal with
                                       % LINAC axis along the x-axis
    
    % Include arrows to indicate directions of magnetic fields in the chicane
    mArrow3([8.5 1 0],[8.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([10.5 -1 1.25],[10.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([12.5 -1 1.25],[12.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([14.5 1 0],[14.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    
    for ii = 1:nBunches 
        % Create coordinates for scatter curve, fresh for each frame - this
        % provides small changes to the appearance of the electrons,
        % thereby adding an element of verisimilitude, dontchaknow 
        ebunchx = 0.01*randn(M,1); 
        ebunchy = 0.01*randn(M,1);
        ebunchz = 0.01*randn(M,1); 
        % Determine the rgb components the electrons should have at the end
        % of the cavity 
        for kk = 1:ceil(nBunches/2)
            rcomp(kk) = 0;
            gcomp(kk) = 2*kk/nBunches;
            bcomp(kk) = 1 - 2*kk/nBunches;
        end
        for kk = ceil(nBunches/2)+1:nBunches
            rcomp(kk) = 2*(kk-ceil(nBunches/2))/nBunches;
            gcomp(kk) = 1 - 2*(kk-ceil(nBunches/2))/nBunches;
            bcomp(kk) = 0;
        end
        % Determine the rgb components as the electrons move through the
        % cavity 
        rcompnow = rcomp*jj/cavLength;
        gcompnow = 1 - (1-gcomp)*jj/cavLength;
        bcompnow = bcomp*jj/cavLength;
        
        % Plot out each bunch/electron 
        ebunch(ii) = scatter3(ebunchx+ebunchxPos(ii)+jj,...
            ebunchy+ebunchyPos(ii),ebunchz+ebunchzPos(ii),2,...
            'filled','MarkerFaceColor',[rcompnow(ii) gcompnow(ii) bcompnow(ii)]);
        alpha = 0.7;
        set(ebunch(ii), 'MarkerEdgeAlpha',0,'MarkerFaceAlpha',alpha);
    end
    
    lp = [7 5 4];
    light('Position',lp,'Style','infinite');
    set(gca, 'Projection','perspective');
    set(gca,'View',[viewAng,10]);
    xlim([-1 LL]);
    ylim([-LL/2.5 LL/2.5]);
    zlim([-LL/2.5 LL/2.5]);
    axis off
    axis equal
    hold off
    frame = getframe(gcf);
    writeVideo(vid,frame);
end

% Move forward another 1 unit past end of cavity

for jj = cavStep:cavStep:1
    viewAng = viewAng+viewStep;
    newplot
    hold on
    plot3([-1 cavLength+9.1], [-4 -4], [-4 -4],'LineWidth',0.1,'Color',[1 1 1]);
    plot3([-1 -1], [-2 -2], [-2 3],'LineWidth',0.1,'Color',[1 1 1]);
    cavity = surf(x,y,z*cavLength,'FaceAlpha',1,'FaceColor',myCopper,'LineStyle',...
        'none','FaceLighting','gouraud','DiffuseStrength',1);
    rotate(cavity,[0 1 0],90,[0,0,0]);
    
    mArrow3([8.5 1 0],[8.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([10.5 -1 1.25],[10.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([12.5 -1 1.25],[12.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([14.5 1 0],[14.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    
    % Time to plot those babes 
    for ii = 1:nBunches
        ebunch(ii) = scatter3(ebunchx+ebunchxPos(ii)+cavLength+jj,...
            ebunchy+ebunchyPos(ii),ebunchz+ebunchzPos(ii),2,...
            'filled','MarkerFaceColor',[rcompnow(ii) gcompnow(ii) bcompnow(ii)]);
        alpha = 0.7;
        set(ebunch(ii), 'MarkerEdgeAlpha',0,'MarkerFaceAlpha',alpha);
    end 
    
    % Visual/image fluff 
    lp = [7 5 4];
    light('Position',lp,'Style','infinite');
    set(gca, 'Projection','perspective');
    set(gca,'View',[viewAng,10]);
    xlim([-1 LL]);
    ylim([-LL/2.5 LL/2.5]);
    zlim([-LL/2.5 LL/2.5]);
    axis off
    axis equal
    hold off
    frame = getframe(gcf);
    writeVideo(vid,frame);
end

% Execute curved path through chicane

ampRed = 0.9; % Amplitude of curved path for lowest-energy bunch
ampBlue = 0.64; % Amplitude of curved path for highest-energy bunch 
% Trace that follows electrons for guide to the eye 
xTrace = 0:cavStep:lChic; 
yTrace = 0.0*xTrace-0.001; 
redTrace = -ampRed*cos(2*pi*xTrace/lChic)+ampRed/2; 
blueTrace = -ampBlue*cos(2*pi*xTrace/lChic)+ampBlue/2; 

for  jj = cavStep:cavStep:lChic
    viewAng = viewAng+viewStep;
    newplot
    hold on
    plot3([-1 cavLength+9.1], [-4 -4], [-4 -4],'LineWidth',0.1,'Color',[1 1 1]);
    plot3([-1 -1], [-2 -2], [-2 3],'LineWidth',0.1,'Color',[1 1 1]);
    cavity = surf(x,y,z*cavLength,'FaceAlpha',1,'FaceColor',myCopper,'LineStyle',...
        'none','FaceLighting','gouraud','DiffuseStrength',1);
    rotate(cavity,[0 1 0],90,[0,0,0]);
    
    % Magnetic-field arrows 
    mArrow3([8.5 1 0],[8.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([10.5 -1 1.25],[10.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([12.5 -1 1.25],[12.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([14.5 1 0],[14.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5); 
    
    % Plot out path of low- and high-energy electrons through chicane 
    if (jj>=0.25)
        redTracePlot = plot3(cavLength+1+xTrace(1:round(1+(jj-0.25)/cavStep)),...
            yTrace(1:round(1+(jj-0.25)/cavStep)),...
            ampRed/2+redTrace(1:round(1+(jj-0.25)/cavStep)),...
            'LineWidth',4,'Color',[1 0 0]);
        redTracePlot.Color(4) = 0.25;
        blueTracePlot = plot3(cavLength+1+xTrace(1:round(1+(jj-0.25)/cavStep)),...
            yTrace(1:round(1+(jj-0.25)/cavStep)),...
            ampBlue/2+blueTrace(1:round(1+(jj-0.25)/cavStep)),...
            'LineWidth',4,'Color',[0 0 1]);
        blueTracePlot.Color(4) = 0.25;

    end 
    for ii=1:nBunches
        posCavEnd = cavLength + ebunchxPos(ii); % Position of bunch at end of cavity
        % Amplitude of curve for each electron/bunch according to its
        % energy/position 
        amp(ii) = ampBlue + (ampRed-ampBlue)*ii/nBunches; 
        % Length of period from perspective of the electron. Close but not
        % identical to lChic 
        Ltot(ii) = lChic-ebunchxPos(ii)+ebunchxEndPos(ii);
        % Now let's plot those suckers out! 
        ebunch(ii) = scatter3(ebunchx+cavLength+1+... % Scatter coords + cavity length + 1 unit 
            ebunchxPos(ii)*(1-jj/lChic)+ ... % Reduce initial scatter
            ebunchxEndPos(ii)*jj/lChic ... % Grow in final small variation in axial positions
            +jj, ... % Move one step forward
            ebunchy+ebunchyPos(ii), ...
            ebunchz+ebunchzPos(ii)-amp(ii)*cos(2*pi*jj/Ltot(ii))+amp(ii), ...
            2,'filled','MarkerFaceColor',...
            [rcompnow(ii) gcompnow(ii) bcompnow(ii)]);
        alpha = 0.7;
        set(ebunch(ii), 'MarkerEdgeAlpha',0,'MarkerFaceAlpha',alpha);
    end
    lp = [7 5 4];
    light('Position',lp,'Style','infinite');
    set(gca, 'Projection','perspective');
    set(gca,'View',[viewAng,10]);
    xlim([-1 LL]);
    ylim([-LL/2.5 LL/2.5]);
    zlim([-LL/2.5 LL/2.5]);
    axis off
    axis equal
    hold off
    frame = getframe(gcf);
    writeVideo(vid,frame);
end

% Move forward one final unit

for jj = cavStep:cavStep:1
    viewAng = viewAng+viewStep;
    newplot
    hold on
    plot3([-1 cavLength+9.1], [-4 -4], [-4 -4],'LineWidth',0.1,'Color',[1 1 1]);
    plot3([-1 -1], [-2 -2], [-2 3],'LineWidth',0.1,'Color',[1 1 1]);
    cavity = surf(x,y,z*cavLength,'FaceAlpha',1,'FaceColor',myCopper,'LineStyle',...
        'none','FaceLighting','gouraud','DiffuseStrength',1);
    rotate(cavity,[0 1 0],90,[0,0,0]);
    
    mArrow3([8.5 1 0],[8.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([10.5 -1 1.25],[10.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([12.5 -1 1.25],[12.5 1 1.25],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    mArrow3([14.5 1 0],[14.5 -1 0],'color',myBlue,'stemWidth',0.1,...
        'tipWidth',0.2,'facealpha',0.5);
    % Finish off trace curves 
    if (jj<0.25)
        endIndex = round((lChic+(jj-0.25))/cavStep);
        redTracePlot = plot3(cavLength+1+xTrace(1:endIndex),...
            yTrace(1:endIndex),...
            ampRed/2+redTrace(1:endIndex),...
            'LineWidth',4,'Color',[1 0 0]);
        redTracePlot.Color(4) = 0.25;
        blueTracePlot = plot3(cavLength+1+xTrace(1:endIndex),...
            yTrace(1:endIndex),...
            ampBlue/2+blueTrace(1:endIndex),...
            'LineWidth',4,'Color',[0 0 1]);
        blueTracePlot.Color(4) = 0.25;
    else
        redTracePlot = plot3(cavLength+1+xTrace,...
            yTrace,...
            ampRed/2+redTrace,...
            'LineWidth',4,'Color',[1 0 0]);
        redTracePlot.Color(4) = 0.25;
        blueTracePlot = plot3(cavLength+1+xTrace,...
            yTrace,...
            ampBlue/2+blueTrace,...
            'LineWidth',4,'Color',[0 0 1]);
        blueTracePlot.Color(4) = 0.25;
    end
    
    for ii = 1:nBunches
        ebunch(ii) = scatter3(ebunchx+cavLength+1+ebunchxEndPos(ii)+lChic+jj,...
            ebunchy+ebunchyPos(ii),ebunchz+ebunchzPos(ii),2,...
            'filled','MarkerFaceColor',[rcompnow(ii) gcompnow(ii) bcompnow(ii)]);
        alpha = 0.7;
        set(ebunch(ii), 'MarkerEdgeAlpha',0,'MarkerFaceAlpha',alpha);
    end
    lp = [7 5 4];
    light('Position',lp,'Style','infinite');
    set(gca, 'Projection','perspective');
    set(gca,'View',[viewAng,10]);
    xlim([-1 LL]);
    ylim([-LL/2.5 LL/2.5]);
    zlim([-LL/2.5 LL/2.5]);
    axis off
    axis equal
    hold off
    frame = getframe(gcf);
    writeVideo(vid,frame);
end

% Output the movie as an mpg file
close(vid);
