function f = jforecast(THIS,data,range,varargin)
% jforecast  Forecast with judgmental adjustments (conditional forecasts).
%
% Syntax
% =======
%
%     F = jforecast(M,D,RANGE,...)
%
% Input arguments
% ================
%
% * `M` [ model ] - Solved model object.
%
% * `D` [ struct ] - Input data from which the initial condition is taken.
%
% * `RANGE` [ numeric ] - Forecast range.
%
% Output arguments
% =================
%
% * `F` [ struct ] - Output struct with the judgmentally adjusted forecast.
%
% Options
% ========
%
% * `'anticipate='` [ *`true`* | `false` ] - If true, real future shocks are
% anticipated, imaginary are unanticipated; vice versa if false.
%
% * `'currentOnly='` [ *`true`* | `false` ] - If true, MSE matrices will be
% computed only for current-dated variables, not for their lags or leads.
%
% * `'deviation='` [ `true` | *`false`* ] - Treat input and output data as
% deviations from balanced-growth path.
%
% * `'dtrends='` [ *`'auto'`* | `true` | `false` ] - Measurement data contain
% deterministic trends.
%
% * `'initCond='` [ *`'data'`* | `'fixed'` ] - Use the MSE for the initial
% conditions if found in the input data or treat the initical conditions as
% fixed.
%
% * `'meanOnly='` [ `true` | *`false`* ] - Return only mean data, i.e. point
% estimates.
%
% * `'plan='` [ plan ] - Simulation plan specifying the exogenised variables
% and endogenised shocks.
%
% * `'vary='` [ struct | *empty* ] - Database with time-varying std
% deviations or cross-correlations of shocks.
%
% Description
% ============
%
% When adjusting the mean and/or std devs of shocks, you can use real and
% imaginary numbers ot distinguish between anticipated and unanticipated
% shocks:
%
% * any shock entered as an imaginary number is treated as an
% anticipated change in the mean of the shock distribution;
%
% * any std dev of a shock entered as an imaginary number indicates that
% the shock will be treated as anticipated when conditioning the forecast
% on the reduced-form tunes.
%
% * the same shock or its std dev can have both the real and the imaginary
% part.
%
% Description
% ============
%
% Example
% ========
%
%

% -IRIS Toolbox.
% -Copyright (c) 2007-2012 Jaromir Benes.

p = inputParser();
p.addRequired('m',@ismodel);
p.addRequired('data',@(x) isstruct(x) || iscell(x));
p.addRequired('range',@isnumeric);
p.parse(THIS,data,range);
range = range(1) : range(end);

if ~isempty(varargin) && ~ischar(varargin{1})
    cond = varargin{1};
    varargin(1) = [];
    iscond = true;
else
    cond = [];
    iscond = false;
end

opt = passvalopt('model.jforecast',varargin{:});

isplancond = isa(opt.plan,'plan') && ~isempty(opt.plan,'cond');
iscond = iscond || isplancond;

if isequal(opt.dtrends,'auto')
    opt.dtrends = ~opt.deviation;
end

% Tunes.
isswap = isplan(opt.plan) && ~isempty(opt.plan,'tunes');

% Create real and imag `stdcorr` vectors from user-supplied databases.
[opt.stdcorrreal,opt.stdcorrimag] = mytune2stdcorr(THIS,range,cond,opt);

% TODO: Remove 'missing', 'contributions' options from jforecast,
% 'anticipate' scalar.

%**************************************************************************

ny = size(THIS.solution{4},1);
nx = size(THIS.solution{1},1);
nb = size(THIS.solution{7},1);
nf = nx - nb;
ne = size(THIS.solution{2},2);
nalt = size(THIS.Assign,3);
nper = length(range);
xrange = range(1)-1 : range(end);
nxper = length(xrange);

% Current-dated variables in the original state vector.
if opt.currentonly
    xcurri = imag(THIS.solutionid{2}) == 0;
else
    xcurri = true(size(THIS.solutionid{2}));
end
nxcurr = sum(xcurri);
fcurri = xcurri(1:nf);
bcurri = xcurri(nf+1:end);

% Get initial condition for the alpha vector. The `datarequest` function
% always expands the `alpha` vector to match `nalt`. The `ainitmse` and
% `xinitmse` matrices can be empty.
[ainit,xinit,naninit,ainitmse,xinitmse] = datarequest('init',THIS,data,range);
% Check for availability of all initial conditions.
dochkinitcond();
ninit = size(ainit,3);
ninitmse = size(ainitmse,4);

% Get input data for y, current dates of [xf;xb], and e. The size of all
% data is equalised in 3rd dimensin in the `datarequest` function.
[yinp,xinp,einp] = datarequest('y,x,e',THIS,data,range);
ndata = size(xinp,3);

% Determine the total number of cycles.
nloop = max([nalt,ninit,ninitmse,ndata]);

lastorzerofunc = @(x) max([0,find(any(x,1),1,'last')]);
vecfunc = @(x) x(:);

if isswap || isplancond
    [ya,xa,ea,ua,Ya,Xa] = myanchors(THIS,opt.plan,range);
end

if isswap
    % Load positions (anchors) of exogenised and endogenised data points.
    if ~opt.anticipate
        [ea,ua] = deal(ua,ea);
    end
    xa = xa(xcurri,:);
    % Check for NaNs in exogenised variables, and check the number of
    % exogenised and endogenised data points.
    dochkexogenised();
    lastea = lastorzerofunc(ea);
    lastua = lastorzerofunc(ua);
    lastya = lastorzerofunc(ya);
    lastxa = lastorzerofunc(xa);
else
    lastea = 0;
    lastua = 0;
    lastya = 0;
    lastxa = 0;
end

if iscond
    % Load conditioning data.
    if isplancond
        Y = yinp;
        X = xinp;
        E = zeros(ne,nper);
        Xa = Xa(xcurri,:);
        X = X(xcurri,:,:);
    else
        Y = datarequest('y',THIS,cond,range);
        X = datarequest('x',THIS,cond,range);
        E = datarequest('e',THIS,cond,range);
        Y = Y(:,:,1);
        X = X(:,:,1);
        E = E(:,:,1);
        X = X(xcurri,:);
        Ya = ~isnan(Y);
        Xa = ~isnan(X);
    end
    lastYa = lastorzerofunc(Ya);
    lastXa = lastorzerofunc(Xa);
    iscond = lastYa > 0 || lastXa > 0;
    % Check for overlaps between shocks from input data and shocks from
    % conditioning data, and add up the overlapping shocks.
    dochkoverlap();
else
    lastYa = 0;
    lastXa = 0;
end

if opt.anticipate
    laste = lastorzerofunc(any(real(einp) ~= 0,3));
    lastu = lastorzerofunc(any(imag(einp) ~= 0,3));
else
    lastu = lastorzerofunc(any(real(einp) ~= 0,3));
    laste = lastorzerofunc(any(imag(einp) ~= 0,3));
end

last = max([lastxa,lastya,laste,lastea,lastu,lastua,lastYa,lastXa]);

if isswap
    ya = ya(:,1:last);
    xa = xa(:,1:last);
    ea = ea(:,1:last);
    ua = ua(:,1:last);
    % Indices of exogenised data points and endogenised shocks.
    exi = [ya(:).',xa(:).'];
    endi = [false,false(1,nb),ua(:).',ea(:).'];
else
    exi = false(1,(ny+nxcurr)*last);
    endi = false(1,1+nb+2*ne*last);
end

if iscond
    Ya = Ya(:,1:last,:);
    Xa = Xa(:,1:last,:);
    Y = Y(:,1:last,:);
    X = X(:,1:last,:);
    % Index of conditions on measurement and transition variables.
    condi = [Ya(:).',Xa(:).'];
    % Index of conditions on measurement and transition variables excluding
    % exogenised position.
    condinotexi = condi(~exi);
end

% Index of parameterisation with solutions not available.
[~,nansolution] = isnan(THIS,'solution');

% Initialise output data.
OUTP = struct();
OUTP.mean = myhdatainit(THIS,nxper,nloop,opt.precision);
if ~opt.meanonly
    OUTP.std = myhdatainit(THIS,nxper,nloop,opt.precision);
end

%**************************************************************************
% Main loop.

if opt.progress
    % Create progress bar.
    progress = progressbar('IRIS model.solve progress');
end

for iloop = 1 : nloop
    
    if iloop <= nalt
        % Expansion needed to t+k.
        k = max([1,last]) - 1;
        THIS = expand(THIS,k);
        Tf = THIS.solution{1}(1:nf,:,iloop);
        Ta = THIS.solution{1}(nf+1:end,:,iloop);
        R = THIS.solution{2}(:,:,iloop);
        Rf = R(1:nf,1:ne);
        Ra = R(nf+1:end,1:ne);
        Kf = THIS.solution{3}(1:nf,:,iloop);
        Ka = THIS.solution{3}(nf+1:end,:,iloop);
        Z = THIS.solution{4}(:,:,iloop);
        H = THIS.solution{5}(:,:,iloop);
        D = THIS.solution{6}(:,:,iloop);
        U = THIS.solution{7}(:,:,iloop);
        Ut = U.';
        % Compute deterministic trends if requested.
        if opt.dtrends
            W = mydtrendsrequest(THIS,'range',range,iloop);
        end
        % Expand solution forward.
        if opt.meanonly
            [M,Ma] = swapsystem(THIS,iloop,exi,endi,last);
        else
            [M,Ma,N,Na] = swapsystem(THIS,iloop,exi,endi,last);
            Nt = N.';
            Nat = Na.';
        end
        stdcorre = [];
        stdcorru = [];
        dostdcorr();
    end
    
    % Solution not available.
    if nansolution(min(iloop,end));
        continue
    end
    
    % Initial condition.
    a0 = ainit(:,1,min(iloop,end));
    x0 = xinit(:,1,min(end,iloop));
    if isempty(ainitmse) || isequal(opt.initcond,'fixed')
        Pa0 = zeros(nb);
        Dxinit = zeros(nb,1);
    else
        Pa0 = ainitmse(:,:,1,min(iloop,end));
        Dxinit = diag(xinitmse(:,:,min(iloop,end)));
    end
    
    % Expected and unexpected shocks.
    if opt.anticipate
        e = real(einp(:,:,min(end,iloop)));
        u = imag(einp(:,:,min(end,iloop)));
    else
        e = imag(einp(:,:,min(end,iloop)));
        u = real(einp(:,:,min(end,iloop)));
    end
    
    if isswap
        % Tunes on measurement variables.
        y = yinp(:,1:last,min(end,iloop));
        if opt.dtrends
            y = y - W(:,1:last);
        end
        % Tunes on transition variables.
        x = xinp(:,1:last,min(end,iloop));
        x = x(xcurri,:);
    else
        y = nan(ny,last);
        x = nan(nxcurr,last);
    end
    
    % Pre-allocate mean arrays.
    xcurr = nan(nxcurr,nper);
    
    % Pre-allocate variance arrays.
    if ~opt.meanonly
        Dy = nan(ny,nper);
        Dxcurr = nan(nxcurr,nper);
        Du = nan(ne,nper);
        De = nan(ne,nper);
    end
    
    % Solve the swap system.
    if last > 0
        % input := [const;a0;u;e].
        input = [+(~opt.deviation);a0(:); ...
            vecfunc(u(:,1:last));vecfunc(e(:,1:last))];
        % output := [y;x].
        output = [y(:);x(:)];
        
        % Swap exogenised outputs and endogenised inputs.
        % rhs := [input(~endi);output(exi)].
        % lhs := [output(~exi);input(endi)].
        rhs = [input(~endi);output(exi)];
        lhs = M*rhs;
        a = Ma*rhs;
        
        if ~opt.meanonly || iscond
            % Prhs is the MSE/Cov matrix of the RHS.
            Prhs = zeros(1+nb+2*ne*last);
            Prhs(1+(1:nb),1+(1:nb)) = Pa0;
            Pu = covfun.stdcorr2cov(stdcorru(:,1:last),ne);
            Pe = covfun.stdcorr2cov(stdcorre(:,1:last),ne);
            index = 1+nb+(1:ne);
            for i = 1 : last
                Prhs(index,index) = Pu(:,:,i);
                index = index + ne;
            end
            for i = 1 : last
                Prhs(index,index) = Pe(:,:,i);
                index = index + ne;
            end
            Prhs = Prhs(~endi,~endi);
            % Add zeros for the std errors of exogenised data points.
            if any(exi)
                Prhs = blkdiag(Prhs,zeros(sum(exi)));
            end
        end
        
        if ~opt.meanonly
            % Plhs is the cov matrix of the LHS.
            Plhs = N*Prhs*Nt;
            Pa = Na*Prhs*Nat;
            Plhs = (Plhs+Plhs.')/2;
            Pa = (Pa+Pa.')/2;
        end
        
        if iscond
            Yd = Y(:,:,min(end,iloop));
            Yd(~Ya) = NaN;
            if opt.dtrends
                Yd = Yd - W(:,1:last);
            end
            Xd = X(:,:,min(end,iloop));
            Xd(~Xa) = NaN;
            output = [Yd(:);Xd(:)];
            z = M(condinotexi,:);
            Pzt = Prhs*z.';
            F = z*Pzt;
            G = Pzt / F;
            % Update the RHS with conditioning information.
            rhs = rhs + G*(output(condi) - lhs(condi));
            % Re-run the forecast with conditioning information.
            lhs = M*rhs;
            a = Ma*rhs;
            if ~opt.meanonly
                % Update the MSE/cov mat of the RHS.
                z = N(condinotexi,:);
                Pzt = Prhs*z.';
                F = z*Pzt;
                G = Pzt / F;
                Prhs = (eye(size(Prhs)) - G*z)*Prhs;
                Prhs = (Prhs+Prhs.')/2;
                Plhs = N*Prhs*Nt;
                Pa = Na*Prhs*Nat;
                Plhs = (Plhs+Plhs.')/2;
                % Covariance of lhs and rhs.
                % Plhsrhs = N*Prhs;
            end
        end
        
        dolhsrhs2yxuea();
        
    else
        u = zeros(ne,last);
        e = zeros(ne,last);
        a = a0;
        if ~opt.meanonly
            Pa = Pa0;
        end
    end
    
    % Forecast between `last+1` and `nper`.
    dobeyond();
    
    % Free memory.
    a = [];
    Pa = [];
    
    % Add measurement detereministic trends.
    if opt.dtrends
        y = y + W;
    end
    
    % Store final results.
    doassignsmooth();
    
    if opt.progress
        % Update progress bar.
        update(progress,iloop/nloop);
    end
end
% End of main loop.

% Report parameterisation with solutions not available.
dochknansolution();

% Prepare final database.
doretsmooth();

% Nested functions follow.

%**************************************************************************
    function dochkinitcond()
        if ~isempty(naninit)
            naninit = unique(naninit);
            utils.error('model', ...
                'Initial condition not available for this variable: ''%s''.', ...
                naninit{:});
        end
    end
% dochkinitcond().

%**************************************************************************
    function dochkexogenised()
        % Check for NaNs in exogenised variables, and check the number of
        % exogenised and endogenised data points.
        index1 = [ya;xa];
        index2 = [any(isnan(yinp),3); ...
            any(isnan(xinp(xcurri,:,:)),3)];
        index = any(index1 & index2,2);
        if any(index)
            yvector = THIS.solutionvector{1};
            xvector = THIS.solutionvector{2};
            xvector = xvector(xcurri);
            vector = [yvector,xvector];
            % Some of the variables are exogenised to NaNs.
            utils.error('model', ...
                'This variable is exogenised to NaN: ''%s''.', ...
                vector{index});
        end
        % Check number of exogenised and endogenised data points.
        if nnzexog(opt.plan) ~= nnzendog(opt.plan)
            utils.warning('model', ...
                ['The number of exogenised data points (%g) does not match ', ...
                'the number of endogenised data points (%g).'], ...
                nnzexog(opt.plan),nnzendog(opt.plan));
        end
    end
% dochkexogenised().

%**************************************************************************
    function dochkoverlap()
        if any(E(:) ~= 0)
            if any(einp(:) ~= 0)
                utils.warning('model', ...
                    ['Both input data and conditioning data include ', ...
                    'structural shock adjustments, and will be added up.']);
            end
            einp = bsxfun(einp,E);
        end
    end
% dochkoverlap().

%**************************************************************************
    function dolhsrhs2yxuea()
        output = zeros((ny+nxcurr)*last,1);
        input = zeros((ne+ne)*last,1);
        output(~exi) = lhs(1:sum(~exi));
        output(exi) = rhs(sum(~endi)+1:end);
        input(~endi) = rhs(1:sum(~endi));
        input(endi) = lhs(sum(~exi)+1:end);
        y = reshape(output(1:ny*last),[ny,last]);
        output(1:ny*last) = [];
        xcurr(:,1:last) = reshape(output,[nxcurr,last]);
        output(1:nxcurr*last) = [];
        
        input(1) = [];
        x0 = U*input(1:nb);
        input(1:nb) = [];
        u = reshape(input(1:ne*last),[ne,last]);
        input(1:ne*last) = [];
        e = reshape(input(1:ne*last),[ne,last]);
        input(1:ne*last) = [];
        
        if opt.meanonly
            return
        end
        
        Poutput = zeros((ny+nxcurr)*last);
        Pinput = zeros((ne+ne)*last);
        Poutput(~exi,~exi) = Plhs(1:sum(~exi),1:sum(~exi));
        Poutput(exi,exi) = Prhs(sum(~endi)+1:end,sum(~endi)+1:end);
        Pinput(~endi,~endi) = Prhs(1:sum(~endi),1:sum(~endi));
        Pinput(endi,endi) = Plhs(sum(~exi)+1:end,sum(~exi)+1:end);
        
        index = 1 : ny;
        for t = 1 : last
            Dy(:,t) = diag(Poutput(index,index));
            index = index + ny;
        end
        Poutput(1:ny*last,:) = [];
        Poutput(:,1:ny*last) = [];
        index = 1 : nxcurr;
        for t = 1 : last
            Dxcurr(:,t) = diag(Poutput(index,index));
            index = index + nxcurr;
        end
        % Poutput(1:nxcurr*last,:) = [];
        % Poutput(:,1:nxcurr*last) = [];
        
        Pinput(1,:) = [];
        Pinput(:,1) = [];
        Pxinit = U*Pinput(1:nb,1:nb)*Ut;
        Dxinit = diag(Pxinit);
        Pinput(1:nb,:) = [];
        Pinput(:,1:nb) = [];
        index = 1 : ne;
        for t = 1 : last
            Du(:,t) = diag(Pinput(index,index));
            index = index + ne;
        end
        Pinput(1:ne*last,:) = [];
        Pinput(:,1:ne*last) = [];
        index = 1 : ne;
        for t = 1 : last
            De(:,t) = diag(Pinput(index,index));
            index = index + ne;
        end
        % Pinput(1:ne*last,:) = [];
        % Pinput(:,1:ne*last) = [];
    end
% dolhsrhs2yxue().

%**************************************************************************
    function dobeyond()
        % Simulate from last to nper.
        xcurr(:,last+1:nper) = 0;
        y(:,last+1:nper) = 0;
        e(:,last+1:nper) = 0;
        u(:,last+1:nper) = 0;        
        Ucurr = U(bcurri,:);
        Tfcurr = Tf(fcurri,:);
        Kfcurr = Kf(fcurri,:);
        for t = last+1 : nper
            xfcurr = Tfcurr*a;
            a = Ta*a;
            if ~opt.deviation
                xfcurr = xfcurr + Kfcurr;
                a = a + Ka;
            end
            xcurr(:,t) = [xfcurr;Ucurr*a];
            if ny > 0
                y(:,t) = Z*a;
                if ~opt.deviation
                    y(:,t) = y(:,t) + D;
                end
            end
        end
        if opt.meanonly
            return
        end
        Du(:,last+1:nper) = stdcorru(1:ne,last+1:nper).^2;
        De(:,last+1:nper) = stdcorre(1:ne,last+1:nper).^2;
        Tfcurrt = Tfcurr.';
        Tat = Ta.';
        Rfcurr = Rf(fcurri,:);
        Rfcurrt = Rfcurr.';
        Rat = Ra.';
        Ht = H.';
        Ucurrt = Ucurr.';
        Zt = Z.';
        for t = last+1 : nper
            Pue = covfun.stdcorr2cov(stdcorru(:,t),ne) ...
                + covfun.stdcorr2cov(stdcorre(:,t),ne);
            Pxfcurr = Tfcurr*Pa*Tfcurrt + Rfcurr*Pue*Rfcurrt;
            Pa = Ta*Pa*Tat + Ra*Pue*Rat;
            Pxbcurr = Ucurr*Pa*Ucurrt;
            Dxcurr(:,t) = [diag(Pxfcurr);diag(Pxbcurr)];
            if ny > 0
                Py = Z*Pa*Zt + H*Pue*Ht;
                Dy(:,t) = diag(Py);
            end
        end
    end
% dobeyond().

%**************************************************************************
    function dochknansolution()
        % Report parameterisations with solutions not available.
        if any(nansolution)
            utils.warning('model', ...
                ['No solution available, no forecast computed ', ...
                'for this parameterisation: #%g.'], ...
                num2cell(find(nansolution)));
        end
    end
% dochknansolution().

%**************************************************************************
    function dostdcorr()
        % TODO: use `mycombinestdcorr` here.
        % Combine `stdcorr` from the current parameterisation and the
        % `stdcorr` supplied through the tune database.
        stdcorre = THIS.stdcorr(1,:,iloop).';
        stdcorre = stdcorre(:,ones(1,nper));
        stdcorrixreal = ~isnan(opt.stdcorrreal);
        if any(stdcorrixreal(:))
            stdcorre(stdcorrixreal) = ...
                opt.stdcorrreal(stdcorrixreal);
        end
        
        stdcorru = THIS.stdcorr(1,:,iloop).';
        stdcorru = stdcorru(:,ones(1,nper));
        stdcorriximag = ~isnan(opt.stdcorrimag);
        if any(stdcorriximag(:))
            stdcorru(stdcorriximag) = ...
                opt.stdcorrimag(stdcorriximag);
        end
        
        % Set the std devs of the endogenised shocks to zero. Otherwise an
        % anticipated endogenised shock would have a non-zero unanticipated
        % std dev, and vice versa.
        if isswap
            tempstd = stdcorre(1:ne,1:last);
            tempstd(ea) = 0;
            tempstd(ua) = 0;
            stdcorre(1:ne,1:last) = tempstd;
            tempstd = stdcorru(1:ne,1:last);
            tempstd(ea) = 0;
            tempstd(ua) = 0;
            stdcorru(1:ne,1:last) = tempstd;
        end
        
        if ~opt.anticipate
            [stdcorru,stdcorre] = deal(stdcorre,stdcorru);
        end
    end
% dostdcorr().

%**************************************************************************
    function doassignsmooth()
        % Final point forecast.
        tempx = nan(nx,nxper);
        tempx(xcurri,2:end) = xcurr;
        tempx(nf+1:end,1) = x0;
        if opt.anticipate
            eu = complex(e,u);
        else
            eu = complex(u,e);
        end
        myhdataassign(THIS,OUTP.mean,iloop, ...
            [nan(ny,1),y], ...
            tempx, ...
            [nan(ne,1),eu]);
        % Final std forecast.
        if ~opt.meanonly
            Dx = nan(nx,nper);
            Dx(xcurri,:) = Dxcurr;            
            if opt.anticipate
                Deu = complex(De,Du);
            else
                Deu = complex(Du,De);
            end
            myhdataassign(THIS,OUTP.std,iloop, ...
                [nan(ny,1),Dy], ...
                [[nan(nf,1);Dxinit],Dx], ...
                [nan(ne,1),Deu], ...
                '-std');
        end
    end
% doassignsmooth().

%**************************************************************************
    function doretsmooth()
        f = struct();
        if opt.meanonly
            f = myhdata2tseries(THIS,OUTP.mean,xrange);
        else
            f.mean = myhdata2tseries(THIS,OUTP.mean,xrange);
            f.std = myhdata2tseries(THIS,OUTP.std,xrange);
        end
    end
% doretsmooth().

end