%% Setup

n = 50; %number of trials
k = 2; %number of cues
s = zeros(n,k); %stimulus matrix (trial x cue)
w = zeros(k,n+1); %initial weight vector
R = zeros(n,1); %outcome sequence
p = zeros(n,1); %prediction on each trial
r = zeros(n,1); %response on each trial
e = .1; %learning rate (only free parameter)

paradigm = 3; %set to 1, 2, or 3 to simulate one of the cases below
switch paradigm
    case 1 %blocking
        s(:,1) = 1; %cue 1 always present
        s(n/2+1:n,2) = 1; %cue 2 present in phase 2 only
        R(:) = 1; %reward on every trial
    case 2 %2 partial cues
        s = randi([0 1],n,k); %independent random cues
        R = rand(n,1) < s*[.5;.5]; %cues contribute additively
    case 3 %1 perfect cue, 1 irrelevant
        s = randi([0 1],n,k); %independent random cues
        R = rand(n,1) < s*[1;0]; %cues contribute additively
end

%% Simulate

%run model
for t=1:n %loop through trials
    p(t) = s(t,:)*w(:,t); %expected outcome
    r(t) = rand<p(t); %simulated response: 1 with probability p(t), 0 with probability 1-p(t)
    d = R(t) - p(t); %prediction error
    w(:,t+1) = w(:,t) + e*d*s(t,:)'; %learning update
end

%plots
figure(1),clf
co = get(gca,'ColorOrder');
for c=1:k %loop through cues
    subplot(k+5,1,c)
    plot(s(:,c),'.','Color',co(c,:)); %plot values of cue number c
    set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.5 1.5],'YTick',0:1,'YTickLabel','','FontSize',16); ylabel('Cue');
end
subplot(k+5,1,k+1)
plot(R,'k.'); %plot outcome sequence
set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.5 1.5],'YTick',0:1,'YTickLabel','','FontSize',16); ylabel('Outcome');
subplot(k+5,1,k+2:k+3)
plot(w') %plot weights
set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.05 1.05],'FontSize',16); ylabel('Weights');
subplot(k+5,1,k+4:k+5)
plot(p,'k'),hold on,plot(r,'k.') %plot predictions
set(gca,'XLim',[0 n+1],'YLim',[-.05 1.05],'FontSize',16); xlabel('Trial'); ylabel('Predictions');


%% Evaluate fit
%assume we're given a sequence of stimuli (s), outcomes (r), and responses (R) 

E = .3; %learning rate for model to be evaluated
ll = zeros(n,1); %log-likelihood of data given model

%run model
for t=1:n %loop through trials
    p(t) = s(t,:)*w(:,t); %expected outcome
    ll(t) = log(abs(1-p(t)-r(t))); %log-likelihood of current response given model's prediction
    d = R(t) - p(t); %prediction error
    w(:,t+1) = w(:,t) + E*d*s(t,:)'; %learning update
end

figure(2),clf
subplot(2,1,1)
plot(r,'b.'), hold on, plot(p,'k') %plot data and model predictions
set(gca,'XLim',[0 n+1],'XTick',[],'YLim',[-.05 1.05],'FontSize',16); ylabel('Response');
subplot(2,1,2)
plot(ll,'b.') %plot log-likelihoods
set(gca,'XLim',[0 n+1],'YLim',[1 -.05]*min(ll),'FontSize',16); xlabel('Trial'); ylabel('Log-likelihood');
text(n*.75,min(ll)/2,['L = ' num2str(round(sum(ll),1))],'FontSize',16) %display total log-likelihood of dataset

%% Estimate learning-rate parameter
%assume we're given a sequence of stimuli (s), outcomes (r), and responses (R) 

E = .01:.01:.5; %range of learning rates to be evaluated
L = zeros(length(E),1); %total log-likelihood for each model

for m = 1:length(E) %loop through candidate models
    for t=1:n %loop through trials
        p(t) = s(t,:)*w(:,t); %expected outcome
        L(m) = L(m) + log(abs(1-p(t)-r(t))); %add log-likelihood of current response given model's prediction
        d = R(t) - p(t); %prediction error
        w(:,t+1) = w(:,t) + E(m)*d*s(t,:)'; %learning update
    end
end

M = find(L==max(L),1); %index of best-fitting model
Ms = find(L > max(L)-chi2inv(.95,1)/2); %confidence interval

figure(3),clf,hold on,box
plot(E,L) %fit as a function of parameter value
plot(E(M),L(M),'r.','MarkerSize',20) %mark winning model
plot(E(Ms),L(Ms),'r.','MarkerSize',5) %confidence interval
set(gca,'FontSize',16); xlabel('Learning Rate'); ylabel('Goodness of Fit')
