% written by Allison Shaw, Contact ashaw@umn.edu for assistance.
% started: 6 Oct 2021
% last updated: 12 June 2023
%
clear all

eps1 = 0.1; % how much defines a 'difference' in migration strategy

fs1 = 10;  % axes labels
fs3 = 09;  % axis numbering
lw2 = 1; % fig edges

width = 13.5;
height = 18;
xpos = 3;
ypos = 2;
wi = 0.16;
he = 0.12;
sy = 0.21;
sx = 0.15;
dy = 0.07;
dx = 0.15;

xx = 40;
yy = 40;

figure(1); clf
hh = gcf;
set(hh,'PaperUnits','centimeters');
set(hh,'Units','centimeters');
set(gcf,'Position',[xpos ypos width height])

colormap(jet)


for k = 1:3
    if     k==1; lambda = 10; gamma = 0.8; alpha_s = 5; phi_max = 2; alpha_f = 0;    % decrease alpha_f
    elseif k==2; lambda = 20; gamma = 0.8; alpha_s = 5; phi_max = 2; alpha_f = 0.8;  % increase lambda
    elseif k==3; lambda = 10; gamma = 0.8; alpha_s = 10; phi_max = 2; alpha_f = 0.8;  % increase alpha_s
    end


    name = strcat(['sims_sensitivity_lambda=' num2str(lambda) '_gamma=' num2str(gamma) '_alphaS=' num2str(alpha_s) '_phimax=' num2str(phi_max) '_alphaF=' num2str(alpha_f) '.mat']);
    load(name)
    
    psivec2 = fliplr(psivec);
    m_bounds = linspace(0,1,length(p_by_distance)); % all possible migration strategies

    for j = 1:length(betavec)
        for jj = 1:length(psivec)
            for jjj = 1:length(muvec)
                % find biggest m_i -- longest distance in the migration strategy
                m_max(jj,jjj,j) = max(m_strategy_vec(:,jj,jjj,j));
                % find smallest m_i -- shortest distance in the migration strategy
                m_min(jj,jjj,j) = min(m_strategy_vec(:,jj,jjj,j));

                % find the index of the furthest distance actually traveled
                ind = find(p_by_distance_all(:,jj,jjj,j)>eps,1,'last');
                % save corresonding distance
                d_max(jj,jjj,j) = m_bounds(ind);
                clear ind
                % find the index of the shortest distance actually traveled
                ind = find(p_by_distance_all(:,jj,jjj,j)>eps,1,'first');
                % save corresonding distance
                d_min(jj,jjj,j) = m_bounds(ind);

            end
        end
    end

    for j = 1:length(betavec)
        beta_min = betavec(j);    % scalar of infection dynamics at x=1

        % binary variable for when migration strategy is conditional or not
        conditional = (m_max(:,:,j)-m_min(:,:,j))>0;
        % create binned set of migraion outcomes
        none = (d_min(:,:,j)==0).*(d_max(:,:,j)==0);
        full = (d_min(:,:,j)==1).*(d_max(:,:,j)==1);
        partial = (d_min(:,:,j)==0).*((d_max(:,:,j)-d_min(:,:,j))>0);
        differential = (d_min(:,:,j)>0).*((d_max(:,:,j)-d_min(:,:,j))>0);
        intermediate = (d_min(:,:,j)>0).*((d_max(:,:,j)-d_min(:,:,j))==0);

        axes('position',[sx+(k-1)*(dx+wi) sy+(j-1)*(dy+he) wi he])
        imagesc(flipud(0*none + 1*partial + 2*intermediate + 3*differential + 4*full),[0 4])
        set(gca,'XTick',1:xx:length(muvec))
        set(gca,'XTickLabel',muvec(1:xx:end))
        set(gca,'YTick',1:yy:length(psivec))
        set(gca,'YTickLabel',psivec2(1:yy:end))
        %xlabel({'infection-independent','mortality (\mu)'},'fontsize',fs1)
        %ylabel({'infection-based','mortality (\psi)'},'fontsize',fs1)
        if     k==1; title(strcat(['\beta_{min}=' num2str(beta_min) '; \alpha_f=' num2str(alpha_f)]),'fontsize',fs1) % decrease alpha_f
        elseif k==2; title(strcat(['\beta_{min}=' num2str(beta_min) '; \lambda=' num2str(lambda)]),'fontsize',fs1)  % increase lambda
        elseif k==3; title(strcat(['\beta_{min}=' num2str(beta_min) '; \alpha_s=' num2str(alpha_s)]),'fontsize',fs1) % increase alpha_s
        end
        
        set(gca,'FontSize',fs3,'LineWidth',lw2);


    end

    
end

axes('position',[0.4 0.01 0.03 0.12])
    imagesc([1:5]')
    set(gca,'XTick',[]);  set(gca,'XTickLabel',{});
    set(gca,'YTick',1:5); set(gca,'YAxisLocation','right'); set(gca,'YTickLabel',{'none','partial','intermediate','differential','full'});
    set(gca,'FontSize',fs3,'LineWidth',lw2);

% label subpanels
axes('position',[0 0 1 1],'visible','off')
    text(0.07,0.16+4*(he+dy),'a)','horizontalalignment','center')
    text(0.07,0.16+3*(he+dy),'b)','horizontalalignment','center')
    text(0.07,0.16+2*(he+dy),'c)','horizontalalignment','center')
    text(0.07,0.16+(he+dy),  'd)','horizontalalignment','center')
    text(0.07+wi+dx,0.16+4*(he+dy),'e)','horizontalalignment','center')
    text(0.07+wi+dx,0.16+3*(he+dy),'f)','horizontalalignment','center')
    text(0.07+wi+dx,0.16+2*(he+dy),'g)','horizontalalignment','center')
    text(0.07+wi+dx,0.16+(he+dy),  'h)','horizontalalignment','center')
    text(0.07+2*(wi+dx),0.16+4*(he+dy),'i)','horizontalalignment','center')
    text(0.07+2*(wi+dx),0.16+3*(he+dy),'j)','horizontalalignment','center')
    text(0.07+2*(wi+dx),0.16+2*(he+dy),'k)','horizontalalignment','center')
    text(0.07+2*(wi+dx),0.16+(he+dy),  'l)','horizontalalignment','center')

    text(0.15,0.97,'decrease \alpha_f')
    text(0.47,0.97,'increase \lambda')
    text(0.79,0.97,'increase \alpha_s')

    annotation('rectangle',[0.13 0.95 0.19 0.04],'LineWidth',lw2)
    annotation('rectangle',[0.45 0.95 0.19 0.04],'LineWidth',lw2)
    annotation('rectangle',[0.77 0.95 0.19 0.04],'LineWidth',lw2)
    %
    text(0.5,0.16,'infection-independent migration mortality (\mu)','horizontalalignment','center')
    text(0.03,0.6,'infection-based mortality (\psi)','horizontalalignment','center','rotation',90)
axis([0 1 0 1])

% Backup previous settings
prePaperType = get(hh,'PaperType');
prePaperPosition = get(hh,'PaperPosition');
prePaperSize = get(hh,'PaperSize');

% Make changing paper type possible
set(hh,'PaperType','<custom>');
% Set the page size and position to match the figure's dimensions
position = get(hh,'Position');
set(hh,'PaperPosition',[0,0,position(3:4)]);
set(hh,'PaperSize',position(3:4));

name = strcat(['FigS4.jpg']);
print(name,'-djpeg','-r600')

% Restore the previous settings
set(hh,'PaperType',prePaperType);
set(hh,'PaperPosition',prePaperPosition);
set(hh,'PaperSize',prePaperSize);


