## Estimate Shocks ##########################################

shocks <- function(prmtr,shocks_no){
  
  prmtr = trans(prmtr)
  
  theta = prmtr[1]         # MA parameter
  gamma_ctau = prmtr[2]    # Consumption response to permanent income
  gamma_ceps = prmtr[3]    # Consumption response to transitory income
  sig_u = prmtr[4]         # Standard deviation of permanent consumption shock
  sig_eta = prmtr[5:14]    # Standard deviation of permanent income shock
  sig_eps = prmtr[15:26]   # Standard deviation of transitory income shoc
  sig_v = prmtr[27:36]     # Standard deviation of transitory consumption shock (1989 and 1988 missing consumption]
  
  Fmat = matrix(0,8,8) # Transition matrix
  Fmat[1,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  Fmat[2,] = c(1, 0, 0, 0, 0, 0, 0, 0)
  Fmat[3,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  Fmat[4,] = c(0, 0, 0, 1, 0, 0, 0, 0)
  Fmat[5,] = c(0, 0, 0, 0, 1, 0, 0, 0)
  Fmat[6,] = c(0, 0, 0, 0, 0, 1, 0, 0)
  Fmat[7,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  Fmat[8,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  
  Fstar = matrix(0,5,5) # Transition matrix of I(0) part
  Fstar[1,] = c(0, 0, 0, 0, 0)
  Fstar[2,] = c(1, 0, 0, 0, 0)
  Fstar[3,] = c(0, 0, 0, 0, 0)
  Fstar[4,] = c(0, 0, 0, 0, 0)
  Fstar[5,] = c(0, 0, 0, 0, 0)

  # Measurement equation
  Hy = matrix(c(1, theta, 0, 1, 0, 0, 0, 0),1,8)
  Hc = matrix(c(0, 0, 1, gamma_ctau, gamma_ceps, 1, 0, 0),1,8)
  H = rbind(Hy, Hc)
  
  Q = matrix(0,8,8) # Covariance matrix
  Q[1,] = c(sig_eps[1]^2, 0, 0, 0, sig_eps[1]^2, 0, 0, 0)
  Q[2,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  Q[3,] = c(0, 0, sig_v[1]^2, 0, 0, 0, 0, 0)
  Q[4,] = c(0, 0, 0, sig_eta[1]^2, 0, 0, sig_eta[1]^2, 0)
  Q[5,] = c(sig_eps[1]^2, 0, 0, 0, sig_eps[1]^2, 0, 0, 0)
  Q[6,] = c(0, 0, 0, 0, 0, sig_u^2, 0, sig_u^2)
  Q[7,] = c(0, 0, 0, sig_eta[1]^2, 0, 0, sig_eta[1]^2, 0)
  Q[8,] = c(0, 0, 0, 0, 0, sig_u^2, 0, sig_u^2)
   
  Qstar = matrix(0,5,5) # Covariance matrix of I(0) part
  Qstar[1,] = c(sig_eps[1]^2, 0, 0, 0, 0)
  Qstar[2,] = c(0, 0, 0, 0, 0)
  Qstar[3,] = c(0, 0, sig_v[1]^2, 0, 0)
  Qstar[4,] = c(0, 0, 0, sig_eta[1]^2, 0)
  Qstar[5,] = c(0, 0, 0, 0, sig_u^2)
  
  Ay = 0
  Ac = 0
  A = matrix(c(Ay, Ac),2,1)
  
  # Vectorization of covariance matrix for stationary components
  vecQstar = matrix(Qstar, ncol = 1)
  vecP_LL = solve(diag(nrow(Qstar)^2) - Fstar%x%Fstar)%*%vecQstar 
  # Covariance matrix
  P_LL = matrix(vecP_LL, nrow(Qstar), nrow(Qstar))
  P_LL = cbind(P_LL[1:3,1:3], matrix(0,3,nrow(Qstar)))
  P_LL = rbind(P_LL, c(0, 0, 0, prior, 0, 0, 0, 0),
                c(0, 0, 0, 0, prior, 0, 0, 0),
                c(0, 0, 0, 0, 0, prior, 0, 0),
                c(0, 0, 0, 0, 0, 0, sig_eta[1]^2, 0),
                c(0, 0, 0, 0, 0, 0, 0, sig_u^2))
  
shocks_mat_v = matrix(999,Ts, N)
shocks_mat_u = matrix(999,Ts, N)
shocks_mat_eta = matrix(999,Ts, N)
shocks_mat_eps = matrix(999,Ts, N)
  
  i_iter = 1
  while(i_iter <= N){ # Loop over individuals
  	
    yin = y[1 + (i_iter - 1)*Ts, 1]

    
    for(z in 1:Ts){
      if(yin == 999){
        yin = y[1 + z + (i_iter - 1)*Ts, 1]
      }
    } 
    
    cin = y[1 + (i_iter - 1)*Ts, 2]
         
    for(z in 1:Ts){
      if(cin == 999){
        cin = y[1+ z + (i_iter - 1)*Ts, 2]
      }
    } 
     
      
  beta_LL = matrix(c(0, 0, 0, yin, 0, cin - gamma_ctau*yin, 0, 0),8,1)
  P_LL = matrix(vecP_LL, nrow(Qstar), nrow(Qstar))
  P_LL = cbind(P_LL[1:3,1:3], matrix(0,3,nrow(Qstar)))
  P_LL = rbind(P_LL, c(0, 0, 0, prior, 0, 0, 0, 0),
                c(0, 0, 0, 0, prior, 0, 0, 0),
                c(0, 0, 0, 0, 0, prior, 0, 0),
                c(0, 0, 0, 0, 0, 0, sig_eta[1]^2, 0),
                c(0, 0, 0, 0, 0, 0, 0, sig_u^2))
  mu = matrix(c(0, 0, 0, 0, 0, 0, 0, 0),8,1)
  
    
    eta_iter = 1
    eps_iter = 1
    v_iter = 1
 
 BTT_MAT = matrix(0,Ts, 8)
 PTL_MAT = matrix(0,Ts, 8^2)
 PTT_MAT = matrix(0,Ts, 8^2)
  
    j_iter = 1
    while(j_iter <= Ts){ # Loop over time periods
  Q[1,] = c(sig_eps[eps_iter]^2, 0, 0, 0, sig_eps[eps_iter]^2, 0, 0, 0)
  Q[2,] = c(0, 0, 0, 0, 0, 0, 0, 0)
  Q[3,] = c(0, 0, sig_v[v_iter]^2, 0, 0, 0, 0, 0)
  Q[4,] = c(0, 0, 0, sig_eta[eta_iter]^2, 0, 0, sig_eta[eta_iter]^2, 0)
  Q[5,] = c(sig_eps[eps_iter]^2, 0, 0, 0, sig_eps[eps_iter]^2, 0, 0, 0)
  Q[6,] = c(0, 0, 0, 0, 0, sig_u^2, 0, sig_u^2)
  Q[7,] = c(0, 0, 0, sig_eta[eta_iter]^2, 0, 0, sig_eta[eta_iter]^2, 0)
  Q[8,] = c(0, 0, 0, 0, 0, sig_u^2, 0, sig_u^2)
      
      beta_TL = mu + Fmat%*%beta_LL
      P_TL = Fmat%*%P_LL%*%t(Fmat) + Q
      
      if(y[j_iter + (i_iter - 1)*Ts, 1] == 999 && y[j_iter + (i_iter - 1)*Ts, 2] == 999){
        beta_TT = beta_TL
        P_TT = P_TL
        
      } else if(y[j_iter + (i_iter - 1)*Ts, 1] == 999){
        vt = y[j_iter + (i_iter - 1)*Ts, 2] - Hc%*%beta_TL - Ac # Prediction error
        ft = Hc%*%P_TL%*%t(Hc) # Variance of forecast error
        
        invft = solve(ft)
        
        beta_TT = beta_TL + P_TL%*%t(Hc)%*%invft%*%vt
        P_TT = P_TL - P_TL%*%t(Hc)%*%invft%*%Hc%*%P_TL
        
shocks_mat_v[j_iter, i_iter] = beta_TT[3]/sig_v[v_iter]
shocks_mat_u[j_iter, i_iter] = beta_TT[8]/sig_u
        
      } else if(y[j_iter + (i_iter - 1)*Ts, 2] == 999){
        vt = y[j_iter + (i_iter - 1)*Ts, 1] - Hy%*%beta_TL - Ay # Prediction error
        ft = Hy%*%P_TL%*%t(Hy) # Variance of forecast error 
        
        invft = solve(ft)
        
        beta_TT = beta_TL + P_TL%*%t(Hy)%*%invft%*%vt
        P_TT = P_TL - P_TL%*%t(Hy)%*%invft%*%Hy%*%P_TL

shocks_mat_eps[j_iter, i_iter] = beta_TT[1]/sig_eps[eps_iter]
shocks_mat_eta[j_iter, i_iter] = beta_TT[7]/sig_eta[eta_iter]
       
      } else if(y[j_iter + (i_iter - 1)*Ts, 1] != 999 && y[j_iter + (i_iter - 1)*Ts, 2] != 999){
        vt = y[j_iter + (i_iter - 1)*Ts, 1:2] - H%*%beta_TL - A # Prediction error
        ft = H%*%P_TL%*%t(H) # Variance of forecast error 
        
        invft = solve(ft)
        
        beta_TT = beta_TL + P_TL%*%t(H)%*%invft%*%vt
        P_TT = P_TL - P_TL%*%t(H)%*%invft%*%H%*%P_TL
 
shocks_mat_v[j_iter, i_iter] = beta_TT[3]/sig_v[v_iter]
shocks_mat_u[j_iter, i_iter] = beta_TT[8]/sig_u
shocks_mat_eps[j_iter, i_iter] = beta_TT[1]/sig_eps[eps_iter]
shocks_mat_eta[j_iter, i_iter] = beta_TT[7]/sig_eta[eta_iter]  
        }

vecP_TL = matrix(P_TL, ncol = 1)
vecP_TT = matrix(P_TT, ncol = 1)

BTT_MAT[j_iter,] = t(beta_TT)
PTL_MAT[j_iter,] = t(vecP_TL)
PTT_MAT[j_iter,] = t(vecP_TT)
      
      beta_LL = beta_TT
      P_LL = P_TT
      
      j_iter = j_iter + 1
      
      if(j_iter == 3 || j_iter == 4){
        eps_iter = eps_iter + 1
        v_iter = v_iter + 1
        
      } else if((j_iter > 4 && j_iter < 10) || j_iter == 12 || j_iter == 13){
        eps_iter = eps_iter + 1
        v_iter = v_iter + 1
        eta_iter = eta_iter + 1      
        
      } else if(j_iter == 10 || j_iter == 11){
        eps_iter = eps_iter + 1
        eta_iter = eta_iter + 1
      }
    } # End of loop over time periods


#Kalman Smoother

BETA_PALL = beta_TT
P_PALL = P_TT

    j_iter = Ts-1
    while(j_iter >= 1){ # Loop over time periods


      if(j_iter == 3 || j_iter == 4){
        eps_iter = eps_iter - 1
        v_iter = v_iter - 1
        
      } else if((j_iter > 4 && j_iter < 10) || j_iter == 12 || j_iter == 13){
        eps_iter = eps_iter + 1
        v_iter = v_iter + 1
        eta_iter = eta_iter + 1      
        
      } else if(j_iter == 10 || j_iter == 11){
        eps_iter = eps_iter - 1
        eta_iter = eta_iter - 1
      }

BTT=BTT_MAT[j_iter,]            
PTL=matrix(t(PTL_MAT[j_iter +1,]),8,8)
PTT=matrix(t(PTT_MAT[j_iter,]),8,8)
             
    BETA_TALL = BTT + PTT %*%
                t(Fmat) %*% (solve(PTL)) %*%
                                (BETA_PALL - Fmat %*% BTT)

    P_TALL = PTT +
            PTT %*%
            t(Fmat) %*% solve(PTL) %*%
            (P_PALL - PTL) %*% t(solve(PTL)) %*% Fmat %*% t(PTT)
            

if (smoother == 1){
if (shocks_mat_v[j_iter, i_iter] != 999){shocks_mat_v[j_iter, i_iter] <- BETA_TALL[3]/sig_v[v_iter]}
if (shocks_mat_u[j_iter, i_iter] != 999){shocks_mat_u[j_iter, i_iter] <- BETA_TALL[8]/sig_u}
if (shocks_mat_eps[j_iter, i_iter] != 999){shocks_mat_eps[j_iter, i_iter] <- BETA_TALL[1]/sig_eps[eps_iter]}
if (shocks_mat_eta[j_iter, i_iter] != 999){shocks_mat_eta[j_iter ,i_iter] <- BETA_TALL[7]/sig_eta[eta_iter]}
}

    BETA_PALL=BETA_TALL;   P_PALL=P_TALL;

j_iter = j_iter - 1

    } # End of loop over time periods
    
     i_iter = i_iter + 1
  } # End of loop over individuals    
    
    eta_iter = 1
    eps_iter = 1
    v_iter = 1
  
    j_iter = 1
    while(j_iter <= Ts){ # Loop over time periods

if (sum(shocks_mat_v[j_iter,1:N])==999*N){shocks_mat_v[j_iter,1:N] <- matrix(999,1,N)}   
    else{   
    v_sub=shocks_mat_v[j_iter,shocks_mat_v[j_iter,] != 999]
    v_sub=t(v_sub)
    if (ncol(v_sub)==N){
    	miss_obs <- vector()} 
    else{
    	miss_obs <- matrix(999,1,N-ncol(v_sub))}
    shocks_mat_v[j_iter,1:N]=cbind(sig_v[v_iter]*(v_sub-mean(v_sub))/sd(v_sub),miss_obs)}

if (sum(shocks_mat_u[j_iter,1:N])==999*N){shocks_mat_u[j_iter,1:N] <- matrix(999,1,N)}   
    else{   
    u_sub=shocks_mat_u[j_iter,shocks_mat_u[j_iter,] != 999]
    u_sub=t(u_sub)
    if (ncol(u_sub)==N){
    	miss_obs <- vector()} 
    else{
    	miss_obs <- matrix(999,1,N-ncol(u_sub))}
    shocks_mat_u[j_iter,1:N]=cbind(sig_u*(u_sub-mean(u_sub))/sd(u_sub),miss_obs)}

if (sum(shocks_mat_eps[j_iter,1:N])==999*N){shocks_mat_eps[j_iter,1:N] <- matrix(999,1,N)}   
    else{   
    eps_sub=shocks_mat_eps[j_iter,shocks_mat_eps[j_iter,] != 999]
    eps_sub=t(eps_sub)
    if (ncol(eps_sub)==N){
    	miss_obs <- vector()} 
    else{
    	miss_obs <- matrix(999,1,N-ncol(eps_sub))}
    shocks_mat_eps[j_iter,1:N]=cbind(sig_eps[eps_iter]*(eps_sub-mean(eps_sub))/sd(eps_sub),miss_obs)}    
    
if (sum(shocks_mat_eta[j_iter,1:N])==999*N){shocks_mat_eta[j_iter,1:N] <- matrix(999,1,N)}   
    else{   
    eta_sub=shocks_mat_eta[j_iter,shocks_mat_eta[j_iter,] != 999]
    eta_sub=t(eta_sub)
    if (ncol(eta_sub)==N){
    	miss_obs <- vector()} 
    else{
    	miss_obs <- matrix(999,1,N-ncol(eta_sub))}
    shocks_mat_eta[j_iter,1:N]=cbind(sig_eta[eta_iter]*(eta_sub-mean(eta_sub))/sd(eta_sub),miss_obs)}    
  
    j_iter = j_iter + 1
      
      if(j_iter == 3 || j_iter == 4){
        eps_iter = eps_iter + 1
        v_iter = v_iter + 1
        
      } else if((j_iter > 4 && j_iter < 10) || j_iter == 12 || j_iter == 13){
        eps_iter = eps_iter + 1
        v_iter = v_iter + 1
        eta_iter = eta_iter + 1      
        
      } else if(j_iter == 10 || j_iter == 11){
        eps_iter = eps_iter + 1
        eta_iter = eta_iter + 1
      }
    } # End of loop over time periods
        
        
if (shocks_no==1){
 shocks_mat=shocks_mat_v}
else if (shocks_no==2){
 shocks_mat=shocks_mat_u}
else if (shocks_no==3){
 shocks_mat=shocks_mat_eps}
else if (shocks_no==4){
 shocks_mat=shocks_mat_eta}
  
  return(shocks_mat) # matrices with estimated shocks
  
} # End of shocks estimation