! test_aux.f90 -- main program. Run without arguments for an
! interactive menu (11 = grids, 21 = steady state, 411 = transition);
! run with an integer argument to execute one configuration of the
! cluster policy grid (see src/cluster/jobfile).

PROGRAM TEST_AUX

    !!!!!!!!!!!!!!!!!!!!!!!!
    !!   USE STATEMENTS   !!
    !!!!!!!!!!!!!!!!!!!!!!!!    
    USE PARAMETERS
    USE GLOBAL
    USE AUXILIARY
    !USE OMP_LIB
    
    IMPLICIT NONE  
    
    !! ARGUMENTS FROM COMMAND LINE
    CHARACTER (LEN = 80)    :: cmdline      ! generic command line  
    
    !! ARGUMENTS FOR PROGRAM
    INTEGER                 :: num          ! case selected on menu
    INTEGER                 :: N_tot
    INTEGER                 :: argtot       ! argument count
    INTEGER                 :: ios
    CHARACTER(LEN=12) string, FORMAT
	
    !! MISC ARGUMENTS
    INTEGER             	:: i, iter 
	INTEGER, DIMENSION(4)   :: idx_vec, vec_size

    
    !! MAIN PROGRAM    
    path = '.'
            
    ! Get command count
    argtot      = COMMAND_ARGUMENT_COUNT()
    
    ! Loop across input pairs and store them
    ifcmdargs: IF (argtot .EQ. 0) THEN
        
        doMenu: DO
            WRITE (*,*) ' --***---***-- '
            WRITE (*,*) ' Enter the number of the routine to run'
            WRITE (*,*) ' ** read data, construct grids and shocks **'  
            WRITE (*,*) '  (11)  baseline' 
            WRITE (*,*) ' ** compute steady state **' 
            WRITE (*,*) '  (21)  baseline'  
            WRITE (*,*) '  (22)  Test wfh'  
            WRITE (*,*) ' ** compute transition **'  
            WRITE (*,*) '  (411) baseline'
            WRITE (*,*) '  (500) test cluster run'
            WRITE (*,*) ' ** other **'
            WRITE (*,*) '  (00)  QUIT'
            READ (*,*) num    

            IF (num>10) THEN
                IF (num<20) THEN
                    SELECT CASE (num)
                    CASE(11)
                        CALL read_data(1)
                        CALL create_grids(1)
                        CALL create_zshocks(1)
                    CASE DEFAULT
                        EXIT
                    END SELECT
                ELSE
                    SELECT CASE (num)
                    CASE(21)
                        CALL computeSS(1)
                    CASE(22)
                        CALL computeSS(2)
                    CASE(411)
                        CALL computetransition(1,1)
                    CASE(500)
                        N_tot       = N_sui*N_curv*N_hloc*N_suidur 
                        vec_size    = (/N_sui,N_curv,N_hloc,N_suidur/)
                        iter = 0
                        PRINT *, 'N_tot:', N_tot
                        DO i = 1,N_tot
		                    idx_vec     = ind2sub(i,vec_size)
        
                            IF (idx_vec(1)>1) THEN
                                sui_cbar    = 0.2_dp + 0.05_dp*REAL(idx_vec(1)-2)
                                policy  = 1
                            ELSE
                                sui_cbar    = zerodp
                                policy  = 0
                            ENDIF
        
                            hcut = zerodp
        
                            sui_curv    = 2.0_dp**(idx_vec(2)-1)
                            loc_curv    = sui_curv
        
                            IF (idx_vec(3)>1) THEN
                                hloc        = 0.045_dp*REAL(idx_vec(3))    ! 0.09, 0.135, 0.18
                                lockdown    = 1
                            ELSE
                                hloc        = lmax
                                lockdown    = 0
                            END IF
        
                            suidur      = 120.0_dp + 30.0_dp*REAL(idx_vec(4)-1)  ! 120, 150, ... , 420
                            locdur      = suidur       	
        
                            IF ((policy.EQ.1).OR.(lockdown.EQ.1)) THEN
                                iter = iter + 1
                            ELSEIF ((idx_vec(2).EQ.1).AND.(idx_vec(4).EQ.1)) THEN
                                iter = iter + 1
                            END IF
                        ENDDO
                        PRINT *, 'total cases:', iter
                    CASE DEFAULT
                        EXIT
                    END SELECT
                END IF ! (num<20)
            ELSE
                EXIT
            END IF ! (num>10)
        END DO doMenu
        
    ELSE
        
        CALL GET_COMMAND_ARGUMENT(NUMBER = 1, VALUE = cmdline)
        READ (cmdline,*) runmode
        WRITE (*,*) ' Program run mode is: ', runmode	
        N_tot       = N_sui*N_curv*N_hloc*N_suidur 
        
        vec_size    = (/N_sui,N_curv,N_hloc,N_suidur/)
        tranguess   = 1 ! 1: use transition guess
        itersave    = 0 ! 0: DO NOT SAVE on server
        IF (runmode.LE.N_tot) THEN
		    idx_vec     = ind2sub(runmode,vec_size)
        
            ! Subsidy amount
            IF (idx_vec(1)>1) THEN
                sui_cbar    = 0.3_dp + 0.05_dp*REAL(idx_vec(1)-2)   ! 0.3, 0.35, ..., 1.25, 1.3
                IF (highsubsidy.EQ.1) sui_cbar  = sui_cbar + 1.05   ! 1.35, 1.4, ..., 2.3, 2.35
                policy  = 1
                policy2 = 0 
            ELSE    ! idx_vec(1)==1
                sui_cbar    = zerodp
                policy  = 0
                policy2 = 0 
            ENDIF
            hcut = zerodp
        
            ! Phase-out curvature
            sui_curv    = 2.0_dp**(idx_vec(2)-1)
            loc_curv    = sui_curv
        
            ! Lockdown intensity
            IF (idx_vec(3)>1) THEN
                hloc        = 0.045_dp*REAL(idx_vec(3))    ! 0.09, 0.135, 0.18
                lockdown    = 1
            ELSE    ! idx_vec(3)==1
                hloc        = lmax
                lockdown    = 0
            END IF
        
            ! Subsidy duration
            suidur      = 120.0_dp + 30.0_dp*REAL(idx_vec(4)-1)  ! 120, 150, ... , 420
            locdur      = suidur
        
            PRINT *, 'sui_cbar:',   sui_cbar
            PRINT *, 'hcut:',       hcut
            PRINT *, 'hloc:',       hloc
            PRINT *, 'suidur:',     suidur
            PRINT *, 'phase-out curvature:',    sui_curv
    
		    DATA FORMAT /'(I6.6)'/
		    WRITE(string,FORMAT) runmode
	        path = 'output/config'//TRIM(string)
		    ios = SYSTEM( 'mkdir -p '//path )
		    PRINT *, 'path:',path   
        
            IF ((policy.EQ.1).OR.(lockdown.EQ.1)) THEN
                CALL computetransition(1,1)  ! Mitigation policy
            ELSEIF ((idx_vec(2).EQ.1).AND.(idx_vec(4).EQ.1)) THEN
                CALL computetransition(1,1)  ! No mitigation policy
            ELSE
                PRINT *, 'Redundant case: Exiting...'
            END IF
        ELSE
            
		    DATA FORMAT /'(I6.6)'/
		    WRITE(string,FORMAT) runmode
	        path = 'output/config'//TRIM(string)
            
            IF (runmode.LE.N_tot+3) THEN
                ! US policy
                ! path = 'output/us_policy'
                policy      = 1 ! 1: run COVID-19 mitigation policy (initial response)
                policy2     = 1 ! 1: second round (CAA + ARPA)
                    sui_cbar    = 0.70_dp
                    hcut        = 0.00_dp
	                suidur 		= 120.0_dp
                    sui_curv    = 3.0_dp
                lockdown    = 1 ! 1: lockdown
                    hloc        = 0.132_dp
	                locdur 		= 042.0_dp
                    loc_curv    = 2.0_dp
                taxtype     = 0 ! 0: no taxes 1: tauc only | 2: taus only | 3: Pigouvian
                IF (runmode.EQ.N_tot+2) THEN
                    ! no lockdown
                    path = 'output/us_subsidy'
                    lockdown    = 0
                    hloc        = lmax
                ELSEIF (runmode.EQ.N_tot+3) THEN
                    ! no subsidy
                    path = 'output/us_lockdown'
                    policy      = 0
                    policy2     = 0 ! 1: second round (CAA + ARPA)
                    sui_cbar    = zerodp
                ENDIF
                runmode = 0
            ELSEIF (runmode.LE.N_tot+6) THEN
                ! Breakdown of optimal policy
                path = 'output/optimum'
                policy      = 1 ! 1: run COVID-19 mitigation policy (initial response)
                policy2     = 0 ! ! 1: second round (CAA + ARPA)
                IF (taxtype.EQ.1) THEN ! 0: no taxes | 1: tauc only | 2: taus only | 3: Pigouvian | 4: tax only
                    tauc_init   = 0.055310182_dp
                    taus_init   = zerodp
                    sui_cbar    = 0.80_dp
                    hcut        = 0.00_dp
	                suidur 		= 300.0_dp
                    sui_curv    = 8.0_dp
                ELSEIF (taxtype.EQ.2) THEN
                    tauc_init   = zerodp
                    taus_init   = 0.0738_dp
                    sui_cbar    = 0.65_dp
                    hcut        = 0.00_dp
	                suidur 		= 390.0_dp
                    sui_curv    = 8.0_dp
                ENDIF  
                lockdown    = 0 ! 1: lockdown
                    hloc        = lmax
                    locdur      = suidur
                    loc_curv    = sui_curv   
                    
                IF (runmode.EQ.N_tot+5) THEN
                    ! subsidy only
                    path = 'output/optimum_sub'
                    taxtype     = 0 ! 0: no taxes | 1: tauc only | 2: taus only | 3: both | 4: tax only
                    tauc_init   = zerodp
                    taus_init   = zerodp
                ELSEIF (runmode.EQ.N_tot+6) THEN
                    ! tax only
                    path = 'output/optimum_tax'
                    taxtype     = 4 ! 0: no taxes | 1: tauc only | 2: taus only | 3: both | 4: tax only
                    sui_cbar    = zerodp
                ENDIF
                runmode = 0
            ELSE
                ! Breakdown of output maximizing policy
                path = 'output/maximizing'
                policy      = 1 ! 1: run COVID-19 mitigation policy (initial response)
                policy2     = 0 ! 1: second round (CAA + ARPA)
                IF (taxtype.EQ.1) THEN ! 0: no taxes | 1: tauc only | 2: taus only | 3: both | 4: tax only
                    sui_cbar    = 0.60_dp
                    hcut        = 0.00_dp
	                suidur 		= 300.0_dp
                    sui_curv    = 4.0_dp
                    tauc_init   = 0.021984795_dp
                    taus_init   = zerodp
                ELSEIF (taxtype.EQ.2) THEN ! 0: no taxes | 1: tauc only | 2: taus only | 3: Pigouvian | 4: tax only
                    sui_cbar    = 0.40_dp
                    hcut        = 0.00_dp
	                suidur 		= 390.0_dp
                    sui_curv    = 4.0_dp
                    tauc_init   = zerodp
                    taus_init   = 0.00458_dp
                ENDIF
                    
                lockdown    = 0 ! 1: lockdown
                    hloc        = lmax
                    locdur      = suidur
                    loc_curv    = sui_curv   
                
                IF (runmode.EQ.N_tot+8) THEN
                    ! subsidy only
                    path = 'output/max_sub'
                    taxtype     = 0 ! 0: no taxes | 1: tauc only | 2: taus only | 3: Pigouvian | 4: tax only
                    tauc_init   = zerodp
                    taus_init   = zerodp
                ELSEIF (runmode.EQ.N_tot+9) THEN
                    ! tax only
                    path = 'output/max_tax'
                    taxtype     = 4 ! 0: no taxes | 1: tauc only | 2: taus only | 3: Pigouvian | 4: tax only
                    sui_cbar    = zerodp
                ENDIF
                runmode = 0
            ENDIF 
            PRINT *, 'sui_cbar:',   sui_cbar
            PRINT *, 'hcut:',       hcut
            PRINT *, 'hloc:',       hloc
            PRINT *, 'suidur:',     suidur
            PRINT *, 'phase-out curvature:',    sui_curv
 
		    ios = SYSTEM( 'mkdir -p '//path )
		    PRINT *, 'path:',path  
            CALL computetransition(1,1)
        ENDIF
        
    END IF ifcmdargs    
        
    CONTAINS
    
SUBROUTINE read_data(showdetails)
    
    INTEGER             :: j
    INTEGER, INTENT(IN) :: showdetails
    REAL (DP)           :: dum
    
    !! begin
    WRITE (*,*) 'read_data...'
    
    !psi_j = days/(20.0_dp*365.0_dp)
    psi_j(1:2) = days/(20.0_dp*365.0_dp)
    psi_j(3) = days/(10.0_dp*365.0_dp)
    
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'aging probability'
        WRITE (*,*) psi_j
    END IF  
    
    !mu_j = onedp/REAL(jmax)
    mu_j = (/0.4_dp,0.4_dp,0.2_dp/)
    
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'age distribution'
        WRITE (*,*) mu_j
    END IF  
    
    eta_j(1,1) = onedp
    eta_j(2,1) = 1.35_dp
    eta_j(jret:jmax,:) = zerodp
    eta_j(:,2) = (onedp-sick_eta_loss)*eta_j(:,1)
    eta_j(:,3) = eta_j(:,1)
    
    IF (showdetails.NE.0) THEN    
        WRITE (*,*) 'age-income profile (susceptible)'
        WRITE (*,*) eta_j(:,1)
        WRITE (*,*) 'age-income profile (infected)'
        WRITE (*,*) eta_j(:,2)
        WRITE (*,*) 'age-income profile (recovered)'
        WRITE (*,*) eta_j(:,3)
    END IF
    
    ! initial SIR distribution
    muh(2) = inf_init
    muh(3) = zerodp
    muh(1) = onedp - muh(2) - muh(3)
    
    IF (showdetails.NE.0) THEN  
        PRINT *, 'Initial SIR distribution', muh
    ENDIF

END SUBROUTINE read_data  

SUBROUTINE create_grids(showdetails)
    
    INTEGER, INTENT(IN) :: showdetails
    INTEGER             :: ind
    
    !! begin
    WRITE (*,*) 'create_grids...'
    
    !!!!!!!!!!!!!!!!!!!
    !! create k_grid !!
    !!!!!!!!!!!!!!!!!!!
    CALL creategrid(k_grid,N_k,1,kmin,kmin,kmax,grdpwr_lo)
    ! creategrid(grid, siz, n_med, grid_min, grid_pivot, grid_max, curvature)

    !!!!!!!!!!!!!!!!!!
    !! create kfine !!
    !!!!!!!!!!!!!!!!!!
    CALL creategrid(kfine,N_kfine,1,kmin,kmin,kfinemax,grdpwr_lo)
    ! creategrid(grid, siz, n_med, grid_min, grid_pivot, grid_max, curvature)
    
    !!!!!!!!!!!!!!!!!!!
    !! create s_grid !!
    !!!!!!!!!!!!!!!!!!!
    CALL creategrid(s_grid,N_s,FLOOR(DBLE(N_s)*0.96_dp),smin,0.9965_dp,smax,1.75_dp) 
    ! creategrid(grid, siz, n_med, grid_min, grid_pivot, grid_max, curvature)
    
    !!!!!!!!!!!!!!!!!!!
    !! create l_grid !!
    !!!!!!!!!!!!!!!!!!!
    ind = FLOOR(DBLE(N_L)*0.80_dp)
    CALL creategrid(l_grid(1:ind),ind,1,lmin,lmin,lmid,onedp) 
    CALL creategrid(l_grid(ind:N_L),N_L-ind+1,1,lmid,lmid,lmax,1.5_dp) 
    ! creategrid(grid, siz, n_med, grid_min, grid_pivot, grid_max, curvature)
    
    !!!!!!!!!!!!!!!!!!!!!
    !! create cfh_grid !!
    !!!!!!!!!!!!!!!!!!!!!
    IF (N_cfh>1) THEN
        CALL creategrid(cfh_grid,N_cfh,1,cfhmin,cfhmin,cfhmax,1.0_dp) 
    ELSE
        cfh_grid = zerodp
    ENDIF 
    
    !!!!!!!!!!!!!!!!!!!!!
    !! create ceq_grid !!
    !!!!!!!!!!!!!!!!!!!!!
    CALL creategrid(ceq_grid,N_ceq,FLOOR(DBLE(N_ceq)*0.75_dp),ceqmin,zerodp,ceqmax,1.2_dp) 
    ! creategrid(grid, siz, n_med, grid_min, grid_pivot, grid_max, curvature)
    
    
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'labour grid'    
        WRITE (*,*) l_grid
        
        WRITE (*,*) 'capital grid'    
        WRITE (*,*) k_grid
        
        WRITE (*,*) 'savings rate grid'
        WRITE (*,*) s_grid
        
        WRITE (*,*) 'cfh grid'
        WRITE (*,*) cfh_grid
    END IF
    
END SUBROUTINE create_grids    

SUBROUTINE create_zshocks(showdetails)

    INTEGER, INTENT(IN)             ::  showdetails
    INTEGER                         ::  iter, i, idx_z
    REAL(DP), DIMENSION(N_z)        ::  Erdold    
    REAL(DP)                        ::  agg_dlinc
    
    !! begin
    WRITE (*,*) 'create_zshocks...'

    ! CONVERT ANNUAL TO RELEVANT FREQUENCY
    rho_z = rho_ann**(days/365.0_dp)
    sigma_eps = SQRT(sigma_ann**2*(onedp-rho_z**2)/(onedp-rho_ann**2))
    
    CALL rouwenhorst
    
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'raw income shocks'
        WRITE (*,*) z_grid
    END IF
    
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'transition matrix'
        DO idx_z=1,N_z
            PRINT *, 'idx=', idx_z, ':', ptran(idx_z,:)
        END DO
    END IF
    
    ! Solve for ergodic distribution (Erd*ptran = Erd)    
    Erdold = onedp/DBLE(N_z)
    iter = 0
    
    ERD_LOOP: DO
        iter = iter + 1
        Erd = MATMUL(Erdold,ptran)
        IF (MAXVAL(ABS(Erd-Erdold))/DBLE(N_z) < eps) THEN
            EXIT ERD_LOOP
        END IF
        Erdold = Erd
    END DO ERD_LOOP
    
    IF (showdetails.NE.0) THEN
10      FORMAT (1X, ' number of iterations:', I10)
        WRITE (*,10) iter
    END IF
    
    IF (showdetails.NE.0) THEN
        PRINT *, 'Ergodic distribution'
        PRINT *, Erd
        PRINT *, 'Erd*ptran'
        PRINT *, MATMUL(Erd,ptran)
        PRINT *, 'sum of ergodic z distribution: ', SUM(Erd)
    END IF
    
    ! Normalize z grid
    z_grid = z_grid - LOG(SUM(EXP(z_grid)*Erd))
    z_grid = EXP(z_grid)
    IF (showdetails.NE.0) THEN
        WRITE (*,*) 'normalized income shocks'
        WRITE (*,*) z_grid
    END IF
    
    IF (highpers.EQ.1) THEN
        DO idx_z = 1,N_z
            Nzmin(idx_z) = MAX(1,idx_z - 3)
            Nzmax(idx_z) = MIN(N_z,idx_z + 3)
        END DO
    ELSE
        Nzmin = 1
        Nzmax = N_z
    END IF
    IF (showdetails.NE.0) THEN
        PRINT *, 'Nzmin', Nzmin
        PRINT *, 'Nzmax', Nzmax
    END IF
    
END SUBROUTINE create_zshocks 

SUBROUTINE computeSS(ss_type)
! ss_type
! 1. baseline:'

    ! local vars
    REAL(DP)            ::  capsule_time_start, capsule_time_stop, dum, zstep
    REAL(DP)            ::  initknew, taunew, phinew, ssnew, betanew,   &
                            tmxdif, kmxdif, pmxdif, smxdif, bmxdif
    INTEGER, INTENT(IN) ::  ss_type
    INTEGER             ::  i, ttt, nsing, idx_j, idx_z
    
    REAL*8, PARAMETER   :: h = 0.9D+00
    REAL*8, PARAMETER   :: updatez = 0.5_dp
    REAL*8, DIMENSION(ncalib)    ::fv
    
    ! Set some parameters (sensitivity)
    IF (highgam.EQ.1) THEN
        gam         = 0.61_dp
        sick_gam    = gam
        beta_ann    = 0.9741230772_dp
        phi         = 20.3090670512_dp
        ss          = 1.0155625963_dp
        tau         = 0.0749468387_dp
        initk       = 2302.4415869062_dp
        rhat        = 0.0028855107_dp
        rhat1       = rhat + 0.000001_dp
        rhat2       = rhat - 0.000001_dp
        IF (vsltype.EQ.2) THEN  
            !vsl     = 6226.0_dp ! $7.4m / $30989 * 365 / 14 (2006) bequestmode
            !ubar    = 25.1053893682628_dp     
            !cbar    = 3.19715245_dp   
        ELSE    
            PRINT *, "Configuration not allowed!"
            PRINT *, CHAR(7)
            PAUSE
        ENDIF
    ELSEIF (bequestmode.EQ.2) THEN
        IF (vsltype.EQ.1) THEN  
            vsl     = 4096.0_dp ! 6.9 m / $43919 * 365 / 14 (2019)
            ubar    = 16.7137997278636_dp
            cbar    = 3.21270104741336_dp    
            tau     = 0.0749629059_dp
            phi     = 19.7720497928_dp
            ss      = 1.0213635822_dp
            beta_ann    = 0.9760155018_dp
            rhat    = 0.0028813266_dp
        ELSE    
            PRINT *, "Configuration not allowed!"
            PRINT *, CHAR(7)
            PAUSE
        ENDIF
    ELSE
        IF (vsltype.EQ.1) THEN ! Low VSL
            vsl     = 4096.0_dp ! 6.9 m / $43919 * 365 / 14 (2019)
            ubar    = 17.3064_dp
            cbar    = 3.20355454_dp  
        ELSEIF (vsltype.EQ.2) THEN ! VSL based on representative agent
            vsl     = 6226.0_dp         ! $7.4m / $30989 * 365 / 14 (2006)
            ubar    = 25.90663_dp
            cbar    = 3.20355454_dp  
        ELSEIF (vsltype.EQ.3) THEN  ! High VSL
            vsl     = 4096.0_dp ! 6.9 m / $43919 * 365 / 14 (2019)
            ubar    = 17.3064_dp
            cbar    = 3.20355454_dp    
        ELSE    
            PRINT *, "VSL configuration not allowed!"
            PRINT *, CHAR(7)
            PAUSE
        ENDIF
    ENDIF    
    IF (highrepnum.EQ.1) THEN
        repnum = 2.5_dp
    ENDIF
    
    CALL read_data(0)
    CALL create_grids(0)
    CALL create_zshocks(0)
    
    delta_adj   = onedp - (onedp-delta_ann)**(days/365.0_dp)
    beta_adj    = beta_ann**(days/365.0_dp)         ! adjusted discount factor
    
    IF (N_cfh>1) THEN   ! Utility function modifications with Cobb-Douglas aggregation
        gamgam  = gam**gam*(onedp-gam)**(onedp-gam)
        ! c_i^gamma*c_o^(1-gamma) = (c*gamma)^gamma*(c*(1-gamma))^(1-gamma)
        ! = c*gamma^gamma*(1-gamma)^(1-gamma)
    ELSE
        gamgam  = onedp
        gam = zerodp
        sick_gam = gam
        phi = phi/gamgam**(onedp-sigma)
        sick_phi = sick_phi/gamgam**(onedp-sigma)
        ubar = ubar/gamgam**(onedp-sigma)
    ENDIF
    
    PRINT *, 'gam', gam
    PRINT *, 'sick_gam', sick_gam
    PRINT *, 'gamgam', gamgam
    PRINT *, 'phi', phi
    PRINT *, 'ubar', ubar
    
    IF (readvguess .EQ. 1) THEN 
        PRINT*, 'Reading in prior value function'
444     FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
        OPEN(3,FILE='vss.out',form='formatted')
        DO i = 1,N_k
            READ(3,444) dum,vss(1,:,i),vss(2,:,i),vss(3,:,i)
        END DO
        CLOSE(3)
    ELSE
        vss     =   zerodp      ! initial value
    END IF
    vssh = vss
    vssl = vss
    
    !! begin
    WRITE (*,*) 'Compute steady state ... '
    
    !! Initialize policy and value functions (May want to read in solution)
    muss    =   zerodp		! steady state distribution 
    sss     =   s_grid(1)   ! savings rate
    lss     =   zerodp      ! labor supply
        
    !! Value of life !!
    IF (ubarcal.EQ.1) THEN
        PRINT *, 'ubar guess:', ubar
        
        dum = SOLVESS(rhat)
        
        ubar = zbrent(SOLVE_UBAR,5.0_dp,50.0_dp,1.0D-5)
        
        dum =  SOLVE_UBAR(ubar)
        PRINT *, "vsl(1,1,1):", v_vsl(1,1,1)/aggcinit
        PRINT *, "vsl(2,1,1):", v_vsl(2,1,1)/aggcinit
        PRINT *, "vsl(1,N_z,1):", v_vsl(1,N_z,1)/aggcinit
        PRINT *, "vsl(2,N_z,1):", v_vsl(2,N_z,1)/aggcinit
        PRINT *, "vsl(3,1,1):", v_vsl(3,1,1)/aggcinit
        PRINT *, "vsl(1,1,N_k):", v_vsl(1,1,N_k)/aggcinit
        PRINT *, "vsl(2,1,N_k):", v_vsl(2,1,N_k)/aggcinit
        PRINT *, "vsl(1,N_z,N_k):", v_vsl(1,N_z,N_k)/aggcinit
        PRINT *, "vsl(2,N_z,N_k):", v_vsl(2,N_z,N_k)/aggcinit
        PRINT *, "vsl(3,1,N_k):", v_vsl(3,1,N_k)/aggcinit
        
499 FORMAT(f18.10,7f18.10,7f18.10,7f18.10)
        OPEN(25,file=TRIM(path)//'/v_vsl.out',form='formatted')
        DO i = 1,N_k
            WRITE(25,499) k_grid(i),v_vsl(1,:,i),v_vsl(2,:,i),v_vsl(3,:,i)
        END DO
        CLOSE(25)
        
        PRINT *, "new ubar:", ubar
        
    ENDIF
    
    capsule_time_start    = DCLOCK() 
940 FORMAT(11F18.10)
    IF (calibed.EQ.0) THEN
    
        IF (readcalguess.EQ.1) THEN
            OPEN(25,file='calib_sol.out',form='formatted')
            READ(25,940) tau,taunew,initk,initknew,phi,phinew,ss,ssnew,beta_ann,betanew,rss
            CLOSE(25)
            tau         = (onedp-updatecal)*tau         +updatecal*taunew
            initk       = (onedp-updatecal)*initk       +updatecal*initknew
            phi         = (onedp-updatecal)*phi         +updatecal*phinew
            ss          = (onedp-updatecal)*ss          +updatecal*ssnew
            beta_ann    = (onedp-updatecal)*beta_ann    +updatecal*betanew
            !rhat = rss
            !dum = SOLVESS(rhat)
            !rhat1 = rhat
            !IF (rnew > rhat) THEN
            !    rhat2 = MIN(rnew, rhat + 0.000005_dp)
            !ELSE
            !    rhat2 = MAX(rnew - 0.000001_dp, rhat - 0.000005_dp)
            !ENDIF
        ENDIF
        
        DO ttt = 1,mxiter

            PRINT *, ttt, 'of' , mxiter
            WRITE(6,571) tau, initk, phi, ss, beta_ann
571         FORMAT(f12.8,2f12.4,2f12.8)

            !IF (ttt>1) THEN
            rhat = rss
            rhat1 = rhat
            dum = SOLVESS(rhat)
            IF (rnew > rhat) THEN
                rhat2 = MIN(rnew, rhat + 0.000001_dp)
            ELSE
                rhat2 = MAX(rnew - 0.0000002_dp, rhat - 0.000002_dp)
            ENDIF            
            !END IF
            CALL getf2(fv)
            
            taunew = fv(1)
            initknew = fv(2)
            phinew = fv(3)
            ssnew = fv(4)
            betanew = fv(5)
            
            tmxdif = DABS(taunew-tau)
            kmxdif = DABS(initknew-initk)/initk
            pmxdif = DABS(phinew-phi)/phi
            smxdif = DABS(ssnew-ss)/ss
            bmxdif = DABS(betanew-beta_ann)/beta_ann
        
            OPEN(25,file='calib_sol.out',form='formatted')
            WRITE(25,940) tau,taunew,initk,initknew,phi,phinew,ss,ssnew,beta_ann,betanew,rss
            CLOSE(25)
            
            OPEN(25,file='vss_old.out',form='formatted')
            DO i = 1,N_k
                WRITE(25,499) k_grid(i),vss(1,:,i),vss(2,:,i),vss(3,:,i)
            END DO
            CLOSE(25)

            tau         = (onedp-updatecal)*tau         +updatecal*taunew
            initk       = (onedp-updatecal)*initk       +updatecal*initknew
            phi         = (onedp-updatecal)*phi         +updatecal*phinew
            ss          = (onedp-updatecal)*ss          +updatecal*ssnew
            beta_ann    = (onedp-updatecal)*beta_ann    +updatecal*betanew

391         FORMAT(i5,5f9.5)
            IF (tmxdif+kmxdif+pmxdif+smxdif+bmxdif .lt. crit2) GOTO 12
            
            IF(multpl(ttt,1) == 1) THEN 
                WRITE(6,391) ttt,tmxdif,kmxdif,pmxdif,smxdif,bmxdif
                PRINT *, "--- Calibration targets ---"
                PRINT *, aggk/aggy*days/365.0_dp, aggh/SUM(muss(1:2,:,:)), ss/(wss*aggl/SUM(muss(1:2,:,:)))
            ENDIF
        
            END DO

        12  WRITE(6,391) ttt,tmxdif,kmxdif,pmxdif,smxdif,bmxdif
        
    ELSE
        dum = SOLVESS(rhat)
        rss = rhat
    END IF
    
2   PRINT*, 'done with initial steady state'
    PRINT *, CHAR(7)
    capsule_time_stop     = DCLOCK()        

100 FORMAT ('steady state computation took: ', F12.3, ' seconds')         
    
    WRITE (*,100) capsule_time_stop - capsule_time_start
    
    wss = (onedp-alpha)*(rss/alpha)**(alpha/(alpha-onedp))
    cbar = aggcinit
    
    PRINT *, "K:", aggkinit
    PRINT *, "C:", aggcinit
    PRINT *, "L:", agglinit
    PRINT *, "Y:", aggyinit
    PRINT *, "Yd", aggcinit + aggkinit*delta_adj
    PRINT *, "Y/worker:", aggy/SUM(muss(1:2,:,:))
    PRINT *, "w:", wss
    PRINT *, "r:", rss
    
    PRINT *, "earnings per worker:", wss*agglinit/SUM(muss(1:2,:,:))
    PRINT *, "median K/Y:", medk/aggyinit*days/365.0_dp
    PRINT *, "tax surplus:", wss*agglinit*tau - ss*SUM(muss(jmax,:,:)) 
    IF (bequestmode.EQ.1) THEN
        PRINT *, "bequest surplus:",  bequest-(initk*hikfrac*psi_j(jmax)*SUM(muss(jmax,:,:)))
        PRINT *, "initial endowment/annual per capita cons.", initk/(aggcinit*365.0_dp/days) 
    ENDIF
    PRINT *, "--- Calibration targets ---"
    PRINT *, "K/Y:", aggkinit/aggyinit*days/365.0_dp
    PRINT *, "hours per worker:", agghinit/SUM(muss(1:2,:,:))
    PRINT *, "ss/(earnings per worker):", ss/(wss*agglinit/SUM(muss(1:2,:,:)))
    !dum =  SOLVE_UBAR(ubar)
    !PRINT *, "Average VSL/annual per capita cons.:", SUM(v_vsl*muss)/(aggcinit*365.0_dp/days)    
    
    IF (ubarcal.EQ.1) THEN
        PRINT *, "Verifying ubar guess"
        PRINT *, "old ubar:", ubar
        ubar = zbrent(SOLVE_UBAR,5.0_dp,50.0_dp,1.0D-5)
        !dum =  SOLVE_UBAR(ubar)
        !PRINT *, "vsl(1,1,1):", v_vsl(1,1,1)/aggcinit
        !PRINT *, "vsl(2,1,1):", v_vsl(2,1,1)/aggcinit
        !PRINT *, "vsl(1,N_z,1):", v_vsl(1,N_z,1)/aggcinit
        !PRINT *, "vsl(2,N_z,1):", v_vsl(2,N_z,1)/aggcinit
        !PRINT *, "vsl(3,1,1):", v_vsl(3,1,1)/aggcinit
        !PRINT *, "vsl(1,1,N_k):", v_vsl(1,1,N_k)/aggcinit
        !PRINT *, "vsl(2,1,N_k):", v_vsl(2,1,N_k)/aggcinit
        !PRINT *, "vsl(1,N_z,N_k):", v_vsl(1,N_z,N_k)/aggcinit
        !PRINT *, "vsl(2,N_z,N_k):", v_vsl(2,N_z,N_k)/aggcinit
        !PRINT *, "vsl(3,1,N_k):", v_vsl(3,1,N_k)/aggcinit
        PRINT *, "new ubar:", ubar
    ENDIF
    
    utilsick = -sick_util_loss*SUM(muss                 &
        *( (css*gamgam)**(onedp-sigma)/(onedp-sigma)    &
        - phi*(lss**(onedp+nu))/(onedp+nu)              &
        + ubar) )
        
    IF (unuloss.EQ.1) THEN
        utilsui = -non_util_loss*SUM(muss                    &
            *( (css*gamgam)**(onedp-sigma)/(onedp-sigma)    &
            - phi*(lss**(onedp+nu))/(onedp+nu)              &
            + ubar) )
    ELSE
        utilsui = zerodp
    ENDIF
        
    PRINT *, "util_sick:", utilsick
    PRINT *, "util_nonemployment:", utilsui
    
    ! EPI PARAMETERS
    IF (transmission.EQ.1) THEN
        beta_e  = repnum*pix
        beta_c  = zerodp
        beta_l  = zerodp
    ELSEIF (transmission.EQ.2) THEN
        beta_e  = repnum*pix*0.35_dp
        beta_c  = repnum*pix*0.35_dp/aggocinit**2
        beta_l  = repnum*pix*0.30_dp/aggohinit**2
    ELSE
        beta_e  = repnum*pix/3.0_dp
        beta_c  = repnum*pix/3.0_dp/aggocinit**2
        beta_l  = repnum*pix/3.0_dp/aggohinit**2
    END IF
    PRINT *, 'beta_e, beta_c, beta_l:'
    PRINT *, beta_e, beta_c, beta_l
    
    IF (wfh.GE.1) THEN
        wfh_z(1,1:2) = 0.029_dp
        wfh_z(1,3) = 0.089_dp
        wfh_z(1,4) = 0.282_dp
        wfh_z(1,5) = 0.467_dp
        wfh_z(1,6:7) = 0.664_dp
        wfh_z(2,1:2) = 0.029_dp
        wfh_z(2,3) = 0.209_dp
        wfh_z(2,4) = 0.376_dp
        wfh_z(2,5) = 0.632_dp
        wfh_z(2,6:7) = 0.664_dp
    ELSE
        wfh_z = zerodp
    ENDIF
    
    ! Save results
    CALL SAVESS
    
    IF (ss_type.EQ.2) THEN
        CALL SOLVESSWFH
    ENDIF  
           
END SUBROUTINE computeSS 

SUBROUTINE computetransition(ss_type,trans_type)
! ss_type
! 1. baseline
! trans_type
! 1. baseline: all shocks

    ! local vars
    REAL(DP)                ::  capsule_time_start, capsule_time_stop, dum, &
                                rmxdif, imxdif, cmxdif, hmxdif, bmxdif
    REAL(DP)                ::  tauc_temp, taus_temp
    INTEGER, INTENT(IN)     ::  ss_type, trans_type
    INTEGER                 ::  i, t, dumi, reason
    LOGICAL                 ::  file_exists
        
    CALL computeSS(0)    
    
    PRINT *, 'Preparing transition...'
    
    betaepath = beta_e
    
940 FORMAT(I5,12F18.10)    
    IF (tranguess .EQ. 1) THEN 
        INQUIRE(file=TRIM(path)//'/trans_sol.out',exist=file_exists)
        IF (file_exists) THEN
            OPEN(25,file=TRIM(path)//'/trans_sol.out',form='formatted')
        ELSE
            OPEN(25,file='trans_sol.out',form='formatted')
        END IF
        bigipath = zerodp
        aggocinfpath = zerodp
        aggohinfpath = zerodp
        betaepath = beta_e
        READ(25,940) dumi,dum,dum,   bigipath(1),dum,   aggocinfpath(1),dum,    &
                    aggohinfpath(1),dum,    tauc_temp,taus_temp,    betaepath(1),dum
        DO t = 2,ncapt
            READ(25,940,IOSTAT=reason) &
                    dumi,dum,dum,   bigipath(t),dum,    aggocinfpath(t),dum,    &
                    aggohinfpath(t),dum,    dum,      dum,          betaepath(t),dum
            IF (Reason .NE. 0)  THEN
				EXIT
			END IF
        END DO
        CLOSE(25)
        betaepath(N_death:ncapt) = beta_e
        
        IF (policy.EQ.1) THEN
            !taxtype ! 1: tauc only | 2: taus only | 3: Pigouvian
            IF (taxtype.EQ.0) THEN
                tauc_init = zerodp
                taus_init = zerodp
            ELSEIF (taxtype.EQ.1) THEN
                IF (tauc_temp>zerodp) THEN
                    tauc_init = tauc_temp
                ENDIF
                taus_init = zerodp
            ELSEIF (taxtype.EQ.2) THEN
                IF (taus_temp>zerodp) THEN
                    taus_init = taus_temp
                ENDIF
                tauc_init = zerodp
            ELSEIF (taxtype.GE.3) THEN
                IF (tauc_temp>zerodp) THEN
                    tauc_init = tauc_temp
                ENDIF
                taus_init = zerodp
            ENDIF
        ELSE ! policy==0
            tauc_init = zerodp
            taus_init = zerodp
        ENDIF
        rpath = rss
    ELSE ! tranguess==0
        ! Assume a constant path
        rpath = rss
        bigipath = zerodp
        bigipath(1) = muh(2)
        aggocinfpath = zerodp
        aggocinfpath(1) = bigipath(1)*aggocinit
        aggohinfpath = zerodp
        aggohinfpath(1) = bigipath(1)*agghinit
        
        betaepath = beta_e
        IF (tvbetae.EQ.1) THEN      
            !betaepath(1:N_death-1) = (/ &
            !    1.3506232284_dp,    &
            !    0.4120540521_dp,    &
            !    0.2562632796_dp,    &
            !    0.1572196014_dp,    &
            !    0.1326382219_dp,    &
            !    0.1413287758_dp,    &
            !    0.4380133016_dp,    &
            !    0.6339753344_dp,    &
            !    0.4388432372_dp,    &
            !    0.3571951210_dp,    &
            !    0.2303908267_dp,    &
            !    0.3232428975_dp,    &
            !    0.2647918159_dp     /) 
            !betaepath(1:N_death-1) = betaepath(1:N_death-1) - 0.57037037037_dp + beta_e
        ENDIF
        
        IF (transmission.EQ.1) THEN ! Exogenous tranmission
            bigspath(1) = onedp - bigipath(1)
            DO t=2,ncapt
                bigipath(t) = bigipath(t-1)*(onedp-pix) + betaepath(t-1)*bigipath(t-1)*bigspath(t-1)
                IF (t.GE.tvac) THEN
                    bigspath(t) = MAX(zerodp,bigspath(t-1)*(onedp-vaccrate(2))      &
                                    - betaepath(t-1)*bigipath(t-1)*bigspath(t-1))
                ELSE
                    bigspath(t) = MAX(zerodp,bigspath(t-1)                          &
                                    - betaepath(t-1)*bigipath(t-1)*bigspath(t-1))
                ENDIF
            ENDDO
        ELSEIF (tvbetae.EQ.1) THEN
            DO t=2,N_death
                bigipath(t) = deaths_data(t)/(pix*deathrate(2))
            ENDDO
            aggocinfpath(1:N_death) = bigipath(1:N_death)*aggocinit
            aggohinfpath(1:N_death) = bigipath(1:N_death)*agghinit*SUM(Erd*(onedp-wfh_z(1,:)))
        ENDIF
    ENDIF
        
    sui_init = sui_cbar*cbar
    
    PRINT *, 'Going to transition'
    DO i = 1,mxtriter

        CALL TRANSITION(itersave)
         
        imxdif = MAXVAL(DABS(ipathnew(1:ncapt)-bigipath(1:ncapt)))
        bmxdif = MAXVAL(DABS(betaepathnew(1:ncapt)-betaepath(1:ncapt)))
        IF (transmission == 1) THEN ! Exogenous transmission
            cmxdif = zerodp
            hmxdif = zerodp
        ELSE
            cmxdif = MAXVAL(DABS(aggocinfpathnew(1:ncapt)-aggocinfpath(1:ncapt)))
            hmxdif = MAXVAL(DABS(aggohinfpathnew(1:ncapt)-aggohinfpath(1:ncapt)))
        END IF
    
        OPEN(25,file=TRIM(path)//'/trans_sol.out',form='formatted')
        DO t = 1,ncapt
            WRITE(25,940) t,rpath(t),rpathnew(t),bigipath(t),ipathnew(t),                        &
                        aggocinfpath(t),aggocinfpathnew(t),aggohinfpath(t),aggohinfpathnew(t),  &
                        tauc_init,taus_init,betaepath(t),betaepathnew(t)
        END DO
        CLOSE(25)
    
        IF (transmission == 1) THEN ! Exogenous transmission
            bigipath = ipathnew
            aggocinfpath    = aggocinfpathnew
            aggohinfpath    = aggohinfpathnew
            betaepath       = betaepathnew
        ELSE ! Endogenous tranmission (dampen updates)
            IF (i.LE.5) THEN
                bigipath        = (bigipath + ipathnew)/2.0_dp
                aggocinfpath    = (aggocinfpath + aggocinfpathnew)/2.0_dp
                aggohinfpath    = (aggohinfpath + aggohinfpathnew)/2.0_dp 
            ELSE
                bigipath        = (onedp-updatei)*bigipath      + updatei*ipathnew
                aggocinfpath    = (onedp-updatei)*aggocinfpath  + updatei*aggocinfpathnew
                aggohinfpath    = (onedp-updatei)*aggohinfpath  + updatei*aggohinfpathnew                
            ENDIF
            betaepath       = (onedp-updateb)*betaepath     + updateb*betaepathnew
        END IF
        
        IF (taxtype.EQ.1) THEN
            tauc_init = (onedp-updatet)*tauc_init+updatet*MAX(zerodp,MIN(2.0_dp,(tauc_init + pvdef*1d-2)))
            IF ((tauc_init>=1.00_dp).AND.(pvdef>1.0d-1)) GOTO 13
        ELSEIF (taxtype.EQ.2) THEN
            taus_init = (onedp-updatet)*taus_init+updatet*MAX(zerodp,MIN(0.9_dp,(taus_init + pvdef*1d-2)))
            IF ((taus_init>=0.89_dp).AND.(pvdef>1.0d-2)) GOTO 13
        ELSEIF (taxtype.EQ.3) THEN
            tauc_init = (onedp-updatet)*tauc_init+updatet*MAX(zerodp,MIN(2.0_dp,(tauc_init + pvdef*1d-2)))
            IF ((tauc_init>=1.00_dp).AND.(pvdef>1.0d-1)) GOTO 13
        ENDIF
        
        IF (taxtype>0) THEN
            IF (taxtype<4) THEN
                IF ((imxdif<criti).AND.(cmxdif<critc).AND.(hmxdif<crith).AND.(bmxdif<critb).AND.(ABS(pvdef)<1.0d-3)) GOTO 12
            ELSE
                IF ((imxdif<criti).AND.(cmxdif<critc).AND.(hmxdif<crith).AND.(bmxdif<critb)) GOTO 12
            ENDIF    
            IF(multpl(i,1) == 1) THEN 
                WRITE(6,392) i,imxdif,cmxdif,hmxdif,bmxdif,pvdef
            ENDIF
        ELSE
            IF ((imxdif<criti).AND.(cmxdif<critc).AND.(hmxdif<crith).AND.(bmxdif<critb)) GOTO 12
            IF(multpl(i,1) == 1) THEN 
                WRITE(6,392) i,imxdif,cmxdif,hmxdif,bmxdif,pvdef
            ENDIF
        ENDIF
        
    END DO

13  PRINT *, 'budget cannot be cleared with specified tax instrument'
    
392 FORMAT(i4,4f10.6,f11.6)
12  WRITE(6,392) i,imxdif,cmxdif,hmxdif,bmxdif,pvdef
    PRINT *, CHAR(7)
    
    IF (itersave.EQ.0) THEN
        PRINT*, 'running one last time and saving...'
        CALL TRANSITION(1)
    ENDIF
        
    PRINT*, 'done with transition'
    
    CALL POST_PROCESS
        
END SUBROUTINE computetransition

END PROGRAM TEST_AUX
 
!!!!!!!!!!!!!!!!!!!!!!!!
!!    END PROGRAM     !!
!!!!!!!!!!!!!!!!!!!!!!!!  