%% plotting figure 7 (simulations)
clearvars
currentFolder = pwd;
addpath(genpath('Matlab_functions'))
path = currentFolder(1:end-8);
data_folder = 'data\processed\';
fits_folder = 'data\fits\';
%%
NS = [1:10,12:24];
CW = [-1,-1,-1,1,1,1,-1,1,1,-1,-1,1,-1,1,-1,1,1,1,-1,-1,-1,1,1,-1];        % -1 no/CCW-CW; 1 yes/CW-CCW
ind_angle = [180:360,1:179];                                               % index of angles so that 0 is at center
%% restructuring Model Predictions:
%                pp    trial  angle    process (MB f s t, PB f s t)
ModPredSim = NaN(24,   495,   360,     6);
ThetaUp =    NaN(24,   495);
for s = NS
    file = sprintf('%s%ssim_same_length%02d.MAT',path,'data\simulations\',s);
    load(file)
    
    ModPredSim(s,:,:,1) = modelpredMB.fast_t(:,ind_angle);
    ModPredSim(s,:,:,2) = modelpredMB.slow_t(:,ind_angle);
    ModPredSim(s,:,:,3) = modelpredMB.total_t(:,ind_angle);
    ModPredSim(s,:,:,4) = modelpredPB.fast_t(:,ind_angle);
    ModPredSim(s,:,:,5) = modelpredPB.slow_t(:,ind_angle);
    ModPredSim(s,:,:,6) = modelpredPB.total_t(:,ind_angle);
    
    ThetaUp(s,:) = T.UpdateAngle(1:495);
end
%% trials to plot generalization for
RelTrials = [190,437;...
    439,450;...
    455,490];
%% plotting:
%% Loading in the colour matrix:
% col(shade,RGBcode,colour)
[~,col] =DefColours(5,5);
shade =  [1,5];
% process:f s t  f s t
ColInd = [6,1,19,6,1,19];
% FF 1 2 SR
Phase_Col = [7,2,10];
%% axis
x =[-180:179]; % x-axis (angles)
y =1:495; % y-axis (trials)
ind_angles_plot = 91:271; % angles between -90 : 90
%% limits
Collim = [min(ModPredSim([1:10,12:end],:,:,:),[],'all'),...
    max(ModPredSim([1:10,12:end],:,:,:),[],'all')];
AI_lim = [-1.2,1.2];
angle_lim = x([ind_angles_plot(1),ind_angles_plot(end)]);
trial_lim = [y(150),y(end)];
Angle_ticks = [-35,0,35];
IndAngle_ticks = find(x == Angle_ticks(1)| x == Angle_ticks(2) |x == Angle_ticks(3));

%% subplotindices
PlotInd = {[1,66] [73,138],[7,144]};

%% ticks and labels based on subplot structure 1:6 imagesc then traces then generalization:
%         MB fast       MB slow      MB total    PB fast    PB slow      PB total
%         traces
%
Ylim    = {angle_lim    angle_lim    angle_lim   angle_lim  angle_lim    angle_lim ...
    AI_lim...
    AI_lim  AI_lim  AI_lim  AI_lim  AI_lim  AI_lim  AI_lim  AI_lim  AI_lim} ;

Xlim    = {trial_lim      trial_lim     trial_lim     trial_lim     trial_lim     trial_lim ...
    trial_lim...
    angle_lim         angle_lim         angle_lim         angle_lim          angle_lim        angle_lim         angle_lim         angle_lim         angle_lim} ;

Yticks  = {[-50,0,50],          [-50,0,50],         [-50,0,50],         [-50,0,50],         [-50,0,50],         [-50,0,50], ...
    [-1,0,1],...
    [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1],       [-1,0,1] };

YtiLa = {  [],                 [],                 [-50,0,50],                 [],                 [],         [-50,0,50],...
    [-1,0,1],...
    [-1,0,1],       [],                 [],                 [],                 [],                 [],                 [],                  [],                []};

Xticks ={  [150:100:450]       [150:100:450]       [150:100:450]       [150:100:450]       [150:100:450]       [150:100:450]...
    [150:100:450],...
    [-50,0,50]          [-50,0,50]          [-50,0,50]          [-50,0,50]          [-50,0,50]          [-50,0,50]          [-50,0,50]          [-50,0,50]           [-50,0,50]};

XtiLa = {  []                  []                  [150:100:450]       []                  []                  [] ...
    [150:100:450]  ...
    [-50,0,50]          [-50,0,50]                  []                  []                  []                  []                  []                  []                  []};
YlabelL  =  {[],...
             [],...
             ['Target Direction'],...
             [],...
             [],...
             ['Target Direction'],...
             ['Adaptation Index'],...
             ['Adaptation Index'],[],         [],         [],         [],         [],            [],         [],         []};
YlabelR = {[],                 [],                 [],                 [],                 [],                 [],...
    [],...
    [],                 [],                 [],                 [],                 [],                 [],                 [],                 [],                 []};

YtiLaR = { [],                 [],                 [],                 [],                 [],                 [],...
    [],...
    [],                 [],                 [],                 [],                 [],                 [],                 [],                 [],                 [] ...
    [],                 [],                 [],                 [],                 [],                 [],                 [],                 [],                 []};

Xlabel  = {[]                  []                  ['Trial']                  []                  []                  []...
    ['Trial']...
    ['Target Direction']                  []                  []                  []                  []                  []                  []                  []                  []};



%% plotting
%% imagesc figures:
position = [0 0 18 15];
figure('units','centimeter',...
    'innerposition',position)
set(gca, 'FontName', 'Arial',...
    'FontSize', 8)

for i = 1:6
    %% model prediction per process
    p1{i} = subplot(1,16,i);
    imagesc(y(trial_lim(1):trial_lim(2)),x(ind_angles_plot),squeeze(mean(ModPredSim([1:10,12:24],trial_lim(1):trial_lim(2),ind_angles_plot,i)))',Collim)
    hold on
    if i <4
        %% if motion based plot mean update angle:
        plot(nanmean(ThetaUp),'w','linewidth',2)
    else
        %% if plan based plot update at 0 (at learning i.e. index trialtype 1 or 2):
        indTT12 =find(~isnan(nanmean(ThetaUp)));
        plot(indTT12([1,end]),[1,1],'w','linewidth',2)
    end
    %% plot lines where the generalization slices are taken
    plot([RelTrials(:),RelTrials(:)]',[ones(length(RelTrials(:)),1)*angle_lim]','w--','linewidth',1.5)
end
%% change colourmap:
colormap('jet')
%% now the traces: 
%       MB  PB
lsym = {'-','--'};
for p = 1:2
        p1{7} =  subplot(1,16,7);
        for j = 1:3
            % plotting the learning curve of individual processes and models at the target direction
            plot(y(trial_lim(1):trial_lim(2)),squeeze(mean(ModPredSim([1:10,12:24],trial_lim(1):trial_lim(2),IndAngle_ticks(2),(p-1)*3+j))),lsym{p},'color',col(1,:,ColInd(j)),'LineWidth',2)
            hold on
        end
end

for i = 1:3
    % mark trials where the generalization curves are plotted: 
    plot([RelTrials(i,:);RelTrials(i,:)],[ones(length(RelTrials(i,:)),1)*angle_lim]','-','linewidth',2,'color',col(1,:,Phase_Col(i)))
end
Leg = legend('motion-referenced fast','motion-referenced slow','motion-referenced total','plan-referenced fast','plan-referenced slow','plan-referenced total', 'FontName', 'Arial','FontSize',8,'Location','south','NumColumns',2);
Leg = legend('boxoff');


%% plot generalization curves and add trial numbers
% move trial text around to the correct spot 
TEx = zeros(2,3,3);
TEx(:,:,1) = [0.1,-0.25;-0.2,0.3;-0.4,0.2]';
TEx(:,:,2) = [0.5,-0.1;0.3,-0.3; 0.15,-0.15]';
TEx(:,:,3) = [-0.1,0.2 ;0.2,-0.3;-0.15,0.15]';
 for  phase = 1:3
     for process =1:3
          p1{7+process+(phase-1)*3} = subplot(1,16,7+process+(phase-1)*3);
         for trial = 1:2
             for model = 1:2
                  plot(x(ind_angles_plot),squeeze(nanmean(ModPredSim([1:10,12:24],RelTrials(phase,trial),ind_angles_plot,(model-1)*3+process))),lsym{model},'color',col(shade(trial),:,ColInd(process)),'LineWidth',2)
                  hold on
                  if model == 1
                     [~, IAmax] = max(abs(squeeze(nanmean(ModPredSim([1:10,12:24],RelTrials(phase,trial),ind_angles_plot,(model-1)*3+process)))));
                     
                     text(x(ind_angles_plot(IAmax)),... x position
                         squeeze(nanmean(ModPredSim([1:10,12:24],RelTrials(phase,trial),ind_angles_plot(IAmax),(model-1)*3+process)))+TEx(trial,process,phase),... y position
                         num2str(RelTrials(phase,trial)),...% trial number
                         'FontSize',8) 
                  end
                  
             end
         end
     end
 end
 %% axis labels ticklabels and limits:
 
 for i = 1:16
    
     subplot(1,16,i)
     yp{i} = ylabel(YlabelL{i});
     xp{i} = xlabel(Xlabel{i});
     if i <8
         xlim(trial_lim)
     else
         xlim(angle_lim)
     end
     ylim(Ylim{i})
     yticks(Yticks{i})
     yticklabels(YtiLa{i})
     xticks(Xticks{i})
     xticklabels(XtiLa{i})
     
     
     yyaxis right
     yp{i} = ylabel(YlabelR{i},'Color',[0,0,0],'Rotation',270,'VerticalAlignment','middle')
     ylim(Ylim{i})
     yticks(Yticks{i})
     yticklabels(YtiLaR{i})
     set(gca, 'YColor',[0,0,0])
     
 end

%% for restructuring figure
%figure('units','normalized','outerposition',[0 0 1 1])
nR =    32; % rows
nC =    36; % columns

Tw = [130/1920,20/1920, 130/1920];
Th = [80/1080, 20/1080, 80/1080];
w = (1-sum(Tw.*[1,nC-1,1]))/nC;
h = (1-sum(Th.*[1,nR-1,1]))/nR;

xSPP = Tw(1):(w+Tw(2)):1-w-Tw(3);
ySPP = 1-h-Th(3):-(h+Th(2)):Th(1);%linspace(1-h-n/2*Th,n/2*Th,nR);
%% moving the plots around
% for background color
col(:,:,20) = 1;

%     motion-referenced plan referenced   trac-  generalization
%     f     s     t     f     s     t     es     f     s     t     f     s     t     f     s     t
BG = [20    20    20    20    20    20    20     7     7     7     2     2     2    10    10    10];
Px = [25    25     1    25    25     1     1     1     5     9    13    17    21    25    29    33];
Py = [12    16    16     3     7     7    23    32    32    32    32    32    32    32    32    32];
Sx = [12    12    24    12    12    24    36     4     4     4     4     4     4     4     4     4];
Sy = [ 4     4     8     4     4     8     6     7     7     7     7     7     7     7     7     7];
subplot(p1{1})
ColorBar = colorbar;
for i = 1:16
    
     set(p1{i},'position',[xSPP(Px(i)),ySPP(Py(i)),w*Sx(i)+Tw(2)*(Sx(i)-1),h*Sy(i)+Th(2)*(Sy(i)-1)],...
          'FontName', 'Arial','FontSize',8,...
         'Color',[col(10,:,BG(i))],...
         'TickLength',[0.2, 0.2]/(Sx(i)+Sy(i)))
end
set(ColorBar,'position',[xSPP(end)+Tw(end)/4,ySPP(Py(2)),0.015, h*51],'FontSize',8)
%% text extra
%% Plan referenced
PBlabel = annotation('textbox',[xSPP(1) ySPP(Py(4))+h*(Sy(4)+1)+Th(2)*(Sy(4)) 0 0],...
    'String','Plan-Referenced',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');


%% PB fast 
PBflabel = annotation('textbox',[xSPP(25) ySPP(5) 0 0],...
    'String','fast',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');

%% PB slow
PBslabel = annotation('textbox',[xSPP(25) ySPP(1) 0 0],...
    'String','slow',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');


%% PB total
PBtlabel = annotation('textbox',[xSPP(1) ySPP(1) 0 0],...
    'String','overall',...
    'EdgeColor','none',...
     'FontName', 'Arial', 'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');


%% Motion referenced
MBlabel = annotation('textbox',[xSPP(1) ySPP(8)+h 0 0],...
    'String','Motion-Referenced',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

%% MB fast 
MBflabel = annotation('textbox',[xSPP(25) ySPP(10) 0 0],...
    'String','fast',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');

%% MB slow
MBslabel = annotation('textbox',[xSPP(25) ySPP(14) 0 0],...
    'String','slow',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');


%% MB total
MBtlabel = annotation('textbox',[xSPP(1) ySPP(10) 0 0],...
    'String','overall',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'bottom');


%% %% Titles
%% First Force Field
FF1label = annotation('textbox',[xSPP(7) ySPP(25)+h 0 0],...
    'String','First Force Field',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'center',...
    'verticalalignment', 'top');
%% Second Force Field
FF2label = annotation('textbox',[xSPP(19) ySPP(25)+h 0 0],...
    'String','Second Force Field',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'center',...
    'verticalalignment', 'top');
%% Spontaneous Recovery
SRlabel = annotation('textbox',[xSPP(31) ySPP(25)+h 0 0],...
    'String','Spontaneous Recovery',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',8,...
    'FitBoxToText','on',...
    'horizontalalignment', 'center',...
    'verticalalignment', 'top');
%% A B C D 
Alabel = annotation('textbox',[0 1 0 0],...
    'String','A',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

Blabel = annotation('textbox',[0 ySPP(7) 0 0],...
    'String','B',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

Clabel = annotation('textbox',[0 ySPP(17) 0 0],...
    'String','C',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

Dlabel = annotation('textbox',[0 ySPP(24) 0 0],...
    'String','D',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

Elabel = annotation('textbox',[xSPP(13)-0.01 ySPP(24) 0 0],...
    'String','E',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');

Flabel = annotation('textbox',[xSPP(25)-0.01 ySPP(24) 0 0],...
    'String','F',...
    'EdgeColor','none',...
     'FontName', 'Arial','FontSize',13,...
    'FitBoxToText','on',...
    'horizontalalignment', 'left',...
    'verticalalignment', 'top');
%%xpT = []
%% moving Xlabel
XposXLabel = zeros(16,1);
XposXLabel(1:6)  = Xticks{3}(1)+diff(Xticks{3}(1:2))*(1/2);
XposXLabel(7)  = Xticks{3}(1)+diff(Xticks{3}(1:2))*(1/3);
XposXLabel(8)  = mean(p1{8}.XLim);

YposXLabel = zeros(16,1);
YposXLabel(1:6)  = max(p1{1}.YLim)+diff(p1{i}.YLim)*5;
YposXLabel(7)  = p1{7}.YLim(1)-diff(p1{i}.YLim)*0.1;
YposXLabel(8)  = p1{8}.YLim(1)-diff(p1{i}.YLim)*0.15;


for i = 1:16
    %% moving ylabel position
    yp{i}.Position(1) = p1{i}.XLim(1);
    yp{i}.Position(2) = mean(p1{i}.YLim);
    xp{i}.Position(1) = XposXLabel(i);
    xp{i}.Position(2) = YposXLabel(i);
    p1{i}.XAxis.FontSize = 8;
    for j = 1:2
        p1{i}.YAxis(j).FontSize = 8;
    end
end
%% save figure removes background color
saveas(gcf,sprintf('%s%s.pdf',path,'Figures\Fig_7'))