function oo_ = shock_decomposition(oo_, M_, options_, vname)
% function oo_ = shock_decomposition(oo_, M_, options_, vname)
%
% INPUTS
% - oo_           [structure]     MATLAB's structure containing the results
% - M_            [structure]     MATLAB's structure describing the model
% - options_      [structure]     MATLAB's structure describing the current options
% - vname         [cell]          array of variable names
% OUTPUT
% - oo_           [structure]     MATLAB's structure containing the results
%
% OccBin smoother data consumed from oo_ (produced by occbin.DSGE_smoother
%   and stored by store_smoother_results):
%   oo_.occbin.smoother.alphahat  (m*T)   raw smoothed states in decision-rule order
%                                         (deviations from steady state)
%   oo_.occbin.smoother.etahat    (r*T)   raw smoothed structural shocks
%   oo_.occbin.smoother.T0        {T}     time-varying transition matrices
%   oo_.occbin.smoother.R0        {T}     time-varying shock impact matrices
%   oo_.occbin.smoother.C0        {T}     time-varying constant vectors
%   These are distinct from oo_.SmoothedVariables/oo_.SmoothedShocks, which
%   are named structs in declaration order with trend+constant added.

% Copyright © 2021-2025 Dynare Team
%
% This file is part of Dynare.
%
% Dynare is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License as published by
% the Free Software Foundation, either version 3 of the License, or
% (at your option) any later version.
%
% Dynare is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
% GNU General Public License for more details.
%
% You should have received a copy of the GNU General Public License
% along with Dynare.  If not, see <https://www.gnu.org/licenses/>.


% list of options
% T0, use_shock_groups, vname, file_name, nfrcst, init_names_

shock_decomp_options = options_.occbin.shock_decomp;

if nargin<4 || isempty(vname)
    vname=cellstr(M_.endo_names(1:M_.orig_endo_nbr,:));
end

opts_simul = options_.occbin.simul;
opts_simul.periods = shock_decomp_options.periods;
opts_simul.check_ahead_periods = shock_decomp_options.check_ahead_periods;
opts_simul.maxit = shock_decomp_options.maxit;
use_shock_groups = shock_decomp_options.use_shock_groups;

init_names_ = shock_decomp_options.init_names_;
nfrcst = shock_decomp_options.nfrcst;

%% new Dynare grouping
if isempty(use_shock_groups)
    use_shock_groups = 'ALL';
    ngroups = M_.exo_nbr;
    ex_names_ = cell(ngroups,1);
    for i=1:ngroups
        ex_names_{i} = M_.exo_names(i);
    end
    shock_decomp_options.main_effect=0;
    shock_decomp_options.additive = 0;
else
    shock_groups = M_.shock_groups.(use_shock_groups);
    shock_ind = fieldnames(shock_groups);
    ngroups = length(shock_ind);
    ex_names_ = shock_ind;
    for i=1:ngroups
        ex_names_{i} = shock_groups.(shock_ind{i}).shocks;
    end
end
%%

if iscell(vname{1})
    vname0=vname;
    clear vname
    vname = cell(1,length(vname0));
    decomp_type = cell(1,length(vname0));
    gtrend = cell(1,length(vname0));
    var_type = nan(1,length(vname0));
    for j=1:length(vname0)
        vname{j}=vname0{j}{1};
        decomp_type{j}=vname0{j}{2};
        if strcmpi(decomp_type{j},'aoa')
            gtrend{j}=vname0{j}{3};
            if strcmpi(vname0{j}{4},'flow')
                var_type(j)=1;
            elseif strcmpi(vname0{j}{4},'deflator')
                var_type(j)=2;
            elseif strcmpi(vname0{j}{4},'stock')
                var_type(j)=0;
            else
                error('wrong var type for aoa decomp')
            end
        end
    end
    clear vname0
else
    decomp_type = cell(1,length(vname));
    for j=1:length(vname)
        decomp_type{j}='qoq';
    end
end

% set time-varying state space matrices
TT = oo_.occbin.smoother.T0;
RR = oo_.occbin.smoother.R0;
CC = oo_.occbin.smoother.C0;

load (options_.datafile, 'T');

indx_init=zeros(length(init_names_),1);
for j=1:length(init_names_)
    indx_init(j)=strmatch(init_names_{j},M_.endo_names(oo_.dr.order_var,:),'exact');
end

if shock_decomp_options.debug
    % re-compute smoothed variables from T, R and  smoothed shocks etahat
    a1=oo_.occbin.smoother.alphahat(:,1);
    as=a1;
    as0=zeros(length(a1),1);
    as0(indx_init)=a1(indx_init);
    etahat=[oo_.occbin.smoother.etahat zeros(size(oo_.occbin.smoother.etahat,1),nfrcst)];
    for j=1:size(etahat,2)-1
        TM = TT(:,:,j+1);
        RM = RR(:,:,j+1);
        CONST = CC(:,j+1);
        as(:,j+1)=TM*as(:,j)+RM*etahat(:,j+1)+CONST;
        as0(:,j+1)=TM*as0(:,j)+CONST;
    end
    err = max(max(abs(oo_.occbin.smoother.alphahat-as)));
    if err>1e-8
        disp('WARNING CHECK SMOOTHER:')
        disp(['simulated model with smoothed shocks differs from stored smoother by ' num2str(err)])
    else
        disp('CHECK SMOOTHER OK:')
        disp(['simulated model with smoothed shocks differs from stored smoother by less than 1e-8 (err=' num2str(err) ')'])
    end
else
    etahat=[oo_.occbin.smoother.etahat zeros(size(oo_.occbin.smoother.etahat,1),nfrcst)];
end
gend=size(etahat,2);

%%
TT1= dates('0Q1'):(dates('0Q1')+gend-1);
if exist('T','var')
    TT1=TT1+T(options_.first_obs)*4;
end

shock_decomp_options = set_default_option(shock_decomp_options,'TINIT',TT1(1));
tinit = max([1,find(TT1==shock_decomp_options.TINIT)]);

%%%%%%%%%%%%%%%%%%%% LINEAR shock decomp
as_lin=zeros(M_.endo_nbr,length(tinit:gend)); % linear smoother reconstructed
att=zeros(M_.endo_nbr,length(tinit:gend)); % linear initial condition effect
inn=zeros(M_.endo_nbr,length(tinit:gend)); % linear aggreage shocks effect without att, i.e. as_lin = inn+att;
deco=zeros(M_.endo_nbr,M_.exo_nbr,length(tinit:gend)); % full decomposition into individual shocks
deco_init = zeros(M_.endo_nbr,M_.endo_nbr+2,length(tinit:gend)); % full decomposition into individual initval
deco_init(:,end,:) = oo_.occbin.linear_smoother.alphahat(:,tinit:gend);
for i=1:M_.endo_nbr
    deco_init(i,i,1) = oo_.occbin.linear_smoother.alphahat(i,tinit);
end

att(:,1)=oo_.occbin.linear_smoother.alphahat(:,tinit);
as_lin(:,1)=oo_.occbin.linear_smoother.alphahat(:,tinit);
TM = oo_.occbin.linear_smoother.T0;
RM= oo_.occbin.linear_smoother.R0;
for j=2:length(tinit:gend)
    as_lin(:,j) = TM*as_lin(:,j-1)+RM*oo_.occbin.linear_smoother.etahat(:,j+tinit-1);
    att(:,j) = TM*att(:,j-1);
    inn(:,j) = RM*oo_.occbin.linear_smoother.etahat(:,j+tinit-1);
    if j>1
        inn(:,j) = inn(:,j) +  TM*inn(:,j-1);
    end
    for iexo=1:M_.exo_nbr
        deco(:,iexo,j) = RM(:,iexo)*oo_.occbin.linear_smoother.etahat(iexo,j+tinit-1);
        if j>1
            deco(:,iexo,j) = deco(:,iexo,j) +  TM*deco(:,iexo,j-1);
        end

    end
    if j>1
        deco_init(:,1:M_.endo_nbr,j) = TM*deco_init(:,1:M_.endo_nbr,j-1);
    end
    deco_init(:,M_.endo_nbr+1,j) = deco_init(:,M_.endo_nbr+2,j) - sum(deco_init(:,1:M_.endo_nbr,j),2);

end
as_lin=as_lin(oo_.dr.inv_order_var,:); % linear smoother reconstructed
att=att(oo_.dr.inv_order_var,:); % linear initial condition effect
inn=inn(oo_.dr.inv_order_var,:); % linear aggreage shocks effect without att
deco=deco(oo_.dr.inv_order_var,:,:); % full decomposition into individual shocks
deco(:,M_.exo_nbr+1,:)=att;
deco(:,M_.exo_nbr+2,:)=as_lin;
oo_.occbin.linear_smoother.decomp=deco;
deco_init=deco_init(oo_.dr.inv_order_var,oo_.dr.inv_order_var,:); % full decomposition into individual initval
oo_.occbin.linear_smoother.init_decomp=deco_init;
%%%%%%%%%%%%%%%%%%%% END LINEAR shock decomp

%%

%%%%%%%%%%%%%%%%%%%% piecewise CONDITIONAL shock decomp
as_p=zeros(M_.endo_nbr,length(tinit:gend)); % smoother reconstructed
att_p=zeros(M_.endo_nbr,length(tinit:gend)); % initial condition effect
inn_p=zeros(M_.endo_nbr,length(tinit:gend)); % aggreage shocks effect without att, i.e. as_lin = inn+att;
deco_p=zeros(M_.endo_nbr,M_.exo_nbr,length(tinit:gend)); % full decomposition into individual shocks
deco_p_init = zeros(M_.endo_nbr,M_.endo_nbr+2,length(tinit:gend)); % full decomposition into individual initval
deco_p_init(:,end,:) = oo_.occbin.smoother.alphahat(:,tinit:gend);
for i=1:M_.endo_nbr
    deco_p_init(i,i,1) = oo_.occbin.smoother.alphahat(i,tinit);
end
reg_p=zeros(M_.endo_nbr,length(tinit:gend)); % pure regime effect (CONST 'shocks')

att_p(:,1)=oo_.occbin.smoother.alphahat(:,tinit);
as_p(:,1)=oo_.occbin.smoother.alphahat(:,tinit);
TT = oo_.occbin.smoother.T0;
RR= oo_.occbin.smoother.R0;
CC= oo_.occbin.smoother.C0;

for j=2:length(tinit:gend)
    TM = TT(:,:,j+tinit-1);
    RM = RR(:,:,j+tinit-1);
    CONST = CC(:,j+tinit-1);
    as_p(:,j) = TM*as_p(:,j-1)+RM*oo_.occbin.smoother.etahat(:,j+tinit-1)+CONST;
    att_p(:,j) = TM*att_p(:,j-1);
    inn_p(:,j) = RM*oo_.occbin.smoother.etahat(:,j+tinit-1);
    reg_p(:,j) = TM*reg_p(:,j-1)+CONST;
    if j>1
        inn_p(:,j) = inn_p(:,j) +  TM*inn_p(:,j-1) ;
    end
    for iexo=1:M_.exo_nbr
        deco_p(:,iexo,j) = RM(:,iexo)*oo_.occbin.smoother.etahat(iexo,j+tinit-1);
        if j>1
            deco_p(:,iexo,j) = deco_p(:,iexo,j) +  TM*deco_p(:,iexo,j-1);
        end

    end
    if j>1
        deco_p_init(:,1:M_.endo_nbr,j) = TM*deco_p_init(:,1:M_.endo_nbr,j-1);
    end
    deco_p_init(:,M_.endo_nbr+1,j) = deco_p_init(:,M_.endo_nbr+2,j) - sum(deco_p_init(:,1:M_.endo_nbr,j),2);
end
as_p=as_p(oo_.dr.inv_order_var,:); % occbin smoother reconstructed
att_p=att_p(oo_.dr.inv_order_var,:); % occbin initial condition effect
inn_p=inn_p(oo_.dr.inv_order_var,:); % occbin aggregage shocks effect without att
deco_p=deco_p(oo_.dr.inv_order_var,:,:); % occbin full decomposition into individual shocks
reg_p=reg_p(oo_.dr.inv_order_var,:); % occbin pure regime effect (CONST 'shocks')
i_reg=strmatch('EPS_REGIME',M_.exo_names,'exact');
if ~isempty(i_reg)
    deco_p(:,i_reg,:)=reg_p;
else
    att_p = att_p + reg_p;
end
deco_p(:,M_.exo_nbr+1,:)=att_p;
deco_p(:,M_.exo_nbr+2,:)=as_p;
oo_.occbin.smoother.decomp=deco_p;
if isempty(i_reg)
    deco_p(:,M_.exo_nbr+1,:)=att_p-reg_p;
end
deco_p_init=deco_p_init(oo_.dr.inv_order_var,oo_.dr.inv_order_var,:); % full decomposition into individual initval
oo_.occbin.smoother.init_decomp=deco_p_init;

rr=abs(deco_p(:,1:end-1,:));
if isequal(use_shock_groups,'ALL') && ~isempty(i_reg)
    rr(:,i_reg,:)=0;
end
ww=zeros(size(rr));
for k=1:size(rr,3)
    tmp=sum(rr(:,:,k),2);
    tmp(tmp<1.e-10)=1;
    for g=1:size(rr,2)
        ww(:,g,k) = rr(:,g,k)./tmp;
    end
end
wdeco_p=deco_p;
if ~isempty(i_reg)
    wdeco_p(:,i_reg,:)=0;
end
for k=1:size(rr,3)
    if any(any(reg_p(:,k)))
        for g=1:size(rr,2)
            wdeco_p(:,g,k) = wdeco_p(:,g,k)+reg_p(:,k).*ww(:,g,k);
        end
    end
end
oo_.occbin.smoother.wdecomp=wdeco_p;
%%%%%%%%%%%%%%%%%%%% END CONDITIONAL shock decomp
