function   [smpl_coeff, smpl_tp, smpl_st]=...
    RS_AR_in_error_uni_fit(y, x, nrep, warmup, p)
% Usage: [beta1_keep, beta2_keep, sigma1_keep, sigma2_keep, tp11_keep, tp00_keep, st_keep]=...
%    Regime_switch_uni_regress(y, x, nrep, warmup)
% estimate general regression for univaraite analysis under two regimes cases
% All priors are given below, one can change some initial values, while
% since I use Gibbs-sampling, the results are less sensitive to the intial values
% Note: for the scale y and x, when they represents the return (which is usually the case), it is bettet to use 
% return itself not the return percentage//

% Written by Weijun Xu, 18/07/2008
% School of Banking and Finance, UNSW
% First edition



[T,k]=size(x);
%% Setting the priors: Assume there is only two latent states
%%%%% Transition probability prior %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
a11=3; b10=1;                    % tp11~Beta(a11, b10)                   
a00=3; 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
%%%%%%% Coefficients beta %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
ub=0*ones(k,1); vb=10*eye(k);  % Beta1~N(u1b, v1b)
%%%%%%% Coefficients regime-shift alpha %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
ualpha=1; valpha=10;    
%%%%%%% Coefficients AR coefficients: r %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
ur=0*ones(p,1); vr=5*eye(p);  
%%%%%%% Varianec prior    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
A=3; B=0.25;              % sigma1~IG(A, B)
%%%%%%% Varianec high regime, omeg %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
a=1; b=10;              % sigma1~IG(a, b)

%% Setting initial values

beta_draw=multi_normrnd(ub, vb,1)';
alpha_draw=multi_normrnd(ualpha, valpha,1)';
arcoeff_draw=multi_normrnd(ur, vr,1)';
sigma_draw=gamrnd(A,1);             
sigma_draw=B/sigma_draw;              
omeg_draw=gamrnd(a,1); 
omeg_draw=b/omeg_draw; 
tp11_draw=0.7; tp00_draw=0.7;
st_draw=binornd(1,0.5,T-p,1);

%% Now let's start Gibbs-Sampling
for iter=1:nrep+warmup
    % step 1: draw latent state variables: st
    % first let's do some data transformation to make the model linear    
    temp_st_draw=binornd(1,0.5,p,1);  % temp_st_draw is to extend st_draw in first p observations
    if p>=1
        alpha_st=alpha_draw * [temp_st_draw; st_draw];       % I actually ignore the first p's st_draw, since we will actually drop first p's observations
        for i=p+1:T   % becasue if there is AR(p), we will drop first four lags
            % yhat=(1-R(L))yt+R(L)(xt*beta+alpha*st), where  R(L)=1-r(1)*L -,...,r(p)L^p
            yhat(i,1)=y(i,1)-(y(i,1)-y(i-(1:p),1)'*arcoeff_draw);      
            for j=1:k
                xhat(i,j)=x(i,j)-x(i-(1:p),j)'*arcoeff_draw;   
            end
            alpha_st_hat=alpha_st(i,1)-alpha_st(i-(1:p),1)'*arcoeff_draw;
            yhat(i,1)=yhat(i,1)+xhat(i,:)*beta_draw+alpha_st_hat;
        end
        ehat=y(p+1:end,:)-yhat(p+1:end,:);       % drop first p observations
    end   
    [st_draw]=state_draw_uni(p, ehat, st_draw, beta_draw, sigma_draw, omeg_draw, tp11_draw, tp00_draw); % where st_draw only have T-p observations

    % step 2: draw the posterior covariate coefficients: beta 
    % agian, we need to some data transformation to make it perfect linear

    alpha_st=alpha_draw * [temp_st_draw;st_draw];  % updating the alpha_st using the new st_draw  
    if p>=1
        for i=p+1:T   % becasue if there is AR(p), we will drop first four lags
            ystar(i,1)=(1+omeg_draw*st_draw(i-p))^-0.5 *(y(i,1)-y(i-(1:p),1)'*arcoeff_draw);        % ystar=(1+omeg*st)^1/2 * R(L)*yt
            alpha_star(i,1)=(1+omeg_draw*st_draw(i-p))^-0.5 *(alpha_st(i,1)-alpha_st(i-(1:p),1)'*arcoeff_draw);
            y_new(i,1)=ystar(i,1)- alpha_star(i,1);    % since now we are sampling beta, we treat the (ystar-alpha_star) as new 'y'
            for j=1:k
                xstar(i,j)=(1+omeg_draw*st_draw(i-p))^-0.5 * (x(i,j)-x(i-(1:p),j)'*arcoeff_draw);                                    % xstar=(1+omeg*st)^1/2* R(L)*xt
            end
        end        
        y_new=y_new(p+1:end,:);
        x_new=xstar(p+1:end,:);
    else 
        y_new=y(p+1:end,:)-alpha_st(p+1:end,:);      % when there is no AR in error, we still need to transform ystar as new 'y'
        x_new=x(p+1:end,:);
    end    
    Db=(x_new'*x_new/sigma_draw+inv(vb))^-1;
    db=x_new'*y_new/sigma_draw+inv(vb)*ub;
    beta_draw=multi_normrnd(Db*db, Db,1)';
    clear x_new y_new
    
    % step 3: draw the posterior regime shifts intercept: alpha
    % agian, we need to some data transformation to make it perfect linear
    % Note that, we assume St=1 is the high volatile period, which has high
    % level in mean and variacne, thus alpha>0,  we need to do some truncated sample distribution
    
    if p>=1
        y_new=ystar-xstar*beta_draw;       % now we are sampling alpha, thus we treat (ystar-xstar*beta) as new 'y'
        y_new=y_new(p+1:end,:);
        x_new=st_draw;                     % now new 'x' is just st_draw             
    else 
        y_new=y(p+1:end,:)-x(p+1:end,:)*beta_draw;   % when there is no AR in error, we still need to transform ystar as new 'y'
        x_new=st_draw;
    end
    
    % since alpha are supposed to be >0, we need do truncated normal sample    
    Dalpha=(x_new'*x_new/sigma_draw+inv(valpha))^-1;
    dalpha=x_new'* y_new/sigma_draw+inv(valpha)*ualpha;
    alpha_draw=multi_normrnd(Dalpha*dalpha, Dalpha,1)';
    clear x_new y_new

    % step 4: draw the posterior variance: sigma
    ystar=ystar(p+1:end,:);  % now we are going to use the ystar, xstar, alpha_star to sample simga, we drop first p observations
    xstar=xstar(p+1:end,:);
    alpha_star=alpha_star(p+1:end,:);
    
    obs=size(st_draw,1);  
    shape=A+obs/2;
    scale=B+ sum((ystar-xstar*beta_draw-alpha_star).^2)/2;
    sigma_draw=gamrnd(shape,1);
    sigma_draw=scale/sigma_draw;
    
    % step 5: draw the posterior high volatilty: omeg, where
    % sigma_high=sigma(1+omeg*1)
    obs_1=sum(st_draw,1); 
    ystar=ystar((st_draw==1),:);
    xstar=xstar((st_draw==1),:);
    alpha_star=alpha_star((st_draw==1),:);  
    
    ybar=ystar.*((ones(obs_1,1)+omeg_draw *1).^0.5);          % ybar=ystar*(1+omeg*st_draw), since we are sampling omeg now
    xbar=xstar.*((ones(obs_1,1)+omeg_draw *1).^0.5);
    alpha_bar=alpha_star.*((ones(obs_1,1)+omeg_draw *1).^0.5);
    shape=a+obs_1/2;
    scale=b+sum((ybar-xbar*beta_draw-alpha_bar).^2)/(2*sigma_draw); % where we need to divide sigma_draw since we are sampling omeg now
    omeg_draw=gamrnd(shape,1);
    omeg_draw=scale/sigma_draw;    
    % make sure omeg_draw>1, since we assume state 1 has high volatility
    while omeg_draw<1
        omeg_draw=gamrnd(shape,1);
        omeg_draw=scale/sigma_draw;
    end          
    
    % step 6: draw the posterior AR coefficients: alpha,
    % Note: we need all the roots of R(L) outside unit circle to make sure it is stationary
    % condition, that is: 1-r1*L-,...,-rp*L^p, all roots should be outside unit circle
    if p>=1

        e=y-x*beta_draw-alpha_draw *[temp_st_draw;st_draw];        
        e=(ones(T,1)+omeg_draw*[temp_st_draw; st_draw]).^-0.5;
        for i=p+1:T   % becasue if there is AR(p), we will drop first four lags
            e_x(i,:)=e(i-(1:p),1)';
        end
        e_x=e_x(p+1:end,:); % treat e_x as new 'x'
        e_y=e(p+1:end,:);   % treat e_y as new 'y'
 
        Dr=(e_x'*e_x/sigma_draw+inv(vr))^-1;       
        dr=e_x'*e_y/sigma_draw+inv(vr)*ur;
        % make sure every root is outside the unit circle
        arcoeff_draw=multi_normrnd(Dr*dr, Dr,1)';
        for i=1:p
            ploy(i)=-arcoeff_draw(end-i+1);
        end           
        while min(abs(roots([ploy,1])))<1  % we need all the roots be outside unit circle
            arcoeff_draw=multi_normrnd(Dr*dr, Dr,1)';
            for i=1:p
                ploy(i)=-arcoeff_draw(end-i+1);
            end
        end      
    end

    % step 7: draw the transistion probability: tp11,and tp00
    [n11, n10, n01, n00]=sigma_draw_number_uni(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)

    % collect each iteration's samples
    if iter>warmup
        beta_keep(iter-warmup,:)=beta_draw';
        alpha_keep(iter-warmup,:)=alpha_draw';
        arcoeff_keep(iter-warmup,:)=arcoeff_draw;
        sigma_keep(iter-warmup,:)=sigma_draw;
        omeg_keep(iter-warmup,:)=omeg_draw;
        tp11_keep(iter-warmup,:)=tp11_draw;
        tp00_keep(iter-warmup,:)=tp00_draw;
        st_keep(iter-warmup,:)=st_draw';       
    end   
    fprintf(' Iteration to %3g \n', iter)
end
smpl_coeff=[beta_keep, alpha_keep, arcoeff_keep, sigma_keep, omeg_keep];
smpl_tp=[tp11_keep, tp00_keep];
smpl_st=st_keep;
end


%% Subfunction 1: to generate latent state variable
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [st_draw]=state_draw_uni(p, ehat, st_draw, beta_draw, sigma_draw, omeg_draw, tp11_draw, tp00_draw);
% Usage: to generate the latent state variable for uni-variate regression/
% this function is only for two state case

T=size(ehat,1);
st_draw(end,1)=binornd(1,0.5); 

for i=T-1:-1:2
    %% first let's calculate the likelihood at each time period // already drop first r observations
    likelihood(i,1)=exp(-(sigma_draw*(1+omeg_draw*1))^-1*(ehat(i).^2)/2); % exp(-v(st)^-1 (yt-yhat_t)^2/2), where v(st)=sigma(1+omeg*st)
    likelihood(i,2)=exp(-(sigma_draw*(1+omeg_draw*0))^-1*(ehat(i).^2)/2);  
    %% second let's calculate the P_st= P(St+1|St)P(St|St-1) when St=1 or 0
    % when st_draw(i)=1
    if st_draw(i+1)==1 & st_draw(i-1)==1
        p_st(i,1)=tp11_draw*tp11_draw;
    elseif st_draw(i+1)==1 & st_draw(i-1)==0
        p_st(i,1)=tp11_draw*(1-tp00_draw);
    elseif st_draw(i+1)==0 & st_draw(i-1)==1
        p_st(i,1)=(1-tp11_draw)*tp11_draw;
    elseif st_draw(i+1)==0 & st_draw(i-1)==0
        p_st(i,1)=(1-tp11_draw)*(1-tp00_draw);
    end
    % when st_draw(i)==0
    if st_draw(i+1)==1 & st_draw(i-1)==1
        p_st(i,2)=(1-tp00_draw)*(1-tp11_draw);
    elseif st_draw(i+1)==1 & st_draw(i-1)==0
        p_st(i,2)=(1-tp00_draw)*tp00_draw;
    elseif st_draw(i+1)==0 & st_draw(i-1)==1
        p_st(i,2)=tp00_draw*(1-tp11_draw);
    elseif st_draw(i+1)==0 & st_draw(i-1)==0
        p_st(i,2)=tp00_draw*tp00_draw;
    end
    %% Last, let's calculate the inference probability P(st|Y_T)=p_st * Product_t_t+r (likelihood)   
    if i<=T-p              % equation 3.6 in Chib (1993)
        llf=[1,1];
        for j=1:p
            llf=llf.*likelihood(i+j-1,:);
        end
        ip_t(i,:)=(p_st(i,:).*llf)./sum(p_st(i,:).*llf);   % where sum(p_st.*llf) is the normalized term
    else                  % equation 3.7 in Chib (1993)
        llf=[1,1];   
        j=i;
        while j<=T-1
            llf=llf.*likelihood(j,:);
            j=j+1;
        end        
        ip_t(i,:)=(p_st(i,:).*llf)./sum(p_st(i,:).*llf);   % where sum(p_st.*llf) is the normalized term
    end
    st_draw(i,1)=binornd(1,ip_t(i,1));        % renew the st_draw based on the calculated inference probability
end
st_draw(1,1)=binornd(1,0.5);
end


%% Subfunction 2: calculate the some count number for inverse garma to
%% generate sigmas
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [n11, n10, n01, n00]=sigma_draw_number_uni(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
    % calculate Pr(St+1|St)
    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

   





















