function [pmhist] = modellearning03(b)

% The purpose of this exercise is to see how model learning and information
%  gain at the level of models can be used to drive sampling behavior.

%% Precursors
% Set up sample space and mu values
x = 0:0.5:10;                   % Sample space
xlen = length(x);
mu = 0:0.02:1;                  % The mu distribution
mulen = length(mu);

% Set up sources and samples
Smu = [4 5 6]';                   % Source means
Ssigma = [0.75 0.75 0.75]';       % Source variances
p = 0.005;              % Precision parameter for beta distributions - use 0.005 for 
Amp = 0.6;              % Amplitude parameter for the normal curuve -  use 0.5
nSources = length(Smu);         % Number of sources
source = 2;                     % Source to use
nTrials = 20;                   % Number of rewards/pellets to be found
locations = Ssigma(source)*randn(nTrials,1) + Smu(source);  % Location of pellet
locations(locations < min(x)) = min(x);                     %  make sure the pellet has the potential to be found
locations(locations > max(x)) = max(x);
sdata = interp1(x,1:length(x),locations,'nearest');         % This is in the index of x where the pellet can be found

% Set up p(mu|M) for each model M - We assume that the models are known in
%   advance (i.e. from previous learning)
probMu = zeros(nSources,xlen,mulen);
for iS = 1:nSources,
    pmu = Amp*rbf(x,Smu(iS),Ssigma(iS));
    pmu(pmu > 1) == 1;
    pvar = pmu*p;
    probMu(iS,:,:) = mkModelProb(pmu,pvar,mu,xlen,mulen);
end

fig02 = figure; 
set(fig02,'position',[18   121   560   665]);
for iS = 1:nSources,
subplot(nSources + 1,1,iS);
h = pcolor(x,mu,squeeze(probMu(iS,:,:))');
set(h,'linestyle','none');
set(gca,'xtick',[]);
ylabel('\mu');
axis xy
end
xlabel('x');


% Set up probability distributions
nOutcomes = 2;                  % The number of outcomes (2 will use beta distributions; more than 2 require Dirichlet distributions)
xhist = zeros(xlen,nOutcomes);  % The histogram of the history
alpha = ones(xlen,nOutcomes);   % The hyperparameters for the beta/dirichlet distributions
alpha(:,2) = 3;
p_out = nan(xlen,mulen);        % Buffer for the probability distribution p(mu|D) for each x
pM = ones(nSources,1)/nSources; % Buffer for the probability distribution p(model|D)
pObs = ones(xlen,1)/xlen;       % Start with a uniform distribution
if nargin < 2,
    use_genprior = 1;           % Use generalized prior flag
end

% Set up history and action selection
trial_hist = zeros(0,3);        % Trial history [sample location, outcome, sample index]
a_method = 'value';              % Action selection method ('value','info','switch')
if nargin == 0
    b = 0;                      % Parameter for softmax-based action selection (0 yields random sampling)
end

% Set up figures
plot_on = 1;
pause_time = 0.001;
if plot_on
    fig01 = figure;
    set(fig01,'position',[864    -4   560   797])
end

% return

%% Main loop
sample_num = 1;                             % Sample index
action = softmax(pObs,b);                   % Select action
pmhist = zeros(nTrials,nSources);           % The probability each model (at the end of each trial)
for iT = 1:nTrials,
    
    actions_used = ones(1,xlen);            % History of what actions have been taken (this will remove repeat samping)
    while action ~= sdata(iT),              % This while loop places a new reward once the old one has been found
        
        trial_hist(sample_num,:) = [x(action) 2 action];    % This is a miss (no pellet found)
        xhist(action,2) = xhist(action,2) + 1;              % Increment histogram
        
        if nOutcomes == 2,
%             [p_out pObs] = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen);
            [p_out pObs] = bayesupdateM(pM,probMu,mu,xlen,mulen,nSources);
            expIGm = infogain_models(mu,probMu,xhist,pObs,xlen,mulen,nSources);
        end
        
        actions_used(action) = 0;           % Remove the posibility of selecting previous selected action
        action = action_selection(actions_used,pObs,expIGm,b,a_method);
        sample_num = sample_num + 1;                % Increment sample number
        
%         if plot_on,
%             figure(fig01);
%             plot01(x,mu,p_out,pObs,xhist,expIGm)
%             title(['Trial: ' num2str(iT) ' of ' num2str(nTrials)]);
%             
%             figure(fig02);
%             subplot(3,1,3);
%             pM = bayesModelupdate(mu,probMu,xhist,xlen,mulen,nSources);
%             h = bar(pM);
%             set(h,'facecolor',[0.2 0.2 0.2]);
%             set(gca,'xlim',[0.4 nSources + 0.6],'ylim',[0 1]);
%             ylabel('p(M)');
%             xlabel('Model');
%             
%             pause(pause_time);
%         end

        
    end
    trial_hist(sample_num,:) = [x(action) 1 action];        % This is a hit (pellet found)
    xhist(action,1) = xhist(action,1) + 1;                  % Increment histogram

    if nOutcomes == 2,
%         [p_out pObs] = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen);
        [p_out pObs] = bayesupdateM(pM,probMu,mu,xlen,mulen,nSources);
        expIG = infogain_parameters(mu,alpha,xhist,xlen,mulen,use_genprior);
        expIGm = infogain_models(mu,probMu,xhist,pObs,xlen,mulen,nSources);
        
        
        if plot_on
            figure(fig01);
            plot01(x,mu,p_out,pObs,xhist,expIGm)
            title(['Trial: ' num2str(iT) ' of ' num2str(nTrials) ' : Samples drawn:' num2str(sample_num)]);
            
            figure(fig02);
            subplot(3,1,3);
            pM = bayesModelupdate(mu,probMu,xhist,xlen,mulen,nSources);
            h = bar(pM);
            set(h,'facecolor',[0.2 0.2 0.2]);
            set(gca,'xlim',[0.4 nSources + 0.6],'ylim',[0 1.1]);
            ylabel('p(m|D)');
            xlabel('Model');
            
            pause(pause_time);
        end
        pmhist(iT,:) = pM;
    end
    
%     action = softmax(pObs,b);          % Select action using the predicted reward
%     action = softmax(expIGm,b);          % Select action using the information gain
    action = action_selection(actions_used,pObs,expIGm,b,a_method);
    sample_num = sample_num + 1;                % Increment sample number
    
end

disp('Done!');

%% Functions
function probMu = mkModelProb(pmu,pvar,mu,xlen,mulen)
Emu = pmu;                          % Set expectation
varmu = pvar;                       % Set variance
probMu = zeros(xlen,mulen);         % Allocate probMu
A = [1./Emu - 1];                   % Constant for computing parameters for beta
alpha = (A - (1 + A).^2.*varmu)./(varmu.*(1 + A).^3);
alpha(find(alpha < 1)) = 1;         % Reset alphas that are too low
beta = alpha.*A;
for iX = 1:xlen, 
    probMu(iX,:) = betapdf(mu,alpha(iX),beta(iX)); 
end
probMu = probMu./repmat(nansum(probMu,2),1,mulen);


function action = action_selection(actions_used,pObs,expIGm,b,a_method)
% Function for action selection method
%
% Inputs
%  actions  - a list of actions that have not yet been taken (this trial)
%  pObs     - the probability of the (reward) observation  across position
%  expIGm   - the expected information gain across position
%  b        - the softmax parameter
%  a_method - the method used for action selection


switch a_method
    case 'value'                            % Maximize value
        pObs = pObs.*actions_used;
        action = softmax(pObs,b);
    case 'info'                             % Maximize information gain
        expIGm = expIGm.*actions_used';
        action = softmax(expIGm,10*b);
    case 'switch'
        if max(expIGm) < 0.001,             % Switch if the expect information gain is less than 0.001 bits
            pObs = pObs.*actions_used;
            action = softmax(pObs,b);
        else
            expIGm = expIGm.*actions_used';
            action = softmax(expIGm,10*b);
        end
end

function y = rbf(x,mu,sigma)
y = exp(-(mu - x).^2/2*sigma^2);

function a = softmax(Q,b)
% Softmax action selection
V = exp(b*Q);
V(V == 0) = 1e-10;
V = cumsum(V);
V = V./max(V);              % Scale the values of the potential actions so they sum to one
randnum = rand(1);
for iA = 1:length(V),
    if randnum < V(iA),
        a = iA;
        break
    end
end


function [p_out pObs] = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen)
% Function for updating p(mu|D) and p(Obs|D) - according to observations  
%    only (no model contributes to the probability calculations)
for iX = 1:xlen,
    p_out(iX,:) = betapdf(mu,alpha(iX,1) + xhist(iX,1),alpha(iX,2) + xhist(iX,2));
end

% Make general prior
if use_genprior,
    htmp = sum(xhist');
    genprior = htmp*p_out;
    p_out = p_out.*repmat(genprior,xlen,1);
end

p_out = p_out./repmat(sum(p_out,2),1,mulen);
pObs = mu*p_out';

function [p_out pObs] = bayesupdateM(pM,probMu,mu,xlen,mulen,nSources)
% Function for updating p(mu|D) and p(Obs|D) - using the models and their
%    probabilities (which are driven by the data)
pObs = zeros(xlen,1);
p_out = zeros(xlen,mulen);
for iS = 1:nSources,
    pObs = pObs + pM(iS)*squeeze(probMu(iS,:,:))*mu';
    p_out = p_out + pM(iS)*squeeze(probMu(iS,:,:));
end
pObs = pObs';


function [pM] = bayesModelupdate(mu,probMu,xhist,xlen,mulen,nSources)
% Function for updating the probability of the model pM
tmp = zeros(xlen,mulen);
for iX = 1:xlen,
    tmp(iX,:) = mu.^xhist(iX,1).*(1 - mu).^xhist(iX,2);
end
tmp = tmp./repmat(sum(tmp,2),1,mulen);

pM = zeros(nSources,1);
for iS = 1:nSources,
    tmp2 = tmp.*squeeze(probMu(iS,:,:));
    pMx = nansum(tmp2,2);           % The probability of the data for each x
    pMx(pMx == 0) = 1e-12;          % Make zeros really small so we can take the sum of the logs rather than the product
    pM(iS) = exp(sum(log(pMx)));
end
pM = pM/nansum(pM);

function expIG = infogain_parameters(mu,alpha,xhist,xlen,mulen,use_genprior)
[p_out pObs] = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen);
p_old = p_out;                      % Current probability p(mu|D)
hist_old = xhist;                   % Current histogram
expIG = zeros(xlen,1);              % Information gain (for the parameters) at each sampling location

for iX = 1:xlen,
    infotmp = 0;
    for iO = 1:2,
        xhist(iX,iO) = xhist(iX,iO) + 1;                            % Increment histogram
        p_out = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen);
        KLtmp = nansum(p_out(iX,:).*log(p_out(iX,:)./p_old(iX,:)));    % KL div for a single new observation
        if iO == 1,
            infotmp = infotmp + pObs(iX)*KLtmp;                     % KL div weighted by probility of obs
        else,
            infotmp = infotmp + (1 - pObs(iX))*KLtmp;
        end
        xhist = hist_old;                           % Reset the history
    end
    expIG(iX) = infotmp;
end

function expIG = infogain_models(mu,probMu,xhist,pObs,xlen,mulen,nSources)
pM = bayesModelupdate(mu,probMu,xhist,xlen,mulen,nSources);
p_old = pM;                         % Current probability p(m|D)
hist_old = xhist;                   % Current histogram
expIG = zeros(xlen,1);              % Information gain (for the parameters) at each sampling location
% pObs_old = pObs;
pObs = zeros(xlen,1);
for iS = 1:nSources,
    pObs = pObs + pM(iS)*squeeze(probMu(iS,:,:))*mu';
end

for iX = 1:xlen,
    infotmp = 0;
    for iO = 1:2,
        xhist(iX,iO) = xhist(iX,iO) + 1;                            % Increment histogram
%         p_out = bayesupdate(mu,alpha,xhist,use_genprior,xlen,mulen);
        p_out = bayesModelupdate(mu,probMu,xhist,xlen,mulen,nSources);
        KLtmp = nansum(p_out.*log(p_out./p_old));                   % KL div for a single new observation
        if iO == 1,
            infotmp = infotmp + pObs(iX)*KLtmp;                     % KL div weighted by probility of obs
        else,
            infotmp = infotmp + (1 - pObs(iX))*KLtmp;
        end
        xhist = hist_old;                           % Reset the history
    end
    expIG(iX) = infotmp;
end

function plot01(x,mu,p_out,pObs,xhist,expIG)
subplot(9,1,6:7);
[AX,H1,H2] = plotyy(x,pObs,x,expIG);
set(AX(2),'ylim',[0 1.1*max(expIG)]);
set(get(AX(2),'Ylabel'),'String','Exp Infogain (bits)');
% set(AX(2),'ylim',[0 0.05]);
set(H1,'linewidth',2);
set(H2,'linewidth',2);
ylabel('E(\mu|D)');
subplot(9,1,8:9);
h = bar(x,xhist,'stacked');
set(h(1),'facecolor',[0.5 0.5 0.5]);
set(h(2),'facecolor',[0.2 0.2 0.2]);
set(gca,'xlim',[min(x) - median(diff(x))/2 max(x) + median(diff(x))/2]);
xlabel('Sample location');
ylabel('N samples');
subplot(9,1,1:5);
imagesc(x,mu,p_out');
axis xy
ylabel('p(\mu|D) (in the color axis)');
