% RL_tutorial_NSCI850.m
%
% learn how to use a simple RL mechanisms to fit data
% Dataset: from a gambling task
%           - 2 choices
%           - 70%-30% reward probabilities
%           - high/low volatility (2 separate subjects)
%               i.e. fast or slow switching of probabilities
%
% written by G. Blohm for NSCI 850 (2017)

clear

%% load data
cd('C:\Users\Gunnar\Documents\Teaching\CompNeurosci-NSCI850\Week13-RL\Matlab')
load('alldataRL.mat','d')
N = length(d{1}.choice); % number trials in experiment

%% plot data
for i = 1:2,
    figure(1)
    subplot(2,1,i)
    hold on
    plot(d{i}.prep.feedbackprob,'b-')
    plot(2-d{i}.choice,'*k')
    plot(2-conv(d{i}.choice,ones(6,1)/6,'same'),'k-')
    
    outcome = d{i}.prep.feedback; % choice outcome
    
    %% simulate RL process
    % parameters
    alp = .4; % sensitivity to reward and punishment
    bet = 2; % inverse learning rate
    
    % initialize
    v = .5*ones(1,2); % initial value
    
    % loop through value function computation...
    for t = 1:N,
        % compute likelihood of the each choice option (softmax)
        ev = exp(bet*v(t,:));
        sev = sum(ev);
        p(t,:) = ev/sev;
        
        % Do a weighted coinflip to make a choice: choose stim 1 if random
        % number is in the [0 p(1)] interval, and 2 otherwise
        if rand(1)<p(t,1)
            c(t) = 1; % choice
        else
            c(t) = 2;
        end
        r = outcome(t,c(t)); % select the outcome for this choice
        
        % update values
        dv  = r-v(t,c(t)); % compute prediction error
        v(t,c(t))= v(t,c(t)) + alp*dv;    % update value
        v(t+1,:) = v(t,:);
    end
    v = v(2:end,:);
    
    % plot result
    subplot(2,1,i)
    plot(p(:,1),'r:')
    plot(v(:,1),'g:')
    plot(2-c,'.r')
    axis([0 135 0 1])
    if i == 2,legend('real prob','choice (data)','smoothed choice (data)',...
            'likelihood of choice (model)','choice value (model)','choice (model)'); end
end
drawnow
% return;

%%=========================================================================
%% GRID SEARCH OPTIMAL PARAMETERS
for ii = 1:50,
    for jj = 1:50,
        
        for i = 1:2,
            outcome = d{i}.prep.feedback; % choice outcome
            
            %% simulate RL process
            % parameters
            alp = ii/50; % sensitivity to reward and punishment
            bet = jj/5; % inverse learning rate
            
            % initialize
            v = .5*ones(1,2); % initial value
            
            % loop through value function computation...
            for t = 1:N,
                % compute likelihood of the each choice option (softmax)
                ev = exp(bet*v(t,:));
                sev = sum(ev);
                p(t,:) = ev/sev;
                
                % Do a weighted coinflip to make a choice: choose stim 1 if random
                % number is in the [0 p(1)] interval, and 2 otherwise
                if rand(1)<p(t,1)
                    c(t) = 1; % choice
                else
                    c(t) = 2;
                end
                r = outcome(t,c(t)); % select the outcome for this choice
                
                % update values
                dv  = r-v(t,c(t)); % compute prediction error
                v(t,c(t))= v(t,c(t)) + alp*dv;    % update value
                v(t+1,:) = v(t,:);
            end
            v = v(2:end,:);
            
            % compute the likelihood of p(choice|model)
            c1 = c==1;
            c2 = c==2;
            loglik(ii,jj,i) = sum(log(p(d{i}.choice==1,1)))+sum(log(p(d{i}.choice==2,2)));
            
        end
    end
end
figure(2)
for i = 1:2,
    subplot(2,1,i)
    [xx,yy]=meshgrid((1:50)/50,(1:50)/5);
    surf(xx,yy,exp((loglik(:,:,i)-min(loglik(:)))/10))
    shading interp;
    view([0 0 -1])
    colorbar
end