
clear; %close all;

%% User options

% probability of target appearing on expected side
cue_validity=0.8;

% SET RANDOM SEED - COMMENT THIS OUT TO GET NEW DATA
% BUT KEEP IT FOR ANMY EXERCISES THAT INVOLVE PLOTTING A SUBSET OF TRIALS
rng(22); 

%% First generate the data set that we will fit the model to

% generate underlying (generative) probabilities jump~B(j)
pJ=1/25; %probability of side reversal
pL(1)=cue_validity; % initial probability of target appearing on left (on trial 1)

% generate 500 trials
for i=2:500
    if rand>pJ
        pL(i)=pL(i-1);   % no reversal
    else
        pL(i)=1-pL(i-1); % reversal
    end
end

%generate data: y(i)~B(q(i))
for i=1:length(pL)
    if rand<pL(i)
        y(i)=1;
    else
        y(i)=0;
   end
end


%%
% set up the state space

pL_space=[0.01:0.01:0.99]; % possible values for p(L) - don't allow 0 of 1, as model may explode!
prior_p_pL=NaN(length(y),length(pL_space)); % initiate matrix for prior probability of each value of pL: size = trials x possible values
post_p_pL=NaN(length(y),length(pL_space)); % initiate matrix for posterior probability of each value of pL: size = trials x possible values
prior_p_pL(1,:)=ones(1,length(pL_space))./length(pL_space); % on trial 1, prior probability for each value of pL is uniform ie 1/number of possible values of pL


%% go through data trial by trial updating probabilities
for i=1:length(y)
    
    if y(i)==1  % if target was on Left on this trial
          p_pL_given_yi(i,:) = pL_space;
    else
          p_pL_given_yi(i,:) = 1-pL_space;
    end;
    
    post_p_pL(i,:) = p_pL_given_yi(i,:) .* prior_p_pL(i,:);
    post_p_pL(i,:) = post_p_pL(i,:)./sum(post_p_pL(i,:),2); % normalise the posterior
    
    % set up the prior for next trial - 24/25 chance that p_pL is the same
    % as on this trial, 
    % and 1/25 chance that it now has a random value between 0 and 1
    prior_p_pL(i+1,:) = post_p_pL(i,:)*0.96 + 0.04*ones(1,length(pL_space))./length(pL_space);
    
end

% Work out the model's best guess of pL on each trial by finding the
% expected value of pL

est_pL = sum(prior_p_pL*pL_space',2);


%% Plot the true value of p(Left) and the model's estimate

figure; plot(pL); hold on; plot(est_pL,'r'); % plot the true values of pL against the model's estimate
plot(y,'k.','MarkerSize',8); % plot the individual data points (left and right targets)
set(gca,'YLim',[-0.2 1.2]); legend('true probability of target on left','model estimate','Data points (1=L, 0=R)');
xlabel('trial number','FontSize',16); ylabel('model estimate of p(Left)','FontSize',16); set(gcf,'color','w'); set(gca,'FontSize',16);

% question - given that the model "knows' pL is only ever either 0.8 or 0.2, why is
% its estimate of pL sometimes in between these two values?

% Hint: how did we get the estimate? 

%% Plot the model's beliefs about p(Left) on some example trials
figure; hold on;

for c=1:5; subplot(1,5,c); hold on;
    t=53:57; % representative trials
    plot([est_pL(t(c)) est_pL(t(c))],[0 1],'r-.','MarkerSize',15);
    plot(pL_space,prior_p_pL(t(c),:));
    title(['prior probability of options on trial ' int2str(t(c))]);
    xlabel('value of pL'); ylabel('p(value of pL)'); set(gcf,'color','w'); set(gca,'XLim',[0 1],'YLim',[0 0.1]);
    legend(['estimate of pL']);
    %annotation(gcf,'textarrow',[est_pL(t(c)) est_pL(t(c))], [0.2+0.3*(c-1) 0.11+0.3*(c-1)],'String','estimated value of pL','Color','r');
end;

%% Plot the model's beliefs as a distribution on each trial
figure; hold on; 
imagesc(1:length(y),flipud(pL_space),post_p_pL'); % plot full posterior over all trials
plot(pL,'w'); plot(est_pL,'k');plot(est_pL,'w--'); % add best estimate of p(left)
plot(0.1+y*0.8,'w.','MarkerSize',5); % add data points (Left/RIght on each trial)
set(gcf,'color','w'); set(gca,'XLim',[0 length(y)],'FontSize',14);
xlabel('trial number','FontSize',14); ylabel('possible values of p(Left)','FontSize',14); 
title('colour represents probability of each possible value of p(Left)','FontSize',14);



