function [info] = fun_msiahvar(y, M, p)
%---------------------------------------------------------------------------------------------------
% function for estimating K-variate MSIAH(M)-VAR(p) model via EM algorithm
%---------------------------------------------------------------------------------------------------
% coding based on:
% Krolzig, H.-M., 1997. Markov-switching vector autoregressions : Modelling, statistical inference,
% and application to business cycle analysis. Springer, Berlin.
%---------------------------------------------------------------------------------------------------
% inputs:
%  -> y:       matrix of dimension (T+p)xK containing the time series of length (T+p) 
%              for all K dependent variables in (each regime-specific) VAR(p) system; 
%              each time series shall be sorted such that the oldest observation is 
%              in the first row and the newest observations is in the last row
%  -> M:       number of different regimes
%  -> p:       number of lags for (each regime-specific) VAR(p) model
%---------------------------------------------------------------------------------------------------
% output:
%  -> info:    structure array whose fields contain all results of the model estimation
%---------------------------------------------------------------------------------------------------
% requirements:
% -> MATLAB Statistics Toolbox / Statistics and Machine Learning Toolbox
% -> James P. Lesage's Spatial Econometrics Toolbox (version 7, updated 3/2010)
%    (available for download at https://www.spatial-econometrics.com/html/jplv7.zip)
%---------------------------------------------------------------------------------------------------

% check input parameters
if nargin < 3,                   error('all three input parameters to be specified');  end
if ~isscalar(M) || ~isscalar(p), error('input parameters M and p have to be scalars'); end
if numel(size(y)) >= 3,          error('input parameter y has too many dimensions');   end

% screen dimensions
T = size(y,1)-p; %number of observations for each dependent variable / equation
K = size(y,2);   %number of variables in (each regime-specific) VAR system

% set up RHS variable
x = ones(T,1);
  for i = 1:p
      x = [x y(p-i+1:end-i,:)];
  end
x = kron(x,eye(K));     %scale to size K (add K rows of zeros)
x = kron(eye(M),    x); %scale to size M (stacking of variables for M regimes)
                        %result: RHS matrix of dimension (M*K*T)xR with R as defined below; cf. Krolzig (1997), p. 108/176
                        %(will be transformed before being returned in 'info', see below)
R = size(x,2);          %total number of autoregressive parameters (incl. intercepts) to be estimated for MS-VAR model
                        %(R = (1+K*P)*K*M) for MSIAH-VAR model type with M regimes)

% set up LHS variable
y_base = y(p+1:end,:);         %LHS matrix of dimension (TxK); cf. Krolzig (1997), p. 191
y      = y_base';              %LHS matrix of dimension (KxT);
y      = kron(ones(M,1),y(:)); %LHS vector of dimension (M*K*Tx1); cf. equation (9.5) in Krolzig (1997), p. 176

% define (random) starting values
b     = randn(R,1)/30;                  %random        starting solution for parameter vector gamma = [gamma_1' gamma_2' ... gamma_M')' with all M regime-specific parameter sets
% b     = (x'*x)\x'*y;                  %OLS           starting solution for parameter vector gamma = [gamma_1' gamma_2' ... gamma_M')' with all M regime-specific parameter sets
% b     = [ones(R/2,1); 2*ones(R/2,1)]; %deterministic starting solution for parameter vector gamma = [gamma_1' gamma_2' ... gamma_M')' with all M regime-specific parameter sets
info.b_start = b;
b_old = ones(size(b));
P = rand(M);
P = P ./ kron(sum(P,2),ones(1,M));      %random starting solution for matrix of transition probabilities
info.P_start = P;
P_old = ones(size(P)) ./ 2;
e      = rand(M,1);
e      = e./sum(e);                     %random starting solution for regime probabilities in first row of matrix e_t1_t
% e      = repmat(1/M,M,1);               %naive guess using equal probabilites for all regimes
e_t_T  = ones(T,M);

% slight "optimization" of starting values
if M > 1
    differentiated_regimes = 0; counter = 0;
    while differentiated_regimes == 0 && counter < 100
        counter = counter +1;
        residuals = reshape(y - x*b,K,M*T)';
        n = calc_n(residuals); %see auxiliary function 'calc_n' on bottom of this script
        if all(sum(n) > 1e-10)
            differentiated_regimes = 1;
        else
            b = randn(R,1)/3;
        end
    end
end

%---------------------------------------------------------------------------------------------------------------------------------------------------------------------
% start of EM algorithm

counter_EM = 0;
while( mean(mean( abs(P-P_old)./P_old )) > 0.000001 && counter_EM < 200)

    counter_EM  = counter_EM + 1;
    residuals = reshape(y - x*b,K,M*T)';
    
    % set up starting values for (scaled) cov matrices of residuals
      if counter_EM == 1
         multIdx  = 0.75; %controls the factor for starting values of cov matrix
         myFactor = 1.5;  %controls how to factor each cov matrix for states
      else
         multIdx  = 1;    %controls the factor for starting values of cov matrix
         myFactor = 1;    %controls how to factor each cov matrix for states
      end
    for i = 1:M
        index = T*(i-1)+1 : T*i;                                              %index for picking regime-specific residual matrix of dimension (TxK) for each of the M regimes
        covmatrix_residuals_scaled = (1/sum(e_t_T(:,i))) * ( residuals(index,:)' * diag(e_t_T(:,i)) * residuals(index,:) ) .* multIdx; %slightly compress starting values for cov matrix
        multIdx = multIdx*myFactor;                                                                                                    %increase compression factor when moving to next regime
        n(:,i) = mvnpdf( residuals(index,:), 0, covmatrix_residuals_scaled ); %stack all vectors nt of size (Mx1) in transposed version for all t; result is matrix of dimension (TxM)
                                                                              %(each vector nt contains the M regime-specific distribution function values for time t)
                                                                              %(NOTE: function 'mvnpdf' from Statistics Toolbox / Statistics and Machine Learning Toolbox)
    end

    % compute regime probabilities via BLHK filter
    e_t_t  = zeros(T,M); %initialize matrix
    e_t_T  = zeros(T,M); %initialize matrix
    e_t1_t = zeros(T,M); %initialize matrix
    y_likl = zeros(T,1); %initialize matrix

    %-----------------------------------------------------------------------------------------------------------------------------------------------------------------
    % start of expectation step

    % calculate e_t1_t and e_t_t by iterating on equations (5.4) and (5.5) in Krolzig (1997), p. 82
    e_t1_t(1,:) = e;
    for i = 1:1:T
        % use (5.4) for e_t_t
        numerator      = n(i,:)' .* e_t1_t(i,:)';
        denominator    = ones(M,1)' * ( n(i,:)' .* e_t1_t(i,:)' );
        y_likl(i,:) = denominator;
        e_t_t(i,:)  = numerator / denominator;
        % use (5.5) for e_t1_t
        if i < T
           e_t1_t(i+1,:) = P' * e_t_t(i,:)';
        end
    end

    % calculate e_t_T as in equation (5.13) in Krolzig (1997), p. 85
    % (Kim's filter for calculating smoothed inferences, the 'K' in the BLHK filter)
    e_t_T(T,:) = e_t_t(T,:);
    for i = T-1:-1:1
        e_t_T(i,:) = ( P * ( e_t_T(i+1,:)' ./ e_t1_t(i+1,:)' ) ) .* e_t_t(i,:)';
    end

    % end of expectation step
    %-----------------------------------------------------------------------------------------------------------------------------------------------------------------
    % start of maximization step
    
    e1 = sum(e_t_T(1:(end-1),:))';               %vector of dimension (Mx1)
    e2 = e_t_T(1:(end-1),:)' * e_t_T(2:(end),:); %matrix of dimension (MxM) containing transition probabilities P(s_t = i, s_t-1 = j | T)
    e2 = reshape(e2,M*M,1);
    vectorized_P_estimate = e2 ./ ( kron(ones(M,1),e1) ); %cf. Krolzig (1997), p. 104 (also equation (6.14) in Krolzig (1997), p. 100)
    P_old = P;
    P = reshape(vectorized_P_estimate,M,M);
    P = P ./ kron(ones(1,M),sum(P,2));
    e = e_t_T(1,:)';
    
    % GLS estimation
    residuals = reshape(y - x*b,K,M*T)'; %matrix of dimension (M*TxK)
    xi = zeros(T,T,M);
    u  = zeros(T,K,M);
    for i = 1:M 
        xi(:,:,i) = diag(e_t_T(:,i));           %diagonal matrices of dimension (TxT); see XI_hat_m in second equation on p. 108 in Krolzig (1997)
        u(:,:,i)  = residuals(T*(i-1)+1:i*T,:); %residual matrices of dimension (TxK); see U_tilde_m on p. 191 in Krolzig (1997)
    end

    temp = zeros(K,K,M);
    for i = 1:M
        temp(:,:,i) = (u(:,:,i)' * xi(:,:,i) * u(:,:,i)) / sum(e_t_T(:,i)); %cov matrices of dimension (KxK); see Sigma_tilde_m on p. 191 in Krolzig (1997)
    end
    sigma_ml = inv(sum(temp,3)/T);

    W = zeros(M*T*K,M*T*K);
    for i = 1:M
        if all(all(temp(:,:,i) == 0))
            W(T*K*(i-1)+1:i*T*K,T*K*(i-1)+1:i*T*K) = kron(diag(e_t_T(:,i)),   (temp(:,:,i)));
        else
            W(T*K*(i-1)+1:i*T*K,T*K*(i-1)+1:i*T*K) = kron(diag(e_t_T(:,i)),inv(temp(:,:,i))); %weighting matrix of dimension (MTKxMTK); see W^(-1) on p. 108 in Krolzig (1997)
        end
    end

    b_old = b;
    b     = inv(x'*W*x)*x'*W*y; %cf. equation (9.6) in Krolzig (1997)

    % end of maximization step
    %-----------------------------------------------------------------------------------------------------------------------------------------------------------------

end %of EM algorithm
%---------------------------------------------------------------------------------------------------------------------------------------------------------------------


% compute 'conditional residuals'; cf. matrix u_tilde_mt on p. 135 in Krolzig (1997)
res = y - x*b;                      %vector of dimension (M*T*K x 1)
info.res1 = reshape(res,K,M*T)';    %matrix of dimension (M*T   x K)
% compute 'smoothed residuals'; cf. matrix u_tilde_t on p. 135 in Krolzig (1997)
info.res2 = zeros(T,K);
for i = 1:M
    info.res2 = info.res2 + (e_t_T(:,i)*ones(1,K)) .* info.res1( (i-1)*T+1:i*T, : );
end
% compute 'one-step prediction errors'; cf. matrix (e_hat_t|t-1) on p. 135 in Krolzig (1997)
info.res3 = zeros(T,K);
for i = 1:M
    info.res3 = info.res3 + (e_t1_t(:,i)*ones(1,K)) .* info.res1( (i-1)*T+1:i*T, : );
end

% compute fitted values corresponding to 'conditional residuals'
info.y_hat1 = x*b;                         % vector of dimension (M*T*K x 1) containing the M regime-specific vectors y_hat of dimension (T*K x 1)
info.y_hat1 = reshape(info.y_hat1,K,M*T)'; % matrix of dimension (M*T   x K) stacking all M regime-specific matrices of dimension (T x K)
% compute fitted values corresponding to 'smoothed residuals'
info.y_hat2 = zeros(T,K);
for i = 1:M
    info.y_hat2 = info.y_hat2 + (e_t_T(:,i)*ones(1,K)) .* info.y_hat1( (i-1)*T+1:i*T, : );
end
% compute fitted values corresponding to 'one-step prediction errors'
info.y_hat3 = zeros(T,K);
for i = 1:M
    info.y_hat3 = info.y_hat3 + (e_t1_t(:,i)*ones(1,K)) .* info.y_hat1( (i-1)*T+1:i*T, : );
end

% save regime-specific residual cov matrices
Sigma_regimes = temp;
info.res_covmatrices = Sigma_regimes;

info.b_clear = b;

% reconstruct all matrices Bm of dimension (K*P+1)x(K); cf. Krolzig (1997), p. 182-183 & p. 188-197
% (named "b_new" here)
b_new = zeros(K,p*K+1,M);
for i = 1:M
    b_new(:,:,i) = reshape(b((i-1)*(K^2*p+K)+1:i*(K^2*p+K)),K,p*K+1);
end
% transform x
x = x(1:K:(K*T),1:p*K+1);

% sort regimes (based on estimated variance for first variable / equation in system)
sort_vec = Sigma_regimes(1,1,:);
[~, sort_vec] = sort(sort_vec(:),'DESCEND');
b                    = b_new(               :,:,sort_vec);
e_t_T                = e_t_T(               :,  sort_vec);
e_t1_t               = e_t1_t(              :,  sort_vec);
e_t_t                = e_t_t(               :,  sort_vec);
Sigma_regimes        = Sigma_regimes(       :,:,sort_vec);
info.res_covmatrices = info.res_covmatrices(:,:,sort_vec);
info.b_new           = b;
info.x               = x;
info.y               = y_base;
info.Sigma           = Sigma_regimes;

% compute transition probabilities
e2 = zeros(T,M^2);
e1 = kron(ones(1,M),e_t_T);
for i = 1:M
    e2(2:end,(i-1)*M+1:i*M) = kron(ones(1,M),e_t_T(2:end,i)) .* e_t_T(1:end-1,:);
end
P = reshape(sum(e2(2:end,:)) ./ sum(e1(2:end,:)),M,M);
P = P./kron(ones(1,M),sum(P,2));
info.P         = P;
info.P_ergodic = info.P^(1000);

% compute regime durations
regime_durations = zeros(1,1,M);
for i = 1:M
     regime_durations(1,1,i) = 1/(1-P(i,i));
end
info.durations = regime_durations;

% save regime probabilities
info.t_T  = e_t_T;
info.t1_t = e_t1_t;
info.t_t  = e_t_t;

% compute and save additional statistics
info.y_likl       = sum(log(y_likl)); % logarithm of the likelihood function evaluated for the parameter estimates
                                      % (cf. equation (6.2) in Krolzig (1997), p. 93, for the likelihood function)
info.aic          = 2*R - 2*info.y_likl;
  rss             = sum( ( reshape(info.res3',K*T,1)       ).^2 );
  y_base_demeaned = y_base - ones(length(y_base),1)*mean(y_base);
  tss             = sum( ( reshape(y_base_demeaned',K*T,1) ).^2 );
info.r2           = 1 - (rss / tss); %cf. Krolzig (1997), p. 137
info.r2_equ_wise  = zeros(1,K); %individual R2 for each of the K dependent variables / equations in (regime-specific) VAR system
for i = 1:K
    info.r2_equ_wise(1,i) = 1 - ( sum(info.res3(:,i).^2) / sum(y_base_demeaned(:,i).^2) );
end
info.model = ['MSIAH(' num2str(M) ')-VAR(' num2str(p) ')']; %store model type/specification

% sort structure fields alphabetically
info = orderfields(info);


end %of 'fun_msiahvar' function



% --------------------------------------------------------------------------------------------------------------------------
% auxiliary function 'calc_n' (based on function 'mvnpdf' from Statistics Toolbox / Statistics and Machine Learning Toolbox)
function [n] = calc_n(residuals)
    M = evalin('caller','M');
    T = evalin('caller','T');    
    n = zeros(T,M);
    for i = 1:M
        index = T*(i-1)+1 : T*i;
        n(:,i) = mvnpdf( residuals(index,:), 0, cov(residuals(index,:)) );
    end
end