function [traj, infStates] = tapas_hgf_jget(r, p, varargin)
% Calculates the trajectories of the agent's representations under the HGF model of the jumping
% Gaussian estimation task (JGET)
%
% This function can be called in two ways:
% 
% (1) tapas_hgf_jget(r, p)
%   
%     where r is the structure generated by fitModel and p is the parameter vector in native space;
%
% (2) tapas_hgf_jget(r, ptrans, 'trans')
% 
%     where r is the structure generated by fitModel, ptrans is the parameter vector in
%     transformed space, and 'trans' is a flag indicating this.
%
% --------------------------------------------------------------------------------------------------
% Copyright (C) 2013 Christoph Mathys, TNU, UZH & ETHZ
%
% This file is part of the HGF toolbox, which is released under the terms of the GNU General Public
% Licence (GPL), version 3. You can redistribute it and/or modify it under the terms of the GPL
% (either version 3 or, at your option, any later version). For further details, see the file
% COPYING or <http://www.gnu.org/licenses/>.


% Transform paramaters back to their native space if needed
if ~isempty(varargin) && strcmp(varargin{1},'trans');
    p = tapas_hgf_jget_transp(r, p);
end

% Number of levels
try
    l = r.c_prc.n_levels;
catch
    l = length(p)/8;
    
    if l ~= floor(l)
        error('tapas:hgf:UndetNumLevels', 'Cannot determine number of levels');
    end
end

% Unpack parameters
mux_0 = p(1:l);
sax_0 = p(l+1:2*l);
mua_0 = p(2*l+1:3*l);
saa_0 = p(3*l+1:4*l);
kau   = p(4*l+1);
kax   = p(4*l+2:5*l);
kaa   = p(5*l+1:6*l-1);
omu   = p(6*l);
omx   = p(6*l+1:7*l-1);
oma   = p(7*l+1:8*l-1);
thx   = exp(p(7*l));
tha   = exp(p(8*l));

% Add dummy "zeroth" trial
u = [0; r.u(:,1)];

% Number of trials (including prior)
n = length(u);

% Construct time axis
t = ones(n,1);

% Initialize updated quantities

% Representations
mux = NaN(n,l);
pix = NaN(n,l);
mua = NaN(n,l);
pia = NaN(n,l);

% Other quantities
muuhat = NaN(n,1);
piuhat = NaN(n,1);
muxhat = NaN(n,l);
pixhat = NaN(n,l);
muahat = NaN(n,l);
piahat = NaN(n,l);
daux   = NaN(n,1);
daua   = NaN(n,1);
wx     = NaN(n,l-1);
dax    = NaN(n,l-1);
wa     = NaN(n,l-1);
daa    = NaN(n,l-1);

% Representation priors
% Note: first entries of the other quantities remain
% NaN because they are undefined and are thrown away
% at the end; their presence simply leads to consistent
% trial indices.
mux(1,:) = mux_0;
pix(1,:) = 1./sax_0;
mua(1,:) = mua_0;
pia(1,:) = 1./saa_0;

% Representation update loop
% Pass through trials 
for k = 2:1:n
    if not(ismember(k-1, r.ign))
        
        %%%%%%%%%%%%%%%%%%%%%%
        % Effect of input u(k)
        %%%%%%%%%%%%%%%%%%%%%%
        
        % Input level
        % ~~~~~~~~~~~
        % Prediction (same as prediction of x_1, see below)
        muuhat(k) = mux(k-1,1);
        
        % Precision of prediction
        piuhat(k) = 1/exp(kau *mua(k-1,1) +omu);
        
        % Mean prediction error
        daux(k) = u(k) -muuhat(k);
        
        % 1st level
        % ~~~~~~~~~
        % Predictions
        muxhat(k,1) = mux(k-1,1);
        muahat(k,1) = mua(k-1,1);
        
        % Precisions of predictions
        pixhat(k,1) = 1/(1/pix(k-1,1) +t(k) *exp(kax(1) *mux(k-1,2) +omx(1)));
        piahat(k,1) = 1/(1/pia(k-1,1) +t(k) *exp(kaa(1) *mua(k-1,2) +oma(1)));
        
        % x-updates
        pix(k,1) = pixhat(k,1) +piuhat(k);
        mux(k,1) = muxhat(k,1) +piuhat(k)/pix(k,1) *daux(k);

        % Prediction error of input precision
        daua(k) = (1/pix(k,1) +(mux(k,1) -u(k))^2) *piuhat(k) -1;

        % alpha-updates
        pia(k,1) = piahat(k,1) +1/2 *kau^2 *(1 +daua(k));
        mua(k,1) = muahat(k,1) +1/2 *1/pia(k,1) *kau *daua(k);

        % Volatility prediction errors
        dax(k,1) = (1/pix(k,1) +(mux(k,1) -muxhat(k,1))^2) *pixhat(k,1) -1;
        daa(k,1) = (1/pia(k,1) +(mua(k,1) -muahat(k,1))^2) *piahat(k,1) -1;
        
        if l > 2
            % Pass through higher levels
            % ~~~~~~~~~~~~~~~~~~~~~~~~~~
            for j = 2:l-1
                % Predictions
                muxhat(k,j) = mux(k-1,j);
                muahat(k,j) = mua(k-1,j);
                
                % Precisions of predictions
                pixhat(k,j) = 1/(1/pix(k-1,j) +t(k) *exp(kax(j) *mux(k-1,j+1) +omx(j)));
                piahat(k,j) = 1/(1/pia(k-1,j) +t(k) *exp(kaa(j) *mua(k-1,j+1) +oma(j)));

                % Weighting factors
                wx(k,j-1) = t(k) *exp(kax(j-1) *mux(k-1,j) +omx(j-1)) *pixhat(k,j-1);
                wa(k,j-1) = t(k) *exp(kaa(j-1) *mua(k-1,j) +oma(j-1)) *piahat(k,j-1);

                % Updates
                pix(k,j) = pixhat(k,j) +1/2 *kax(j-1)^2 *wx(k,j-1) *(wx(k,j-1) +(2 *wx(k,j-1) -1) *dax(k,j-1));
                pia(k,j) = piahat(k,j) +1/2 *kaa(j-1)^2 *wa(k,j-1) *(wa(k,j-1) +(2 *wa(k,j-1) -1) *daa(k,j-1));

                if pix(k,j) <= 0 || pia(k,j) <= 0
                    error('tapas:hgf:NegPostPrec', 'Negative posterior precision. Parameters are in a region where model assumptions are violated.');
                end

                mux(k,j) = muxhat(k,j) +1/2 *1/pix(k,j) *kax(j-1) *wx(k,j-1) *dax(k,j-1);
                mua(k,j) = muahat(k,j) +1/2 *1/pia(k,j) *kaa(j-1) *wa(k,j-1) *daa(k,j-1);
    
                % Volatility prediction errors
                dax(k,j) = (1/pix(k,j) +(mux(k,j) -muxhat(k,j))^2) *pixhat(k,j) -1;
                daa(k,j) = (1/pia(k,j) +(mua(k,j) -muahat(k,j))^2) *piahat(k,j) -1;
            end
        end

        % Last level
        % ~~~~~~~~~~
        % Predictions
        muxhat(k,l) = mux(k-1,l);
        muahat(k,l) = mua(k-1,l);
        
        % Precision of prediction
        pixhat(k,l) = 1/(1/pix(k-1,l) +t(k) *thx);
        piahat(k,l) = 1/(1/pia(k-1,l) +t(k) *tha);

        % Weighting factor
        wx(k,l-1) = t(k) *exp(kax(l-1) *mux(k-1,l) +omx(l-1)) *pixhat(k,l-1);
        wa(k,l-1) = t(k) *exp(kaa(l-1) *mua(k-1,l) +oma(l-1)) *piahat(k,l-1);
        
        % Updates
        pix(k,l) = pixhat(k,l) +1/2 *kax(l-1)^2 *wx(k,l-1) *(wx(k,l-1) +(2 *wx(k,l-1) -1) *dax(k,l-1));
        pia(k,l) = piahat(k,l) +1/2 *kaa(l-1)^2 *wa(k,l-1) *(wa(k,l-1) +(2 *wa(k,l-1) -1) *daa(k,l-1));

        if pix(k,l) <= 0 || pia(k,l) <= 0
            error('tapas:hgf:NegPostPrec', 'Negative posterior precision. Parameters are in a region where model assumptions are violated.');
        end

        mux(k,l) = muxhat(k,l) +1/2 *1/pix(k,l) *kax(l-1) *wx(k,l-1) *dax(k,l-1);
        mua(k,l) = muahat(k,l) +1/2 *1/pia(k,l) *kaa(l-1) *wa(k,l-1) *daa(k,l-1);
    
        % Volatility prediction error
        dax(k,l) = (1/pix(k,l) +(mux(k,l) -muxhat(k,l))^2) *pixhat(k,l) -1;
        daa(k,l) = (1/pia(k,l) +(mua(k,l) -muahat(k,l))^2) *piahat(k,l) -1;
    else

        mux(k,:) = mux(k-1,:); 
        mua(k,:) = mua(k-1,:); 
        pix(k,:) = pix(k-1,:);
        pia(k,:) = pia(k-1,:);

        muuhat(k) = muuhat(k-1);
        piuhat(k) = piuhat(k-1);

        muxhat(k,:) = muxhat(k-1,:);
        muahat(k,:) = muahat(k-1,:);
        pixhat(k,:) = pixhat(k-1,:);
        piahat(k,:) = piahat(k-1,:);

        daux(k) = daux(k-1);
        daua(k) = daua(k-1);

        wx(k,:)  = wx(k-1,:);
        wa(k,:)  = wa(k-1,:);
        dax(k,:) = dax(k-1,:);
        daa(k,:) = daa(k-1,:);
        
    end
end

% Remove representation priors
mux(1,:)  = [];
mua(1,:)  = [];
pix(1,:)  = [];
pia(1,:)  = [];

% Check validity of trajectories
if any(isnan(mux(:))) || any(isnan(pix(:))) || any(isnan(mua(:))) || any(isnan(pia(:)))
    error('tapas:hgf:VarApproxInvalid', 'Variational approximation invalid. Parameters are in a region where model assumptions are violated.');
else
    % Check for implausible jumps in trajectories
    dmux = diff(mux);
    dmua = diff(mua);
    dpix = diff(pix);
    dpia = diff(pia);
    rmdmux = repmat(sqrt(mean(dmux.^2)),length(dmux),1);
    rmdmua = repmat(sqrt(mean(dmua.^2)),length(dmua),1);
    rmdpix = repmat(sqrt(mean(dpix.^2)),length(dpix),1);
    rmdpia = repmat(sqrt(mean(dpia.^2)),length(dpia),1);

    jumpTol = 256;
    if any(abs(dmux(:)) > jumpTol*rmdmux(:)) || any(abs(dmua(:)) > jumpTol*rmdmua(:)) || any(abs(dpix(:)) > jumpTol*rmdpix(:)) || any(abs(dpia(:)) > jumpTol*rmdpia(:))
        error('tapas:hgf:VarApproxInvalid', 'Variational approximation invalid. Parameters are in a region where model assumptions are violated.');
    end
end

% Remove other dummy initial values
muuhat(1)   = [];
piuhat(1)   = [];
muxhat(1,:) = [];
muahat(1,:) = [];
pixhat(1,:) = [];
piahat(1,:) = [];
wx(1,:)     = [];
wa(1,:)     = [];
daux(1)     = [];
daua(1)     = [];
dax(1,:)    = [];
daa(1,:)    = [];

% Extract learning rates
lrx = NaN(n-1,l);
lra = NaN(n-1,l);

lrx(:,1) = piuhat./pix(:,1);
lrx(:,2:end) = kax./2 *wx./pix(:,2:end);
lra(:,1) = 1/2 *kau./pia(:,1);
lra(:,2:end) = kaa./2 *wa./pia(:,2:end);

% Create result data structure
traj = struct;

traj.mux     = mux;
traj.mua     = mua;
traj.sax     = 1./pix;
traj.saa     = 1./pia;

traj.muuhat  = muuhat;
traj.muxhat  = muxhat;
traj.muahat  = muahat;
traj.sauhat  = 1./piuhat;
traj.saxhat  = 1./pixhat;
traj.saahat  = 1./piahat;

traj.wx      = wx;
traj.wa      = wa;

traj.daux    = daux;
traj.daua    = daua;
traj.dax     = dax;
traj.daa     = daa;

traj.lrx     = lrx;
traj.lra     = lra;

% Create matrices for use by the observation model
infStates = NaN(n-1,1,10);
infStates(:,1,1)  = traj.muuhat;
infStates(:,1,2)  = traj.sauhat;
infStates(:,1,3)  = traj.muxhat(:,1);
infStates(:,1,4)  = traj.saxhat(:,1);
infStates(:,1,5)  = traj.muahat(:,1);
infStates(:,1,6)  = traj.saahat(:,1);
infStates(:,1,7)  = traj.mux(:,1);
infStates(:,1,8)  = traj.sax(:,1);
infStates(:,1,9)  = traj.mua(:,1);
infStates(:,1,10) = traj.saa(:,1);

return;
