%This script will create an automated animation of a particle evolution
%from an Impact simulation by reading-in text files containing particle
%data at various time slices. Also, this script can create an interactive
%plot where the user can adjust the time slice with a slider object. 
%
%The particle data for each time slice must be in 6-column format where the
%columns represent (x,px,y,py,z,pz) which are generated from Impact. Each
%file must be placed in a folder specified by ifolder in the user section.
%Next, the files must be named 'fort.xxxx' where xxxx is a 4 digit number
%representing the time slice.
%
%The particles can be plotted as filled circles (faster) or shaded 3d
%bubbles (slower) using the included script: bubbleplot3.m written by
%Peter (PB) Bodin ( pbodin@kth.se ). It is recommended to use the bubble
%plotting method for animation mode and the circles plotting method for
%interactive mode for improved responsiveness.
%
%Written by David Bizzozero
%April 27, 2021

%% User inputs section
%Plot type as scatter3 'circles' or 3d 'bubbles'. For best results, use
%the plot type 'circles' in interactive mode and 'bubbles' for a video
%output of an animation.

ptype = 'circles';  %Ptype can be: 'circles', 'bubbles'
mode = 'single';    %Mode can be: 'animation', 'interactive', 'single'

ifolder = 'fort_data';	%Folder name for particles if mode is not 'single'
ifile = 'fort.50';  %File name if mode is set to 'single'
nparts = 9984;      %Maximum number of particles to plot (not to exceed 
                    %particle data file length)

%Set up video parameters for video animation recording
v_flag = false;     %Toggle recording to video
v_name = 'particle_test';    %Video filename (filename must be free)
v_fps = 30;         %Video framerate (default 30)

%Optional plotting variables (can be omitted for defaults or automatic)
bc = 'k';	%Plot backround color or rgb 3-vector (default: 'k')
tc = 'w';	%Plot text color or rgb 3-vector (default: 'w')
pc = 'r';	%Particle color code or rgb 3-vector (e.g. default: 'r')
xlims = [-0.005, 0.005];	%X limits for plotting box (default: 'auto')
ylims = [-0.005, 0.005];	%Y limits for plotting box (default: 'auto')
zlims = [-0.024, 0.024];	%Z limits for plotting box (default: 'auto')
iview = [30,20];           %Initial viewing angle (default: [30,20])
xplot = 'x';    %Variable to plot along x-axis: 'x', 'px', 'y', or 'py'
yplot = 'y';   %Variable to plot along y-axis: 'x', 'px', 'y', or 'py'

%Optional plotting variables for bubble plots
rb = 'auto';    %Bubble plot particle radius (default: 'auto')

%Optional plotting variables for animation mode only
cflag = true;	%Flag to toggle plotting backgr. cyl. (default: true)
rc = 0.012;     %Backgr. cyl. radius for anim. plots (default: 'auto')
tmin = -3*pi/5;	%Backgr. cyl. angle min for anim. plots (default: -pi/2)
tmax = 3*pi/5;	%Backgr. cyl. angle max for anim. plots (default: pi/2)
dt = pi/12;     %Backgr. cyl. angle step for anim. plots (default: pi/12)
sc = 'interp';	%Backgr. cyl. face color/shading for anim. plots
sa = 1.0;       %Backgr. cyl. face opacity for anim. plots
ec = 'k';       %Backgr. cyl. edge color for anim. plots
ea = 1.0;       %Backgr. cyl. edge opacity for anim. plots
xtick = 'auto';     %Sets plot tick marks along x axis for anim. plots
ytick = 'auto';     %Sets plot tick marks along y axis for anim. plots
ztick = 0:0.02:14;  %Sets plot tick marks along z axis for anim. plots

%% Check for common errors

if ~exist('xplot','var'),	xplot = 'x';	end
if ~exist('yplot','var'),	yplot = 'y';	end
if (~strcmpi(xplot,'x') || ~strcmpi(yplot,'y')) && cflag == true
    disp('Cylinder plot only available if xplot = ''x'' and yplot = ''y''')
    cflag = false;
end
if (~strcmpi(xplot,'x') || ~strcmpi(yplot,'y')) && strcmpi(ptype,'bubbles')
    disp('Bubble plot only available if xplot = ''x'' and yplot = ''y''')
    ptype = 'circles';
end
if strcmpi(xplot,yplot)
    disp('Error: xplot and yplot axes variables must be different.')
    return
end
if ~exist('bubbleplot3.m','file') && strcmpi(ptype,'bubbles')
    disp('Error: bubbleplot3.m not found, switching ptype to ''circles''')
    ptype = 'circles';
end

%% Load particle files and set up particle arrays

switch mode
    case {'animation','interactive'}
        if ~exist('load_flag','var')	%Read-in Impact particle data once
            load_flag = 0;	%Create flag to save on unnessary reloading
        end
        if load_flag ~= 1
            ifiles = dir(fullfile(ifolder, 'fort.????'));	%File list
            ind_start = eval(strrep(ifiles(1).name,'fort.',''));  %1st file
            ind_end = eval(strrep(ifiles(end).name,'fort.',''));  %End file
            nsamps = length(ifiles);	%Total number of samples

            fortdata = zeros(nparts,6,nsamps);
            pflag = 10;
            disp('Loading particle data files...')
            for i = 1:nsamps	%Loop over samples and read-in data by file
                fortdata(:,:,i) = dlmread(ifiles(i).name,'',[0,0,nparts-1,5]);
                if i/nsamps > pflag/100	%Print output every 10%
                    disp(string(pflag)+'% complete')
                    pflag = pflag+10;
                end
            end
            disp('Particle data loading complete.')
            load_flag = 1;	%Flag set to 1 if time-sliced data
        end
        
    case 'single'
        fortdata = dlmread(ifile,'',[0,0,nparts-1,5]);
        disp('Particle data loading complete.')
        load_flag = 2;	%Flag set to 2 if single data
end

switch xplot
    case 'x'
        xp = squeeze(fortdata(:,1,:));
    case 'px'
        xp = squeeze(fortdata(:,2,:));
    case 'y'
        xp = squeeze(fortdata(:,3,:));
    case 'py'
        xp = squeeze(fortdata(:,4,:));
end
switch yplot
    case 'x'
        yp = squeeze(fortdata(:,1,:));
    case 'px'
        yp = squeeze(fortdata(:,2,:));
    case 'y'
        yp = squeeze(fortdata(:,3,:));
    case 'py'
        yp = squeeze(fortdata(:,4,:));
end
zp = squeeze(fortdata(:,5,:));  %Z-axis always 'z' coordinate

%Set defaults for various plotting options if nonexistent
if ~exist('bc','var'),	bc = 'k'; end
if ~exist('tc','var'),	tc = 'w'; end
if ~exist('pc','var'),  pc = 'r';    end
if ~exist('fntsz','var'),	fntsz = 20;     end
if ~exist('fntszn','var'),	fntszn = 14;	end
if ischar(pc)    %Convert color code to rgb vector (used with bubbles)
    pc = rem(floor((strfind('kbgcrmyw',pc)-1)*[0.25 0.5 1]),2);
end
if ~exist('rb','var'),	rb = 'auto';	end
if strcmpi(rb,'auto')
    rb = max(max([abs(xp(:)),abs(yp(:))]))/50;
end
if ~exist('cflag','var'),	cflag = true;	end
if ~exist('iview','var'),	iview = [30,20];	end
if ~exist('xlims','var'),	xlims = 'auto';	end
if strcmpi(xlims,'auto'),	xlims = [min(xp(:)), max(xp(:))];	end
if ~exist('ylims','var'),	ylims = 'auto';	end
if strcmpi(ylims,'auto'),	ylims = [min(yp(:)), max(yp(:))];	end
if ~exist('zlims','var'),	zlims = 'auto';	end
if strcmpi(zlims,'auto')
    zlims = [min(min(zp-mean(zp,1))), max(max(zp-mean(zp,1)))];
end
if ~exist('xtick','var'), xtick = 'auto'; end
if ~exist('ytick','var'), ytick = 'auto'; end
if ~exist('ztick','var'), ztick = 'auto'; end

%Pack all plotting options into struct
opts = struct;
opts.xplot = xplot;
opts.yplot = yplot;
opts.bc = bc;
opts.tc = tc;
opts.pc = pc;
opts.iview = iview;
opts.fntsz = fntsz;
opts.fntszn = fntszn;
opts.lims = [zlims;xlims;ylims];
opts.rb = rb;
opts.ptype = ptype;
opts.mode = mode;
opts.cflag = cflag;
opts.xtick = xtick;
opts.ytick = ytick;
opts.ztick = ztick;

%Reset values when rerunning script
clear xplot yplot pc fntsz fntszn bc tc pc rb cflag iview ...
    xlims ylims zlims xtick ytick ztick


%% Plotting section by mode option

f = figure('position',get(0,'screensize'),'color',opts.bc);
ax = axes('box','on','parent',f,'position',[0.1 0.25 0.8 0.68],...
    'color',opts.bc,'Xcolor',opts.tc,'Ycolor',opts.tc,...
    'Zcolor',opts.tc,...
    'Xgrid','on','Ygrid','on','Zgrid','on',...
    'fontname','times','fontsize',opts.fntszn);
hold(ax,'on')
view(opts.iview)
plot_parts(xp,yp,zp,1,ax,opts)
fix_aspect_ratio(ax,opts)

switch mode
case 'animation'
    
    if opts.cflag
        if ~exist('rc','var'),	rc = max(abs([xp(:);yp(:)]))*1.5;	end
        if ~exist('tmin','var'),	tmin = pi/2;    end
        if ~exist('tmax','var'),	tmax = -pi/2;   end
        if ~exist('dt','var'),	dt = pi/12;	end
        if ~exist('sc','var'),	sc = 'interp';	end
        if ~exist('sa','var'),	sa = 1.0;	end
        if ~exist('ec','var'),	ec = 'k';	end
        if ~exist('ea','var'),	ea = 1.0;	end
        opts.sc = sc;
        opts.ec = ec;
        [Tris,X,Y,Z,opts.fc] = make_cylinder(rc,min(zp(:)),max(zp(:)),tmin,tmax,dt);
        ntris = size(Tris,1);	%Total number of triangles in cylinder
        opts.nt = sum(Z==min(Z));	%Number vertices per slice
        if opts.fc
            opts.nz = ntris/(4*opts.nt);	%Number of full dual-layer slices
            opts.dz = Z(2*opts.nt+1)-Z(1);	%Width of full dual-layer slice
            opts.sa = max(sa,0.75);
            opts.ea = max(ea,0.75);
        else
            opts.nz = ntris/(4*opts.nt-4);	%Number of full dual-layer slices
            opts.dz = Z(2*opts.nt)-Z(1);	%Width of full dual-layer slice
            opts.sa = sa;
            opts.ea = ea;
        end
        clear rc tmin tmax dt sc sa ec ea
    end
    
    if v_flag	%Set up video writer object
        video_obj = VideoWriter(v_name,'MPEG-4');
        video_obj.FrameRate = v_fps;
        open(video_obj)
    end
    
    xticks(ax,opts.ztick)
    yticks(ax,opts.xtick)
    zticks(ax,opts.ytick)
    xtickformat('%.3f')
    ytickformat('%.3f')
    ztickformat('%.3f')
    
    for i = 1:nsamps
        plot_parts(xp,yp,zp,i,ax,opts)
        if opts.cflag
            hold on
            plot_cylinder(Tris,X,Y,Z,mean(zp(:,i)),opts)
        end
        drawnow
        if v_flag
            writeVideo(video_obj,getframe(gcf))
        end
    end
    
    if v_flag; close(video_obj); end
    
case 'interactive'
    
    %Create slider object for callback to change sample
    hsl = uicontrol('Parent',f,'Style','slider','Min',1,'Max',nsamps,...
                    'SliderStep',[1 10]/(nsamps-1),'Value',1,...
                    'Position',[680,120,700,23],...
                    'FontName','Times','FontSize',opts.fntsz);
    bl1 = uicontrol('Parent',f,'Style','text','Position',[620,90,60,60],...
                    'String','1','BackgroundColor',f.Color,...
                    'FontName','Times','FontSize',opts.fntsz,'ForeGroundColor',opts.tc);
    bl2 = uicontrol('Parent',f,'Style','text','Position',[1400,90,60,60],...
                    'String',num2str(nsamps),'BackgroundColor',f.Color,...
                    'FontName','Times','FontSize',opts.fntsz,'ForeGroundColor',opts.tc);
    bl3 = uicontrol('Parent',f,'Style','text','Position',[880,60,300,60],...
                    'String','Sample','BackgroundColor',f.Color,...
                    'FontName','Times','FontSize',opts.fntsz,'ForeGroundColor',opts.tc);
    
    %Use selected plotting routine with callback listeners
    set(hsl,'Callback',@(hObject,eventdata) ...
        plot_parts(xp,yp,zp,round(get(hObject,'Value')),ax,opts))
    addlistener(hsl,'ContinuousValueChange',@(hObject,eventdata) ...
        plot_parts(xp,yp,zp,round(get(hObject,'Value')),ax,opts));
    
end

%% Plotting and other routines (called by plotting section)

function plot_parts(xp,yp,zp,n,ax,opts)
%This function plots particles given in the arrays xp, yp, and zp which are
%indexed by time slice (array column) and plots time slice n using 2D
%filled cicles or 3D bubbles for the particles. Depending on the mode
%{'animation', 'single', or 'interactive'}, the routine will adjust the
%title and axes accordingly.

[az,el] = view;
x = xp(:,n);
y = yp(:,n);
z = zp(:,n);
mz = mean(z);
switch opts.mode
    case 'animation'
        lims = [(mz+opts.lims(1,1)),(mz+opts.lims(1,2));opts.lims(2:3,1:2)];
    otherwise
        z = z-mz;
        lims = opts.lims;
end
switch opts.ptype
    case 'circles'
        set(ax,'NextPlot','replacechildren');
        scatter3(z,x,y,'Marker','o','MarkerFaceColor',opts.pc,...
            'MarkerEdgeColor','k')
    case 'bubbles'
        cla
        rb = opts.rb*ones(size(z));
        bubbleplot3(z,x,y,rb,repmat(opts.pc,length(z),1));
        camlight right
        lighting phong
end
xlim(lims(1,:)), ylim(lims(2,:)), zlim(lims(3,:))
view([az,el])
switch opts.mode
    case 'animation'
        title(ax,'$\bar z$ = '+string(num2str(mean(zp(:,n)),'%.4f'))+' m',...
            'interpreter','latex','fontsize',opts.fntsz,'color',opts.tc)
    case 'single'
        title(ax,'Particles at $\bar z$ = '+string(mean(zp(:,n)))+' m',...
            'interpreter','latex','fontsize',opts.fntsz,'color',opts.tc)
        drawnow
    case 'interactive'
        title(ax,'Particle sample '+string(n)+...
            ' at $\bar z$ = '+string(mean(zp(:,n)))+' m',...
            'interpreter','latex','fontsize',opts.fntsz,'color',opts.tc)
        drawnow
end
end

function plot_cylinder(Tris,X,Y,Z,mz,opts)
%This function plots a triangular surface plot of the cylinder defined by
%the vertices X,Y,Z and connectivity matrix Tris. The argument mz is the
%shift in the z coordinate to limit the number of triangles displayed only
%to the plotting box (for faster rendering).

%Set window limits centered about mz from limits defined in opts
lims = [(mz+opts.lims(1,1)),(mz+opts.lims(1,2))];

dz = opts.dz;   nz = opts.nz;   nt = opts.nt;
%Indices for triangles which are visible in window range
if opts.fc
    i1 = max(floor(lims(1)/dz),0)*4*nt+1;
    i2 = min(ceil(lims(2)/dz),nz)*4*nt;
else
    i1 = max(floor(lims(1)/dz),0)*(4*nt-4)+1;
    i2 = min(ceil(lims(2)/dz),nz)*(4*nt-4);
end
trisurf(Tris(i1:i2,:),Z,X,Y,...
    'EdgeColor',opts.ec,'EdgeAlpha',opts.ea,...
    'FaceColor',opts.sc,'FaceAlpha',opts.sa);
lighting phong
    
end

function [Tris,X,Y,Z,fc] = make_cylinder(r,z1,z2,t1,t2,dt,dz)
%This script will generate a cylinder or arc of a cylinder of vertices
%given by the arrays X, Y, Z, and the connectivity matrix Tris. The flag
%fc denotes a full circle cylinder; if false, then the triangulation
%represents an arc of a circle cross-section.

if nargin<4 %Default to full cylinder (for surfaces with 0<facealpha<1)
    t1 = 0;
    t2 = 2*pi;
end
if nargin<6 %Default dt to pi/12 angular resolution
    dt = pi/12;
end
if nargin<7 %Default dz for approx. equilateral triangles
    dz = sqrt(3)*r*dt;
end

if (t2-t1)>=(2*pi-dt)   %If arc angle is too large, make full circle
    t1 = 0;
    t2 = 2*pi;
    fc = true;  %Flag for full circle (vertices wrap around)
else
    fc = false; %Flag for arc of a circle (vertices do not wrap around)
end

nt = ceil((t2-t1)/dt);	%Number of vertices around circle or arc
dt = (t2-t1)/nt;        %Angle between two consecutive vertices
if fc
    ts1 = (t1:dt:t2-dt)';       %Angles for vertices in circle
    ts2 = ts1+dt/2;             %Shifted angles for mid-layer by dt/2
else
    nt = nt+1;
    ts1 = (t1:dt:t2)';          %Angles for vertices in circle
    ts2 = ts1(1:end-1)+dt/2;	%Shifted angles for mid-layer by dt/2
end

nz = ceil((z2-z1)/dz);	%Number of full dual-layers in z from z1 to z2
dz = (z2-z1)/nz;        %Thickness of a full dual-layer
zs = (z1:dz/2:z2)';     %Z-coords. of layers (odd num. of layers)

xs1 = r*cos(ts1);   ys1 = r*sin(ts1);   %Coords. of circle vertices
xs2 = r*cos(ts2);   ys2 = r*sin(ts2);   %Coords. of shifted vertices

%Set up triangle connectivity matrix for up to (4*nt*nz) triangles
if fc
    Tris = zeros(4*nt*nz,3);
else
    Tris = zeros(4*(nt-1)*nz,3);
end

%Define all X, Y, Z corrdinates of triangle vertices
X = [xs1;repmat([xs2;xs1],nz,1)];
Y = [ys1;repmat([ys2;ys1],nz,1)];
Z = repelem(zs,nt);
if ~fc	%Remove extra vertices if not full circle cross section
    Z((2*nt):(2*nt):end)=[];    
end

%Create triangles based on vertices for circle or arc of circle
for n = 1:nz
    if fc
        %Offsets for given dual-layer of triangles
        zoff = (n-1)*4*nt;
        toff = (n-1)*2*nt;
        
        %Indices of vertices for layer 1: (p1,2), 2: (p3,4), 3: (p5,6)
        p1 = toff + (1:nt);
        p2 = toff + ([2:nt,1]);
        p3 = toff + ((nt+1):(2*nt));
        p4 = toff + ([2*nt,(nt+1):(2*nt-1)]);
        p5 = toff + ((2*nt+1):(3*nt));
        p6 = toff + ([(2*nt+2):(3*nt),(2*nt+1)]);
        
        %Define triangles based on layers of vertices
        Tris((1+zoff):(nt+zoff),:) = [p1',p2',p3'];
        Tris((nt+1+zoff):(2*nt+zoff),:) = [p3',p4',p1'];
        Tris((2*nt+1+zoff):(3*nt+zoff),:) = [p4',p3',p5'];
        Tris((3*nt+1+zoff):(4*nt+zoff),:) = [p6',p5',p3'];
    else
        %Offsets for given dual-layer of triangles
        zoff = (n-1)*4*(nt-1);
        toff = (n-1)*(2*nt-1);
        
        %Indices of vertices for layer 1: (p1,2), 2: (p3,4), 3: (p5,6)
        p1 = toff + (1:(nt-1));
        p2 = toff + (2:nt);
        p3 = toff + ((nt+1):(2*nt-2));
        p4 = toff + ((nt+2):(2*nt-1));
        p5 = toff + ((2*nt):(3*nt-2));
        p6 = toff + ((2*nt+1):(3*nt-1));
        
        %Define triangles based on layers of vertices
        Tris((1+zoff):((nt-1)+zoff),:) = [p1',p2',[p3,p4(end)]'];
        Tris((nt+zoff):((2*nt-3)+zoff),:) = [p4',p3',p1(2:end)'];
        Tris((2*nt-2+zoff),:) = [p1(1),p3(1),p5(1)];        %Edge triangle
        Tris((2*nt-1+zoff):((3*nt-4)+zoff),:) = [p3',p4',p5(2:end)'];
        Tris((3*nt-3+zoff):((4*nt-5)+zoff),:) = [p6',p5',[p3,p4(end)]'];
        Tris((4*nt-4+zoff),:) = [p6(end),p4(end),p2(end)];	%Edge triangle
    end
end

end

function fix_aspect_ratio(ax,opts)

xplot = opts.xplot;
yplot = opts.yplot;
axsort = sort({xplot,yplot});

xstring = '$z\,[{\rm m}]$';
if sum(startsWith(axsort,'p')) == 0
    axis(ax,'equal')
    ystring = horzcat('$',xplot,'\,[{\rm m}]$');
    zstring = horzcat('$',yplot,'\,[{\rm m}]$');
    
elseif  sum(startsWith(axsort,'p')) == 1
    da = get(ax,'DataAspectRatio');
    if startsWith(xplot,'p')
        ystring = horzcat('$p_',strip(xplot,'p'),'/{m_e c}$');
        zstring = horzcat('$',yplot,'\,[{\rm m}]$');
        da(1) = 1;  da(3) = 1;
    else
        ystring = horzcat('$',xplot,'\,[{\rm m}]$');
        zstring = horzcat('$p_',strip(yplot,'p'),'/{m_e c}$');
        da(1) = 1;  da(2) = 1;
    end
    set(ax,'DataAspectRatio',da)
    
elseif  sum(startsWith(axsort,'p')) == 2
    da = get(ax,'DataAspectRatio');
    ystring = horzcat('$p_',strip(xplot,'p'),'/{m_e c}$');
    zstring = horzcat('$p_',strip(yplot,'p'),'/{m_e c}$');
    da(2) = da(3);
    set(ax,'DataAspectRatio',da)
end

xlabel(ax,xstring,'interpreter','latex',...
        'fontsize',opts.fntsz,'color',opts.tc)
ylabel(ax,ystring,'interpreter','latex',...
        'fontsize',opts.fntsz,'color',opts.tc)
zlabel(ax,zstring,'interpreter','latex',...
        'fontsize',opts.fntsz,'color',opts.tc)

end