%setup
n = 50; %number of trials
k = 2; %number of cues
N = 1000; %number of datasets to generate and fit
e = .1; %learning rate for data-generating model
E = .01:.01:.5; %range of learning rates to be evaluated in fitting models

s = zeros(n,k); %stimulus matrix (trial x cue)
R = zeros(n,1); %outcome sequence
p = zeros(n,1); %prediction on each trial
r = zeros(n,1); %response on each trial
diff = zeros(1,N); %for tracking results: difference in fit between models

for i=1:N
    
    %create cue-outcome schedule (using 2-partial-cue paradigm)
    s = randi([0 1],n,k); %independent random cues
    R = rand(n,1) < s*[.5;.5]; %cues contribute additively
    
    %simulate common-e model to create test data
    w = zeros(k,n+1); %initialize weight vector
    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
    
    %fit common-e model to the data
    Lcommon = zeros(length(E),1); %total log-likelihood for each model
    for m = 1:length(E) %loop through candidate models
        w = zeros(k,n+1); %initialize weight vector
        for t=1:n %loop through trials
            p(t) = s(t,:)*w(:,t); %expected outcome
            Lcommon(m) = Lcommon(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) + d*E(m)*s(t,:)'; %learning update
        end
    end
    fitCommon = max(Lcommon); %loglikelihood of best-fitting common-e model
    
    %fit separate-e model to the data
    Lseparate = zeros(length(E),length(E)); %total log-likelihood for each model
    for e1 = 1:length(E) %loop through values for first learning rate
        for e2 = 1:length(E) %loop through values for second learning rate
            w = zeros(k,n+1); %initialize weight vector
            for t=1:n %loop through trials
                p(t) = s(t,:)*w(:,t); %expected outcome
                Lseparate(e1,e2) = Lseparate(e1,e2) + 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) + d*[E(e1);E(e2)].*s(t,:)'; %learning update
            end
        end
    end
    fitSeparate = max(max(Lseparate)); %loglikelihood of best-fitting common-e model
    
    diff(i) = fitSeparate - fitCommon; %difference in fit between models
    disp(i)
    if fitSeparate<fitCommon,break,end
end