clear; close all; 

%% user options

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

% SET RANDOM SEED - YOU COULD COMMENT THIS OUT TO GET NEW DATA
% BUT KEEP IT FOR ANY 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


%% now set up the state space
% possible values for p and j
pL_space=(0.01:0.01:0.99)';
pJ_space=exp(log(0.01):(log(0.2)-log(0.01))/20:log(0.2))';

% grids
[pL2d,pJ2d]=ndgrid(pL_space,pJ_space);

% transition function
transfunc=  (reshape(repmat(eye(length(pL_space)),1,length(pJ_space)),length(pL_space),length(pL_space),length(pJ_space))... % p(pL(t)| no jump occurred) * p(no jump occurred)
            .* permute(reshape(repmat(1-pJ_space,length(pL_space),length(pL_space)),length(pJ_space),length(pL_space),length(pL_space)),[2 3 1]))...
            + ((ones(length(pL_space),length(pL_space),length(pJ_space))./length(pL_space))... % + p(pL(t)| jump occurred) * p(jump_occurred)
            .* permute(reshape(repmat(pJ_space,length(pL_space),length(pL_space)),length(pJ_space),length(pL_space),length(pL_space)),[2 3 1]));

prior_p_pLpJ=NaN(size(pL2d,1),size(pL2d,2),100); %initiate prior and post matrices
post_p_pLpJ=NaN(size(pL2d,1),size(pL2d,2),100);
% uniform prior to start with
prior_p_pLpJ(:,:,1)=ones(size(pL2d))./length(pL2d(:));

%% fit the model
for i=1:length(y)

    % 'leak' or apply transition function
    if i>1
        unpacked_post=permute(reshape(repmat(post_p_pLpJ(:,:,i-1),1,length(pL_space)),length(pL_space),length(pJ_space),length(pL_space)),[1 3 2]);
        unpacked_prior=unpacked_post.*transfunc;
        prior_p_pLpJ(:,:,i)=squeeze(sum(unpacked_prior,1));
        prior_p_pLpJ(:,:,i)=prior_p_pLpJ(:,:,i)./sum(sum(prior_p_pLpJ(:,:,i)));
    end
        
    % p(y|pL,pJ)
    if y(i)==1
        % if y(i) is a hit, p(y(i)|pL) is simply pL
        py_g_pLpJ=pL2d;
    else
        % if y(i) is a miss, p(y(i)|pL) is 1-pL
        py_g_pLpJ=1-pL2d;
    end
    
    % Bayes' theorem: p(pL,pJ|y)=p(y|pL,pJ)p(pL,pJ);
    post_p_pLpJ(:,:,i)=py_g_pLpJ.*prior_p_pLpJ(:,:,i);
    
    %normalise posterior
    post_p_pLpJ(:,:,i)=post_p_pLpJ(:,:,i)./sum(sum(post_p_pLpJ(:,:,i)));
    
    % get joint maximum of prior for pL and pJ
    [mx,ix]=max(myvect(prior_p_pLpJ(:,:,i)));
    [ixq,ixj]=ind2sub(size(pL2d),ix);
    joint_mode(i,:)=[pL_space(ixq), pJ_space(ixj)];
    
    % get marginal expected values for pL and pJ
    EpL(i)=sum(myvect(prior_p_pLpJ(:,:,i).*pL2d));
    EpJ(i)=sum(myvect(prior_p_pLpJ(:,:,i).*pJ2d));
    
end


%% plot true p(Left) and model's estimate
figure; hold on;

% plot estimates based on mode/peak of distribution
%plot(joint_mode(:,1),'r'); plot([0 500],[1/25 1/25],'c--'); plot(joint_mode(:,2),'g'); plot(pL,'b'); 

% plot estimates based on marginal means of distribution
plot(pL,'b'); plot(EpL,'r'); plot([0 500],[1/25 1/25],'c--'); plot(EpJ,'g');  

plot(y,'k.','MarkerSize',8); % plot the individual data points (left and right targets)
set(gca,'YLim',[-0.2 1.2]); legend('true p(Left)','estimated p(Left)','true p(Jump)','estimated p(Jump)','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);

%% Plot joint probability distribution over pL and pJ
figure; hold on;
for i=1:50
        subplot(5,10,i);
        imagesc(pJ_space,pL_space,prior_p_pLpJ(:,:,i)); ylabel('value of p(Left)'); xlabel('value of p(Jump)'); 
end
set(gcf,'color','w');

%% plot the prior over pL on each trial (marginalizing over pJ)
figure; hold on;
prior_p_pL=squeeze(sum(prior_p_pLpJ,2)); % marginalize over pJ
imagesc(1:length(y),pL_space,prior_p_pL); %plot probability distribution over pL on each trial
plot(pL,'w'); % plot actual value of pL on each trial as a white line
plot((joint_mode(:,1)),'k'); plot(joint_mode(:,1),'w--');% plot peak value for pL on each trial as black and white dashed line


plot(y,'k.','MarkerSize',20);
set(gcf,'Color','w');

