%% Create paradigms
m = [50 50 5]; %phase lengths
R = [ones(m(1)+m(2),1);zeros(m(3),1)]; %reward in phases 1:2 only
sBlock = [[ones(m(1),1);ones(m(2),1);zeros(m(3),1)],[zeros(m(1),1);ones(m(2),1);ones(m(3),1)]]; %blocking paradigm
sBack = [[ones(m(1),1);ones(m(2),1);zeros(m(3),1)],[ones(m(1),1);zeros(m(2),1);ones(m(3),1)]]; %backward-blocking paradigm

%demo models
runRW(.1,sBlock,R);
runKalman(25,1,sBlock,R);
runRW(.1,sBack,R);
runKalman(25,1,sBack,R);

%% fit both models to data from Kalman filter

%create data from each model
N = 10; %number of (homogenous) subjects to simulate under each model
e = .3; %parameter for correct RW model
r_rw = zeros(sum(m),N); %responses for RW-generated data, subject x trial
for s=1:N %loop subjects
    p = runRW(e,sBack,R,0); %response probabilities for one subject (plot suppressed)
    r_rw(:,s) = rand(sum(m),1)<p; %probability matching
end
se = 25; sh = 1; %parameters for correct Kalman model
r_k = zeros(sum(m),N); %responses for Kalman-generated data, subject x trial
for s=1:N %loop subjects
    p = runKalman(se,sh,sBack,R,0); %response probabilities for one subject (plot suppressed)
    r_k(:,s) = rand(sum(m),1)<p; %probability matching
end

%fit RL model to both datasets
Es = (.01:.01:.5); %parameter values to test
LRW_rw = zeros(length(Es),1); %loglikelihoods for all parameter values (rw data)
LRW_k = zeros(length(Es),1); %loglikelihoods for all parameter values (kalman data)
for i=1:length(Es)
    p = runRW(Es(i),sBack,R,0); %get model predictions
    LRW_rw(i) = sum(sum(log(1-abs(repmat(p,1,N)-r_rw)))); %loglikelihood for RW-generated data, summed over subjects
    LRW_k(i) = sum(sum(log(1-abs(repmat(p,1,N)-r_k)))); %loglikelihood for Kalman-generated data, summed over subjects
end
AICRW_rw = -2*max(max(LRW_rw))+2; %minimum deviance plus 2*(1 free parameter)
bestE_rw = Es(find(LRW_rw==max(LRW_rw),1)); %best parameter for fitting rw data
AICRW_k = -2*max(max(LRW_k))+2; %minimum deviance plus 2*(1 free parameter)
bestE_k = Es(find(LRW_k==max(LRW_k),1)); %best parameter for fitting kalman data

%fit Kalman model
SEs = (.5:.5:10).^2; SHs = (.1:.1:2).^2; %sets of parameter values to test
LKalman_rw = zeros(length(SEs),length(SHs)); %loglikelihoods for all parameter combinations (rw data)
LKalman_k = zeros(length(SEs),length(SHs)); %loglikelihoods for all parameter combinations (kalman data)
for i=1:length(SEs),for j=1:length(SHs)
    p = runKalman(SEs(i),SHs(j),sBack,R,0); %get model predictions
    LKalman_rw(i,j) = sum(sum(log(1-abs(repmat(p,1,N)-r_rw)))); %loglikelihood summed over subjects
    LKalman_k(i,j) = sum(sum(log(1-abs(repmat(p,1,N)-r_k)))); %loglikelihood summed over subjects
end,end
AICKalman_rw = -2*max(max(LKalman_rw))+4; %minimum deviance plus 2*(2 free parameters)
bestSE_rw = SEs(find(max(LKalman_rw')==max(max(LKalman_rw)),1)); bestSH_rw = SHs(find(max(LKalman_rw)==max(max(LKalman_rw)),1)); %best parameters for fitting rw data
AICKalman_k = -2*max(max(LKalman_k))+4; %minimum deviance plus 2*(2 free parameters)
bestSE_k = SEs(find(max(LKalman_k')==max(max(LKalman_k)),1)); bestSH_k = SHs(find(max(LKalman_k)==max(max(LKalman_k)),1)); %best parameters for fitting kalman data

%print results
disp(' ');disp(['Fitting data generated by RW model'])
disp(['RW: e = ' num2str(bestE_rw) ', AIC = ' num2str(AICRW_rw)])
disp(['Kalman: se = ' num2str(bestSE_rw) ', sh = ' num2str(bestSH_rw) ', AIC = ' num2str(AICKalman_rw)])
disp(' ');disp(['Fitting data generated by Kalman model'])
disp(['RW: e = ' num2str(bestE_k) ', AIC = ' num2str(AICRW_k)])
disp(['Kalman: se = ' num2str(bestSE_k) ', sh = ' num2str(bestSH_k) ', AIC = ' num2str(AICKalman_k)])

%plot fit surfaces
figure(3)
[X Y] = meshgrid(SEs,SHs);
subplot(2,2,1)
plot(Es,LRW_rw); xlabel('e');ylabel('loglikelihood (rw data)')
subplot(2,2,2)
surf(X,Y,LKalman_rw'); xlabel('se');ylabel('sh')
subplot(2,2,3)
plot(Es,LRW_k); xlabel('e');ylabel('loglikelihood (k data)')
subplot(2,2,4)
surf(X,Y,LKalman_k'); xlabel('se');ylabel('sh')
