function [smpl_coeff, Sigma1_keep, Sigma0_keep, smpl_st, smooth_t_keep,  filter_t_keep, in_sample_results, out_sample_results, residual]=...
    Bay_RS_Multi_intercept_VAR_fit(y, x, nrep, warmup, out_sample_forecast_step)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Model: Y1t=Xt*B + St*Alpha + V(st)^1/2(ut), where V(st)=Sigma_low when st=0 and Sigma_high when st=1
%% Bayesian Gibbs-Sampling for MRS-VAR model, Allow regime shift in intercept and cov-var matrix
%% Constraints on ALpha < 0, and Variance in regime 1 > variance in regime 0
% Written By Weijun Xu, School of Banking & Finance, UNSW
% First edition, 01/08/08
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if nargin<=4 | isempty(out_sample_forecast_step)
    out_sample_forecast_step=0;
else
    out_sample_forecast_step=out_sample_forecast_step;
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% First let's do some data transformation
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
[T,k]=size(x);
num_eq=size(y,2);
num_cov=num_eq*(num_eq+1)/2;   %num of covariance parameters need to be estimated

yt=y';          % Yt:  num_eq x T
for i=1:T       % Xt: num_eq x k x T, three dimension 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

% Un-regime shift beta coefficients
OLS_coeff=x\y;
OLS_sigma=(y-x*OLS_coeff)'*(y-x*OLS_coeff)/(T-k);
ub=OLS_coeff(:);
vb=10*eye(num_eq*k);

% Regime shift intercept, only for state 1, i.e. this term will not exist if State=0
ua=-1e-3*ones(num_eq,1);                      % we assume regime 1 is high volatility regime with low intercept
va=10*eye(num_eq);

% 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
alpha_draw=multi_normrnd(ua, va,1)';      % regime shift intercetpts, High volatility has low intercept
beta_draw=multi_normrnd(ub, vb,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,alpha_draw, beta_draw,  Sigma1_draw, Sigma0_draw, tp11_draw, tp00_draw,T,num_eq);
    for i=1:T
        st(:,:,i)=kron(eye(num_eq),st_draw(i));
    end

    % step 2: draw the posterior coefficients: beta // no regime shifts
    for i=1:T
        yt_star(:,i)=yt(:,i)-st(:,:,i)*alpha_draw;  % treat new {yt-st*alpha} as new 'y' since we are sampling beta
    end
    xt_star=xt; % new 'x'

    num_1=sum(st_draw); % num of states been 1
    num_0=T-num_1;      % num of states been 0

    yt_star1=yt_star(:,(st_draw==1));      yt_star0=yt_star(:,(st_draw==0));     % yt_star1=yt_star when st==1
    xt_star1=xt_star(:,:,(st_draw==1));    xt_star0=xt_star(:,:,(st_draw==0));   % xt_star0=xt_star when st==0

    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
    vb_bar1=vb_bar1+inv(vb);
    b_bar1=b_bar1+inv(vb)*ub;

    vb_bar2=zeros(num_eq*k); b_bar2=zeros(num_eq*k,1);
    for i=1:num_0   % when st==0
        vb_bar2=vb_bar2+ xt_star0(:,:,i)'*inv(Sigma0_draw)*xt_star0(:,:,i);
        b_bar2=b_bar2+xt_star0(:,:,i)'*inv(Sigma0_draw)*yt_star0(:,i);
    end

    Db=inv(vb_bar1+vb_bar2);        % Db=(x1'inv(sigma1)x1+x0'inv(sigma0)x0+vb^-1)^-1
    db=(b_bar2+b_bar1);             % db=(x1'inv(sigma1)y1+x0'inv(sigma0)y0+vb^-1*ub)
    beta_draw=multi_normrnd(Db*db, Db,1)';
    clear yt_star xt_star yt_star1 xt_star1 yt_star0 xt_star0 

    % step 3: draw the posterior regime shifts coefficients: alpha < 0,
    % truncated normal sampling skill is used
    for i=1:T
        yt_star(:,i)=yt(:,i)-xt(:,:,i)*beta_draw;  % treat new {yt-xt*beta} as new 'y', since we are sampling alpha now
    end
    xt_star=st; % new 'x'
    yt_star1=yt_star(:,(st_draw==1));      yt_star0=yt_star(:,(st_draw==0));             % yt_star1=yt_star when st==1
    xt_star1=xt_star(:,:,(st_draw==1));    xt_star0=xt_star(:,:,(st_draw==0));           % xt_star1=xt_star when st==1

    va_bar1=zeros(num_eq); a_bar1=zeros(num_eq,1);
    for i=1:num_1   % when st==1
        va_bar1=va_bar1+xt_star1(:,:,i)'*inv(Sigma1_draw)*xt_star1(:,:,i);
        a_bar1=a_bar1+xt_star1(:,:,i)'*inv(Sigma1_draw)*yt_star1(:,i);
    end
    va_bar1=va_bar1+inv(va);
    a_bar1=a_bar1+inv(va)*ua;

    va_bar2=zeros(num_eq); a_bar2=zeros(num_eq,1);
    for i=1:num_0   % when st==0
        va_bar2=va_bar2+ xt_star0(:,:,i)'*inv(Sigma0_draw)*xt_star0(:,:,i);
        a_bar2=a_bar2+xt_star0(:,:,i)'*inv(Sigma0_draw)*yt_star0(:,i);
    end
    Da=inv(va_bar1+va_bar2);   % Da=(x1'inv(sigma1)x1+x0'inv(sigma0)x0+va^-1)^-1
    da=(a_bar2+a_bar1);        % da=(x1'inv(sigma1)y1+x0'inv(sigma0)y0+va^-1*ua)

    % sample alpha<0, since we say regime 1 is high volatility with low intercept regime
    ll=-1e+6*ones(num_eq,1); ul=-1e-6*ones(num_eq,1);
    u=unifrnd(0,1,num_eq,1);
    p1=normcdf(inv(chol(Da)')*(ll-Da*da));
    p2=normcdf(inv(chol(Da)')*(ul-Da*da));

    inv_CDF=norminv(u.*(p2-p1)+p1,0,1);
    for i=1:num_eq             % sometimes, inv_CDF may be "-inf or inf", which brings problem in our later sampling, thus we do like follows
        if inv_CDF(i,:)==-inf;
            inv_CDF(i,:)=-3;   % since less than -3 times sigmas only has 0.001 probability
        elseif inv_CDF(i,:)==inf
            inv_CDF(i,:)=3;    % since greater than 3 times sigmas only has 0.001 probability            
        end
    end
    alpha_draw=Da*da+chol(Da)'*inv_CDF;
    clear yt_star xt_star yt_star1 xt_star1 yt_star0 xt_star0 

    % step 4: draw the posterior volatility variance: Sigma0_draw & Sigma1_draw
    % Note that, diag(Sigma1_draw)>diag(Sigma0_draw) since regime 1 is high volatility regime
    yt_star1=yt(:,(st_draw==1));      yt_star0=yt(:,(st_draw==0));
    xt_star1=xt(:,:,(st_draw==1));    xt_star0=xt(:,:,(st_draw==0));
    
    post_omeg1=zeros(num_eq);    
    for i=1:num_1  % when st=1
        post_omeg1=post_omeg1+(yt_star1(:,i)- xt_star1(:,:,i)*beta_draw-eye(num_eq)*alpha_draw)*...
            (yt_star1(:,i)- xt_star1(:,:,i)*beta_draw-eye(num_eq)*alpha_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)*beta_draw)*...
            (yt_star0(:,i)-xt_star0(:,:,i)*beta_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>=100
            Sigma1_draw=inv(wishrnd(omeg1,v1));  % if after 100 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)

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    %% One can have his own in-sample or out-sample functions to be drawn in each iteration
    %% the following two functions are used to calculate the hedge ratio in-sample and out-sample forecast
   
    % step 6: draw the in-sample hedge-ratio function
    if nargout > 6
        hedge1=Sigma1_draw(1,2)/Sigma1_draw(2,2);
        hedge0=Sigma0_draw(1,2)/Sigma0_draw(2,2);           % hedge ratio in regime 1 and regime 0
        hedge_insample=filter_t*[hedge1,hedge0]';           % in-sample hedge ration calculation
    end
    
    % step 7: draw the Out-sample forecast function, if out-sapmle forecast >0, here we draw the out-sample forecast hedge ratio
    if nargout > 7
        hedge1=Sigma1_draw(1,2)/Sigma1_draw(2,2);
        hedge0=Sigma0_draw(1,2)/Sigma0_draw(2,2);           % hedge ratio in regime 1 and regime 0
        P=[tp11_draw, 1-tp11_draw; 1-tp00_draw, tp00_draw]; % Transistion probability Matrix
        state_probability(1,:)=filter_t(end,:)*P;

        % when out_sample_forecast_step=1
        hedge_outsample(1)=state_probability(1,:)*[hedge1,hedge0]';
        if out_sample_forecast_step>1
            % when out_sample_forecast_step>1, we use state_probability(1,:) as new filter probability
            for j=2:out_sample_forecast_step
                state_probability(j,:)=state_probability(j-1,:)*P;
                hedge_outsample(j)=state_probability(j,:)*[hedge1,hedge0]';
            end
        end
    end
    
     % step 8: Draw the insample residuals
    if nargout > 8
        for i=1:T
            if st_draw(i)==1                
                residual1_draw(:,i)=yt(:,i)-xt(:,:,i)*beta_draw-eye(num_eq)*alpha_draw;
                residual2_draw(:,i)=inv(chol(Sigma1_draw)')*(yt(:,i)-xt(:,:,i)*beta_draw-eye(num_eq)*alpha_draw); % adjusted cov-matrix
                residual3_draw(:,i)=residual1_draw(:,i)./sqrt(diag(Sigma1_draw));
            elseif st_draw(i)==0
                residual1_draw(:,i)=yt(:,i)-xt(:,:,i)*beta_draw;
                residual2_draw(:,i)=inv(chol(Sigma0_draw)')*(yt(:,i)-xt(:,:,i)*beta_draw);                        % adjusted cov-matrix
                residual3_draw(:,i)=residual1_draw(:,i)./sqrt(diag(Sigma0_draw));
            end
        end
    end
    
    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   
    if iter>warmup
        alpha_keep(iter-warmup,:)=alpha_draw';
        beta_keep(iter-warmup,:)=beta_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(:,1)';
        smooth_t_keep(iter-warmup,:)=smooth_t(:,1)';
        if nargout > 6
            in_sample_results(iter-warmup,:)=hedge_insample';
            if nargout > 7
                out_sample_results(iter-warmup,:)=hedge_outsample;
                if nargout > 8
                    residual{1}(:,:,iter-warmup)=residual1_draw;
                    residual{2}(:,:,iter-warmup)=residual2_draw; 
                    residual{3}(:,:,iter-warmup)=residual3_draw;
                end
            end
        end
    end
    %fprintf(' Iteration to %3g \n', iter)
end
smpl_coeff=[alpha_keep, beta_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,alpha_draw, beta_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)*beta_draw-eye(num_eq)*alpha_draw)'*inv(Sigma1_draw)*...
        (yt(:,i)-xt(:,:,i)*beta_draw-eye(num_eq)*alpha_draw));      % when st=1, we use sigma1_draw, and beta_draw and alpha_draw
    likelihood(i,2)=(2*pi)^(-0.5*num_eq) *(det(Sigma0_draw)).^(-0.5)*...
        exp(-0.5*(yt(:,i)-xt(:,:,i)*beta_draw)'*inv(Sigma0_draw)*...
        (yt(:,i)-xt(:,:,i)*beta_draw));                             % when st=0, we use sigma0_draw, and beta_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

smooth_t(end,:)=filter_t(end,:);
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



