! auxiliary.f90 -- core model routines: steady-state solver,
! transition solver (backward value iteration + forward simulation of the
! SIR-macro dynamics), welfare calculations, and all file output.

MODULE AUXILIARY

!!!!!!!!!!!!!!!!!!!!!!!!
!!   USE STATEMENTS   !!
!!!!!!!!!!!!!!!!!!!!!!!! 

USE PARAMETERS
USE GLOBAL
USE TOOLKIT
!USE IFPORT  ! Intel-only; not needed with gfortran

IMPLICIT NONE

    CONTAINS

FUNCTION DCLOCK() RESULT(t)
    ! Portable replacement for Intel IFPORT's DCLOCK (wall-clock seconds)
    REAL(8) :: t
    INTEGER(8) :: cnt, rate
    CALL SYSTEM_CLOCK(cnt, rate)
    t = REAL(cnt,8)/REAL(rate,8)
END FUNCTION DCLOCK

SUBROUTINE getf2(f)
    real*8, dimension(ncalib), intent(out) :: f
    
    rss = zbrent(SOLVESS,rhat1,rhat2,1.0D-14)
    wss = (onedp-alpha)*(rss/alpha)**(alpha/(alpha-onedp))
      
    f(1) = (ss*SUM(muss(jmax,:,:)))/(wss*aggl)
    IF (bequestmode.EQ.1) THEN
        f(2) = bequest/(hikfrac*psi_j(jmax)*SUM(muss(jmax,:,:)))
    ELSE
        f(2) = initk
    ENDIF
    f(3) = phi + 50.0_dp*( aggh/SUM(muss(1:2,:,:)) - 0.30705_dp ) ! 34.39/112
    f(4) = (wss*aggl/SUM(muss(1:2,:,:)))*0.3_dp
    f(5) = beta_ann - 5d-3*(aggk/aggy*days/365.0_dp-4.8_dp)
    
    PRINT *, '       x'
    WRITE(6,334) tau, initk, phi, ss, beta_ann
    PRINT *, '       implied x'
    WRITE(6,334) f 
334 FORMAT(f12.8,2f12.4,2f12.8)
    !PRINT *, CHAR(7)
    !PAUSE
      
END SUBROUTINE getf2

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

DOUBLE PRECISION FUNCTION SOLVESS(rate)
	!! Solves the model
        
    ! LOCAL VARIABLES USED WITHIN LOOPS
    REAL(DP), INTENT(IN)    :: rate
    INTEGER                 :: showdetails = 0
    INTEGER     :: idx_j, idx_k, idx_kprime
	INTEGER     :: idx_z, idx_zprime ! index for income shock
	INTEGER     :: idx_s            ! index for savings rate
    INTEGER     :: idx_l            ! index for labor supply
    INTEGER     :: i, j             ! dummy index
    INTEGER     :: iter, i_s_star, i_l_star
    INTEGER, DIMENSION(N_z) :: Nsmin
    INTEGER :: Nlmin      
    INTEGER, DIMENSION(N_z) :: Nlmax
    
	REAL(DP)    :: wealth		    ! incoming wealth
    REAL(DP)    :: cash             ! cash-on-hand   
    REAL(DP)    :: srate            ! savings rate        
    REAL(DP)    :: kprime		    ! next days saving
    REAL(DP)    :: cons, uu, hours, wagetilde, fu
    
    ! AUXILIARY VALUE FUNCTION & POLICY FUNCTIONS
    REAL(DP)    :: evp, evo, gapv, gaps, mudif
    REAL(DP), DIMENSION(N_k)        :: vjmax,vjmax_old,sjmax,sjmax_old
    REAL(DP), DIMENSION(N_z,N_k)    :: vjwork,vjwork_old,sjwork,sjwork_old,nextv
    
    REAL(DP), ALLOCATABLE   :: kssfine(:,:,:), cssfine(:,:,:)
    
    INTEGER, ALLOCATABLE    :: jlos1(:,:,:)
    REAL(DP), ALLOCATABLE   :: wghts1(:,:,:)
    INTEGER                 :: jlo1
    REAL(DP)                :: wght1, fsum
    
    ALLOCATE( kssfine(jmax,N_z,N_kfine) )
    ALLOCATE( cssfine(jmax,N_z,N_kfine) )
    ALLOCATE( jlos1(jmax,N_z,N_kfine) )
    ALLOCATE( wghts1(jmax,N_z,N_kfine) )
    
    r = rate
    wage = (onedp-alpha)*(r/alpha)**(alpha/(alpha-onedp))
    wagetilde = wage*(onedp-tau)
    R_adj = onedp + r - delta_adj
    beta_adj    = beta_ann**(days/365.0_dp)         ! adjusted discount factor
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! Value function computation for j=jmax !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    idx_j = jmax
    
    IF (rate>rhat+1.0d-6) THEN
        vjmax = vssh(idx_j,1,:)
    ELSEIF (rate<rhat-1.0d-6) THEN
        vjmax = vssl(idx_j,1,:)
    ELSE
        vjmax = vss(idx_j,1,:)
    ENDIF
    sjmax = sss(idx_j,1,:)
    
    JMAX_LOOP: DO iter = 1,vfitmax

        vjmax_old = vjmax
	    sjmax_old = sjmax
	
	    ! Compute optimal savings (jmax)
	    K_LOOP: DO idx_k = 1,N_k
    
		    evo = -infty
		    i_s_star = 1
		    cash = ss + k_grid(idx_k)*R_adj
        
            S_LOOP: DO idx_s = 1,N_s
        
                kprime = cash*s_grid(idx_s)
                cons = cash - kprime
                CALL utility_ret(cons,gam,uu)
			    evp = uu + beta_adj*(onedp-psi_j(idx_j))*interp1(k_grid,vjmax_old,N_k,kprime)
            
			    IF ( evp > evo ) THEN
				    evo = evp
				    i_s_star = idx_s
			    END IF	

		    END DO S_LOOP
        
            vjmax(idx_k) = evo
            sjmax(idx_k) = s_grid(i_s_star)
        
        END DO K_LOOP
    
        gapv = SUM(ABS(vjmax - vjmax_old))/(N_k)
        gaps = SUM(ABS(sjmax - sjmax_old))/(N_k)
    
        !IF (showdetails.EQ.1) THEN
        IF(multpl(iter,vfitmax) == 1) THEN 
            PRINT *, "avg abs val of v",  gapv
            PRINT *, "avg abs val of saving", gaps
            PRINT *, "iter, age:", iter, idx_j
        ENDIF
    
        IF (gapv<critv .AND. gaps<crits) EXIT
    
    END DO JMAX_LOOP

    DO idx_z=1,N_z
        vss(jmax,idx_z,:) = vjmax
        sss(jmax,idx_z,:) = sjmax
    END DO    

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! Value function computation for j=2 !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    idx_j = 2

    IF (readvguess.EQ.1) THEN
        IF (rate>rhat+1.0d-6) THEN
                vjwork = vssh(idx_j,:,:)
            ELSEIF (rate<rhat-1.0d-6) THEN
                vjwork = vssl(idx_j,:,:)
            ELSE
                vjwork = vss(idx_j,:,:)
            ENDIF
    ELSE
        vjwork = vss(idx_j+1,:,:) 
    END IF
    sjwork = sss(idx_j,:,:)
    
    J2_LOOP: DO iter = 1,vfitmax

        vjwork_old = vjwork
	    sjwork_old = sjwork
	
        ! Compute conitinuation value vnext(z,kprime)
        DO idx_z = 1,N_z
        DO idx_kprime = 1,N_k
            evp = beta_adj*psi_j(idx_j)*vss(idx_j+1,1,idx_kprime)
            DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                evp = evp + beta_adj*(onedp-psi_j(idx_j))   &
                    *ptran(idx_z,idx_zprime)*vjwork_old(idx_zprime,idx_kprime)
            END DO
            nextv(idx_z,idx_kprime) = evp
        END DO
        END DO
    
        ! Compute optimal savings and hours (j=2)
        Nsmin = 1
        Nlmax = N_l
	    K2_LOOP: DO idx_k = 1,N_k
        Nlmin = 1
        Z2_LOOP: DO idx_z = 1,N_z
		
            evo = -infty
		    i_s_star = 1
		    i_l_star = 1
        
            S2_LOOP: DO idx_s = Nsmin(idx_z),N_s
            L2_LOOP: DO idx_l = Nlmin,Nlmax(idx_z)
        
                hours = l_grid(idx_l)
                cash = wagetilde*eta_j(idx_j,1)*z_grid(idx_z)*hours + k_grid(idx_k)*R_adj
                kprime = cash*s_grid(idx_s)
                cons = cash - kprime
            
                CALL utility(cons,gam,hours,uu)
			    evp = uu + interp1(k_grid,nextv(idx_z,:),N_k,kprime)
            
			    IF ( evp > evo ) THEN
				    evo = evp
				    i_s_star = idx_s
                    i_l_star = idx_l
			    END IF	
        
            END DO L2_LOOP
            END DO S2_LOOP
        
            IF (i_s_star.EQ.Nsmin(idx_z)) THEN
                IF (i_s_star>1) THEN
                    PRINT *, 'nsmin binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmin) THEN
                IF (i_l_star>1) THEN
                    PRINT *, 'nlmin binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmax(idx_z)) THEN
                IF (i_l_star<N_l) THEN
                    PRINT *, 'nlmax binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            vjwork(idx_z,idx_k) = evo
            sjwork(idx_z,idx_k) = s_grid(i_s_star)
            lss(idx_j,idx_z,idx_k) = l_grid(i_l_star)
            Nsmin(idx_z) = MAX(1,i_s_star - 5)  ! Given z, savings rate increasing in wealth
            Nlmin = MAX(1,i_l_star - 5)         ! Given k, labor supply increasing in productivity
            Nlmax(idx_z) = MIN(N_l,i_l_star + 4)! Given z, labor supply decreasing in wealth
        END DO Z2_LOOP
        END DO K2_LOOP
    
        gapv = SUM(ABS(vjwork - vjwork_old))/(N_k*N_z)
        gaps = SUM(ABS(sjwork - sjwork_old))/(N_k*N_z)
    
        IF(multpl(iter,vfitmax) == 1) THEN 
            PRINT *, "avg abs val of v",  gapv
            PRINT *, "avg abs val of saving", gaps
            PRINT *, "iter, age:", iter, idx_j
        ENDIF
        IF (gapv<critv .AND. gaps<crits) EXIT    
    
    END DO J2_LOOP

    vss(idx_j,:,:) = vjwork
    sss(idx_j,:,:) = sjwork

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! Value function computation for j=1 !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    idx_j = 1

    IF (readvguess.EQ.1) THEN
        IF (rate>rhat+1.0d-6) THEN
                vjwork = vssh(idx_j,:,:)
            ELSEIF (rate<rhat-1.0d-6) THEN
                vjwork = vssl(idx_j,:,:)
            ELSE
                vjwork = vss(idx_j,:,:)
            ENDIF
    ELSE
        vjwork = vss(idx_j+1,:,:) 
    END IF
    sjwork = sss(idx_j,:,:)
    
    J1_LOOP: DO iter = 1,vfitmax

        vjwork_old = vjwork
	    sjwork_old = sjwork
    
        ! Compute conitinuation value vnext(z,kprime)
        DO idx_z = 1,N_z
        DO idx_kprime = 1,N_k
            evp = zerodp
            DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                evp = evp + beta_adj*ptran(idx_z,idx_zprime)                    &
                    *( psi_j(idx_j)*vss(idx_j+1,idx_zprime,idx_kprime)          &
                    + (onedp-psi_j(idx_j))*vjwork_old(idx_zprime,idx_kprime) )
            END DO
            nextv(idx_z,idx_kprime) = evp
        END DO
        END DO
    
	    ! Compute optimal savings and hours (j=1)
        Nlmax = N_l
        Nsmin = 1
	    K1_LOOP: DO idx_k = 1,N_k
        Nlmin = 1
        Z1_LOOP: DO idx_z = 1,N_z
		
            evo = -infty
		    i_s_star = 1
		    i_l_star = 1
        
            S1_LOOP: DO idx_s = Nsmin(idx_z),N_s
            L1_LOOP: DO idx_l = Nlmin,Nlmax(idx_z)
        
                hours = l_grid(idx_l)
                cash = wagetilde*eta_j(idx_j,1)*z_grid(idx_z)*hours + k_grid(idx_k)*R_adj
                kprime = cash*s_grid(idx_s)
                cons = cash - kprime
                CALL utility(cons,gam,hours,uu)
			    evp = uu + interp1(k_grid,nextv(idx_z,:),N_k,kprime)
            
			    IF ( evp > evo ) THEN
				    evo = evp
				    i_s_star = idx_s
                    i_l_star = idx_l
			    END IF	
        
            END DO L1_LOOP
		    END DO S1_LOOP
            IF (i_s_star.EQ.Nsmin(idx_z)) THEN
                IF (i_s_star>1) THEN
                    PRINT *, 'nsmin binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmin) THEN
                IF (i_l_star>1) THEN
                    PRINT *, 'nlmin binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmax(idx_z)) THEN
                IF (i_l_star<N_l) THEN
                    PRINT *, 'nlmax binding:', idx_j,idx_z,idx_k
                ENDIF
            ENDIF
            vjwork(idx_z,idx_k) = evo
            sjwork(idx_z,idx_k) = s_grid(i_s_star)
            lss(idx_j,idx_z,idx_k) = l_grid(i_l_star)
            Nsmin(idx_z) = MAX(1,i_s_star - 5)  ! Given z, savings rate increasing in wealth
            Nlmin = MAX(1,i_l_star - 5)         ! Given k, labor supply increasing in productivity
            Nlmax(idx_z) = MIN(N_l,i_l_star + 4)! Given z, labor supply decreasing in wealth
        END DO Z1_LOOP
        END DO K1_LOOP
    
        gapv = SUM(ABS(vjwork - vjwork_old))/(N_k*N_z)
        gaps = SUM(ABS(sjwork - sjwork_old))/(N_k*N_z)
    
        !IF (showdetails.EQ.1) THEN
        IF(multpl(iter,vfitmax) == 1) THEN 
            PRINT *, "avg abs val of v",  gapv
            PRINT *, "avg abs val of saving", gaps
            PRINT *, "iter, age:", iter, idx_j
        ENDIF
    
        IF (gapv<critv .AND. gaps<crits) EXIT
    
    END DO J1_LOOP

    vss(idx_j,:,:) = vjwork
    sss(idx_j,:,:) = sjwork

    IF (rate>rhat) THEN
        vssh = vss
    ELSE
        vssl = vss
    ENDIF
    
    ! Store policy functions
    DO idx_j=1,jmax
    DO idx_k=1,N_k
        IF (idx_j.EQ.3) THEN
            cash = ss + k_grid(idx_k)*R_adj
            kprime = cash*sss(idx_j,1,idx_k)
            cons = cash - kprime
            kss(idx_j,:,idx_k) = kprime
            css(idx_j,:,idx_k) = cons
        ELSE
            DO idx_z=1,N_z
                hours = lss(idx_j,idx_z,idx_k)
                cash = wagetilde*eta_j(idx_j,1)*z_grid(idx_z)*hours + k_grid(idx_k)*R_adj
                kprime = cash*sss(idx_j,idx_z,idx_k)
                cons = cash - kprime
                kss(idx_j,idx_z,idx_k) = kprime
                css(idx_j,idx_z,idx_k) = cons
            END DO
        END IF
    END DO
    END DO

    IF (sswrite.EQ.1) THEN
    499 FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
        OPEN(25,file='vss.out',form='formatted')
        DO idx_k = 1,N_k
            WRITE(25,499) k_grid(idx_k),vss(1,:,idx_k),vss(2,:,idx_k),vss(3,:,idx_k)
        END DO
        CLOSE(25)
    ENDIF
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! Compute stationary distribution !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    kssfine = zerodp
    cssfine = zerodp
    lssfine = zerodp
    mussfine = zerodp

    ! Compute fine policy functions
    DO idx_j=1,jmax
    DO idx_k=1,N_kfine
        wealth = kfine(idx_k)
        IF (idx_j.EQ.3) THEN
            kssfine(idx_j,:,idx_k) = interp1(k_grid,kss(idx_j,1,:),N_k,wealth)
            !FUNCTION interp1(X,Y,siz,xi)
            cssfine(idx_j,:,idx_k) = interp1(k_grid,css(idx_j,1,:),N_k,wealth)
        ELSE
            DO idx_z=1,N_z
                kssfine(idx_j,idx_z,idx_k) = interp1(k_grid,kss(idx_j,idx_z,:),N_k,wealth)
                cssfine(idx_j,idx_z,idx_k) = interp1(k_grid,css(idx_j,idx_z,:),N_k,wealth)
                lssfine(idx_j,idx_z,idx_k) = interp1(k_grid,lss(idx_j,idx_z,:),N_k,wealth)
            END DO
        END IF
    END DO
    END DO
    
    ! Initialize distribution at zero wealth
    DO idx_j=1,jmax
    DO idx_z=1,N_z
        mussfine(idx_j,idx_z,1) = Erd(idx_z)*mu_j(idx_j)
    END DO
    END DO

    ! Construct weights for interpolating policy functions over kfine grid
    jlo1 = 1
    wght1 = zerodp
    jlos1 = 1
    wghts1 = zerodp
    DO idx_j=1,jmax
    DO idx_z=1,N_z
    DO idx_k=1,N_kfine
        CALL hunt(kfine,N_kfine,kssfine(idx_j,idx_z,idx_k),jlo1)  
        IF (jlo1 < 1) THEN
            jlo1 = 1
            wght1 = onedp
        ELSEIF (jlo1 >= N_kfine) THEN
            jlo1 = N_kfine - 1
            wght1 = zerodp
        ELSE   
            wght1 = onedp-(kssfine(idx_j,idx_z,idx_k)-kfine(jlo1))/(kfine(jlo1+1)-kfine(jlo1))   !weight/share going to lower grid point
        ENDIF
        wghts1(idx_j,idx_z,idx_k) = wght1
	    jlos1(idx_j,idx_z,idx_k) = jlo1
        IF (wght1 .lt. 0.0D+00 .or. wght1 .gt. 1.0D+00) THEN
            PRINT*, 'wgth1 is messed up',wght1
            PAUSE
        ENDIF
    END DO
    END DO
    END DO

    MU_LOOP: DO iter = 1,mufitmax  

        mu_old = mussfine
    
        DO idx_j=1,jmax
        DO idx_z=1,N_z
        DO idx_k=1,N_kfine
            IF (mu_old(idx_j,idx_z,idx_k)>zerodp) THEN
                jlo1 = jlos1(idx_j,idx_z,idx_k)
                wght1 = wghts1(idx_j,idx_z,idx_k)
                kprime = kssfine(idx_j,idx_z,idx_k)
                IF (idx_j<jmax) THEN
                    mussfine(idx_j+1,:,jlo1) = mussfine(idx_j+1,:,jlo1) &
				        + mu_old(idx_j,idx_z,idx_k)*psi_j(idx_j)  &
				        *ptran(idx_z,:)*wght1
                    mussfine(idx_j+1,:,jlo1+1) = mussfine(idx_j+1,:,jlo1+1) &
				        + mu_old(idx_j,idx_z,idx_k)*psi_j(idx_j)  &
				        *ptran(idx_z,:)*(onedp-wght1)    
		        END IF
                mussfine(idx_j,:,jlo1) = mussfine(idx_j,:,jlo1) &
				        + mu_old(idx_j,idx_z,idx_k)*(onedp-psi_j(idx_j))  &
				        *ptran(idx_z,:)*wght1
                mussfine(idx_j,:,jlo1+1) = mussfine(idx_j,:,jlo1+1) &
				        + mu_old(idx_j,idx_z,idx_k)*(onedp-psi_j(idx_j))  &
				        *ptran(idx_z,:)*(onedp-wght1)
                mussfine(idx_j,idx_z,idx_k) = mussfine(idx_j,idx_z,idx_k)  &
                        - mu_old(idx_j,idx_z,idx_k)
            END IF
        END DO
        END DO
        END DO

        IF (bequestmode.EQ.1) THEN
            ! Construct weight for interpolating policy functions over kfine grid (initial wealth of newborns that receive a positive endowment)
            CALL hunt(kfine,N_kfine,initk,jlo1)  
            IF (jlo1 < 1) THEN
                jlo1 = 1
                wght1 = onedp
            ELSEIF (jlo1 >= N_kfine) THEN
                jlo1 = N_kfine - 1
                wght1 = zerodp
            ELSE   
                wght1 = onedp-(initk-kfine(jlo1))/(kfine(jlo1+1)-kfine(jlo1))   !weight/share going to lower grid point
            ENDIF
    
            mussfine(1,:,1)         = mussfine(1,:,1)       &
                + (onedp-hikfrac)*Erd*psi_j(jmax)*SUM(mussfine(jmax,:,:)) 
            mussfine(1,:,jlo1)      = mussfine(1,:,jlo1)    &
                + wght1*hikfrac*Erd*psi_j(jmax)*SUM(mussfine(jmax,:,:)) 
            mussfine(1,:,jlo1+1)    = mussfine(1,:,jlo1+1)  &
                + (onedp-wght1)*hikfrac*Erd*psi_j(jmax)*SUM(mussfine(jmax,:,:)) 
        ELSE
            mussfine(1,:,:)         = mussfine(1,:,:)       &
                + psi_j(jmax)*mussfine(jmax,:,:) 
        ENDIF
        
    
        mudif = SUM(DABS(mussfine-mu_old))/(jmax*N_z*N_kfine)
        
        fsum = SUM(mussfine(:,:,:))                 !integrating the pmf
    
        IF (DABS(fsum-onedp) > 1.0D-06) THEN
            PRINT *, "distribution does not integrate to 1!"
        ENDIF
        
        IF (mudif <= 1.0D-12) THEN
            IF (showdetails.EQ.1) THEN
                PRINT *, "DONE WITH STEADY STATE DISTRIBUTION"
            ENDIF
            EXIT
        ELSE
            IF (showdetails.EQ.1) THEN
                WRITE(6,703) iter,mudif
    703         FORMAT(I10,2F18.12)
            ENDIF
        ENDIF
    END DO MU_LOOP

    IF (DABS(fsum-onedp) > 1.0D-06) THEN
        PRINT *, "distribution does not integrate to 1!"
        PAUSE
    ENDIF

    IF (iter.GE.mufitmax) THEN
        PRINT *, "iteration max reached!"
        PAUSE
    END IF

    IF (ABS(SUM(mussfine(jmax,:,:))-mu_j(3))>1.0d-05) THEN
        PRINT *, "Retired mass not correct!"
        PAUSE
    END IF
    
    ! COMPUTE AGGREGATES
    aggk = zerodp
    aggc = zerodp
    aggh = zerodp
    aggl = zerodp
    bequest = zerodp
    DO idx_j=1,jmax
    DO idx_z=1,N_z
    DO idx_k=1,N_kfine
        aggk = aggk + mussfine(idx_j,idx_z,idx_k)*kfine(idx_k)
        aggc = aggc + mussfine(idx_j,idx_z,idx_k)*cssfine(idx_j,idx_z,idx_k)
        aggh = aggh + mussfine(idx_j,idx_z,idx_k)*lssfine(idx_j,idx_z,idx_k)
        aggl = aggl + mussfine(idx_j,idx_z,idx_k)*lssfine(idx_j,idx_z,idx_k)*eta_j(idx_j,1)*z_grid(idx_z)
        IF (idx_j.EQ.3) THEN
            bequest = bequest + psi_j(idx_j)*mussfine(idx_j,idx_z,idx_k)*kfine(idx_k)
        END IF
    END DO
    END DO
    END DO
    aggy = aggk**alpha*aggl**(onedp-alpha)

    aggkinit = aggk
    aggcinit = aggc
    aggocinit = aggc*(onedp-gam)
    agghinit = aggh
    aggohinit = aggh*(onedp-sswfhfractot)
    agglinit = aggl
    aggyinit = aggy

    ! Distribution with coarse kgrid
    muss = zerodp
    DO idx_j=1,jmax
    DO idx_z=1,N_z
    DO idx_k=1,N_kfine
        IF (mussfine(idx_j,idx_z,idx_k)>zerodp) THEN
            ! Create weights
            CALL hunt(k_grid,N_k,kfine(idx_k),jlo1)  
            IF (jlo1 < 1) THEN
                jlo1 = 1
                wght1 = onedp
            ELSEIF (jlo1 >= N_k) THEN
                jlo1 = N_k - 1
                wght1 = zerodp
            ELSE   
                wght1 = onedp-(kfine(idx_k)-k_grid(jlo1))/(k_grid(jlo1+1)-k_grid(jlo1))   !weight/share going to lower grid point
            ENDIF
            IF (wght1 .lt. 0.0D+00 .or. wght1 .gt. 1.0D+00) THEN
                PRINT*, 'wgth1 is messed up',wght1
                PAUSE
            ENDIF
        
            muss(idx_j,idx_z,jlo1) = muss(idx_j,idx_z,jlo1)   &
                + mussfine(idx_j,idx_z,idx_k)*wght1
            muss(idx_j,idx_z,jlo1+1) = muss(idx_j,idx_z,jlo1+1)   &
                + mussfine(idx_j,idx_z,idx_k)*(onedp-wght1)
        ENDIF
    END DO
    END DO
    END DO

    rnew = alpha*(aggk/aggl)**(alpha-onedp)

    SOLVESS  = rnew - r

    IF (showdetails.EQ.1) THEN
        PRINT *, "K:", aggk
        PRINT *, "C:", aggc
        PRINT *, "hours per worker:", aggh/SUM(mussfine(1:2,:,:))
        PRINT *, "L:", aggl
        PRINT *, "Y:", aggy
        PRINT *, "w:", wage
        PRINT *, "Y/worker:", aggy/SUM(mussfine(1:2,:,:))
        PRINT *, "r:", r
        PRINT *, "mu_j:", SUM(mussfine(1,:,:)), SUM(mussfine(2,:,:)), SUM(mussfine(3,:,:)) 
    END IF

    WRITE(6,299) r,rnew,SOLVESS 
    299 FORMAT(3f16.10)

    IF (ALLOCATED (kssfine)) DEALLOCATE (kssfine)
    IF (ALLOCATED (cssfine)) DEALLOCATE (cssfine)
    IF (ALLOCATED (jlos1)) DEALLOCATE (jlos1)
    IF (ALLOCATED (wghts1)) DEALLOCATE (wghts1)
    
END ! SOLVESS

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

DOUBLE PRECISION FUNCTION SOLVE_UBAR(ubarcand)
	!! ubar calibration
        
    ! LOCAL VARIABLES USED WITHIN LOOPS
    REAL(DP), INTENT(IN)    :: ubarcand
	!! Computes VSL
    REAL(DP)    :: v_u(jmax,N_z,N_k), v_uold(jmax,N_z,N_k)
        
    ! LOCAL VARIABLES USED WITHIN LOOPS
    INTEGER     :: idx_j, idx_k
	INTEGER     :: idx_z, idx_zprime ! index for income shock
    INTEGER     :: iter
 
    REAL(DP)    :: kprime		    ! next days saving
    REAL(DP)    :: evp, v_c, v_psi, gapv

    v_u = vss    
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! VSL computation for j=jmax !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    idx_j = jmax
    
    DO iter = 1,vfitmax
        v_uold = v_u
        
        ! c_i^gamma*c_o^(1-gamma) = (c*gamma)^gamma*(c*(1-gamma))^(1-gamma)
        ! = c*gamma^gamma*(1-gamma)^(1-gamma) = c*gamgam
        DO idx_k = 1,N_k
            v_u(idx_j,:,idx_k) = ubarcand                                       &
                + ( css(idx_j,1,idx_k)*gamgam )**(onedp-sigma)/(onedp-sigma)    &
                + beta_adj*(onedp-psi_j(idx_j))*interp1(k_grid,                 &
                v_uold(idx_j,1,:),N_k,kss(idx_j,1,idx_k))
        END DO
        gapv = SUM(ABS(v_u - v_uold))/(N_k*N_z)

        IF (gapv<critv) EXIT    
    
    END DO

	K_LOOP: DO idx_k = 1,N_k
       
        v_c = css(idx_j,1,idx_k)**(-sigma)*gamgam**(onedp-sigma)

        v_psi = beta_adj*(                          &
            (onedp-psi_j(idx_j))*interp1(k_grid,    &
            v_u(idx_j,1,:),N_k,kss(3,1,idx_k))      &
            )
        
        v_vsl(idx_j,:,idx_k) = v_psi/v_c

    END DO K_LOOP

    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! VSL computation for j=2 !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    idx_j = 2
    
    DO iter = 1,vfitmax
        v_uold = v_u
        
        DO idx_k = 1,N_k
        DO idx_z = 1,N_z
            evp = ubarcand                                              &
                - phi*(lss(idx_j,idx_z,idx_k)**(onedp+nu))/(onedp+nu)   &
                + ( css(idx_j,idx_z,idx_k)*gamgam )**(onedp-sigma)/(onedp-sigma)
            kprime = kss(idx_j,idx_z,idx_k)
            
            evp = beta_adj*psi_j(idx_j)*interp1(k_grid,         &
                v_u(idx_j+1,1,:),N_k,kprime)
            DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                evp = evp + beta_adj*(onedp-psi_j(idx_j))       &
                    *ptran(idx_z,idx_zprime)*interp1(k_grid,    &
                    v_uold(idx_j,idx_zprime,:),N_k,kprime)
            END DO
            v_u(idx_j,idx_z,idx_k) = evp
        END DO
        END DO  
        
        gapv = SUM(ABS(v_u - v_uold))/(N_k*N_z)
    
        IF (gapv<critv) EXIT    
    
    END DO
    
    K2_LOOP: DO idx_k = 1,N_k
    Z2_LOOP: DO idx_z = 1,N_z
        
        kprime = kss(idx_j,idx_z,idx_k)
        
        v_c = css(idx_j,idx_z,idx_k)**(-sigma)*gamgam**(onedp-sigma)
        
		v_psi = beta_adj*psi_j(idx_j)*interp1(k_grid,       &
            v_u(idx_j+1,1,:),N_k,kprime)
        DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
            v_psi = v_psi + beta_adj*(onedp-psi_j(idx_j))   &
                *ptran(idx_z,idx_zprime)*interp1(k_grid,    &
                v_u(idx_j,idx_zprime,:),N_k,kprime)
        END DO
       
        v_vsl(idx_j,idx_z,idx_k) = v_psi/v_c

    END DO Z2_LOOP
    END DO K2_LOOP
    
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !! VSL computation for j=1 !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    idx_j = 1
    
    DO iter = 1,vfitmax
        v_uold = v_u
        
        DO idx_k = 1,N_k
        DO idx_z = 1,N_z
            evp = ubarcand                                              &
                - phi*(lss(idx_j,idx_z,idx_k)**(onedp+nu))/(onedp+nu)   &
                + ( css(idx_j,idx_z,idx_k)*gamgam )**(onedp-sigma)/(onedp-sigma)
            kprime = kss(idx_j,idx_z,idx_k)
            
            evp = zerodp
            DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                evp = evp + beta_adj*ptran(idx_z,idx_zprime)            &
                    *interp1(k_grid,                                    &
                    psi_j(idx_j)*v_u(idx_j+1,idx_zprime,:)              &
                    + (onedp-psi_j(idx_j))*v_uold(idx_j,idx_zprime,:),  &
                    N_k,kprime)
            END DO
            v_u(idx_j,idx_z,idx_k) = evp
        END DO
        END DO  
        
        gapv = SUM(ABS(v_u - v_uold))/(N_k*N_z)
    
        IF (gapv<critv) EXIT    
    
    END DO
    
	K1_LOOP: DO idx_k = 1,N_k
    Z1_LOOP: DO idx_z = 1,N_z
        
        kprime = kss(idx_j,idx_z,idx_k)
        
        v_c = css(idx_j,idx_z,idx_k)**(-sigma)*gamgam**(onedp-sigma)
        
		v_psi = zerodp
        DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
            v_psi = v_psi + beta_adj*ptran(idx_z,idx_zprime)            &
                *interp1(k_grid,                                        &
                psi_j(idx_j)*v_u(idx_j+1,idx_zprime,:)                  &
                + (onedp-psi_j(idx_j))*v_u(idx_j,idx_zprime,:),N_k,kprime)
        END DO
        
        v_vsl(idx_j,idx_z,idx_k) = v_psi/v_c
            
    END DO Z1_LOOP
    END DO K1_LOOP
    
    SOLVE_UBAR  = SUM(v_vsl*muss)/aggcinit - vsl
    
    WRITE(*,299) ubarcand, SUM(v_vsl*muss)/aggcinit
299 FORMAT(f12.6,f16.8)
    
END ! SOLVE_UBAR

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

SUBROUTINE SAVESS
    
    INTEGER :: idx_k

    ! Policy and value functions
499 FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
    OPEN(25,file=TRIM(path)//'/kss.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),kss(1,:,idx_k),kss(2,:,idx_k),kss(3,:,idx_k)
    END DO
    CLOSE(25)
    
    OPEN(25,file=TRIM(path)//'/css.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),css(1,:,idx_k),css(2,:,idx_k),css(3,:,idx_k)
    END DO
    CLOSE(25)
    
    OPEN(25,file=TRIM(path)//'/lss.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),lss(1,:,idx_k),lss(2,:,idx_k),lss(3,:,idx_k)
    END DO
    CLOSE(25)
    
    OPEN(25,file=TRIM(path)//'/sss.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),sss(1,:,idx_k),sss(2,:,idx_k),sss(3,:,idx_k)
    END DO
    CLOSE(25)
    
    OPEN(25,file=TRIM(path)//'/muss.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),muss(1,:,idx_k),muss(2,:,idx_k),muss(3,:,idx_k)
    END DO
    CLOSE(25)
    
    OPEN (UNIT=9, FILE=TRIM(path)//'/eta_j.out', POSITION = "REWIND")
    WRITE (9, *) eta_j
    CLOSE (UNIT=9)    
    OPEN (UNIT=9, FILE=TRIM(path)//'/k_grid.out', POSITION = "REWIND")
    WRITE (9, *) k_grid
    CLOSE (UNIT=9)    
    OPEN (UNIT=9, FILE=TRIM(path)//'/z_grid.out', POSITION = "REWIND")
    WRITE (9, *) z_grid
    CLOSE (UNIT=9)
    OPEN (UNIT=9, FILE=TRIM(path)//'/ptran.out', POSITION = "REWIND")
    WRITE (9, *) ptran
    CLOSE (UNIT=9)
    OPEN (UNIT=9, FILE=TRIM(path)//'/Erd.out', POSITION = "REWIND")
    WRITE (9, *) Erd
    CLOSE (UNIT=9)
    OPEN (UNIT=9, FILE=TRIM(path)//'/wfh_z.out', POSITION = "REWIND")
    WRITE (9, *) wfh_z
    CLOSE (UNIT=9)
    
    !!!!!!!!!!!!!!!!!!!!
    ! Model parameters !
    !!!!!!!!!!!!!!!!!!!!    
   
100 FORMAT ('     days',    T16,    &
            '      N_k',    T32,    &
            '      N_z',    T48,    &
            '     jmax',    T64,    &
            '     beta',    T80,    &
            '   beta_e',    T96,    &
            '   beta_c',    T112,   &
            '   beta_l',    T128,   &
            '    pi_ir',    T144,   &
            '   hoscap',    T160,   &
            '    psi_j',    T176,   &
            '    initk',    T192,   &
            '  hikfrac',    T208,   &
            ' sui_cbar',    T224,   &
            '     hcut',    T240,   &
            '   suidur',    T256,   &
            ' lockdown',    T272,   &
            '     hloc',    T288,   &
            '    gamma',    T304,   &
            ' sui_curv',    T320    &
                )
                
102 FORMAT (4I16, 12F16.10, I16, 3F16.10)  
    
    OPEN (UNIT=11, FILE=TRIM(path)//'/PARAMS_MATLAB.out', POSITION="REWIND")
    WRITE (11,100)
    WRITE (11,102)  FLOOR(days), N_k, N_z, jmax,                    &
                    beta_adj, beta_e, beta_c, beta_l, pix,          &
                    hoscap, psi_j(1), initk, hikfrac,               &
                    sui_cbar, hcut, suidur, lockdown, hloc,         &
                    gam, sui_curv
    CLOSE (UNIT=11)

    !!!!!!!!!!!!!!!!!!!!!!
    ! EQUILIBRIUM OUTPUT !
    !!!!!!!!!!!!!!!!!!!!!!  
    
104 FORMAT ('      wss',    T16,    &
            '      rss',    T32,    &
            '        K',    T48,    &
            '        C',    T64,    &
            '        Y',    T80,    &
            '        L',    T96,    &
            '    hours',    T112,   &
            ' ret. inc',    T128,   &
            ' ret. tax',    T144    &
                )

106 FORMAT (9F16.10)  
    
    OPEN (UNIT=11, FILE=TRIM(path)//'/equilibrium_init.out', POSITION="REWIND")
    WRITE (11,104)
    WRITE (11,106)  wss, rss,                       &
                    aggkinit, aggcinit, aggyinit, agglinit, agghinit,   &
                    ss, tau
            
    CLOSE (UNIT=11)
    
END SUBROUTINE SAVESS

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

SUBROUTINE SOLVESSWFH
    
    REAL(DP)    :: sswfhhours
    INTEGER     :: idx_k
    
    sswfhhours = zerodp
    DO idx_k=1,N_kfine
        sswfhhours = sswfhhours + SUM(mussfine(1:jret-1,:,idx_k)*lssfine(1:jret-1,:,idx_k)*wfh_z*sswfhfrac)
    END DO
    
    PRINT *, "Percent of steady state WFH:", sswfhhours/SUM(mussfine(1:jret-1,:,:)*lssfine(1:jret-1,:,:))*100.0_dp
    
    IF (ABS(sswfhhours/SUM(mussfine(1:jret-1,:,:)*lssfine(1:jret-1,:,:))*100.0_dp-14.4_dp)>1D-1) THEN
        sswfhfrac = zbrent(SSWFH,zerodp,onedp,1.0D-14)
        PRINT *, "sswfhfrac:", sswfhfrac
    ENDIF
    
    CONTAINS

    DOUBLE PRECISION FUNCTION SSWFH(x)
    !! Solves the model
        
    ! LOCAL VARIABLES USED WITHIN LOOPS
    REAL(DP), INTENT(IN)    :: x
    REAL(DP)                :: y
    INTEGER                 :: idx_k
    
    y = zerodp
    DO idx_k=1,N_kfine
        y = y + SUM(mussfine(1:jret-1,:,idx_k)*lssfine(1:jret-1,:,idx_k)*wfh_z*x)
    END DO
    y = y/SUM(mussfine(1:jret-1,:,:)*lssfine(1:jret-1,:,:))
    SSWFH = sswfhfractot - y
    
    END FUNCTION SSWFH  

END SUBROUTINE SOLVESSWFH
    
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

SUBROUTINE TRANSITION(runsave)
    !! Solves the transition, given a sequence of prices
    
    INTEGER, INTENT(IN) :: runsave
    INTEGER     :: showdetails = 0
	REAL(DP)    ::  capsule_time_start, capsule_time_stop
    
    ! LOCAL VARIABLES USED WITHIN LOOPS
    INTEGER     :: idx_j
    INTEGER     :: idx_k, idx_kprime, idx_kfine    ! index for wealth
    INTEGER     :: idx_z, idx_zprime
    INTEGER	 	:: idx_h, idx_hprime
    INTEGER     :: idx_l            ! index for labor supply
	INTEGER     :: idx_s            ! index for savings rate
	INTEGER     :: idx_c
    INTEGER     :: t, i, j          ! dummy index
    INTEGER     :: iter, i_s_star, i_l_star, i_c_star, ik50, ik90, new_ik90
    INTEGER, DIMENSION(N_z,N_h)	    :: Nsmin    ! savings rate increases in wealth
    INTEGER, DIMENSION(N_z,N_h)     :: Nlmax    ! labor decreases in wealth
    INTEGER, DIMENSION(N_h)         :: Nlmin    ! labor increases in productivity
    INTEGER, DIMENSION(N_z)         :: Ncmin    ! inside cons frac increases in wealth
    
	REAL(DP)    :: wealth		    ! incoming wealth
    REAL(DP)    :: cash             ! cash-on-hand   
    REAL(DP)    :: srate            ! savings rate        
    REAL(DP)    :: kprime		    ! next days saving
    REAL(DP)    :: cons, uu, hours, wagetilde, cfhrat
    REAL(DP)    :: evp, evo
    REAL(DP), DIMENSION(N_z,N_k,N_h)    :: nextv
    REAL(DP), DIMENSION(jret-1,N_z)     :: hcap
    INTEGER, DIMENSION(jret-1,N_z)      :: Nlcap
    REAL(DP), DIMENSION(N_h,jmax)       :: aggc_hj, aggoc_hj, aggh_hj, aggoh_hj
    REAL(DP), DIMENSION(N_h,jmax,ncaptdist) ::  &
                    aggc_hjpath, aggoc_hjpath, aggh_hjpath, aggoh_hjpath
    
    REAL(DP)    :: taus, tauc, sui, ohcap
    REAL(DP)    :: aggsui, aggtax, aggoc, aggoh, aggocinf, aggohinf, mu_takeup
    
    REAL(DP), DIMENSION(N_kfine)        :: cumpop, ktemp
    REAL(DP), DIMENSION(jmax,N_z,N_kfine)   :: new_muss

    !! VALUE FUNCTION & POLICY FUNCTIONS (Transition)
    REAL(DP), ALLOCATABLE   :: vt(:,:,:,:,:)
    REAL(DP), ALLOCATABLE   :: st(:,:,:,:,:), lt(:,:,:,:,:), ct(:,:,:,:,:), kt(:,:,:,:,:), cfht(:,:,:,:,:)
    REAL(DP), ALLOCATABLE   :: ktfine(:,:,:,:), ctfine(:,:,:,:), ltfine(:,:,:,:), cfhtfine(:,:,:,:)
    REAL(DP), ALLOCATABLE   :: mu_curr(:,:,:,:), mu_next(:,:,:,:)
    
    INTEGER, ALLOCATABLE    :: jlos1(:,:,:,:)
    REAL(DP), ALLOCATABLE   :: wghts1(:,:,:,:)
    
    INTEGER                 :: jlo1
    REAL(DP)                :: wght1
    
    !! ALLOCATE LARGE ARRAYS
    ALLOCATE( vt(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( st(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( lt(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( ct(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( kt(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( cfht(jmax,N_z,N_k,N_h,ncapt+1) )
    ALLOCATE( ltfine(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( ctfine(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( ktfine(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( cfhtfine(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( mu_curr(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( mu_next(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( jlos1(jmax,N_z,N_kfine,N_h) )
    ALLOCATE( wghts1(jmax,N_z,N_kfine,N_h) )
    
    ! Initialize functions
    vt = zerodp
    st = zerodp
    lt = zerodp
    ct = eps
    kt = zerodp
    cfht = gam
    
    ! T+1 value function
    DO idx_h = 1,N_h
    DO idx_k = 1,N_k
    DO idx_z = 1,N_z
    DO idx_j = 1,jmax
        vt(idx_j,idx_z,idx_k,idx_h,ncapt+1) =  vss(idx_j,idx_z,idx_k)
    END DO
    END DO
    END DO
    END DO
    
    betaepathnew = betaepath
    wpath = (onedp-alpha)*(rpath/alpha)**(alpha/(alpha-onedp))
    
    T_LOOP: DO t = ncapt,1,-1
        
        capsule_time_start    = DCLOCK() 
        
        IF (policy.EQ.1) THEN
            IF (t.GE.365.0_dp/days*yrs_cov) THEN 
                taus = zerodp
                tauc = zerodp
                sui = zerodp
            ELSEIF (t.GE.suidur/days) THEN
                taus = taus_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                tauc = tauc_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                sui = sui_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                IF ((t.GE.280.0_dp/days).AND.(t.LE.530_dp/days).AND.(policy2.EQ.1)) THEN
                    sui = MAX(sui,sui_init*0.5_dp)
                END IF
            ELSE
                taus = taus_init
                tauc = tauc_init
                sui = sui_init
            END IF
        ELSE ! policy==0
            taus = zerodp
            tauc = zerodp
            sui = zerodp
        END IF
        
        IF (lockdown.EQ.1) THEN
            IF (t.GE.365.0_dp/days*yrs_cov) THEN ! after covid
                Nlcap = N_l
                hcap = lmax
                ohcap = lmax
            ELSE
                IF (t.GE.locdur/days) THEN ! lockdown winding down
                    ohcap = hloc*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-locdur))**loc_curv      &
                        + (onedp - ((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-locdur))**loc_curv)
                ELSE    ! lockdown
                    ohcap = hloc
                ENDIF
                DO idx_z=1,N_z
                    hcap(1,idx_z) = MIN(lmax,ohcap/(onedp-wfh_z(1,idx_z)))
                    Nlcap(1,idx_z) = MAXVAL( MINLOC( ABS( l_grid - hcap(1,idx_z) ) ) )
                    hcap(2,idx_z) = MIN(lmax,ohcap/(onedp-wfh_z(2,idx_z)))
                    Nlcap(2,idx_z) = MAXVAL( MINLOC( ABS( l_grid - hcap(2,idx_z) ) ) )
                END DO
            END IF
        ELSE ! no lockdown
            Nlcap = N_l
            hcap = lmax
            ohcap = lmax
        ENDIF
    
        bigi = bigipath(t)
        betaet = betaepath(t)
        aggocinf = aggocinfpath(t)
        aggohinf = aggohinfpath(t)
        
        r = rpath(t)
        R_adj = onedp + r - delta_adj
        wage = wpath(t) 
        wagetilde = wage*(onedp-tau-taus)
		
        ! Initialize pih, pid to zero
        pih = zerodp
        pid = zerodp
		IF (t.GE.365.0_dp/days*yrs_cov) THEN ! vaccine/cure
            
			pih(:,:,3) = onedp  ! deterministic recovery with cure
            
        ELSE ! during pandemic
            
            IF ((capcon.EQ.1).AND.(bigi>hoscap)) THEN
                pid = (hoscap/bigi)*deathrate       &
                    + (onedp - hoscap/bigi)*capdr
            ELSE
                pid = deathrate     
            END IF
            
			! pih(:,1,:) determined invidividually below
            
            pih(:,2,2) = onedp - pix
			pih(:,2,3) = pix*(onedp-pid)
            pih(:,3,3) = onedp
            
        END IF
        
        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        !! Value function computation for retired !!
        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        
        idx_j = jmax

        ! inside cons frac increases in wealth (Ncmin(z))
        
        Ncmin = 1
	    
        ! Compute optimal savings (jmax)
	    K_LOOP: DO idx_k = 1,N_k
            
            !!!!!!!!!!!
            idx_h = 1 ! Susceptible
            !!!!!!!!!!!
		    evo = -infty
		    i_s_star = 1
            i_c_star = 1
		    cash = ss + k_grid(idx_k)*R_adj
        
            S_LOOP: DO idx_s = 1,N_s
            C_LOOP: DO idx_c = Ncmin(1),N_cfh
        
                kprime = cash*s_grid(idx_s)
                IF (taxtype.EQ.3) THEN ! Pigouvian tax on outside consumption
                    cons = (cash - kprime)/(onedp+(onedp-cfh_grid(idx_c))*tauc)
                ELSE ! Tax on total consumption
                    cons = (cash - kprime)/(onedp+tauc)
                ENDIF
                CALL utility_ret(cons,cfh_grid(idx_c),uu)
                evp = uu

                IF (t<365.0_dp/days*yrs_cov) THEN ! During pandemic
					IF (t.GE.tvac) THEN
                        pih(idx_j,1,3) = vaccrate(idx_j)
                    ELSE
                        pih(idx_j,1,3) = zerodp
                    ENDIF
                    pih(idx_j,1,2) = MIN(onedp - pih(idx_j,1,3),        &
                        beta_c*aggocinf*cons*(onedp - cfh_grid(idx_c))  &
                        + betaet*bigi)
                ELSE ! End of pandemic
                    pih(idx_j,1,2) = zerodp
					pih(idx_j,1,3) = onedp
                END IF
                pih(idx_j,1,1) = onedp - pih(idx_j,1,2) - pih(idx_j,1,3)

                evp = evp + beta_adj*(onedp-psi_j(idx_j))*    				    &
                        interp1(k_grid,pih(idx_j,idx_h,1)*vt(idx_j,1,:,1,t+1)	&
                        +pih(idx_j,idx_h,2)*vt(idx_j,1,:,2,t+1)				    &
                        +pih(idx_j,idx_h,3)*vt(idx_j,1,:,3,t+1),N_k,kprime)
                
			    IF ( evp > evo ) THEN
				    evo = evp
				    i_s_star = idx_s
				    i_c_star = idx_c
			    END IF	

		    END DO C_LOOP
		    END DO S_LOOP
        
            IF (i_c_star.EQ.Ncmin(1)) THEN
                IF (i_c_star>1) THEN
                    PRINT *, 'ncmin binding', idx_j,1,idx_k,t
                    !PRINT *, CHAR(7)
                ENDIF
            ENDIF
            
            vt(idx_j,:,idx_k,idx_h,t) = evo ! vt(jmax,N_z,N_k,N_h,ncapt+1)
            st(idx_j,:,idx_k,idx_h,t) = s_grid(i_s_star)
            cfht(idx_j,:,idx_k,idx_h,t) = cfh_grid(i_c_star)
            Ncmin(1) = MAX(1,i_c_star - 2) ! inside cons frac increasing in wealth (Ncmin(z))
            
            !!!!!!!!!!!!!!!!!!!!!!!!!
            ! Infected or Recovered !
            !!!!!!!!!!!!!!!!!!!!!!!!!
            
            H_LOOP: DO idx_h = 2,N_h
		    
                evo = -infty
		        i_s_star = 1
		        cash = ss + k_grid(idx_k)*R_adj
            
                DO idx_s = 1,N_s
            
                    kprime = cash*s_grid(idx_s)
                
                    IF (idx_h.EQ.2) THEN ! Infected
                        
                        IF (taxtype.EQ.3) THEN ! Pigouvian tax
                            ! INSIDE CONS. SHARE 
                            ! c_i/(c_i + c_o)
                            ! = c_i/(c_i + (1+tauc)c_o)*(c_i + (1+tauc)*c_o)/(c_i + c_o)
                            ! = sick_gam*(1 + tauc*c_o/(c_i + c_o))
                            ! = sick_gam + sick_gam*tauc*(1 - c_i/(c_i + c_o))
                            ! --> c_i/(c_i + c_o) = sick_gam*(1+tauc)/(1+sick_gam*tauc)
                            cfhrat = sick_gam*(onedp+tauc)/(onedp+sick_gam*tauc)
                            cons = (cash - kprime)/(onedp+(onedp-cfhrat)*tauc)
                        ELSE
                            cfhrat = sick_gam
                            cons = (cash - kprime)/(onedp+tauc)
                        ENDIF
                        CALL utility_sick_ret(cons,cfhrat,uu)
                        evp = uu + beta_adj*(onedp-psi_j(idx_j))*                   &
                                interp1(k_grid,pih(idx_j,2,2)*vt(idx_j,1,:,2,t+1)   &
                                +pih(idx_j,2,3)*vt(idx_j,1,:,3,t+1),N_k,kprime)
                    ELSE ! (idx_h.EQ.3) Recovered
                        
                        IF (taxtype.EQ.3) THEN ! Pigouvian tax
                            cfhrat = gam*(onedp+tauc)/(onedp+gam*tauc)
                            cons = (cash - kprime)/(onedp+(onedp-cfhrat)*tauc)
                        ELSE
                            cfhrat = gam
                            cons = (cash - kprime)/(onedp+tauc)
                        ENDIF
                        CALL utility_ret(cons,cfhrat,uu)
                        evp = uu + beta_adj*(onedp-psi_j(idx_j))*                   &
                                interp1(k_grid,vt(idx_j,1,:,3,t+1),N_k,kprime)
                        
                    ENDIF            
                
			        IF ( evp > evo ) THEN
				        evo = evp
				        i_s_star = idx_s
			        END IF	

		        END DO
        
                vt(idx_j,:,idx_k,idx_h,t) = evo ! vt(jmax,N_z,N_k,N_h,ncapt+1)
                st(idx_j,:,idx_k,idx_h,t) = s_grid(i_s_star)
            
            END DO H_LOOP
        END DO K_LOOP
      
        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
        !! Value function computation for working age !!
        !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

        J_LOOP: DO idx_j = jret-1,1,-1
        
        ! Compute conitinuation value vnext(z,hprime,kprime)
        ! Initialize to zero
        nextv = zerodp 
        
        DO idx_z = 1,N_z
        DO idx_hprime = 1,N_h
        DO idx_kprime = 1,N_k
            IF (idx_j.EQ.jret-1) THEN
                evp = beta_adj*psi_j(idx_j)*vt(idx_j+1,1,idx_kprime,idx_hprime,t+1)
                DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                    evp = evp + beta_adj*ptran(idx_z,idx_zprime)*(onedp-psi_j(idx_j))   &
                        *vt(idx_j,idx_zprime,idx_kprime,idx_hprime,t+1)
                END DO
            ELSE
                evp = zerodp
                DO idx_zprime=Nzmin(idx_z),Nzmax(idx_z)
                    evp = evp + beta_adj*ptran(idx_z,idx_zprime)*(                      &
                        psi_j(idx_j)*vt(idx_j+1,idx_zprime,idx_kprime,idx_hprime,t+1)   &
                        +(onedp-psi_j(idx_j))*vt(idx_j,idx_zprime,idx_kprime,idx_hprime,t+1) )
                END DO
            ENDIF
            nextv(idx_z,idx_kprime,idx_hprime) = evp
        END DO
        END DO
        END DO
        
        ! Given (z,h), savings rate increases in wealth (Nsmin(z,h))
        ! Given (z,h), labor decreases in wealth (Nlmax(z,h))        
        ! Given (k,h), labor increases in productivity (Nlmin(h))
        ! Given z and h=S, inside cons frac increases in wealth (Ncmin(z))
        
        Nsmin = 1
		DO idx_h=1,N_h
			Nlmax(:,idx_h) = Nlcap(idx_j,:)
        END DO
        Ncmin = 1
        
        K2_LOOP: DO idx_k = 1,N_k
        Nlmin = 1
        Z2_LOOP: DO idx_z = 1,N_z
        
            !!!!!!!!!!!
            idx_h = 1 ! Susceptible
            !!!!!!!!!!!
            
            evo = -infty
		    i_s_star = 1
		    i_l_star = 1
            i_c_star = 1
        
            C2_LOOP: DO idx_c = Ncmin(idx_z),N_cfh
            S2_LOOP: DO idx_s = Nsmin(idx_z,idx_h),N_s
            L2_LOOP: DO idx_l = 1,Nlmax(idx_z,idx_h)
                
            hours = l_grid(idx_l)
            IF ((hours.LE.hcut).OR.(idx_l.GE.Nlmin(idx_h))) THEN    
                
                cash = wagetilde*eta_j(idx_j,idx_h)*z_grid(idx_z)*hours     &
                    + k_grid(idx_k)*R_adj
                IF (hours.LE.hcut) cash = cash + sui
                
                kprime = cash*s_grid(idx_s)
                IF (taxtype.EQ.3) THEN ! Pigouvian tax
                    cons = (cash - kprime)/(onedp+(onedp-cfh_grid(idx_c))*tauc)
                ELSE ! Tax on total consumption
                    cons = (cash - kprime)/(onedp+tauc)
                ENDIF
                
                CALL utility(cons,cfh_grid(idx_c),hours,uu)
                
                IF (t<365.0_dp/days*yrs_cov) THEN
                    IF (t.GE.tvac) THEN
                        pih(idx_j,1,3) = vaccrate(idx_j)
                    ELSE
                        pih(idx_j,1,3) = zerodp
                    ENDIF
                    pih(idx_j,1,2) = MIN(onedp - pih(idx_j,1,3),            &
                        beta_c*aggocinf*cons*(onedp-cfh_grid(idx_c))        &
                        + beta_l*aggohinf*hours*(onedp-wfh_z(idx_j,idx_z))  &
                        + betaet*bigi)
					
                ELSE
                    pih(idx_j,1,2) = zerodp
					pih(idx_j,1,3) = onedp
                ENDIF
                pih(idx_j,1,1) = onedp - pih(idx_j,1,2) - pih(idx_j,1,3)

                IF (hours.LE.hcut) uu = uu + utilsui
                
                evp = uu + interp1(k_grid,                  &
                    pih(idx_j,idx_h,1)*nextv(idx_z,:,1)     &
                    + pih(idx_j,idx_h,2)*nextv(idx_z,:,2)   &
                    + pih(idx_j,idx_h,3)*nextv(idx_z,:,3),N_k,kprime)

                IF ( evp > evo ) THEN
				    evo = evp
				    i_s_star = idx_s
                    i_l_star = idx_l
                    i_c_star = idx_c
                END IF	
            END IF  ! ((hours.LE.hcut).OR.(idx_l.GE.Nlmin(idx_h)))
            END DO L2_LOOP
            END DO S2_LOOP
            END DO C2_LOOP
        
            IF (i_s_star.EQ.Nsmin(idx_z,idx_h)) THEN
                IF (i_s_star>1) THEN
                    PRINT *, 'nsmin binding', idx_j,idx_z,idx_k,idx_h, i_s_star
                    !PRINT *, CHAR(7)
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmin(idx_h)) THEN
                IF (i_l_star>1) THEN
                    PRINT *, 'nlmin binding', idx_j,idx_z,idx_k,idx_h, i_l_star
                    !PRINT *, CHAR(7)
                ENDIF
            ENDIF
            IF (i_l_star.EQ.Nlmax(idx_z,idx_h)) THEN
                IF (i_l_star<Nlcap(idx_j,idx_z)) THEN
                    PRINT *, 'nlmax binding', idx_j,idx_z,idx_k,idx_h,i_l_star
                    !PRINT *, CHAR(7)
                ENDIF
            ENDIF
            IF (i_c_star.EQ.Ncmin(idx_z)) THEN
                IF (i_c_star>1) THEN
                    PRINT *, 'ncmin binding', idx_j,idx_z,idx_k,t
                    !PRINT *, CHAR(7)
                ENDIF
            ENDIF
            
            vt(idx_j,idx_z,idx_k,idx_h,t) = evo ! vt(jmax,N_z,N_k,N_h,ncapt+1)
            st(idx_j,idx_z,idx_k,idx_h,t) = s_grid(i_s_star)
            lt(idx_j,idx_z,idx_k,idx_h,t) = l_grid(i_l_star)
            cfht(idx_j,idx_z,idx_k,idx_h,t) = cfh_grid(i_c_star)
            
            ! Given (z,h), savings rate increases in wealth (Nsmin(z,h))
            IF (idx_k>1) Nsmin(idx_z,idx_h) = MAX(1,i_s_star-8)
            ! Given (k,h), labor increases in productivity (Nlmin(h))
            Nlmin(idx_h) = MAX(1,i_l_star-8)
            ! Given z and h=S, inside cons frac increases in wealth (Ncmin(z))
			Ncmin(idx_z) = MAX(1,i_c_star - 5)
            IF (l_grid(i_l_star)>hcut) THEN
                ! Given (z,h), labor decreases in wealth (Nlmax(z,h)) 
                Nlmax(idx_z,idx_h) = MIN(Nlcap(idx_j,idx_z),i_l_star + 3)
            END IF            
            
            !!!!!!!!!!!!!!!!!!!!!!!!!
            ! Infected or Recovered !
            !!!!!!!!!!!!!!!!!!!!!!!!!
            
            H2_LOOP: DO idx_h = 2,N_h
                
                evo = -infty
		        i_s_star = 1
		        i_l_star = 1
        
                DO idx_s = Nsmin(idx_z,idx_h),N_s
                DO idx_l = 1,Nlmax(idx_z,idx_h)

                hours = l_grid(idx_l)
                IF ((hours.LE.hcut).OR.(idx_l.GE.Nlmin(idx_h))) THEN      
                    
                    cash = wagetilde*eta_j(idx_j,idx_h)*z_grid(idx_z)*hours     &
                        + k_grid(idx_k)*R_adj
                    IF (hours.LE.hcut) cash = cash + sui
                    kprime = cash*s_grid(idx_s)
                    
                    IF (idx_h.EQ.2) THEN    ! Infected
                        
                        IF (taxtype.EQ.3) THEN  ! Pigouvian tax
                            cfhrat = sick_gam*(onedp+tauc)/(onedp+sick_gam*tauc)
                            cons = (cash - kprime)/(onedp+(onedp-cfhrat)*tauc)
                        ELSE ! Tax on total consumption
                            cfhrat = sick_gam
                            cons = (cash - kprime)/(onedp+tauc)
                        ENDIF
                        CALL utility_sick(cons,cfhrat,hours,uu)
                        
                    ELSE    ! Recovered
                        
                        IF (taxtype.EQ.3) THEN  ! Pigouvian tax
                            cfhrat = gam*(onedp+tauc)/(onedp+gam*tauc)
                            cons = (cash - kprime)/(onedp+(onedp-cfhrat)*tauc)
                        ELSE ! Tax on total consumption
                            cfhrat = gam
                            cons = (cash - kprime)/(onedp+tauc)
                        ENDIF
                        CALL utility(cons,cfhrat,hours,uu)
                        
                    ENDIF
                    
                    IF (hours.LE.hcut) uu = uu + utilsui
                    
                    evp = uu + interp1(k_grid,                  &
                        + pih(idx_j,idx_h,2)*nextv(idx_z,:,2)   &
                        + pih(idx_j,idx_h,3)*nextv(idx_z,:,3),N_k,kprime)
                    
			        IF ( evp > evo ) THEN
				        evo = evp
				        i_s_star = idx_s
                        i_l_star = idx_l
                    END IF	
                    
                END IF
                END DO
                END DO
        
                IF (i_s_star.EQ.Nsmin(idx_z,idx_h)) THEN
                    IF (i_s_star>1) THEN
                        PRINT *, 'nsmin binding', idx_j,idx_z,idx_k,idx_h, i_s_star
                        !PRINT *, CHAR(7)
                    ENDIF
                ENDIF
                IF (i_l_star.EQ.Nlmin(idx_h)) THEN
                    IF (i_l_star>1) THEN
                        PRINT *, 'nlmin binding', idx_j,idx_z,idx_k,idx_h, i_l_star
                        !PRINT *, CHAR(7)
                    ENDIF
                ENDIF
                IF (i_l_star.EQ.Nlmax(idx_z,idx_h)) THEN
                    IF (i_l_star<Nlcap(idx_j,idx_z)) THEN
                        PRINT *, 'nlmax binding', idx_j,idx_z,idx_k,idx_h, i_l_star
                        !PRINT *, CHAR(7)
                    ENDIF
                ENDIF
            
                vt(idx_j,idx_z,idx_k,idx_h,t) = evo ! vt(jmax,N_z,N_k,N_h,ncapt+1)
                st(idx_j,idx_z,idx_k,idx_h,t) = s_grid(i_s_star)
                lt(idx_j,idx_z,idx_k,idx_h,t) = l_grid(i_l_star)
                
                ! Given (z,h), savings rate increases in wealth (Nsmin(z,h))
                IF (idx_k>1) Nsmin(idx_z,idx_h) = MAX(1,i_s_star-10)
                ! Given (k,h), labor increases in productivity (Nlmin(h))
                Nlmin(idx_h) = MAX(1,i_l_star-8)
			    IF (l_grid(i_l_star)>hcut) THEN
                    ! Given (z,h), labor decreases in wealth (Nlmax(z,h)) 
                    Nlmax(idx_z,idx_h) = MIN(Nlcap(idx_j,idx_z),i_l_star + 3)
                END IF
            
            ENDDO H2_LOOP
            
        END DO Z2_LOOP
        END DO K2_LOOP 
        
        END DO J_LOOP
  
        IF (taxtype.EQ.3) THEN ! Pigouvian tax
            cfht(:,:,:,2,t) = sick_gam*(onedp+tauc)/(onedp+sick_gam*tauc)
            cfht(:,:,:,3,t) = gam*(onedp+tauc)/(onedp+gam*tauc)
        ELSE
            cfht(:,:,:,2,t) = sick_gam
            cfht(:,:,:,3,t) = gam
        ENDIF  
        
        ! Store policy functions
        DO idx_j=1,jmax
        DO idx_k=1,N_k
        DO idx_h=1,N_h
            IF (idx_j.GE.jret) THEN
                cash = ss + k_grid(idx_k)*R_adj
                kprime = cash*st(idx_j,1,idx_k,idx_h,t)
                IF (taxtype.EQ.3) THEN ! Pigouvian tax
                    cons = (cash - kprime)/(onedp+(onedp-cfht(idx_j,1,idx_k,idx_h,t))*tauc)
                ELSE
                    cons = (cash - kprime)/(onedp+tauc)
                ENDIF
                kt(idx_j,:,idx_k,idx_h,t) = kprime
                ct(idx_j,:,idx_k,idx_h,t) = cons
            ELSE
                DO idx_z=1,N_z
                    hours = lt(idx_j,idx_z,idx_k,idx_h,t)
                    cash = wagetilde*eta_j(idx_j,idx_h)*z_grid(idx_z)*hours     &
                        + k_grid(idx_k)*R_adj
                    IF (hours.LE.hcut) cash = cash + sui
                    kprime = cash*st(idx_j,idx_z,idx_k,idx_h,t)
                    IF (taxtype.EQ.3) THEN ! Pigouvian tax
                        cons = (cash - kprime)/(onedp+(onedp-cfht(idx_j,idx_z,idx_k,idx_h,t))*tauc)
                    ELSE
                        cons = (cash - kprime)/(onedp+tauc)
                    ENDIF
                    kt(idx_j,idx_z,idx_k,idx_h,t) = kprime
                    ct(idx_j,idx_z,idx_k,idx_h,t) = cons
                END DO
            END IF
        END DO
        END DO
        END DO
    
        capsule_time_stop     = DCLOCK()   
        IF (showdetails.EQ.1) THEN
        IF(multpl(t,10) == 1) THEN 
    	    PRINT *, 't:', t, '| computation:', capsule_time_stop - capsule_time_start , 'seconds'
        END IF
        END IF
        
    END DO T_LOOP
      
    DO idx_j=1,jmax
    DO idx_z=1,N_z
    DO idx_k=1,N_k
    DO idx_h=1,N_h
        st(idx_j,idx_z,idx_k,idx_h,ncapt+1) = sss(idx_j,idx_z,idx_k)
        lt(idx_j,idx_z,idx_k,idx_h,ncapt+1) = lss(idx_j,idx_z,idx_k)
        kt(idx_j,idx_z,idx_k,idx_h,ncapt+1) = kss(idx_j,idx_z,idx_k)
        ct(idx_j,idx_z,idx_k,idx_h,ncapt+1) = css(idx_j,idx_z,idx_k)
    END DO
    END DO
    END DO
    END DO
    cfht(:,:,:,:,ncapt+1) = gam
                
    !!!!!!!!!!!!!!!!!!!!!!!!!!!
    !!! Compute distribution !!
    !!!!!!!!!!!!!!!!!!!!!!!!!!!
        
    IF (showdetails.EQ.1) THEN
        PRINT *, 'Computing distributions ...'
    ENDIF

    ! Initial distribution
    DO idx_j = 1,jmax
    DO idx_z = 1,N_z
    DO idx_kfine = 1,N_kfine
    DO idx_h = 1,N_h
        mu_next(idx_j,idx_z,idx_kfine,idx_h) = mussfine(idx_j,idx_z,idx_kfine)*muh(idx_h)
    END DO
    END DO
    END DO
    END DO 
    
    DO t = 1,ncaptdist   
        
        bigi = bigipath(t)
        betaet = betaepath(t)
        aggocinf = aggocinfpath(t)
        aggohinf = aggohinfpath(t)
        wage = wpath(t) 
        
        IF (policy.EQ.1) THEN
            IF (t.GE.365.0_dp/days*yrs_cov) THEN 
                taus = zerodp
                tauc = zerodp
                sui = zerodp
            ELSEIF (t.GE.suidur/days) THEN
                taus = taus_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                tauc = tauc_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                sui = sui_init*((365.0_dp*yrs_cov-DBLE(t)*days)/(365.0_dp*yrs_cov-suidur))**sui_curv
                IF ((t.GE.280.0_dp/days).AND.(t.LE.530_dp/days).AND.(policy2.EQ.1)) THEN
                    sui = MAX(sui,sui_init*0.5_dp)
                END IF
            ELSE
                taus = taus_init
                tauc = tauc_init
                sui = sui_init
            END IF
        ELSE
            taus = zerodp
            tauc = zerodp
            sui = zerodp
        END IF
        
        ! Initialize pih, pid to zero
        pih = zerodp
        pid = zerodp
		IF (t.GE.365.0_dp/days*yrs_cov) THEN ! vaccine/cure
                        
			pih(:,:,3) = onedp  ! deterministic recovery with cure
            
        ELSE ! During pandemic
            
            IF ((capcon.EQ.1).AND.(bigi>hoscap)) THEN
                pid = hoscap/bigi*deathrate         &
                    + (onedp - hoscap/bigi)*capdr
            ELSE
                pid = deathrate     
            END IF
            
			! pih(:,1,:) determined invidividually below
            
            pih(:,2,2) = onedp - pix
			pih(:,2,3) = pix*(onedp-pid)
            pih(:,3,3) = onedp
            
		END IF
        
        mu_curr = mu_next
        mu_next = zerodp   
        
        ktfine = zerodp
        ctfine = zerodp
        ltfine = zerodp
        cfhtfine = zerodp
        
        DO idx_j=1,jmax
        DO idx_kfine=1,N_kfine
            wealth = kfine(idx_kfine)
            DO idx_z=1,N_z
            DO idx_h=1,N_h
                
                CALL hunt(k_grid,N_k,wealth,idx_k)  
                IF ((idx_k < 1).OR.(idx_k >= N_k)) THEN
                    ktfine(idx_j,idx_z,idx_kfine,idx_h)     = kt(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                    ctfine(idx_j,idx_z,idx_kfine,idx_h)     = ct(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                    ltfine(idx_j,idx_z,idx_kfine,idx_h)     = lt(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                    cfhtfine(idx_j,idx_z,idx_kfine,idx_h)   = cfht(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                ELSE
                    IF (sui>zerodp) THEN ! Policy
                        IF ((lt(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1)).LE.hcut).AND.(lt(idx_j,idx_z,idx_k+1,idx_h,MIN(t,ncapt+1))>hcut)) THEN
                            IF (ABS(wealth-k_grid(idx_k)).LE.0.5_dp*(k_grid(idx_k+1)-k_grid(idx_k))) THEN
                                ktfine(idx_j,idx_z,idx_kfine,idx_h)     = kt(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                                ctfine(idx_j,idx_z,idx_kfine,idx_h)     = ct(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                                ltfine(idx_j,idx_z,idx_kfine,idx_h)     = lt(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                                cfhtfine(idx_j,idx_z,idx_kfine,idx_h)   = cfht(idx_j,idx_z,idx_k,idx_h,MIN(t,ncapt+1))
                            ELSE
                                ktfine(idx_j,idx_z,idx_kfine,idx_h)     = kt(idx_j,idx_z,idx_k+1,idx_h,MIN(t,ncapt+1))
                                ctfine(idx_j,idx_z,idx_kfine,idx_h)     = ct(idx_j,idx_z,idx_k+1,idx_h,MIN(t,ncapt+1))
                                ltfine(idx_j,idx_z,idx_kfine,idx_h)     = lt(idx_j,idx_z,idx_k+1,idx_h,MIN(t,ncapt+1))
                                cfhtfine(idx_j,idx_z,idx_kfine,idx_h)   = cfht(idx_j,idx_z,idx_k+1,idx_h,MIN(t,ncapt+1))
                            ENDIF
                        ELSE
                            ktfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,kt(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                            ctfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,ct(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                            ltfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,lt(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                            cfhtfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,cfht(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                        ENDIF
                    ELSE ! No policy
                        ktfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,kt(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                        ctfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,ct(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                        ltfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,lt(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                        cfhtfine(idx_j,idx_z,idx_kfine,idx_h) = interp1(k_grid,cfht(idx_j,idx_z,:,idx_h,MIN(t,ncapt+1)),N_k,wealth)
                    ENDIF
                ENDIF
            END DO
            END DO
        END DO
        END DO

        jlo1 = 1
        wght1 = zerodp
        jlos1 = 1
        wghts1 = zerodp
        DO idx_j=1,jmax
        DO idx_z=1,N_z
        DO idx_kfine=1,N_kfine
        DO idx_h=1,N_h
            CALL hunt(kfine,N_kfine,ktfine(idx_j,idx_z,idx_kfine,idx_h),jlo1)  
            IF (jlo1 < 1) THEN
                jlo1 = 1
                wght1 = onedp
            ELSEIF (jlo1 >= N_kfine) THEN
                jlo1 = N_kfine - 1
                wght1 = zerodp
            ELSE   
                wght1 = onedp-(ktfine(idx_j,idx_z,idx_kfine,idx_h)-kfine(jlo1))/(kfine(jlo1+1)-kfine(jlo1))   !weight/share going to lower grid point
            ENDIF
            wghts1(idx_j,idx_z,idx_kfine,idx_h) = wght1
	        jlos1(idx_j,idx_z,idx_kfine,idx_h) = jlo1
            IF (wght1 .lt. 0.0D+00 .or. wght1 .gt. 1.0D+00) THEN
                PRINT*, 'wgth1 is messed up',wght1
                PAUSE
            ENDIF
        END DO
        END DO
        END DO
        END DO
        
        ! Initialize aggregate supplemental income
        mu_takeup   = zerodp
        
        DO idx_j=1,jmax
        DO idx_z=1,N_z
        DO idx_kfine=1,N_kfine
        DO idx_h=1,N_h
            
            IF (mu_curr(idx_j,idx_z,idx_kfine,idx_h)>zerodp) THEN
                
                jlo1 = jlos1(idx_j,idx_z,idx_kfine,idx_h)
                wght1 = wghts1(idx_j,idx_z,idx_kfine,idx_h)
                kprime = ktfine(idx_j,idx_z,idx_kfine,idx_h)
                cons = ctfine(idx_j,idx_z,idx_kfine,idx_h)
                hours = ltfine(idx_j,idx_z,idx_kfine,idx_h)

                IF (idx_j<jmax) THEN
                    IF ((hours.LE.hcut).AND.(sui>zerodp)) THEN
                        mu_takeup = mu_takeup + mu_curr(idx_j,idx_z,idx_kfine,idx_h)
                    END IF    
                ENDIF
                
				IF (idx_h.EQ.1) THEN
                    IF (t<365.0_dp/days*yrs_cov) THEN
					    IF (t.GE.tvac) THEN
                            pih(idx_j,1,3) = vaccrate(idx_j)
                        ELSE
                            pih(idx_j,1,3) = zerodp
                        ENDIF
                        pih(idx_j,1,2) = MIN(onedp - pih(idx_j,1,3),            &
                            beta_c*aggocinf*cons*(onedp-cfhtfine(idx_j,idx_z,idx_kfine,idx_h))  &
                            + beta_l*aggohinf*hours*(onedp-wfh_z(idx_j,idx_z))  &
                            + betaet*bigi)
                    ELSE ! After pandemic
                        pih(idx_j,1,2) = zerodp
						pih(idx_j,1,3) = onedp
                    ENDIF
                    pih(idx_j,1,1) = onedp - pih(idx_j,1,2) - pih(idx_j,1,3)
                END IF
                				
                DO idx_hprime=idx_h,N_h
                DO idx_zprime=1,N_z
                    IF (idx_j<jmax) THEN
                        mu_next(idx_j+1,idx_zprime,jlo1,idx_hprime) =               	&
                            mu_next(idx_j+1,idx_zprime,jlo1,idx_hprime)             	&
				            + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*psi_j(idx_j)         	&
				            *ptran(idx_z,idx_zprime)*wght1*pih(idx_j,idx_h,idx_hprime)
                        mu_next(idx_j+1,idx_zprime,jlo1+1,idx_hprime) =             	&
                            mu_next(idx_j+1,idx_zprime,jlo1+1,idx_hprime)           	&
				            + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*psi_j(idx_j)         	&
				            *ptran(idx_z,idx_zprime)*(onedp-wght1)*pih(idx_j,idx_h,idx_hprime)
                    END IF
                    mu_next(idx_j,idx_zprime,jlo1,idx_hprime) =                     &
                        mu_next(idx_j,idx_zprime,jlo1,idx_hprime)                   &
				        + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*(onedp-psi_j(idx_j))     &
				        *ptran(idx_z,idx_zprime)*wght1*pih(idx_j,idx_h,idx_hprime)
                    mu_next(idx_j,idx_zprime,jlo1+1,idx_hprime) =                   &
                        mu_next(idx_j,idx_zprime,jlo1+1,idx_hprime)                 &
				        + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*(onedp-psi_j(idx_j))     &
				        *ptran(idx_z,idx_zprime)*(onedp-wght1)*pih(idx_j,idx_h,idx_hprime)
                END DO
                END DO
            END IF
        END DO
        END DO
        END DO
        END DO
        
        aggsui = sui*mu_takeup
        
        ! newborns enter at the constant steady state level
        IF (bequestmode.EQ.1) THEN
            CALL hunt(kfine,N_kfine,initk,jlo1)  
            IF (jlo1 < 1) THEN
                jlo1 = 1
                wght1 = onedp
            ELSEIF (jlo1 >= N_kfine) THEN
                jlo1 = N_kfine - 1
                wght1 = zerodp
            ELSE   
                wght1 = onedp-(initk-kfine(jlo1))/(kfine(jlo1+1)-kfine(jlo1))   !weight/share going to lower grid point
            ENDIF
            mu_next(1,:,1,1)        = mu_next(1,:,1,1)      &
                + Erd*psi_j(jmax)*SUM(muss(jmax,:,:))*(onedp-hikfrac)
            mu_next(1,:,jlo1,1)     = mu_next(1,:,jlo1,1)   &
                + Erd*psi_j(jmax)*SUM(muss(jmax,:,:))*hikfrac*wght1
            mu_next(1,:,jlo1+1,1)   = mu_next(1,:,jlo1+1,1) &
                + Erd*psi_j(jmax)*SUM(muss(jmax,:,:))*hikfrac*(onedp-wght1)
        ELSE
            mu_next(1,:,:,1)        = mu_next(1,:,:,1)      &
                + psi_j(jmax)*mussfine(jmax,:,:) 
        ENDIF
        
        ! COMPUTE AGGREGATES
        aggk = zerodp
        aggl = zerodp
        aggc_hj = zerodp
        aggoc_hj = zerodp
        aggh_hj = zerodp
        aggoh_hj = zerodp
        DO idx_j=1,jmax
        DO idx_z=1,N_z
        DO idx_kfine=1,N_kfine
            aggk    = aggk  + SUM(mu_curr(idx_j,idx_z,idx_kfine,:)*kfine(idx_kfine))
            aggl    = aggl  + SUM(mu_curr(idx_j,idx_z,idx_kfine,:)*ltfine(idx_j,idx_z,idx_kfine,:)*eta_j(idx_j,:)*z_grid(idx_z))
            DO idx_h=1,N_h
                aggc_hj(idx_h,idx_j)    = aggc_hj(idx_h,idx_j)  &
                    + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*ctfine(idx_j,idx_z,idx_kfine,idx_h)
                aggoc_hj(idx_h,idx_j)   = aggoc_hj(idx_h,idx_j) &
                    + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*ctfine(idx_j,idx_z,idx_kfine,idx_h)*(onedp-cfhtfine(idx_j,idx_z,idx_kfine,idx_h))
                IF (idx_j<jret) THEN
                    aggh_hj(idx_h,idx_j)    = aggh_hj(idx_h,idx_j)  &
                        + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*ltfine(idx_j,idx_z,idx_kfine,idx_h)
                    IF (idx_h<3.OR.wfh<2) THEN
                        aggoh_hj(idx_h,idx_j)   = aggoh_hj(idx_h,idx_j) &
                            + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*ltfine(idx_j,idx_z,idx_kfine,idx_h)*(onedp-wfh_z(idx_j,idx_z))
                    ELSE ! (idx_h==3.AND.wfh==2, i.e. recovered do not work from home)
                        aggoh_hj(idx_h,idx_j)   = aggoh_hj(idx_h,idx_j) &
                            + mu_curr(idx_j,idx_z,idx_kfine,idx_h)*ltfine(idx_j,idx_z,idx_kfine,idx_h)*(onedp-wfh_z(idx_j,idx_z)*sswfhfrac)
                    END IF
                END IF
            END DO
        END DO
        END DO
        END DO
        aggc = SUM(aggc_hj)
        aggoc = SUM(aggoc_hj)
        aggh = SUM(aggh_hj)
        aggoh = SUM(aggoh_hj)
        aggy = aggk**alpha*aggl**(onedp-alpha)

        aggc1path(t) = SUM(aggc_hj(:,1))
        aggc2path(t) = SUM(aggc_hj(:,2))
        aggc3path(t) = SUM(aggc_hj(:,3))
        aggoc1path(t) = SUM(aggoc_hj(:,1))
        aggoc2path(t) = SUM(aggoc_hj(:,2))
        aggoc3path(t) = SUM(aggoc_hj(:,3))
        aggh1path(t) = SUM(aggh_hj(:,1))
        aggh2path(t) = SUM(aggh_hj(:,2))
        aggh3path(t) = SUM(aggh_hj(:,3))
        aggoh1path(t) = SUM(aggoh_hj(:,1))
        aggoh2path(t) = SUM(aggoh_hj(:,2))
        aggoh3path(t) = SUM(aggoh_hj(:,3))
        
        aggocinfpathnew(t) = SUM(aggoc_hj(2,:))
        aggohinfpathnew(t) = SUM(aggoh_hj(2,1:2))
        
        DO idx_j=1,jmax
            DO idx_h=1,N_h
                aggc_hjpath(idx_h,idx_j,t) = aggc_hj(idx_h,idx_j)
                aggoc_hjpath(idx_h,idx_j,t) = aggoc_hj(idx_h,idx_j)
                aggh_hjpath(idx_h,idx_j,t) = aggh_hj(idx_h,idx_j)
                aggoh_hjpath(idx_h,idx_j,t) = aggoh_hj(idx_h,idx_j)
            END DO
        END DO
        
        bigspath(t) = SUM(mu_curr(:,:,:,1))
        ipathnew(t) = SUM(mu_curr(:,:,:,2))
        bigrpath(t) = SUM(mu_curr(:,:,:,3))
        
        s1path(t) = SUM(mu_curr(1,:,:,1))
        i1path(t) = SUM(mu_curr(1,:,:,2))
        r1path(t) = SUM(mu_curr(1,:,:,3))
        s2path(t) = SUM(mu_curr(2,:,:,1))
        i2path(t) = SUM(mu_curr(2,:,:,2))
        r2path(t) = SUM(mu_curr(2,:,:,3))
        s3path(t) = SUM(mu_curr(3,:,:,1))
        i3path(t) = SUM(mu_curr(3,:,:,2))
        r3path(t) = SUM(mu_curr(3,:,:,3))
        
        d1path(t)   = SUM(mu_curr(1,:,:,2))*pix*pid(1)
		d2path(t)   = SUM(mu_curr(2,:,:,2))*pix*pid(2)
		d3path(t)   = SUM(mu_curr(3,:,:,2))*pix*pid(3)
        
        dpath(t)    = d1path(t) + d2path(t) + d3path(t)
        
        IF (tvbetaecal.GE.1) THEN
            IF ((t>1).AND.(t.LE.N_death)) THEN 
                betaepathnew(t-1) = MIN(2.00_dp,MAX(0.00_dp,( deaths_data(t)/pix                &
                    - pid(1)*(                                                                  &
                        i1path(t-1)*(onedp-psi_j(1))*(onedp-pix)                                &
                        + beta_c*aggoc_hjpath(1,1,t-1)*(onedp-psi_j(1))*aggocinfpathnew(t-1)    &
                        + beta_l*aggoh_hjpath(1,1,t-1)*(onedp-psi_j(1))*aggohinfpathnew(t-1)    &
                        )                                                                       &
                    - pid(2)*(                                                                  &
                        ( i2path(t-1)*(onedp-psi_j(2))                                          &
                        + i1path(t-1)*psi_j(1) )*(onedp-pix)                                    &
                        + beta_c*( aggoc_hjpath(1,2,t-1)*(onedp-psi_j(2))                       &
                            + aggoc_hjpath(1,1,t-1)*psi_j(1) )*aggocinfpathnew(t-1)             &
                        + beta_l*( aggoh_hjpath(1,2,t-1)*(onedp-psi_j(2))                       &
                            + aggoh_hjpath(1,1,t-1)*psi_j(1) )*aggohinfpathnew(t-1))            &
                    - pid(3)*(                                                                  &
                        ( i3path(t-1)*(onedp-psi_j(3))                                          &
                        + i2path(t-1)*psi_j(2) )*(onedp-pix)                                    &
                        + beta_c*( aggoc_hjpath(1,3,t-1)*(onedp-psi_j(3))                       &
                            + aggoc_hjpath(1,2,t-1)*psi_j(2) )*aggocinfpathnew(t-1)) )          &
                    /( ipathnew(t-1)*(                                                          &
                        s1path(t-1)*(onedp-psi_j(1))*pid(1)                                     &
                        + (s2path(t-1)*(onedp-psi_j(2))+s1path(t-1)*psi_j(1))*pid(2)            &
                        + (s3path(t-1)*(onedp-psi_j(3))+s2path(t-1)*psi_j(2))*pid(3)) ) ) )
            END IF
        END IF
        
        IF (taxtype.EQ.3) THEN ! Pigvouivan tax
            aggtax = aggoc*tauc +  wage*aggl*taus
        ELSE ! Tax on total consumption
            aggtax = aggc*tauc +  wage*aggl*taus
        ENDIF
        
        aggkpath(t) = aggk
        aggcpath(t) = aggc
        aggocpath(t) = aggoc
        agghpath(t) = aggh
        aggohpath(t) = aggoh
        agglpath(t) = aggl
        aggypath(t) = aggy
        aggtaxpath(t) = aggtax
        aggsuipath(t) = aggsui
        mu_takeuppath(t) = mu_takeup/SUM(mu_curr(1:2,:,:,:))
        rpathnew(t) = alpha*(aggk/aggl)**(alpha-onedp)
    
    END DO
    
    pvdef = zerodp
    DO t=ncapt,1,-1
        pvdef = pvdef/(onedp+rpath(t)-delta_adj) + aggsuipath(t) - aggtaxpath(t)
    END DO    
    
    IF (showdetails.EQ.1) THEN
        PRINT *, 'peak infection rate:', MAXVAL(ipathnew)
        PRINT *, 'Infected labor decline: ', SUM(muss(1:2,:,:)*(    &
            lt(1:2,:,:,2,ncapt)/lss(1:2,:,:)-onedp))/SUM(muss(1:2,:,:))
    ENDIF
    PRINT *, 'max takeup:', MAXVAL(mu_takeuppath)
    PRINT *, 'output:', SUM(aggypath(1:52))/(aggyinit*52.0_dp)
	    
    vhat_s = vt(:,:,:,1,1)

499 FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
498 FORMAT(i5, f18.10,7f18.10,7f18.10,7f18.10)
    
    IF ((runsave.EQ.1).AND.(runmode.EQ.0)) THEN
        ! Policy and value functions
    
        OPEN(25,file=TRIM(path)//'/vt_s.out',form='formatted')
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),vt(1,:,idx_k,1,t),vt(2,:,idx_k,1,t),vt(3,:,idx_k,1,t)
        END DO
        END DO
        CLOSE(25)    
        OPEN(25,file=TRIM(path)//'/vt_i.out',form='formatted')
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),vt(1,:,idx_k,2,t),vt(2,:,idx_k,2,t),vt(3,:,idx_k,2,t)
        END DO
        END DO
        CLOSE(25)       

        OPEN (25,FILE=TRIM(path)//'/lt_s.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),lt(1,:,idx_k,1,t),lt(2,:,idx_k,1,t),lt(3,:,idx_k,1,t)
        END DO
        END DO
        CLOSE(25)
        OPEN (25,FILE=TRIM(path)//'/ct_s.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),ct(1,:,idx_k,1,t),ct(2,:,idx_k,1,t),ct(3,:,idx_k,1,t)
        END DO
        END DO
        CLOSE (25)
        OPEN (25,FILE=TRIM(path)//'/cfht_s.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),cfht(1,:,idx_k,1,t),cfht(2,:,idx_k,1,t),cfht(3,:,idx_k,1,t)
        END DO
        END DO
        CLOSE (25) 
        OPEN (25,FILE=TRIM(path)//'/lt_i.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),lt(1,:,idx_k,2,t),lt(2,:,idx_k,2,t),lt(3,:,idx_k,2,t)
        END DO
        END DO
        CLOSE(25)
        OPEN (25,FILE=TRIM(path)//'/ct_i.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),ct(1,:,idx_k,2,t),ct(2,:,idx_k,2,t),ct(3,:,idx_k,2,t)
        END DO
        END DO
        CLOSE (25)
        OPEN (25,FILE=TRIM(path)//'/ct_r.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),ct(1,:,idx_k,3,t),ct(2,:,idx_k,3,t),ct(3,:,idx_k,3,t)
        END DO
        END DO
        CLOSE (25)
        OPEN (25,FILE=TRIM(path)//'/lt_r.out',FORM="FORMATTED")
        DO t = 1,ncapt+1
        DO idx_k = 1,N_k
            WRITE(25,498) t, k_grid(idx_k),lt(1,:,idx_k,3,t),lt(2,:,idx_k,3,t),lt(3,:,idx_k,3,t)
        END DO
        END DO
        CLOSE (25)    
    END IF   
            
    IF (ALLOCATED (vt)) DEALLOCATE (vt)
    IF (ALLOCATED (st)) DEALLOCATE (st)
    IF (ALLOCATED (lt)) DEALLOCATE (lt)
    IF (ALLOCATED (ct)) DEALLOCATE (ct)
    IF (ALLOCATED (cfht)) DEALLOCATE (cfht)
    IF (ALLOCATED (kt)) DEALLOCATE (kt)
    IF (ALLOCATED (cfht)) DEALLOCATE (cfht)
    IF (ALLOCATED (ltfine)) DEALLOCATE (ltfine)
    IF (ALLOCATED (ctfine)) DEALLOCATE (ctfine)
    IF (ALLOCATED (cfhtfine)) DEALLOCATE (cfhtfine)
    IF (ALLOCATED (ktfine)) DEALLOCATE (ktfine)
    IF (ALLOCATED (mu_curr)) DEALLOCATE (mu_curr)
    IF (ALLOCATED (mu_next)) DEALLOCATE (mu_next)
    IF (ALLOCATED (jlos1)) DEALLOCATE (jlos1)
    IF (ALLOCATED (wghts1)) DEALLOCATE (wghts1)
    
    !!!!!!!!!!!!!!!!!!!!!!
    ! EQUILIBRIUM OUTPUT !
    !!!!!!!!!!!!!!!!!!!!!!  
    
104 FORMAT ('   period',    T5,     &
            '     wage',    T20,    &
            '        r',    T35,    &
            '        K',    T50,    &
            '        C',    T65,    &
            '        Y',    T80,    &
            '        L',    T95,    &
            '    hours',    T110,   &
            '        S',    T125,   &
            '        I',    T140,   &
            '        R',    T155,   &
            ' S(young)',    T170,   &
            ' I(young)',    T185,   &
            ' R(young)',    T200,   &
            '   S(mid)',    T215,   &
            '   I(mid)',    T230,   &
            '   R(mid)',    T245,   &
            '   S(old)',    T260,   &
            '   I(old)',    T275,   &
            '   R(old)',    T290,   &
            ' C(young)',    T305,   &
            ' h(young)',    T320,   &
            '   C(mid)',    T335,   &
            '   h(mid)',    T350,   &
            '   C(old)',    T365,   &
            '   h(old)',    T380,   &
            '   C(y,s)',    T395,   &
            '   C(m,s)',    T410,   &
            '   C(o,s)',    T425,   &
            '   C(y,i)',    T440,   &
            '   C(m,i)',    T455,   &
            '   C(o,i)',    T470,   &
            '   C(y,r)',    T485,   &
            '   C(m,r)',    T500,   &
            '   C(o,r)',    T515,   &
            '  oh(y,s)',    T530,   &
            '  oh(m,s)',    T545,   &
            '  oh(o,s)',    T560,   &
            '  oh(y,i)',    T575,   &
            '  oh(m,i)',    T590,   &
            '  oh(o,i)',    T605,   &
            '  oh(y,r)',    T620,   &
            '  oh(m,r)',    T635,   &
            '  oh(o,r)',    T650,   &
            '   h(y,s)',    T665,   &
            '   h(m,s)',    T680,   &
            '   h(o,s)',    T695,   &
            '   h(y,i)',    T710,   &
            '   h(m,i)',    T725,   &
            '   h(o,i)',    T740,   &
            '   h(y,r)',    T755,   &
            '   h(m,r)',    T770,   &
            '   h(o,r)',    T785,   &
            '   deaths',    T800,   &
            '  tax rev',    T815,   &
            '  agg sui',    T830,   &
            '   pv def',    T845,   &
            '  oc(y,s)',    T860,   &
            '  oc(m,s)',    T875,   &
            '  oc(o,s)',    T890,   &
            '  oc(y,i)',    T905,   &
            '  oc(m,i)',    T920,   &
            '  oc(o,i)',    T935,   &
            '  oc(y,r)',    T950,   &
            '  oc(m,r)',    T965,   &
            '  oc(o,r)',    T980,   &
            'mu_takeup',    T995,   &
            '  beta_et',    T1010,  &
            ' deaths_y',    T1025,  &
            ' deaths_m',    T1040,  &
            ' deaths_o',    T1055   &
                )

106 FORMAT (I5, 70f15.10)
    
    OPEN (UNIT=11, FILE=TRIM(path)//'/equilibrium_tran.out', POSITION="REWIND")
    WRITE (11,104)
    WRITE (11,106)  0, wss, rss,        aggkinit,   aggcinit,   &
                    aggyinit,           agglinit,   agghinit,   &
                    onedp,              zerodp,     zerodp,     &
                    SUM(muss(1,:,:)),   zerodp,     zerodp,     &
                    SUM(muss(2,:,:)),   zerodp,     zerodp,     &
                    SUM(muss(3,:,:)),   zerodp,     zerodp,     &
                    SUM(css(1,:,:)*muss(1,:,:)),                &
                    SUM(lss(1,:,:)*muss(1,:,:)),                &
                    SUM(css(2,:,:)*muss(2,:,:)),                &
                    SUM(lss(2,:,:)*muss(2,:,:)),                &
                    SUM(css(3,:,:)*muss(3,:,:)),                &
                    SUM(lss(3,:,:)*muss(3,:,:)),                &
                    SUM(css(1,:,:)*muss(1,:,:)),                &
                    SUM(css(2,:,:)*muss(2,:,:)),                &
                    SUM(css(3,:,:)*muss(3,:,:)),                &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp,             zerodp,     zerodp,     &
                    SUM(lss(1,:,:)*muss(1,:,:))*(onedp-sswfhfractot),   &
                    SUM(lss(2,:,:)*muss(2,:,:))*(onedp-sswfhfractot),   &
                    SUM(lss(3,:,:)*muss(3,:,:))*(onedp-sswfhfractot),   &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp,             zerodp,     zerodp,     &
                    SUM(lss(1,:,:)*muss(1,:,:)),                &
                    SUM(lss(2,:,:)*muss(2,:,:)),                &
                    SUM(lss(3,:,:)*muss(3,:,:)),                &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp, tauc_init, taus_init, pvdef,        &
                    SUM(css(1,:,:)*muss(1,:,:))*(onedp-gam),    &
                    SUM(css(2,:,:)*muss(2,:,:))*(onedp-gam),    &
                    SUM(css(3,:,:)*muss(3,:,:))*(onedp-gam),    &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp,             zerodp,     zerodp,     &
                    zerodp, beta_e,                             &
                    zerodp,             zerodp,     zerodp
    
    DO t=1,ncaptdist
        WRITE (11,106)  t, wpath(t), rpath(t), aggkpath(t),     aggcpath(t),            &
                        aggypath(t),        agglpath(t),        agghpath(t),            &
                        bigspath(t),        bigipath(t),        bigrpath(t),            &
                        s1path(t),          i1path(t),          r1path(t),              &
                        s2path(t),          i2path(t),          r2path(t),              &
                        s3path(t),          i3path(t),          r3path(t),              &
                        aggc1path(t),       aggh1path(t),       aggc2path(t),           &
                        aggh2path(t),       aggc3path(t),       aggh3path(t),           &
                        aggc_hjpath(1,1,t), aggc_hjpath(1,2,t), aggc_hjpath(1,3,t),     &
                        aggc_hjpath(2,1,t), aggc_hjpath(2,2,t), aggc_hjpath(2,3,t),     &
                        aggc_hjpath(3,1,t), aggc_hjpath(3,2,t), aggc_hjpath(3,3,t),     &
                        aggoh_hjpath(1,1,t),aggoh_hjpath(1,2,t),aggoh_hjpath(1,3,t),    &
                        aggoh_hjpath(2,1,t),aggoh_hjpath(2,2,t),aggoh_hjpath(2,3,t),    &
                        aggoh_hjpath(3,1,t),aggoh_hjpath(3,2,t),aggoh_hjpath(3,3,t),    &
                        aggh_hjpath(1,1,t), aggh_hjpath(1,2,t), aggh_hjpath(1,3,t),     &
                        aggh_hjpath(2,1,t), aggh_hjpath(2,2,t), aggh_hjpath(2,3,t),     &
                        aggh_hjpath(3,1,t), aggh_hjpath(3,2,t), aggh_hjpath(3,3,t),     &
                        dpath(t), aggtaxpath(t), aggsuipath(t), pvdef,                  &
                        aggoc_hjpath(1,1,t), aggoc_hjpath(1,2,t), aggoc_hjpath(1,3,t),  &
                        aggoc_hjpath(2,1,t), aggoc_hjpath(2,2,t), aggoc_hjpath(2,3,t),  &
                        aggoc_hjpath(3,1,t), aggoc_hjpath(3,2,t), aggoc_hjpath(3,3,t),  &
                        mu_takeuppath(t), betaepath(t),                                 &
                        d1path(t),           d2path(t),         d3path(t) 
    END DO  
    CLOSE (UNIT=11)

END SUBROUTINE TRANSITION

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

SUBROUTINE POST_PROCESS

    INTEGER     :: idx_j, idx_z, idx_k, idx_zprime
    INTEGER     :: i             ! dummy index
    INTEGER     :: iter
    
	REAL(DP)    :: kprime		    ! next days saving
    REAL(DP)    :: cons, uu, hours
    REAL(DP)    :: evp, dum
    
    PRINT *, "prepare welfare calculation ..."
    IF (welfguess.EQ.1) THEN
        OPEN(25,file='vhati.out',form='formatted')
        DO idx_k = 1,N_k
        DO i = 1,N_ceq
            READ(25,497) dum,dum,vhati(1,:,idx_k,i),vhati(2,:,idx_k,i),vhati(3,:,idx_k,i)
        END DO
        END DO
        CLOSE(25) 
497     FORMAT(2f18.10,7f18.10,7f18.10,7f18.10)
    ELSE
        DO i = 1,N_ceq
            vhati(:,:,:,i) = vss
        END DO
    END IF
    
    DO i = 1,N_ceq
        iterloop: DO iter = 1,mxiter
            vhati_old(:,:,:,i) = vhati(:,:,:,i)
            DO idx_j = jmax,1,-1
                !!$OMP PARALLEL DEFAULT(SHARED) PRIVATE(cons,hours,kprime,uu,evp)
                !!$OMP DO COLLAPSE(2) 
                DO idx_z = 1,N_z
                DO idx_k = 1,N_k
            
                    cons = (onedp+ceq_grid(i))*css(idx_j,idx_z,idx_k)
                    kprime = kss(idx_j,idx_z,idx_k)
                
                    IF (idx_j<jret) THEN
                
                        hours = lss(idx_j,idx_z,idx_k)
                        CALL utility(cons,gam,hours,uu)
                    
                        IF (idx_j.EQ.jret-1) THEN
                
                            evp = uu + beta_adj*psi_j(idx_j)*interp1(k_grid,vhati(idx_j+1,1,:,i),N_k,kprime)
                            DO idx_zprime = Nzmin(idx_z), Nzmax(idx_z)
                                evp = evp + beta_adj*(onedp-psi_j(idx_j))*ptran(idx_z,idx_zprime)       &
                                    *interp1(k_grid,vhati_old(idx_j,idx_zprime,:,i),N_k,kprime)
                            END DO
                
                        ELSE ! (idx_j < jret-1)
                
                            evp = uu
                            DO idx_zprime = Nzmin(idx_z), Nzmax(idx_z)
                                evp = evp + beta_adj*ptran(idx_z,idx_zprime)                            &
                                    *interp1(k_grid,psi_j(idx_j)*vhati(idx_j+1,idx_zprime,:,i)          &
                                    + (onedp-psi_j(idx_j))*vhati_old(idx_j,idx_zprime,:,i),N_k,kprime)
                            END DO
                        
                        END IF
                    ELSE ! (j.GE.jret)
                        CALL utility_ret(cons,gam,uu)
                        evp = uu + beta_adj*(onedp-psi_j(idx_j))*interp1(k_grid,vhati_old(idx_j,1,:,i),N_k,kprime) 
                    
                    ENDIF
                    vhati(idx_j,idx_z,idx_k,i) = evp
                END DO
                END DO
                !!$OMP END DO
                !!$OMP END PARALLEL
            END DO !idx_j
                
            IF (MAXVAL(DABS(vhati_old(:,:,:,i)-vhati(:,:,:,i))) < 1.0D-08) THEN
                EXIT iterloop
            ENDIF
            IF (multpl(iter,mxiter) .EQ. 1) THEN
                IF (multpl(i,10) .EQ. 1) THEN
                    PRINT*, iter, i, MAXVAL(DABS(vhati_old(:,:,:,i)-vhati(:,:,:,i)))
                ENDIF
            ENDIF
        END DO iterloop
    END DO
    
13  PRINT *, "computing welfare"

    IF (sswrite.EQ.1) THEN
        OPEN(25,file=TRIM(path)//'/vhati.out',form='formatted')
        DO idx_k = 1,N_k
        DO i = 1,N_ceq
            WRITE(25,497) k_grid(idx_k),ceq_grid(i),vhati(1,:,idx_k,i),vhati(2,:,idx_k,i),vhati(3,:,idx_k,i)
        END DO
        END DO
        CLOSE(25) 
    ENDIF
    

    DO idx_j=1,jmax
    DO idx_z=1,N_z
    DO idx_k=1,N_k
        welf(idx_j,idx_z,idx_k) = zbrent(WELFARE,ceqmin,ceqmax,1.0D-14)
    END DO
    END DO
    END DO

    OPEN(25,file=TRIM(path)//'/welf.out',form='formatted')
    DO idx_k = 1,N_k
        WRITE(25,499) k_grid(idx_k),welf(1,:,idx_k),welf(2,:,idx_k),welf(3,:,idx_k)
499     FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
    END DO
    CLOSE(25)  
    
    CONTAINS

    DOUBLE PRECISION FUNCTION WELFARE(ceq)
    !! Solves the model
        
    ! LOCAL VARIABLES USED WITHIN LOOPS
    REAL(DP), INTENT(IN)    :: ceq
    REAL(DP)                :: vc
    
    vc = INTERP1(ceq_grid,vhati(idx_j,idx_z,idx_k,:),N_ceq,ceq)
    !FUNCTION interp1(X,Y,siz,xi)
    WELFARE = vc - vhat_s(idx_j,idx_z,idx_k)
    
    END FUNCTION WELFARE
    
END SUBROUTINE POST_PROCESS
   
END MODULE AUXILIARY