function [smpl_coeff, Sigma1_keep, Sigma0_keep, smpl_st, smooth_t_keep, filter_t_keep]=...
    Bay_RS_Multi_all_fit(y, x, nrep, warmup)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Model: Y1t=Xt*B(St) +  V(st)^1/2(ut), where V(st)=Sigma_low when st=0 and Sigma_high when st=1
%% Allow regime shift everything: intercept, coefficients, and cov-var matrix
%% This function can be used to test Timmerman(2008)'s paper in a Bayesian Gibbs-sampling approach
% Written By Weijun Xu, School of Banking & Finance, UNSW
% First edition, 01/08/08
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% First let's do some data transformation
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[T,k]=size(x1);
num_eq=size(y,2);
yt=y';                         % Yt: num_eq x T           
for i=1:T                      % Xt: num_eq x k x T, three-D matrix
    xt(:,:,i)=kron(eye(num_eq,x(i,:)); 
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Setting the priors: Assume there is only two latent states
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Transition probability prior
a11=10; b10=1;                    % tp11~Beta(a11, b10)                   
a00=10; b01=1;                    % tp00~Beta(a00, b01), tp10->St=0|St-1=1
% the tranition matrix is like:          St=1            St=0
%                               St-1=1   tp11            (1-tp11)
%                               St-1=0   (1-tp00)         tp00
% Coefficient priors
OLS_coeff=x\y;  
OLS_sigma=(y-x*OLS_coeff)'*(y-x*OLS_coeff)/(T-k);

ub1=[-1e-3*ones(1,k); OLS_coeff(2:end,:)];
ub1=ub1(:);       vb1=10*eye(num_eq*k);     % we want high volitliity regime has low intercept
ub0=OLS_coeff(:); vb0=10*eye(num_eq*k);

% Cov-Var Matrix priors: Sigma0_draw & Sigma1_draw
v1=10; omeg1=(1.2*v1*OLS_sigma)^-1;         % Sigma1_draw ~ Inv(Wishrnd(omeg1, v1))
v0=10; omeg0=(0.8*v0*OLS_sigma)^-1;         % Sigma0_draw ~ Inv(Wishrnd(omeg0, v0))

% Setting initial values
beta1_draw=multi_normrnd(ub1, vb1,1)';      % regime shift intercetpts, High volatility has low intercept
beta0_draw=multi_normrnd(ub0, vb0,1)';
Sigma1_draw=inv(wishrnd(omeg1,v1));
Sigma0_draw=inv(wishrnd(omeg0,v0));
tp11_draw=0.9; tp00_draw=0.9;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Running Gibbs-Sampling
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
for iter=1:nrep+warmup
    % step 1: draw latent state variables: st
    [st_draw, smooth_t, filter_t]=state_draw_Multi(yt,xt, beta1_draw, beta0_draw,  Sigma1_draw, Sigma0_draw, tp11_draw, tp00_draw,T,num_eq);
  
    % step 2: draw the posterior coefficients: beta1_draw & beta0_draw
    num_1=sum(st_draw); % num of states been 1
    num_0=T-num_1;      % num of states been 0
    yt_star1=yt(:,(st_draw==1));      yt_star0=yt(:,(st_draw==0));     % yt_star1=yt when st==1
    xt_star1=xt(:,:,(st_draw==1));    xt_star0=xt(:,:,(st_draw==0));   % xt_star0=xt when st==0
    
    % sample beta1_draw    
    vb_bar1=zeros(num_eq*k); b_bar1=zeros(num_eq*k,1);
    for i=1:num_1   % when st==1
        vb_bar1=vb_bar1+xt_star1(:,:,i)'*inv(Sigma1_draw)*xt_star1(:,:,i);
        b_bar1=b_bar1+xt_star1(:,:,i)'*inv(Sigma1_draw)*yt_star1(:,i);
    end
    Db1=(vb_bar1+inv(vb1))^-1;
    db1=b_bar1+inv(vb1)*ub1;

    % sample beta0_draw    
    vb_bar0=zeros(num_eq*k); b_bar0=zeros(num_eq*k,1);
    for i=1:num_0   % when st==0
        vb_bar0=vb_bar0+xt_star0(:,:,i)'*inv(Sigma0_draw)*xt_star0(:,:,i);
        b_bar0=b_bar0+xt_star0(:,:,i)'*inv(Sigma0_draw)*yt_star0(:,i);
    end
    Db0=(vb_bar0+inv(vb0))^-1;
    db0=b_bar0+inv(vb0)*ub0;
        
    beta1_draw=multi_normrnd(Db1*db1, Db1,1)';    
    beta0_draw=multi_normrnd(Db0*db0, Db0,1)';
   
    % step 3: draw the posterior volatility variance: Sigma0_draw & Sigma1_draw
    % Note that, diag(Sigma1_draw)>diag(Sigma0_draw) since regime 1 is high volatility regime    
    post_omeg1=zeros(num_eq);    
    for i=1:num_1  % when st=1
        post_omeg1=post_omeg1+(yt_star1(:,i)- xt_star1(:,:,i)*beta1_draw)*...
            (yt_star1(:,i)- xt_star1(:,:,i)*beta1_draw)';
    end    
    post_omeg1=(post_omeg1+inv(omeg1))^-1;
    post_v1=num_1+v1;
    Sigma1_draw= inv(wishrnd(post_omeg1,post_v1));

    post_omeg0=zeros(num_eq);
    for i=1:num_0  % when st=0
        post_omeg0=post_omeg0+(yt_star0(:,i)- xt_star0(:,:,i)*beta0_draw)*...
            (yt_star0(:,i)-xt_star0(:,:,i)*beta0_draw)';
    end
    post_omeg0=(post_omeg0+inv(omeg0))^-1;
    post_v0=num_0+v0;
    Sigma0_draw= inv(wishrnd(post_omeg0,post_v0));

    % make sure each element of diag(Sigma1_draw)> diag(Sigma0_draw), since we assume state 1 has high volatility
    j=1;
    while min(diag(Sigma1_draw)-diag(Sigma0_draw))<0
        Sigma0_draw=inv(wishrnd(post_omeg0,post_v0));
        Sigma1_draw=inv(wishrnd(post_omeg1,post_v1));
        j=j+1;
        if j>=50
            Sigma1_draw=inv(wishrnd(omeg1,v1));  % if after 50 iterations, we still not get results, we sample directly from the priors
            Sigma0_draw=inv(wishrnd(omeg0,v0));  
            break
        end
    end
    
    % step 5: draw the transistion probability: tp11,and tp00
    [n11, n10, n01, n00]=sigma_draw_number_multi(st_draw);  % n10 is number of states that St=0|St-1=1
    tp11_draw=betarnd(a11+n11, b10+n10);   % tp11~Beta(a11+n11, b10+n10), where n10=number of times that St-1=1|St=0
    tp00_draw=betarnd(a00+n00, b01+n01);   % tp00~Beta(a00+n00, b01+n01)

    
    if iter>warmup
        beta1_keep(iter-warmup,:)=beta1_draw';
        beta0_keep(iter-warmup,:)=beta0_draw';
        Sigma1_keep(:,:,iter-warmup)=Sigma1_draw;
        Sigma0_keep(:,:,iter-warmup)=Sigma0_draw;    
        tp11_keep(iter-warmup,:)=tp11_draw;
        tp00_keep(iter-warmup,:)=tp00_draw;
        smpl_st(iter-warmup,:)=st_draw';
        filter_t_keep(iter-warmup,:)=filter_t';
        smooth_t_keep(iter-warmup,:)=smooth_t';
    end
    fprintf(' Iteration to %3g \n', iter)
end
smpl_coeff=[beta1_keep, beta0_keep, tp11_keep, tp00_keep];

end


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Subfunction 1: Generate latent state variable, follows Chib(1996)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [st_draw, smooth_t, filter_t]=state_draw_Multi(yt,xt,beta1_draw, beta0_draw,  Sigma1_draw, Sigma0_draw, tp11_draw, tp00_draw,T,num_eq)
% Usage: to generate the latent state variable for Multivariate regression
% this function is only for two state case

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% first let's define some initial settings
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
st_draw=zeros(T,1);
P=[tp11_draw, 1-tp11_draw; 1-tp00_draw, tp00_draw]; % which is the transistion matrix
filter_t=zeros(T,2);       % which is the filter proability: filter_t=Pr(St|Y_t)
conditional_t=zeros(T,2);  % conditional_t=Pr(St|Y_t-1)
smooth_t=zeros(T,2);       % smooth_t=Pr(St|Y_T, S_t+1,...,S_T)
filter_t(1,:)=[0.5, 0.5];  % assign the initial filter_t value
conditional_t(1,:)=[0.5, 0.5];
U=unifrnd(0,1,T,1);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Calculate the filter probability, then calculate  posterior conditional probability: P(st|Y_T)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
for i=2:T
    likelihood(i,1)=(2*pi)^(-0.5*num_eq) *(det(Sigma1_draw)).^(-0.5)*...
        exp(-0.5*(yt(:,i)-xt(:,:,i)*beta1_draw)'*inv(Sigma1_draw)*...
        (yt(:,i)-xt(:,:,i)*beta1_draw));      % when st=1, we use sigma1_draw, and beta1_draw
    likelihood(i,2)=(2*pi)^(-0.5*num_eq) *(det(Sigma0_draw)).^(-0.5)*...
        exp(-0.5*(yt(:,i)-xt(:,:,i)*beta0_draw)'*inv(Sigma0_draw)*...
        (yt(:,i)-xt(:,:,i)*beta0_draw));      % when st=0, we use sigma0_draw, and beta0_draw

    if sum(likelihood(i,:)==0)>1        % i.e. two likelihood are both zero, avoid 0/0 problem
        likelihood(i,:)=1e-30.*ones(2,1);
    end
    conditional_t(i,:)=filter_t(i-1,:)*P;
    joint_likelihood(i,1)=sum(conditional_t(i,:).*likelihood(i,:));
    filter_t(i,:)=(conditional_t(i,:).*likelihood(i,:))./joint_likelihood(i,1);   % Normalized filter_t
end

st_draw(end)=(U(end)<=filter_t(end,1));  % if random uniform number is less than filter_t, we set it as 1, otherwise set it as 0
% caculate the posterior conditional probability, i.e. the smooth probability: Pr(St|Y_T, S_t+1,...,S_T)
for i=T-1:-1:1
    if st_draw(i+1)==1
        smooth_t(i,:)=filter_t(i,:).*P(:,1)';   % -->[filter_t_1*p11, filter_t_0* p01]
        smooth_t(i,:)=smooth_t(i,:)./sum(smooth_t(i,:),2);   % Normalize p_st1
    elseif st_draw(i+1)==0
        smooth_t(i,:)=filter_t(i,:).*P(:,2)';   % -->[filter_t_1*p10, filter_t_0* p00]
        smooth_t(i,:)=smooth_t(i,:)./sum(smooth_t(i,:),2);
    end
    st_draw(i)=(U(i)<=smooth_t(i,1));
end
conditional_t=conditional_t(:,1);       % conditional_t probabiilty:  P(St=1|Y_t-1)
filter_t=filter_t(:,1);                 % filter_t probability: P(St=1|Y_t)
smooth_t(end,:)=filter_t(end,:);
smooth_t=smooth_t(:,1);                 % smooth_t density:     P(St=1|Y_T, S_T+1)

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Subfunction 2: calculate the st count number for generate the transition probability
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [n11, n10, n01, n00]=sigma_draw_number_multi(st_draw)
% to calculate the count number n11, n10, n01, n00;
% generally, one count will be missing in the total sum of counts
T=size(st_draw,1);
n11=0; n10=0; n01=0; n00=0;
for i=2:T
    if  st_draw(i)==1
        if st_draw(i-1)==1
            n11=n11+1;     % n10 is number of times that St=0|St-1=1
        else
            n01=n01+1;
        end
    elseif  st_draw(i)==0
        if st_draw(i-1)==0
            n00=n00+1;
        else
            n10=n10+1;
        end
    end
end
end



