MODULE PS6PARA
    IMPLICIT NONE
    INTEGER:: I
    INTEGER, PARAMETER:: N=66, RETIRE_AGE=46, T=30, CHANGE=2 ! LIFE EXPECTANCY, RETIREMENT, PERIODS
    REAL(KIND=8), PARAMETER:: A_MAX=20.1, A_MIN=0.1, STEP=1.0
    INTEGER, PARAMETER:: NZ=2, NA=(A_MAX-A_MIN)/STEP+1
    REAL(KIND=8):: ALPHA=0.36, BETA=0.97
    REAL(KIND=8), PARAMETER:: GAMMA=0.42, SIGMA=2.0, DELTA=0.06, POP_GROWTH= 0.011
    REAL(KIND=8), DIMENSION(2):: PROD_STATES = (/3.0, 0.5/)
    REAL(KIND=8), DIMENSION(2,2):: PROD_MARKOV = TRANSPOSE(RESHAPE((/0.9261,(1-0.9261),(1-0.9811),0.9811/),(/NZ,NZ/)))
    REAL(KIND=8), DIMENSION(NA):: A_GRID = (/(I*STEP, I=1,NA)/) + A_MIN - STEP
    REAL(KIND=8), DIMENSION(N):: COHORT_POP
    REAL(KIND=8), DIMENSION(N,NZ):: AGE_EFF=0.
    REAL(KIND=8), DIMENSION(45):: EF
    REAL(KIND=8)::THETA, THE_OLD=0.11, THE_NEW=0.0
    CONTAINS
        SUBROUTINE INIT_ARRAY()
            COHORT_POP = (/((1.0/((1.0+POP_GROWTH)**I)), I=1,N)/)
            COHORT_POP = COHORT_POP/SUM(COHORT_POP)
            OPEN(UNIT=13, FILE='ef.txt', STATUS='OLD')
            DO I=1,RETIRE_AGE-1
                READ(UNIT=13, FMT=*) EF(I)
            ENDDO
            CLOSE(UNIT=13)
            DO I=1,RETIRE_AGE-1
            AGE_EFF(I,1) = EF(I)*PROD_STATES(1)
            AGE_EFF(I,2) = EF(I)*PROD_STATES(2)
            ENDDO
        END SUBROUTINE
END MODULE

MODULE PS6RES
    USE PS6PARA
    IMPLICIT NONE
    REAL(KIND=8):: ERROR_MAIN=100., ERROR_EP
    REAL(KIND=8), DIMENSION(T):: WAGE, AGG_LABOR=1.0, AGG_CAP=2.0, BEN, INTEREST
    REAL(KIND=8), DIMENSION(T):: ACAP_NEW, ALAB_NEW
    REAL(KIND=8), DIMENSION(NA,NZ,N,T):: VFUNC, LFUNC, STAT_DIST, VFUNC_OLD
    INTEGER, DIMENSION(NA,NZ,N,T):: PFUNC=1
END MODULE

PROGRAM PS6
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    REAL(KIND=8):: CRIT_MAIN=1e-2, DEL_LAB, DEL_CAP
    INTEGER:: IDX, PERIOD, DEBUG=0.

    CALL INIT_ARRAY()
    CALL COMPUTE_EP() ! CALCULATE STEADY STATE EQ UNDER NEW THETA
    DO IDX=2,T-1 ! INITIAL GUESS
        DEL_CAP = AGG_CAP(T)-AGG_CAP(1)
        DEL_LAB = AGG_LABOR(T)-AGG_LABOR(1)
        AGG_LABOR(IDX) =AGG_LABOR(1) + DEL_LAB*(FLOAT(IDX)-1.)/FLOAT(T)
        AGG_CAP(IDX) = AGG_CAP(1)+DEL_CAP*(FLOAT(IDX)-1.)/FLOAT(T)
    ENDDO
       ! START MAIN LOOP
    DO WHILE (ERROR_MAIN>CRIT_MAIN)
        DO IDX=2, T-1
            CALL CAL_POINT_MOM(IDX)
        ENDDO
        VFUNC(:,:,:,2:T-1)=0.
        LFUNC(:,:,:,2:T-1) =0.
        PFUNC(:,:,:,2:T-1)=1

        !$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(THETA,IDX, PERIOD)
        DO IDX=1, T-1
            IF (IDX<CHANGE) THEN
                THETA = THE_OLD
            ELSE
                THETA = THE_NEW
            ENDIF
            PERIOD = T-IDX
            CALL BACKWARD_INDUCTION(PERIOD)

        ENDDO
        !$OMP END PARALLEL DO
        CALL FIND_STAT_DIST()
        ! COMPUTE ERROR UPDATE PRICE
        CALL COMPUTE_ERROR()
        ! END CONDITIONS FOR DEBUG
        DEBUG = DEBUG+1
        IF (DEBUG>1000) THEN
            ERROR_MAIN=0.00000001
            PRINT*, "OK, I AM SCREWEDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDDD"
        ENDIF
        PRINT*, "STATDIST, SUM, NONZERO", SUM(STAT_DIST), COUNT(STAT_DIST/=0)
        PRINT*, "ERROR: MAIN, ENDPOINT", ERROR_MAIN, ERROR_EP
        PRINT*, "R, B,W:", INTEREST(T-1), SUM(BEN(2:T))/FLOAT(T-1), SUM(WAGE)/FLOAT(T)
        PRINT*, "VFUNC: NONZERO, SIZE", COUNT(VFUNC/=0), SIZE(VFUNC)
        PRINT*, "======================================="
    ENDDO
    DO IDX=1, T
        CALL CAL_POINT_MOM(IDX)
    ENDDO

    CALL WRITE_ALL()
    CALL CONSUM_EQ()
END PROGRAM PS6

SUBROUTINE COMPUTE_EP()
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    REAL(KIND=8):: CRIT_EP=1e-3
    INTEGER, DIMENSION(2):: EPS=(/1,T/)
    INTEGER:: TIDX, IDX

    DO IDX=1,2
        TIDX=EPS(IDX)
        PRINT*, NA
        ERROR_EP=100.
        DO WHILE(ERROR_EP>CRIT_EP)
            CALL CAL_POINT_MOM(TIDX)
            PRINT*, "=================================================="
            PRINT*, "R, B,W:", INTEREST(TIDX), BEN(TIDX), WAGE(TIDX)
            PRINT*,  "L, K:",AGG_LABOR(TIDX), AGG_CAP(TIDX)
            ! BACKWARD INDUCTION USING NEW VARS (BELLMAN)
            CALL BACKWARD_INDUCTION(TIDX)
            ! UPDATE STAT_DIST
            CALL EP_STAT_DIST(TIDX)
            ! OBSERVE ERROR AND UPDATE PRICE
            CALL EP_PRICE(TIDX, .TRUE.)
        ENDDO

        CALL CAL_POINT_MOM(TIDX)
    ENDDO

END SUBROUTINE

SUBROUTINE FIND_STAT_DIST()
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER:: SROWIDX, SCOLIDX, AGEIDX, LAST_AGE, LAST_CHOICE, TIDX,IDX
    REAL(KIND=8), DIMENSION(NA,NZ,N,T):: STAT_DIST_NEW
    REAL(KIND=8):: OLD_POP

    STAT_DIST_NEW(:,:,:,1) = STAT_DIST(:,:,:,1)
    STAT_DIST_NEW(:,:,:,T) = STAT_DIST(:,:,:,T)
    DO TIDX=2,T-1 ! INITIAL EXCEPT ENDPOINTS
        ! EVERYONE HAVE ZERO ASSET AT STARTING AGE
        STAT_DIST_NEW(:,:,:,TIDX) = 0.
        STAT_DIST_NEW(1,1,1,TIDX) = 0.2037*COHORT_POP(1) ! ALWAYS NO ASSET AT FIRST AGE, BUT 2 PROD STATES
        STAT_DIST_NEW(1,2,1,TIDX) = 0.7963*COHORT_POP(1)
    ENDDO
    DO TIDX=2,T-1 ! OMIT END POINTS
        DO AGEIDX=2, N
        DO SROWIDX=1, NA ! LAST PERIOD ASSET
        DO SCOLIDX=1, NZ !LAST PERIOD PRODUCTIVITY
            LAST_AGE = AGEIDX-1
            LAST_CHOICE = PFUNC(SROWIDX, SCOLIDX, LAST_AGE, TIDX-1) ! TWO POSSIBLE
            OLD_POP= STAT_DIST_NEW(SROWIDX,SCOLIDX,LAST_AGE,TIDX-1)/(1.+POP_GROWTH)
            STAT_DIST_NEW(LAST_CHOICE,1,AGEIDX,TIDX) = STAT_DIST_NEW(LAST_CHOICE,1, AGEIDX,TIDX)+OLD_POP*PROD_MARKOV(SCOLIDX,1)
            STAT_DIST_NEW(LAST_CHOICE,2,AGEIDX,TIDX)= STAT_DIST_NEW(LAST_CHOICE,2, AGEIDX,TIDX)+OLD_POP*PROD_MARKOV(SCOLIDX,2)
        ENDDO
        ENDDO
        ENDDO
    ENDDO ! END BIG T LOOP
    DO TIDX=2, T-1
        STAT_DIST(:,:,:,TIDX) = STAT_DIST_NEW(:,:,:,TIDX)
    ENDDO
    PRINT*, COUNT(STAT_DIST/=0), SUM(STAT_DIST(:,:,:,1)), SUM(STAT_DIST(:,:,:,2)), SUM(STAT_DIST(:,:,:,T)), SUM(STAT_DIST)

END SUBROUTINE

SUBROUTINE COMPUTE_ERROR() ! FOR LARGE LOOP
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    REAL(KIND=8), DIMENSION(T-2):: ERR_CAP, ERR_LAB
    INTEGER:: TIDX

    DO TIDX=2,T-1
        CALL EP_PRICE(TIDX,.FALSE.)
    ENDDO
    ERR_CAP = ABS(ACAP_NEW(2:T-1) - AGG_CAP(2:T-1))
    ERR_LAB = ABS(ALAB_NEW(2:T-1) - AGG_LABOR(2:T-1))
    ERROR_MAIN = ABS(MAXVAL((/MAXVAL(ERR_CAP), MAXVAL(ERR_LAB)/)))
    PRINT*, "CAP LAB ERROR:", MAXVAL(ERR_CAP), MAXVAL(ERR_LAB)
    AGG_CAP(2:T-1) = 0.01*ACAP_NEW(2:T-1)+0.99  *AGG_CAP(2:T-1)
    AGG_LABOR(2:T-1) = 0.01*ALAB_NEW(2:T-1) + 0.99*AGG_LABOR(2:T-1)
END SUBROUTINE

SUBROUTINE BACKWARD_INDUCTION(TIDX)
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: TIDX
    INTEGER:: AGEIDX, AGE, SROWIDX, NEXTP, OLDP
    REAL(KIND=8):: CONSUM, UTIL, ERROR_VFI, CRIT_VFI=1e-3

    IF (TIDX/=1 .AND. TIDX/=T) THEN
        NEXTP = TIDX+1
        OLDP = TIDX-1
    ELSE
        NEXTP = TIDX
        OLDP = TIDX
    ENDIF

    ERROR_VFI=100.
    DO WHILE(ERROR_VFI> CRIT_VFI)
        ! FOR LAST AGE, ! SHOULD BE DETERMINISTIC, DO ONCE IS OK (NO CHOICE)
        DO SROWIDX=1,NA ! THIS IS THE CHOICE FROM LAST PERIOD/ CURRENT ASSET
                CONSUM = (1+INTEREST(OLDP))*A_GRID(SROWIDX) + BEN(TIDX) ! SPEND ALL IN LAST PERIOD
                PFUNC(:,:,N, TIDX) = 1 ! ACTUALLY DOES NOT MATTER
                UTIL = CONSUM**((1-SIGMA)*GAMMA)/(1-SIGMA) ! NO NEXT PERIOD VALUE
                VFUNC(SROWIDX,1,N, TIDX) = UTIL
                VFUNC(SROWIDX,2,N, TIDX) = UTIL
                ! ASSET MUST BE ZERO FOR AGE
        ENDDO

        DO AGEIDX = 2, N ! FOR AGE 20-65
            AGE = N- AGEIDX+1
            CALL BELLMAN(AGE, TIDX)
        ENDDO
        ERROR_VFI = MAXVAL(ABS(VFUNC-VFUNC_OLD))
        VFUNC_OLD = VFUNC
    ENDDO
END SUBROUTINE

SUBROUTINE BELLMAN(AGE, TIDX)
    USE OMP_LIB
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: AGE, TIDX
    INTEGER:: NEXTP, OLDP
    REAL(KIND=8):: COND_MAX_UTIL, CONSUM, UTIL, WORKING, PROD, NEXTU
    INTEGER:: SROWIDX, SCOLIDX, CHOICEIDX

    IF (TIDX/=1 .AND. TIDX/=T) THEN
        NEXTP = TIDX+1
        OLDP = TIDX-1
    ELSE
        NEXTP = TIDX
        OLDP = TIDX
    ENDIF
    IF (AGE< RETIRE_AGE) THEN ! WOKRING BELLMAN
        !$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(SROWIDX, SCOLIDX, COND_MAX_UTIL, WORKING, PROD, UTIL, NEXTU, CHOICEIDX, CONSUM)
        DO SROWIDX=1, NA
            DO SCOLIDX=1, NZ
                COND_MAX_UTIL = -1e12
                PROD = AGE_EFF(AGE, SCOLIDX)
                DO CHOICEIDX=1,NA ! LOOP OVER CHOICE OF ASSET PRIME
                    CALL FIND_LABOR(CHOICEIDX, SROWIDX, SCOLIDX, AGE, WORKING, TIDX)
                    CONSUM = (1+INTEREST(OLDP))*A_GRID(SROWIDX) + (1-THETA)*PROD*WORKING*WAGE(TIDX) - A_GRID(CHOICEIDX)
                    IF (CONSUM > 0.) THEN
                        NEXTU = BETA*SUM(PROD_MARKOV(SCOLIDX,:)*VFUNC(CHOICEIDX,:,AGE+1, NEXTP))
                        UTIL = (CONSUM**GAMMA*(1-WORKING)**(1-GAMMA))**(1-SIGMA)/(1-SIGMA)+NEXTU
                        IF (UTIL>COND_MAX_UTIL) THEN
                            PFUNC(SROWIDX, SCOLIDX, AGE,TIDX) = CHOICEIDX
                            LFUNC(SROWIDX,SCOLIDX, AGE,TIDX) = WORKING
                            COND_MAX_UTIL = UTIL
                        ENDIF
                    ENDIF
                ENDDO ! END LOOP CHOICE SPACE FOR ONE STATE
                    IF (COND_MAX_UTIL<-1E10) THEN
                    STOP
                    ENDIF
                VFUNC(SROWIDX, SCOLIDX, AGE,TIDX) = COND_MAX_UTIL
            ENDDO
        ENDDO
        !$OMP END PARALLEL DO
    ELSE ! BELLMAN FOR RETIRED PEOPLE
        !$OMP PARALLEL DEFAULT(SHARED) PRIVATE(SROWIDX, SCOLIDX, COND_MAX_UTIL, UTIL, CHOICEIDX, CONSUM )
        !$OMP DO
        DO SROWIDX=1, NA
            COND_MAX_UTIL = -1e12
            DO CHOICEIDX=1, NA
                CONSUM = (1+INTEREST(OLDP))*A_GRID(SROWIDX) + BEN(TIDX) - A_GRID(CHOICEIDX)
                IF (CONSUM>0.) THEN
                    UTIL = CONSUM**((1-SIGMA)*GAMMA)/(1-SIGMA)+ BETA*(VFUNC(CHOICEIDX,1,AGE+1,NEXTP))
                    IF (UTIL>COND_MAX_UTIL) THEN
                        PFUNC(SROWIDX,:, AGE, TIDX) = CHOICEIDX ! NO PRODUCTIVITY DIFFERENCE
                        COND_MAX_UTIL = UTIL
                    ENDIF
                ENDIF
            ENDDO
                IF (COND_MAX_UTIL<-1E10) THEN
                    PRINT*, COND_MAX_UTIL
                    STOP
                ENDIF
            VFUNC(SROWIDX, 1, AGE, TIDX) = COND_MAX_UTIL ! PROD_STATE DO NOT AFFECT RETIRED PEOPLE
            VFUNC(SROWIDX, 2, AGE, TIDX) = COND_MAX_UTIL ! PROD_STATE DO NOT AFFECT RETIRED PEOPLE
        ENDDO
        !$OMP END DO
        !$OMP END PARALLEL
    ENDIF
END SUBROUTINE

SUBROUTINE EP_STAT_DIST(TIDX)
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: TIDX
    INTEGER:: SROWIDX, SCOLIDX, AGEIDX, LAST_AGE, LAST_CHOICE
    REAL(KIND=8), DIMENSION(NA,NZ,N, T):: STAT_DIST_NEW

    ! EVERYONE HAVE ZERO ASSET AT STARTING AGE
    STAT_DIST_NEW = 0.
    STAT_DIST_NEW(1,1,1, TIDX) = 0.2037 ! ALWAYS NO ASSET AT FIRST AGE,
    STAT_DIST_NEW(1,2,1, TIDX) = 0.7963 !  BUT 2 PROD STATES
    DO AGEIDX=2, N
        DO SROWIDX=1, NA ! LAST PERIOD ASSET
            DO SCOLIDX=1, NZ !LAST PERIOD PRODUCTIVITY
                LAST_CHOICE = PFUNC(SROWIDX, SCOLIDX, AGEIDX-1, TIDX) ! TWO POSSIBLE
                LAST_AGE = AGEIDX-1
                STAT_DIST_NEW(LAST_CHOICE, 1, AGEIDX, TIDX) = &
                STAT_DIST_NEW(LAST_CHOICE, 1, AGEIDX, TIDX)+STAT_DIST_NEW(SROWIDX,SCOLIDX,LAST_AGE, TIDX)*PROD_MARKOV(SCOLIDX,1)
                STAT_DIST_NEW(LAST_CHOICE, 2, AGEIDX,TIDX) = &
                STAT_DIST_NEW(LAST_CHOICE, 2, AGEIDX, TIDX) + STAT_DIST_NEW(SROWIDX,SCOLIDX,LAST_AGE, TIDX)*PROD_MARKOV(SCOLIDX,2)
           ENDDO
        ENDDO
    ENDDO
    DO AGEIDX=1,N
        STAT_DIST(:,:,AGEIDX,TIDX) = STAT_DIST_NEW(:,:,AGEIDX,TIDX)*COHORT_POP(AGEIDX)
    ENDDO
END SUBROUTINE

SUBROUTINE EP_PRICE(TIDX, EP)
    ! COMPUTE ERROR AND UPDATE PRICE
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: TIDX
    LOGICAL, INTENT(IN):: EP
    INTEGER:: SROWIDX, SCOLIDX, AGEIDX
    REAL(KIND=8):: EPL_NEW=0., EPK_NEW=0., PORTION, ERROR_EPLAB, ERROR_EPCAP

    EPL_NEW=0.
    EPK_NEW=0.
    DO AGEIDX=1, N ! CAL ERROR
        DO SROWIDX=1,NA
            DO SCOLIDX=1,NZ
                PORTION = STAT_DIST(SROWIDX, SCOLIDX, AGEIDX,TIDX)
                    EPK_NEW = EPK_NEW+ A_GRID(SROWIDX)*PORTION
                IF (AGEIDX<RETIRE_AGE) THEN
                    EPL_NEW = EPL_NEW+LFUNC(SROWIDX,SCOLIDX,AGEIDX, TIDX)*PORTION*AGE_EFF(AGEIDX,SCOLIDX)
                ENDIF
            ENDDO
        ENDDO
    ENDDO

    IF (EP) THEN
        ERROR_EPLAB = ABS(EPL_NEW-AGG_LABOR(TIDX))
        ERROR_EPCAP = ABS(EPK_NEW-AGG_CAP(TIDX))
        PRINT*, "UPDATE_PRICE ERROR", ERROR_EPLAB, ERROR_EPCAP
        ERROR_EP = MAX(ERROR_EPLAB, ERROR_EPCAP)
        AGG_LABOR(TIDX) = 0.05*EPL_NEW + 0.95*AGG_LABOR(TIDX)
        AGG_CAP(TIDX) = 0.05*EPK_NEW + 0.95*AGG_CAP(TIDX)
    ELSE
        ACAP_NEW(TIDX)= EPK_NEW
        ALAB_NEW(TIDX)= EPL_NEW
    ENDIF
END SUBROUTINE

SUBROUTINE FIND_LABOR(CHOICEIDX, SROWIDX, SCOLIDX, AGE, WORKING, TIDX)
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: CHOICEIDX, SROWIDX, SCOLIDX, AGE, TIDX
    INTEGER:: NEXTP, OLDP
    REAL(KIND=8), INTENT(OUT):: WORKING
    REAL(KIND=8):: LAST_CHOICE, CUR_CHOICE, OPT_LABOR, NOM

    IF (TIDX<CHANGE) THEN
        THETA = THE_OLD
    ELSE
        THETA = THE_NEW
    ENDIF

    IF (TIDX/=1 .AND. TIDX/=T) THEN
        NEXTP = TIDX+1
        OLDP = TIDX-1
    ELSE
        NEXTP = TIDX
        OLDP = TIDX
    ENDIF

    CUR_CHOICE = A_GRID(CHOICEIDX)
    IF (AGE/=1) THEN
        LAST_CHOICE = (1+INTEREST(OLDP))*A_GRID(SROWIDX) ! DISCOUNTED
    ELSE
        LAST_CHOICE = 0.
    ENDIF
    NOM = (GAMMA*(1-THETA)*AGE_EFF(AGE,SCOLIDX)*WAGE(TIDX)-(1-GAMMA)*(LAST_CHOICE-CUR_CHOICE))
    OPT_LABOR = NOM/((1-THETA)*WAGE(TIDX)*AGE_EFF(AGE,SCOLIDX))
    IF (OPT_LABOR > 1.0) THEN
        OPT_LABOR = 1.0
    ELSE IF (OPT_LABOR < 0.) THEN
        OPT_LABOR = 0.0
    ENDIF
    WORKING = OPT_LABOR
END SUBROUTINE

SUBROUTINE CAL_POINT_MOM(TIDX)
    USE PS6PARA
    USE PS6RES
    IMPLICIT NONE
    INTEGER, INTENT(IN):: TIDX

    IF (TIDX<CHANGE) THEN
        THETA = THE_OLD
    ELSE
        THETA = THE_NEW
    ENDIF

    WAGE(TIDX) = (1-ALPHA)*(AGG_CAP(TIDX)**(ALPHA))*(AGG_LABOR(TIDX)**(-ALPHA))! FOC WRT LABOR
    INTEREST(TIDX) = ALPHA*(AGG_CAP(TIDX)**(ALPHA-1))*(AGG_LABOR(TIDX)**(1-ALPHA))-DELTA! FOC WRT CAPITAL
    BEN(TIDX) = THETA*WAGE(TIDX)*AGG_LABOR(TIDX)/SUM(COHORT_POP(RETIRE_AGE:N)) ! TAX REV

END SUBROUTINE

SUBROUTINE WRITE_ALL()
    USE PS6RES
    USE PS6PARA
    IMPLICIT NONE
    INTEGER:: SROWIDX
    CHARACTER(LEN=130):: PATH="/Users/chek_choi/Downloads/fortran/"
    CHARACTER(LEN=150):: FILE_NAME

        FILE_NAME = TRIM(PATH)//"VFUNC"
        OPEN(UNIT=1, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, NA
            WRITE(UNIT=1,FMT=*) VFUNC(SROWIDX,:, 50,1)
        ENDDO
        CLOSE(UNIT=1)

        FILE_NAME = TRIM(PATH)//"PFUNC"
        OPEN(UNIT=2, FILE=FILE_NAME, STATUS='REPLACE') ! ALSO SAVE POLICY FUNCTIONS
        DO SROWIDX=1, NA
            WRITE(UNIT=2,FMT=*) A_GRID(PFUNC(SROWIDX,1, 20,1)), A_GRID(PFUNC(SROWIDX,2, 20,1))
        ENDDO
        CLOSE(UNIT=2)

        FILE_NAME = TRIM(PATH)//"STATDIST"
        OPEN(UNIT=3, FILE=FILE_NAME, STATUS='REPLACE') ! ALSO SAVE POLICY FUNCTIONS
        DO SROWIDX=1, NA
            WRITE(UNIT=3,FMT=*) STAT_DIST(SROWIDX, :, 50,1)
        ENDDO
        CLOSE(UNIT=3)

        FILE_NAME = TRIM(PATH)//"CAPPATH"
        OPEN(UNIT=4, FILE=FILE_NAME, STATUS='REPLACE') ! FOR HAVING THE X-AXIS OF PLOT
        DO SROWIDX=1, T
            WRITE(UNIT=4,FMT=*) AGG_CAP(SROWIDX)
        ENDDO
        CLOSE(UNIT=4)

        FILE_NAME = TRIM(PATH)//"AGRID"
        OPEN(UNIT=5, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, NA
            WRITE(UNIT=5,FMT=*) A_GRID(SROWIDX)
        ENDDO
        CLOSE(UNIT=5)

        FILE_NAME = TRIM(PATH)//"LABPATH"
        OPEN(UNIT=32, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, T
            WRITE(UNIT=32,FMT=*) AGG_LABOR(SROWIDX)
        ENDDO
        CLOSE(UNIT=32)

        FILE_NAME = TRIM(PATH)//"RPATH"
        OPEN(UNIT=33, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, T
            WRITE(UNIT=33,FMT=*) INTEREST(SROWIDX)
        ENDDO
        CLOSE(UNIT=33)

        FILE_NAME = TRIM(PATH)//"WPATH"
        OPEN(UNIT=34, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, T
            WRITE(UNIT=34,FMT=*) WAGE(SROWIDX)
        ENDDO
        CLOSE(UNIT=34)

END SUBROUTINE

SUBROUTINE CONSUM_EQ()
    USE PS6RES
    USE PS6PARA
    IMPLICIT NONE
    INTEGER:: RIDX, CIDX, IDX
    REAL(KIND=8), DIMENSION(NA,NZ,N):: VFUNC_BEN
    REAL(KIND=8), DIMENSION(NA,NZ,N,2):: VFUNC_NOBEN=0.
    REAL(KIND=8), DIMENSION(N,2):: CONS_EQ

    VFUNC_BEN = VFUNC(:,:,:,1)
    VFUNC_NOBEN(:,:,:,1) = VFUNC(:,:,:,2)
    VFUNC_NOBEN(:,:,:,2) = VFUNC(:,:,:,2)
    !$OMP PARALLEL DO DEFAULT(SHARED) PRIVATE(RIDX)
    DO RIDX=1, N
        CONS_EQ(RIDX,1) = SUM(VFUNC_NOBEN(:,:,RIDX,1))-SUM(VFUNC_BEN(:,:,RIDX))
        CONS_EQ(RIDX,2) = SUM(VFUNC_NOBEN(:,:,RIDX,2))-SUM(VFUNC_BEN(:,:,RIDX))
    ENDDO
    !$OMP END PARALLEL DO

    OPEN(UNIT=28, FILE='CONSUM_EQ', STATUS='REPLACE')
        DO RIDX=1,N
            WRITE(UNIT=28, FMT=*) CONS_EQ(RIDX, :)
        ENDDO
    CLOSE(UNIT=28)
    PRINT*, "========================================"
END SUBROUTINE


      SUBROUTINE MB03OD( JOBQR, M, N, A, LDA, JPVT, RCOND, SVLMAX, TAU,
     $                   RANK, SVAL, DWORK, LDWORK, INFO )
C
C     SLICOT RELEASE 5.7.
C
C     Copyright (c) 2002-2020 NICONET e.V.
C
C     PURPOSE
C
C     To compute (optionally) a rank-revealing QR factorization of a
C     real general M-by-N matrix  A,  which may be rank-deficient,
C     and estimate its effective rank using incremental condition
C     estimation.
C
C     The routine uses a QR factorization with column pivoting:
C        A * P = Q * R,  where  R = [ R11 R12 ],
C                                   [  0  R22 ]
C     with R11 defined as the largest leading submatrix whose estimated
C     condition number is less than 1/RCOND.  The order of R11, RANK,
C     is the effective rank of A.
C
C     MB03OD  does not perform any scaling of the matrix A.
C
C     ARGUMENTS
C
C     Mode Parameters
C
C     JOBQR   CHARACTER*1
C             = 'Q':  Perform a QR factorization with column pivoting;
C             = 'N':  Do not perform the QR factorization (but assume
C                     that it has been done outside).
C
C     Input/Output Parameters
C
C     M       (input) INTEGER
C             The number of rows of the matrix A.  M >= 0.
C
C     N       (input) INTEGER
C             The number of columns of the matrix A.  N >= 0.
C
C     A       (input/output) DOUBLE PRECISION array, dimension
C             ( LDA, N )
C             On entry with JOBQR = 'Q', the leading M by N part of this
C             array must contain the given matrix A.
C             On exit with JOBQR = 'Q', the leading min(M,N) by N upper
C             triangular part of A contains the triangular factor R,
C             and the elements below the diagonal, with the array TAU,
C             represent the orthogonal matrix Q as a product of
C             min(M,N) elementary reflectors.
C             On entry and on exit with JOBQR = 'N', the leading
C             min(M,N) by N upper triangular part of A contains the
C             triangular factor R, as determined by the QR factorization
C             with pivoting.  The elements below the diagonal of A are
C             not referenced.
C
C     LDA     INTEGER
C             The leading dimension of the array A.  LDA >= max(1,M).
C
C     JPVT    (input/output) INTEGER array, dimension ( N )
C             On entry with JOBQR = 'Q', if JPVT(i) <> 0, the i-th
C             column of A is an initial column, otherwise it is a free
C             column. Before the QR factorization of A, all initial
C             columns are permuted to the leading positions; only the
C             remaining free columns are moved as a result of column
C             pivoting during the factorization.  For rank determination
C             it is preferable that all columns be free.
C             On exit with JOBQR = 'Q', if JPVT(i) = k, then the i-th
C             column of A*P was the k-th column of A.
C             Array JPVT is not referenced when JOBQR = 'N'.
C
C     RCOND   (input) DOUBLE PRECISION
C             RCOND is used to determine the effective rank of A, which
C             is defined as the order of the largest leading triangular
C             submatrix R11 in the QR factorization with pivoting of A,
C             whose estimated condition number is less than 1/RCOND.
C             RCOND >= 0.
C             NOTE that when SVLMAX > 0, the estimated rank could be
C             less than that defined above (see SVLMAX).
C
C     SVLMAX  (input) DOUBLE PRECISION
C             If A is a submatrix of another matrix B, and the rank
C             decision should be related to that matrix, then SVLMAX
C             should be an estimate of the largest singular value of B
C             (for instance, the Frobenius norm of B).  If this is not
C             the case, the input value SVLMAX = 0 should work.
C             SVLMAX >= 0.
C
C     TAU     (output) DOUBLE PRECISION array, dimension ( MIN( M, N ) )
C             On exit with JOBQR = 'Q', the leading min(M,N) elements of
C             TAU contain the scalar factors of the elementary
C             reflectors.
C             Array TAU is not referenced when JOBQR = 'N'.
C
C     RANK    (output) INTEGER
C             The effective (estimated) rank of A, i.e. the order of
C             the submatrix R11.
C
C     SVAL    (output) DOUBLE PRECISION array, dimension ( 3 )
C             The estimates of some of the singular values of the
C             triangular factor R:
C             SVAL(1): largest singular value of R(1:RANK,1:RANK);
C             SVAL(2): smallest singular value of R(1:RANK,1:RANK);
C             SVAL(3): smallest singular value of R(1:RANK+1,1:RANK+1),
C                      if RANK < MIN( M, N ), or of R(1:RANK,1:RANK),
C                      otherwise.
C             If the triangular factorization is a rank-revealing one
C             (which will be the case if the leading columns were well-
C             conditioned), then SVAL(1) will also be an estimate for
C             the largest singular value of A, and SVAL(2) and SVAL(3)
C             will be estimates for the RANK-th and (RANK+1)-st singular
C             values of A, respectively.
C             By examining these values, one can confirm that the rank
C             is well defined with respect to the chosen value of RCOND.
C             The ratio SVAL(1)/SVAL(2) is an estimate of the condition
C             number of R(1:RANK,1:RANK).
C
C     Workspace
C
C     DWORK   DOUBLE PRECISION array, dimension ( LDWORK )
C             On exit, if  INFO = 0,  DWORK(1) returns the optimal value
C             of LDWORK.
C
C     LDWORK  INTEGER
C             The length of the array DWORK.
C             LDWORK >= 3*N + 1,                 if JOBQR = 'Q';
C             LDWORK >= max( 1, 2*min( M, N ) ), if JOBQR = 'N'.
C             For good performance when JOBQR = 'Q', LDWORK should be
C             larger. Specifically, LDWORK >= 2*N + ( N + 1 )*NB, where
C             NB is the optimal block size for the LAPACK Library
C             routine DGEQP3.
C
C             If LDWORK = -1, then a workspace query is assumed;
C             the routine only calculates the optimal size of the
C             DWORK array, returns this value as the first entry of
C             the DWORK array, and no error message related to LDWORK
C             is issued by XERBLA.
C
C     Error Indicator
C
C     INFO    INTEGER
C             = 0:  successful exit
C             < 0:  if INFO = -i, the i-th argument had an illegal
C                   value.
C
C     METHOD
C
C     The routine computes or uses a QR factorization with column
C     pivoting of A,  A * P = Q * R,  with  R  defined above, and then
C     finds the largest leading submatrix whose estimated condition
C     number is less than 1/RCOND, taking the possible positive value of
C     SVLMAX into account.  This is performed using the LAPACK
C     incremental condition estimation scheme and a slightly modified
C     rank decision test.
C
C     CONTRIBUTOR
C
C     V. Sima, Katholieke Univ. Leuven, Belgium, Nov. 1996.
C
C     REVISIONS
C
C     V. Sima, Research Institute for Informatics, Bucharest, Mar. 2005,
C     Aug. 2011.
C
C    ******************************************************************
C
C     .. Parameters ..
      INTEGER            IMAX, IMIN
      PARAMETER          ( IMAX = 1, IMIN = 2 )
      DOUBLE PRECISION   ZERO, ONE
      PARAMETER          ( ZERO = 0.0D0, ONE = 1.0D0 )
C     .. Scalar Arguments ..
      CHARACTER          JOBQR
      INTEGER            INFO, LDA, LDWORK, M, N, RANK
      DOUBLE PRECISION   RCOND, SVLMAX
C     .. Array Arguments ..
      INTEGER            JPVT( * )
      DOUBLE PRECISION   A( LDA, * ), SVAL( 3 ), TAU( * ), DWORK( * )
C     .. Local Scalars ..
      LOGICAL            LJOBQR, LQUERY
      INTEGER            I, ISMAX, ISMIN, MAXWRK, MINWRK, MN
      DOUBLE PRECISION   C1, C2, S1, S2, SMAX, SMAXPR, SMIN, SMINPR
C     ..
C     .. External Functions ..
      LOGICAL            LSAME
      EXTERNAL           LSAME
C     .. External Subroutines ..
      EXTERNAL           DGEQP3, DLAIC1, XERBLA
C     .. Intrinsic Functions ..
      INTRINSIC          ABS, INT, MAX, MIN
C     ..
C     .. Executable Statements ..
C
      LJOBQR = LSAME( JOBQR, 'Q' )
      MN = MIN( M, N )
      IF( LJOBQR ) THEN
         MINWRK = 3*N + 1
      ELSE
         MINWRK = MAX( 1, 2*MN )
      END IF
      MAXWRK = MINWRK
C
C     Test the input scalar arguments.
C
      INFO = 0
      IF( .NOT.LJOBQR .AND. .NOT.LSAME( JOBQR, 'N' ) ) THEN
         INFO = -1
      ELSE IF( M.LT.0 ) THEN
         INFO = -2
      ELSE IF( N.LT.0 ) THEN
         INFO = -3
      ELSE IF( LDA.LT.MAX( 1, M ) ) THEN
         INFO = -5
      ELSE IF( RCOND.LT.ZERO ) THEN
         INFO = -7
      ELSE IF( SVLMAX.LT.ZERO ) THEN
         INFO = -8
      ELSE
         LQUERY = LDWORK.EQ.-1
         IF ( LJOBQR ) THEN
            CALL DGEQP3( M, N, A, LDA, JPVT, TAU, DWORK, -1, INFO )
            MAXWRK = MAX( MAXWRK, INT( DWORK(1) ) )
         END IF
         IF( LDWORK.LT.MINWRK .AND. .NOT.LQUERY )
     $      INFO = -13
      END IF
C
      IF( INFO.NE.0 ) THEN
         CALL XERBLA( 'MB03OD', -INFO )
         RETURN
      ELSE IF( LQUERY ) THEN
         DWORK( 1 ) = MAXWRK
         RETURN
      END IF
C
C     Quick return if possible
C
      IF( MN.EQ.0 ) THEN
         RANK = 0
         SVAL(  1 ) = ZERO
         SVAL(  2 ) = ZERO
         SVAL(  3 ) = ZERO
         DWORK( 1 ) = ONE
         RETURN
      END IF
C
      IF ( LJOBQR ) THEN
C
C        Compute QR factorization with column pivoting of A:
C           A * P = Q * R
C        Workspace need   3*N + 1;
C                  prefer 2*N + (N+1)*NB.
C        Details of Householder rotations stored in TAU.
C
         CALL DGEQP3( M, N, A, LDA, JPVT, TAU, DWORK, LDWORK, INFO )
      END IF
C
C     Determine RANK using incremental condition estimation
C
      ISMIN = 1
      ISMAX = MN + 1
      DWORK( ISMIN ) = ONE
      DWORK( ISMAX ) = ONE
      SMAX = ABS( A( 1, 1 ) )
      SMIN = SMAX
      IF( SMAX.EQ.ZERO .OR. SVLMAX*RCOND.GT.SMAX ) THEN
         RANK = 0
         SVAL( 1 ) = SMAX
         SVAL( 2 ) = ZERO
         SVAL( 3 ) = ZERO
      ELSE
         RANK = 1
         SMINPR = SMIN
C
   10    CONTINUE
         IF( RANK.LT.MN ) THEN
            I = RANK + 1
            CALL DLAIC1( IMIN, RANK, DWORK( ISMIN ), SMIN, A( 1, I ),
     $                   A( I, I ), SMINPR, S1, C1 )
            CALL DLAIC1( IMAX, RANK, DWORK( ISMAX ), SMAX, A( 1, I ),
     $                   A( I, I ), SMAXPR, S2, C2 )
C
            IF( SVLMAX*RCOND.LE.SMAXPR ) THEN
               IF( SVLMAX*RCOND.LE.SMINPR ) THEN
                  IF( SMAXPR*RCOND.LE.SMINPR ) THEN
                     DO 20 I = 1, RANK
                        DWORK( ISMIN+I-1 ) = S1*DWORK( ISMIN+I-1 )
                        DWORK( ISMAX+I-1 ) = S2*DWORK( ISMAX+I-1 )
   20                CONTINUE
                     DWORK( ISMIN+RANK ) = C1
                     DWORK( ISMAX+RANK ) = C2
                     SMIN = SMINPR
                     SMAX = SMAXPR
                     RANK = RANK + 1
                     GO TO 10
                  END IF
               END IF
            END IF
         END IF
         SVAL( 1 ) = SMAX
         SVAL( 2 ) = SMIN
         SVAL( 3 ) = SMINPR
      END IF
C
      DWORK( 1 ) = MAXWRK
      RETURN
C *** Last line of MB03OD ***
      END


module mod_initflow
  use mpi_f08
  use mod_common_mpi, only: myid_block,ierr,comm_block
  use mod_param     , only: pi
  use mod_types
  implicit none
  private
  public initflow,init_inflow
  contains
  subroutine initflow(inivel,is_wallturb,lo,hi,lo_g,hi_g,l,uref,lref,visc,bforce, &
                      xc,xf,yc,yf,zc,zf,dxc,dxf,dyc,dyf,dzc,dzf,u,v,w,p)
    !
    ! computes initial conditions for the velocity field
    !
    implicit none
    character(len=3), intent(in) :: inivel
    logical , intent(in) :: is_wallturb
    integer , intent(in), dimension(3) :: lo,hi,lo_g,hi_g
    real(rp), intent(in), dimension(3) :: l
    real(rp), intent(inout) :: uref
    real(rp), intent(in) :: lref,visc,bforce
    real(rp), intent(in), dimension(lo(1)-1:) :: xc,xf,dxc,dxf
    real(rp), intent(in), dimension(lo(2)-1:) :: yc,yf,dyc,dyf
    real(rp), intent(in), dimension(lo(3)-1:) :: zc,zf,dzc,dzf
    real(rp), dimension(lo(1)-1:,lo(2)-1:,lo(3)-1:), intent(out) :: u,v,w,p
    real(rp), allocatable, dimension(:) :: u1d
    integer  :: i,j,k
    logical  :: is_noise,is_mean,is_pair
    real(rp) :: reb,retau,lambda
    real(rp) :: xcl,xfl,ycl,yfl,zcl,zfl
    !
    allocate(u1d(lo(3):hi(3)))
    is_noise = .false.
    is_mean  = .false.
    is_pair  = .false.
    select case(trim(inivel))
    case('cou')
      call couette(lo(3),hi(3),zc,l(3),uref,u1d)
    case('poi')
      call poiseuille(lo(3),hi(3),zc,l(3),uref,u1d)
      is_mean=.true.
    case('zer')
      u1d(:) = 0._rp
    case('uni')
      u1d(:) = uref
    case('log')
      call log_profile(lo(3),hi(3),zc,l(3),uref,lref,visc,u1d)
      is_noise = .true.
      is_mean = .true.
    case('hcp')
      call poiseuille(lo(3),hi(3),zc,2.*l(3),uref,u1d)
      is_mean = .true.
    case('hcl')
      call log_profile(lo(3),hi(3),zc,2.*l(3),uref,lref,visc,u1d)
      is_noise = .true.
      is_mean=.true.
    case('tgv')
      do k=lo(3),hi(3)
        zcl = zc(k)/l(3)*2._rp*pi
        do j=lo(2),hi(2)
          ycl = yc(j)/l(2)*2._rp*pi
          yfl = yf(j)/l(2)*2._rp*pi
          do i=lo(1),hi(1)
            xcl = xc(i)/l(1)*2._rp*pi
            xfl = xf(i)/l(1)*2._rp*pi
            u(i,j,k) =  sin(xfl)*cos(ycl)*cos(zcl)
            v(i,j,k) = -cos(xcl)*sin(yfl)*cos(zcl)
            w(i,j,k) = 0._rp
            p(i,j,k) = 0._rp!(cos(2._rp*xc)+cos(2._rp*yc))*(cos(2._rp*zc)+2._rp)/16._rp
          end do
        end do
      end do
    case('kov')
      reb = uref*lref/visc
      lambda = reb/2-sqrt(reb**2/4+4*pi**2)
      do k=lo(3),hi(3)
        zcl = zc(k)
        zfl = zf(k)
        do j=lo(2),hi(2)
          ycl = yc(j)
          yfl = yf(j)
          do i=lo(1),hi(1)
            xcl = xc(i)
            xfl = xf(i)
            u(i,j,k) = 1._rp            -exp(lambda*xfl)*cos(2*pi*zcl)
            v(i,j,k) = 0._rp
            w(i,j,k) = lambda/(2._rp*pi)*exp(lambda*xcl)*sin(2*pi*zfl)
            p(i,j,k) = 0._rp !.5_rp*(1._rp-exp(2._rp*lambda*xcl))
          end do
        end do
      end do
    case('pdc')
      if(is_wallturb) then ! turbulent flow
        retau  = sqrt(bforce*lref)*uref/visc
        reb    = (retau/.09_rp)**(1._rp/.88_rp)
        uref   = (reb/2._rp)/retau
      else                 ! laminar flow
        uref = (bforce*lref**2/(3._rp*visc))
      end if
      call poiseuille(lo(3),hi(3),zc,l(3),uref,u1d)
      is_mean=.true.
    case default
      if(myid_block == 0) write(stderr,*) 'ERROR: invalid name for initial velocity field'
      if(myid_block == 0) write(stderr,*) ''
      if(myid_block == 0) write(stderr,*) '*** Simulation aborted due to errors in the case file ***'
      if(myid_block == 0) write(stderr,*) '    check INFO_INPUT.md'
      call MPI_FINALIZE()
      error stop
    end select
    if(all(inivel /= ['tgv','kov'])) then
      do k=lo(3),hi(3)
        do j=lo(2),hi(2)
          do i=lo(1),hi(1)
            u(i,j,k) = u1d(k)
            v(i,j,k) = 0._rp
            w(i,j,k) = 0._rp
            p(i,j,k) = 0._rp
          end do
        end do
      end do
    end if
    if(is_noise) then
      call add_noise(lo,hi,lo_g,hi_g,123,.05_rp,u)
      call add_noise(lo,hi,lo_g,hi_g,456,.05_rp,v)
      call add_noise(lo,hi,lo_g,hi_g,789,.05_rp,w)
    end if
    if(is_mean) then
      call set_mean(lo,hi,l,dxc,dyf,dzf,uref,u)
    end if
    if(is_wallturb) is_pair = .true.
    if(is_pair) then
      !
      ! initialize a streamwise vortex pair for a fast transition
      ! to turbulence in a pressure-driven channel:
      !        psi(x,y,z)  = f(z)*g(x,y), with
      !        f(z)        = (1-z**2)**2, and
      !        g(x,y)      = y*exp[-(16x**2-4y**2)]
      ! (x,y,z) --> (streamwise, spanwise, wall-normal) directions
      !
      ! see Henningson and Kim, JFM 1991
      !
      do k=lo(3),hi(3)
        zcl = 2._rp*zc(k)/l(3) - 1._rp ! z rescaled to be between -1 and +1
        zfl = 2._rp*zf(k)/l(3) - 1._rp
        do j=lo(2),hi(2)
          ycl = 2._rp*yc(j)/l(2) - 1._rp ! y rescaled to be between -1 and +1
          yfl = 2._rp*yf(j)/l(2) - 1._rp
          do i=lo(1),hi(1)
            xcl = 2._rp*xc(i)/l(1) - 1._rp ! x rescaled to be between -1 and +1
            xfl = 2._rp*xf(i)/l(1) - 1._rp
            v(i,j,k) = -1._rp*gxy(yfl,xcl)*dfz(zcl)*uref
            w(i,j,k) =  1._rp*fz(zfl)*dgxy(ycl,xcl)*uref
            p(i,j,k) =  0._rp
          end do
        end do
      end do
      !
      ! alternatively, using a Taylor-Green vortex
      ! for the cross-stream velocity components
      ! (commented below)
      !
      !do k=lo(3),hi(3)
      !  zcl = zc(k)/l(3)*2._rp*pi
      !  zfl = zf(k)/l(3)*2._rp*pi
      !  do j=lo(2),hi(2)
      !    ycl = yc(j)/l(2)*2._rp*pi
      !    yfl = yf(j)/l(2)*2._rp*pi
      !    do i=lo(1),hi(1)
      !      xcl = xc(i)/l(1)*2._rp*pi
      !      xfl = xf(i)/l(1)*2._rp*pi
      !      v(i,j,k) =  sin(xcl)*cos(yfl)*cos(zcl)
      !      w(i,j,k) = -cos(xcl)*sin(ycl)*cos(zfl)
      !      p(i,j,k) = 0._rp!(cos(2._rp*xcl)+cos(2._rp*ycl))*(cos(2._rp*zcl)+2._rp)/16._rp
      !    end do
      !  end do
      !end do
    end if
    deallocate(u1d)
  end subroutine initflow
  !
  subroutine add_noise(lo,hi,lo_g,hi_g,iseed,norm,p)
    implicit none
    integer , intent(in), dimension(3) :: lo,hi,lo_g,hi_g
    integer , intent(in) :: iseed
    real(rp), intent(in) :: norm
    real(rp), intent(inout), dimension(lo(1)-1:,lo(2)-1:,lo(3)-1:) :: p
    integer(4), allocatable, dimension(:) :: seed
    real(rp) :: rn
    integer :: i,j,k
    allocate(seed(64))
    seed(:) = iseed
    call random_seed( put = seed )
    do k=lo_g(3),hi_g(3)
      do j=lo_g(2),hi_g(2)
        do i=lo_g(1),hi_g(1)
          call random_number(rn)
          if(i>=lo(1).and.i<=hi(1) .and. &
             j>=lo(2).and.j<=hi(2) .and. &
             k>=lo(3).and.k<=hi(3) ) then
             p(i,j,k) = p(i,j,k) + 2._rp*(rn-.5_rp)*norm
          end if
        end do
      end do
    end do
  end subroutine add_noise
  !
  subroutine set_mean(lo,hi,l,dx,dy,dz,mean,p)
    implicit none
    integer , intent(in), dimension(3) :: lo,hi
    real(rp), intent(in), dimension(3) :: l
    real(rp), intent(in), dimension(lo(1)-1:) :: dx
    real(rp), intent(in), dimension(lo(2)-1:) :: dy
    real(rp), intent(in), dimension(lo(3)-1:) :: dz
    real(rp), intent(in) :: mean
    real(rp), intent(inout), dimension(lo(1)-1:,lo(2)-1:,lo(3)-1:) :: p
    real(rp) :: meanold
    integer :: i,j,k
    meanold = 0._rp
    !
    !$OMP PARALLEL DO DEFAULT(none) &
    !$OMP SHARED(lo,hi,l,dx,dy,dz,p) &
    !$OMP PRIVATE(i,j,k) &
    !$OMP REDUCTION(+:meanold)
    do k=lo(3),hi(3)
      do j=lo(2),hi(2)
        do i=lo(1),hi(1)
          meanold = meanold + p(i,j,k)*dx(i)*dy(j)*dz(k)/(l(1)*l(2)*l(3))
        end do
      end do
    end do
    !$OMP END PARALLEL DO
    call mpi_allreduce(MPI_IN_PLACE,meanold,1,MPI_REAL_RP,MPI_SUM,comm_block)
    if(meanold /= 0._rp) then
      !$OMP WORKSHARE
      p(:,:,:) = p(:,:,:)/meanold*mean
      !$OMP END WORKSHARE
    end if
  end subroutine set_mean
  !
  subroutine couette(lo,hi,zc,l,norm,p)
    !
    ! plane couette profile normalized by the wall velocity difference
    !
    implicit none
    integer , intent(in)   :: lo,hi
    real(rp), intent(in), dimension(lo-1:) :: zc
    real(rp), intent(in)   :: l,norm
    real(rp), intent(out), dimension(lo:) :: p
    integer :: k
    real(rp) :: z
    do k=lo,hi
      z    = zc(k)/l
      p(k) = .5_rp*(1._rp-2._rp*z)*norm
    end do
  end subroutine couette
  !
  subroutine poiseuille(lo,hi,zc,l,norm,p)
    implicit none
    integer , intent(in)   :: lo,hi
    real(rp), intent(in), dimension(lo-1:) :: zc
    real(rp), intent(in)   :: l,norm
    real(rp), intent(out), dimension(lo:) :: p
    integer :: k
    real(rp) :: z
    !
    ! plane poiseuille profile normalized by the bulk velocity
    !
    do k=lo,hi
      z    = zc(k)/l
      p(k) = 6._rp*z*(1._rp-z)*norm
    end do
  end subroutine poiseuille
  !
  function poiseuille_square(r,lmin,lmax) result(vel)
    !
    ! below a poiseuille inflow is calculated
    !   it is convinient that the integral a function of this kind
    !   is equal to one, such that it can be easily recycled
    !   for the inflow of a square duct, as done in
    !   in init_inflow below
    !
    real(rp), intent(in) :: r,lmin,lmax ! position where velocity is calculated
                                        ! and locations lower and upper boundary
    real(rp) :: vel,rr
    !
    rr = (r-lmin)/(lmax-lmin)
    vel = 6._rp*(rr*(1._rp-rr))
  end function poiseuille_square
  !
  subroutine init_inflow(inflow_type,periods,lo,hi,lmin,lmax,x1c,x1f,x2c,x2f, &
                         uref,lref,visc,veln,velt1,velt2)
    !
    ! below an inflow is calculated for a 2d plane
    ! later this subroutine can be generalized with other shapes of
    ! inflow velocity; by construction, the profile will not vary along
    ! a direction that does not have Dirichlet BCs; note that
    ! since the inflow is evaluated at the center of a face,
    ! there is no danger of experiencing a singularity '0._rp**0'.
    !
    implicit none
    integer , parameter :: INFLOW_POISEUILLE = 1, &
                           INFLOW_KOVASZNAY  = 2
    !
    integer , intent(in )                      :: inflow_type
    integer , intent(in ), dimension(2       ) :: periods
    integer , intent(in ), dimension(2       ) :: lo,hi
    real(rp), intent(in ), dimension(2       ) :: lmin,lmax
    real(rp), intent(in ), dimension(lo(1)-1:) :: x1c,x1f
    real(rp), intent(in ), dimension(lo(2)-1:) :: x2c,x2f
    real(rp), intent(in ) :: uref,lref,visc ! target bulk velocity
    real(rp), intent(out), dimension(lo(1)-1:,lo(2)-1:) :: veln,velt1,velt2
    integer, dimension(2) :: q
    integer :: i1,i2
    real(rp) :: lambda,reb,x0
    !
    select case(inflow_type)
    case(INFLOW_POISEUILLE)
      q(:) = 1
      where(periods(:) /= 0) q(:) = 0
      do i2 = lo(2)-1,hi(2)+1
        do i1 = lo(1)-1,hi(1)+1
          veln(i1,i2) = uref*poiseuille_square(x1c(i1),lmin(1),lmax(1))**q(1) * &
                             poiseuille_square(x2c(i2),lmin(2),lmax(2))**q(2)
          velt1(i1,i2) = 0._rp
          velt2(i1,i2) = 0._rp
        end do
      end do
    case(INFLOW_KOVASZNAY)
      reb = uref*lref/visc
      lambda = reb/2-sqrt(reb**2/4+4*pi**2)
      x0 = -0.5
      do i2 = lo(2)-1,hi(2)+1
        do i1 = lo(1)-1,hi(1)+1
          veln( i1,i2) = 1._rp-            exp(lambda*x0)*cos(2*pi*x2c(i2))
          velt1(i1,i2) = 0._rp
          velt2(i1,i2) = lambda/(2._rp*pi)*exp(lambda*x0)*sin(2*pi*x2f(i2))
        end do
      end do
    case default
      q(:) = 1
      where(periods(:) /= 0) q(:) = 0
      do i2 = lo(2)-1,hi(2)+1
        do i1 = lo(1)-1,hi(1)+1
          veln(i1,i2) = 0._rp
          velt1(i1,i2) = 0._rp
          velt2(i1,i2) = 0._rp
        end do
      end do
    end select
  end subroutine init_inflow
  !
  subroutine log_profile(lo,hi,zc,l,uref,lref,visc,p)
    implicit none
    integer , intent(in)   :: lo,hi
    real(rp), intent(in), dimension(lo-1:) :: zc
    real(rp), intent(in)   :: l,uref,lref,visc
    real(rp), intent(out), dimension(lo:) :: p
    integer :: k
    real(rp) :: z,reb,retau ! z/lz and bulk Reynolds number
    reb = lref*uref/visc
    retau = 0.09_rp*reb**(0.88_rp) ! from Pope's book
    do k=lo,hi
      z    = zc(k)/l
      if(z>0.5_rp) z = 1._rp-z
      z    = zc(k)*2._rp*retau
      p(k) = 2.5_rp*log(z) + 5.5_rp
      if (z<=11.6_rp) p(k)=z
    end do
  end subroutine log_profile
  !
  ! functions to initialize the streamwise vortex pair
  ! (explained above)
  !
  function fz(zc)
  real(rp), intent(in) :: zc
  real(rp) :: fz
    fz = ((1._rp-zc**2)**2)
  end function
  !
  function dfz(zc)
  real(rp), intent(in) :: zc
  real(rp) :: dfz
    dfz = -4._rp*zc*(1._rp-zc**2)
  end function
  !
  function gxy(xc,yc)
  real(rp), intent(in) :: xc,yc
  real(rp) :: gxy
    gxy = yc*exp(-4._rp*(4._rp*xc**2+yc**2))
  end function
  !
  function dgxy(xc,yc)
  real(rp), intent(in) :: xc,yc
  real(rp) :: dgxy
    dgxy = exp(-4._rp*(4._rp*xc**2+yc**2))*(1._rp-8._rp*yc**2)
  end function
end module mod_initflow


*###[ SPENCE:
	DOUBLE COMPLEX function SPENCE(z)
***#[*comment:***********************************************************
*									*
*	Interface to the FF dilogarithms compatible with the FormF 	*
*	SPENCE function.  All error propagation is lost and the terms	*
*	pi^2/12 are added.						*
*									*
*	Input:	z	complex		cannot lie on the real axis for	*
*					Re(z)>1				*
*	Output:	SPENCE	complex		Sp(z) = Li2(z) = \sum z^n/n^2	*
*					= \int_0^z log(1-x)/x dx	*
*	Calls: 	ffzzdl							*
*									*
***#]*comment:***********************************************************
*  #[ declarations:
	implicit none
*
*	arguments
*
	DOUBLE COMPLEX z
*
*	local variables
*
	integer init,ipi12,ier
	DOUBLE COMPLEX zdilog,zdum
	save init
*
*	common blocks
*
	include 'ff.h'
*
*  #] declarations:
*  #[ initialisations:
	data init /0/
	if ( init .eq. 0 ) then
	    init = 1
	    call ffini
	endif
*  #] initialisations:
*  #[ work:
	ier = 0
	call ffzzdl(zdilog,ipi12,zdum,z,ier)
	SPENCE = zdilog + ipi12*pi12
*  #] work:
*###] SPENCE:
	end


program hmipt
  USE DMFT_IPT
  USE SCIFOR
  USE DMFT_TOOLS
  implicit none
  integer                :: i,iloop,L
  logical                :: converged
  complex(8)             :: zeta
  real(8)                :: n,wband,wmix,dt,tmax,dw
  real(8),allocatable    :: wr(:),time(:)
  complex(8),allocatable :: sigma(:),fg(:),fg0(:),sold(:),sigt(:),sigw(:)

  call parse_input_variable(L,"L",'inputIPT.in',default=10000)
  call parse_input_variable(Wband,'wband','inputIPT.in',default=1d0)
  call parse_input_variable(wmix,'wmix','inputIPT.in',default=1d0)
  call read_input("inputIPT.in")
  allocate(fg(L),sigma(L),fg0(L),wr(L),sold(L))

  wr=linspace(-wmax,wmax,L,mesh=dw)
  dt= pi/wmax
  tmax=dt*L/2
  allocate(time(L),sigt(L))
  time = linspace(-tmax,tmax-dt,L)

  sigma=zero ; iloop=0 ; converged=.false.
  do while (.not.converged)
     iloop=iloop+1
     write(*,"(A,i5)",advance="no")"DMFT-loop",iloop
     do i=1,L
        zeta  = cmplx(wr(i),eps) - sigma(i)
        fg(i) = gfbether(wr(i),zeta,wband)
     enddo
     fg0 = one/(one/fg + sigma)
     sold = sigma
     sigma= ipt_solve_real(fg0,wr)
     sigt = f_fft_rw2rt(sigma)*dw/pi2
     sigw = f_fft_rt2rw(sigt)*dt
     call splot("Sigma_t.ipt",time,sigt)
     call splot("Sigma_w.ipt",wr,sigw)
     sigma = wmix*sigma + (1.d0-wmix)*sold
     converged=check_convergence(sigma,dmft_error,nsuccess,nloop)
  enddo
  call splot("G_realw.ipt",wr,-dimag(fg)/pi,dreal(fg))
  call splot("G0_realw.ipt",wr,fg0)
  call splot("Sigma_realw.ipt",wr,sigma)

end program hmipt


module eigen
    use types
    use utils !assert
    use callbacks
    implicit none
    private
    public lanczos,symetric_eigenvalues,matvecmul,LanczosTimeEvolution

contains

    subroutine lanczos(Dim, StartVec, Evaluator, EigenValue,EigenVector)
        !direct algorithm to find the lowest eigenvalue of a system
        ! http://www.netlib.org/utk/people/JackDongarra/etemplates/node103.html

        integer,intent(in) :: Dim
        real(dp),intent(in) :: StartVec(Dim)
        class(VectorFunctor) :: Evaluator
        real(dp),intent(out) :: EigenVector(Dim)
        real(dp),intent(out) :: EigenValue
        real(dp),allocatable :: Tj(:,:), RitzValues(:),RitzVectors(:,:)
        real(dp),allocatable :: RitzNorms(:)

        !private variables
        real(dp) :: V(Dim,Dim) ! the matrix V that holds the orthonormal basis {v,Av,A^2v,...}
        real(dp) :: r(Dim) ! aux vector for computation
        integer :: j
        real(dp) :: Beta(1:Dim),Alpha(1:Dim) !banded elements of the matrix

        real(dp),parameter :: eps = tiny(1.0_dp), thres = sqrt(eps)
        real(dp) :: NormA !aproximate norm of matrix A

        Beta(1) =  norm2(StartVec)

        V(:,1) = StartVec /Beta(1)

        do j=1,Dim - 1

            r = Evaluator%eval(V(:,j))

            alpha(j) = DOT_PRODUCT(V(:,j),r)
            if (j > 1)r = r - V(:,j-1)*Beta(j)
            r =  r - V(:,j)*alpha(j)

            Beta(j + 1) = NORM2(r)
            V(:,j+1) = r /Beta(j+1)

            ! Reorthogonalize
            ! Here we do full reorthogonalization
            ! http://www.netlib.org/utk/people/JackDongarra/etemplates/node109.html
            call FullReorthogonalization(j,V)
            if(j > 2) then

                call ComputeRitzNorms(j)
                if( RitzNorms(1) < 1e-5) then
                     EigenValue = RitzValues(1)
                     EigenVector = matvecmul(Dim,j,V(:,1:j),RitzVectors(:,1) )

                     print *,"Iters" , j
                    return
                endif
            endif
            ! call SelectiveReorthogonalization(j)

        enddo

        r = Evaluator%eval(V(:,Dim))
        alpha(Dim) =   dot_product(r,V(:,Dim))
        call ComputeRitzNorms(Dim)

        EigenValue = RitzValues(1)

    contains

            subroutine FullReorthogonalization(j, V)
                integer,intent(in) :: j
                real(dp),intent(inout) :: V(:,:)
                real(dp),dimension(j,size(V,1)) :: TransposedV
                real(dp),dimension(size(V,1)) :: x
                real(dp),dimension(j) :: h
                real(dp) :: Norm
                integer :: i,k
                ! print *, "Starting full reorthonormalization"
                x = V(:,j+1)
                do i =1,j
                    do k = 1,size(V,1)
                        TransposedV(i,k) = V(k,i)
                    enddo
                enddo
                Norm = norm2(x)
                ! compute reortogonalization coeff
                h = matmul(TransposedV,x)
                V(:,j+1) = x - matmul(V(:,1:j),h)
            end subroutine FullReorthogonalization

            subroutine SelectiveReorthogonalization(j)
                ! performs selective reortonormalization using
                ! a recurrence relation described here
                ! http://www.netlib.org/utk/people/JackDongarra/etemplates/node110.html
                !
                integer,intent(in) :: j
                integer :: i
                integer :: k !loop variable
                real(dp),save,dimension(:), allocatable :: WCurrent, WNext, WPrev
                logical, save :: FirstCall = .true.
                ! print *, "Starting selective reorthogonalization"

                if (FirstCall) then
                    WPrev = [0.0_dp, 0.0_dp]
                    WCurrent = [0.0_dp, eps, eps]
                    allocate(WNext(4))
                    FirstCall = .false.
                endif

                ! print *, "Calling with j ",j
                !compute WNext using WCurrent && WPrev
                WNext(1) = 0
                do k = 2, j
                    WNext(k) = beta(k)* WCurrent(k+1) + (alpha(k)- alpha(j)) * WCurrent(k) + &
                    beta(k-1)*WCurrent(k-1) - beta(j-1) * WPrev(k)
                enddo

                call UpdateNormA(j)
                WNext(2:j) = (WNext(2:j) + sign([(1.0_dp,i=2,j)],WNext(2:j))* eps * NormA)/Beta( j + 1)
                WNext(j + 1 ) = NormA * eps
                WNext(j + 2 ) = eps

                if (maxval(WNext) > thres) then
                    ! reorthogonalize
                    ! print *, "Reorthogonalize"
                    call FullReorthogonalization(j,V)

                endif

                WPrev = WCurrent
                WCurrent = WNext
                deallocate(WNext)
                allocate(WNext(j+3))

            end subroutine SelectiveReorthogonalization

            subroutine UpdateNormA(j)
                ! Computes an aproximation of the norm of A
                ! using the norm of T
                integer :: j
                ! if (j == 1) then
                !     NormA = abs(alpha(1) + beta(2))
                ! else
                !     NormA = max(NormA, abs(beta(j) + alpha(j) + beta(j+1)))
                ! endif
            end subroutine UpdateNormA

            subroutine ComputeRitzNorms(j)
                integer,intent(in) :: j
                integer :: i
                if (allocated(Tj)) deallocate(Tj)
                if(allocated(RitzValues)) deallocate(RitzValues)
                if(allocated(RitzVectors)) deallocate(RitzVectors)
                allocate(Tj(j,j))
                Tj = 0

                allocate(RitzValues(j))
                allocate(RitzVectors(j,j))
                ! fill T matrix

                do i=1,j
                    Tj(i,i) = alpha(i)
                    if (i < j ) Tj(i,i + 1) = beta(i + 1)
                enddo

                ! solve eigen problem
                call symetric_eigen(j,Tj,RitzValues,RitzVectors)

                if(allocated(RitzNorms)) deallocate(RitzNorms)
                allocate(RitzNorms(j))

                RitzNorms = abs(Beta(1:j)*RitzVectors(j,:))
                ! print *, "RitzNorms", RitzNorms

            end subroutine ComputeRitzNorms

    end subroutine lanczos

    subroutine symetric_eigenvalues(Dim,A, Values)
        integer :: Dim
        real(dp),intent(in) :: A(Dim,Dim)
        real(dp) :: A_(Dim,Dim)
        real(dp),intent(out) :: Values(Dim)
        integer,parameter :: LWORK_MAX = 1000
        real(dp) :: Work(LWORK_MAX)
        integer :: lWork, info
        A_ = A
        !query optimal lwork
        lWork = -1

        call dsyev("N", "L",Dim,A_,Dim,Values,Work,lWork,info)
        lWork = min(LWORK_MAX, int(Work(1)))

        !solve eigenproblem
        call dsyev("N", "U", Dim, A_, Dim, Values,Work,lWork,info)

        call assert(info == 0)!, "symetric_eigenvalues: diagonalization failed")

    end subroutine symetric_eigenvalues

    subroutine symetric_eigen(Dim,A, Values,Vectors)
        ! symetric_eigen solves the eigebproblem for
        ! a symetric real matrix A using LAPACK dsyev
        ! the eignevalues are stored in Values
        ! and the eigenvectors are stored in Vectors
        integer :: Dim
        real(dp) :: A(Dim,Dim)
        real(dp) :: A_(Dim,Dim)
        real(dp),intent(out) :: Values(Dim), Vectors(Dim,Dim)
        integer,parameter :: LWORK_MAX = 1000
        real(dp) :: Work(LWORK_MAX)
        integer :: lWork, info
        integer :: i

        A_ = A
        !query optimal lwork
        lWork = -1
        call dsyev("V", "L",Dim,A_,Dim,Values,Work,lWork,info)
        lWork = min(LWORK_MAX, int(Work(1)))

        !solve eigenproblem
        call DSYEV("V", "U", Dim, A_, Dim, Values,Work,lWork,info)

        call assert(info == 0)!, "symetric_eigenvalues: diagonalization failed")

        Vectors = A_

    end subroutine symetric_eigen

    function matvecmul(M,N,A,v) result(rs)
        implicit none
        integer,intent(in) :: M,N
        real(dp):: A(M,N)
        real(dp):: v(N)
        real(dp):: rs(M)
        integer :: i,j

        do i = 1,M
            rs(i) = 0
            do j= 1,N
                rs(i) =  A(i,j)*v(j) + rs(i)
            enddo
        enddo

    end function matvecmul

    subroutine LanczosTimeEvolution(BasisDim,GroundState,InitialState, Time,Evaluator,Evolution)
        integer,intent(in) :: BasisDim
        real(dp),intent(in) :: InitialState(BasisDim),GroundState(BasisDim)
        real(dp),intent(out) :: Evolution(:)
        real(dp),intent(in) :: Time(:)
        class(VectorFunctor) :: Evaluator
        complex(dp),allocatable :: D(:,:)

        !internal variables
        real(dp),allocatable :: Tj(:,:), RitzValues(:),RitzVectors(:,:)
        real(dp),allocatable :: RitzNorms(:)
        real(dp) :: FinalState(BasisDim)
        !private variables
        real(dp) :: V(BasisDim,BasisDim) ! the matrix V that holds the orthonormal basis {v,Av,A^2v,...}
                               !, lanczos basis
        real(dp) :: r(BasisDim) ! aux vector for computation
        integer :: j, k
        real(dp) :: Beta(1:BasisDim),Alpha(1:BasisDim) !banded elements of the matrix

        real(dp),parameter :: eps = tiny(1.0_dp), thres = sqrt(eps)

        Beta(1) =  norm2(InitialState)
        V(:,1) = InitialState /Beta(1)

        do j=1,BasisDim-1

            r = Evaluator%eval(V(:,j))

            alpha(j) = DOT_PRODUCT(V(:,j),r)
            if (j > 1)r = r - V(:,j-1)*Beta(j)
            r =  r - V(:,j)*alpha(j)

            Beta(j + 1) = NORM2(r)
            V(:,j+1) = r /Beta(j+1)

            ! Reorthogonalize
            ! Here we do full reorthogonalization
            ! http://www.netlib.org/utk/people/JackDongarra/etemplates/node109.html
            !call FullReorthogonalization(j,V)
            if(j > 2) then

                call ComputeRitzNorms(j)
                if( RitzNorms(1) < 1e-8) then
                    if(allocated(D))deallocate(D)
                    print *, "Computing  time evolution"
                    print *, "Original basis: ", BasisDim
                    print *, "Lanczos basis: ", j
                    allocate(D(j,j))
                    do k=1,size(Time)
                        ! print *, Time(k)
                        call exponent(j,RitzValues,D,Time(k))

                        D= matmul(RitzVectors,D)
                        D = matmul(D, transpose(RitzVectors))
                        V =  matmul(V(:,1:j),D)
                        FinalState = V(:,1)
                        Evolution(k) = DOT_PRODUCT(FinalState,GroundState)
                    enddo
                    print *, "Computed evolution"

                    exit

                endif
            endif
            ! call SelectiveReorthogonalization(j)

        enddo
        contains
        !give the eigenvalues, returns
        !the exponential matrix
        subroutine exponent(N,V,A,Time)
            integer,intent(in) :: N
            complex(dp),intent(out) :: A(N,N)
            real(dp),intent(in) :: V(N),Time

            integer :: i
            A = 0.0_dp
            do i=1,N
                A(i,i) = exp(-  (0.0_dp,1.0_dp)*V(i)*Time)
            enddo
        end subroutine exponent
         subroutine ComputeRitzNorms(j)
                integer,intent(in) :: j
                integer :: i
                if (allocated(Tj)) deallocate(Tj)
                if(allocated(RitzValues)) deallocate(RitzValues)
                if(allocated(RitzVectors)) deallocate(RitzVectors)
                allocate(Tj(j,j))
                Tj = 0

                allocate(RitzValues(j))
                allocate(RitzVectors(j,j))
                ! fill T matrix

                do i=1,j
                    Tj(i,i) = alpha(i)
                    if (i < j ) Tj(i,i + 1) = beta(i + 1)
                enddo

                ! solve eigen problem
                call symetric_eigen(j,Tj,RitzValues,RitzVectors)

                if(allocated(RitzNorms)) deallocate(RitzNorms)
                allocate(RitzNorms(j))

                RitzNorms = abs(Beta(1:j)*RitzVectors(j,:))
                ! print *, "RitzNorms", RitzNorms

            end subroutine ComputeRitzNorms

    end subroutine LanczosTimeEvolution

end module eigen


      SUBROUTINE MB03AH( SHFT, K, N, AMAP, S, SINV, A, LDA1, LDA2, C1,
     $                   S1, C2, S2 )
C
C     SLICOT RELEASE 5.7.
C
C     Copyright (c) 2002-2020 NICONET e.V.
C
C     PURPOSE
C
C     To compute two Givens rotations (C1,S1) and (C2,S2) such that the
C     orthogonal matrix
C
C               [ Q  0 ]        [  C1  S1  0 ]   [ 1  0   0  ]
C           Z = [      ],  Q := [ -S1  C1  0 ] * [ 0  C2  S2 ],
C               [ 0  I ]        [  0   0   1 ]   [ 0 -S2  C2 ]
C
C     makes the first column of the real Wilkinson double shift
C     polynomial of the product of matrices in periodic upper Hessenberg
C     form, stored in the array A, parallel to the first unit vector.
C     Only the rotation defined by C1 and S1 is used for the real
C     Wilkinson single shift polynomial (see SLICOT Library routines
C     MB03BE or MB03BF). All factors whose exponents differ from that of
C     the Hessenberg factor are assumed nonsingular. The trailing 2-by-2
C     submatrix and the five nonzero elements in the first two columns
C     of the matrix product are evaluated when a double shift is used.
C
C     ARGUMENTS
C
C     Mode Parameters
C
C     SHFT    CHARACTER*1
C             Specifies the number of shifts employed by the shift
C             polynomial, as follows:
C             = 'D':  two shifts (assumes N > 2);
C             = 'S':  one real shift.
C
C     Input/Output Parameters
C
C     K       (input)  INTEGER
C             The number of factors.  K >= 1.
C
C     N       (input)  INTEGER
C             The order of the factors.  N >= 2.
C
C     AMAP    (input)  INTEGER array, dimension (K)
C             The map for accessing the factors, i.e., if AMAP(I) = J,
C             then the factor A_I is stored at the J-th position in A.
C             AMAP(K) is the pointer to the Hessenberg matrix.
C             Before calling this routine, AMAP returned by SLICOT
C             Library routine MB03BA should be modified as follows:
C             J = AMAP(1), AMAP(I) = AMAP(I+1), I = 1:K-1, AMAP(K) = J.
C
C     S       (input)  INTEGER array, dimension (K)
C             The signature array. Each entry of S must be 1 or -1.
C
C     SINV    (input)  INTEGER
C             Signature multiplier. Entries of S are virtually
C             multiplied by SINV.
C
C     A       (input)  DOUBLE PRECISION array, dimension (LDA1,LDA2,K)
C             The leading N-by-N-by-K part of this array must contain
C             the product (implicitly represented by its K factors)
C             in periodic upper Hessenberg form.
C
C     LDA1    INTEGER
C             The first leading dimension of the array A.  LDA1 >= N.
C
C     LDA2    INTEGER
C             The second leading dimension of the array A.  LDA2 >= N.
C
C     C1      (output)  DOUBLE PRECISION
C     S1      (output)  DOUBLE PRECISION
C             On exit, C1 and S1 contain the parameters for the first
C             Givens rotation.
C
C     C2      (output)  DOUBLE PRECISION
C     S2      (output)  DOUBLE PRECISION
C             On exit, if SHFT = 'D' and N > 2, C2 and S2 contain the
C             parameters for the second Givens rotation. Otherwise,
C             C2 = 1, S2 = 0.
C
C     METHOD
C
C     The necessary elements of the real Wilkinson double/single shift
C     polynomial are computed, and suitable Givens rotations are found.
C     For numerical reasons, this routine should be called when
C     convergence difficulties are encountered. For a double shift, if
C     there are two real eigenvalues of the trailing 2-by-2 part of the
C     product, both shifts are chosen equal to the eigenvalue with
C     minimum modulus. The trailing element of the product is used as a
C     single shift. If SINV is negative, the shift(s) correspond to the
C     reciprocals of the eigenvalues of the product, as required by the
C     SLICOT Library routine MB03BD.
C
C     CONTRIBUTOR
C
C     V. Sima, Research Institute for Informatics, Bucharest, Romania,
C     Aug. 2019, Dec. 2020.
C
C     REVISIONS
C
C     V. Sima, Sep. 2019, Dec. 2019, Jan. 2020, Feb. 2020, Mar. 2020.
C
C     KEYWORDS
C
C     Eigenvalues, QZ algorithm, periodic QZ algorithm, orthogonal
C     transformation.
C
C     ******************************************************************
C
C     .. Parameters ..
      DOUBLE PRECISION  ONE, TWO, ZERO
      PARAMETER         ( ONE = 1.0D0, TWO = 2.0D0, ZERO = 0.0D0 )
C     .. Scalar Arguments ..
      CHARACTER         SHFT
      INTEGER           K, LDA1, LDA2, N, SINV
      DOUBLE PRECISION  C1, C2, S1, S2
C     .. Array Arguments ..
      INTEGER           AMAP(*), S(*)
      DOUBLE PRECISION  A(LDA1,LDA2,*)
C     .. Local Scalars ..
      LOGICAL           SGLE
      INTEGER           I, IC, IND, J, M, MM
      DOUBLE PRECISION  E1, E2, P1, P2, P3, PR, SCL, SM, T
C     .. Local Arrays ..
      INTEGER           IP(3), JP(3)
      DOUBLE PRECISION  DWORK(9), WI(2), WR(2), Y(9), Z(2,2)
C     .. External Functions ..
      LOGICAL           LSAME
      EXTERNAL          LSAME
C     .. External Subroutines ..
      EXTERNAL          DCOPY, DGESC2, DGETC2, DLACPY, DLANV2, DLARTG,
     $                  DLASET, DTRMV
C     .. Intrinsic Functions ..
      INTRINSIC         ABS, MIN
C
C     .. Executable Statements ..
C
C     For efficiency reasons, the parameters are not checked for errors.
C
C     Evaluate the needed part of the matrix product.
C
      SGLE = LSAME( SHFT, 'S' ) .OR. N.EQ.2
C
      M  = MIN( N, 3 )
      MM = M*M
C
      CALL DLASET( 'Full', M, M, ZERO, ONE, DWORK, M )
C
      DO 30  J = K - 1, 1, -1
         I = AMAP(J)
         IF ( S(I).EQ.SINV ) THEN
C
            DO 10  IC = 1, MM, M
               CALL DTRMV( 'Upper', 'NoTran', 'NonUnit', M, A(1,1,I),
     $                     LDA1, DWORK(IC), 1 )
   10       CONTINUE
C
         ELSE
C
C           Complete pivoting is used for triangular factors whose
C           exponents differ from SINV. It is assumed that no overflow
C           could appear when solving linear systems, hence SCL = 1.
C
            CALL DLACPY( 'Upper', M, M, A(1,1,I), LDA1, Y, M )
            CALL DLASET( 'Lower', M-1, M-1, ZERO, ZERO, Y(2), M )
            CALL DGETC2( M, Y, M, IP, JP, IND )
C
            DO 20  IC = 1, MM, M
               CALL DGESC2( M, Y, M, DWORK(IC), IP, JP, SCL )
   20       CONTINUE
C
         END IF
   30 CONTINUE
C
C     Compute successively in Y(1:3) the nonzero elements of the first
C     two columns of the product, and save the results.
C
      I = AMAP(K)
      CALL DCOPY( 2, A(1,1,I), 1, Y, 1 )
      CALL DTRMV( 'Upper', 'NoTran', 'NonUnit', 2, DWORK, M, Y, 1 )
      E1 = Y(1)
      E2 = Y(2)
C
      IF ( SGLE ) THEN
         P1 = ONE
C
C        Compute the (N,N) element of the product and the rotations.
C        This element is used as shift.
C
         DO 40  J = 1, K
            I = AMAP(J)
            IF ( S(I).EQ.SINV ) THEN
               P1 = P1*A(N,N,I)
            ELSE
               P1 = P1/A(N,N,I)
            END IF
   40    CONTINUE
C
         CALL DLARTG( E1 - P1, E2, C1, S1, E1 )
         C2 = ONE
         S2 = ZERO
C
      ELSE
C
         CALL DCOPY( M, A(1,2,I), 1, Y, 1 )
         CALL DTRMV( 'Upper', 'NoTran', 'NonUnit', M, DWORK, M, Y, 1)
         P1 = Y(1)
         P2 = Y(2)
         P3 = Y(3)
C
C        Compute the bottom 2-by-2 part using complete pivoting.
C
         CALL DLASET( 'Full', 2, 2, ZERO, ONE, Z, 2 )
C
         M = N - 1
C
         DO 50  J = K - 1, 1, -1
            I = AMAP(J)
            IF ( S(I).EQ.SINV ) THEN
               Z(1,1) = A(M,M,I)*Z(1,1)
               Z(1,2) = A(M,M,I)*Z(1,2) + A(M,N,I)*Z(2,2)
               Z(2,2) = A(N,N,I)*Z(2,2)
            ELSE
               Y(1) = A(M,M,I)
               Y(2) = ZERO
               CALL DCOPY(  2, A(M,N,I), 1, Y(3), 1 )
               CALL DGETC2( 2, Y, 2, IP, JP, IND )
               CALL DGESC2( 2, Y, 2, Z, IP, JP, SCL )
               CALL DGESC2( 2, Y, 2, Z(1,2), IP, JP, SCL )
            END IF
   50    CONTINUE
C
         I = AMAP(K)
         T      = Z(1,1)*A(M,M,I) + Z(1,2)*A(N,M,I)
         Z(1,2) = Z(1,1)*A(M,N,I) + Z(1,2)*A(N,N,I)
         Z(1,1) = T
         Z(2,1) = Z(2,2)*A(N,M,I)
         Z(2,2) = Z(2,2)*A(N,N,I)
C
C        Compute the eigenvalues of the bottom 2-by-2 part.
C        If there are two real eigenvalues, both shifts are chosen equal
C        to the eigenvalue with minimum modulus. Only the sum and
C        product of the shifts are needed.
C
         CALL DLANV2( Z(1,1), Z(1,2), Z(2,1), Z(2,2), WR(1), WI(1),
     $                WR(2), WI(2), C1, S1 )
         IF ( WI(1).EQ.ZERO ) THEN
            IF ( ABS( WR(1) ).LT.ABS( WR(2) ) ) THEN
               T = WR(1)
            ELSE
               T = WR(2)
            END IF
            SM = TWO*T
            PR = T**2
         ELSE
            SM = TWO*WR(1)
            PR = WR(1)**2 + WI(1)**2
         END IF
C
C        Compute a multiple of the first column of the real Wilkinson
C        double shift polynomial, having only three nonzero elements.
C
         P1 = P1 + ( ( E1 - SM )*E1 + PR )/E2
         P2 = P2 + E1 - SM
C
C        Compute the rotations to annihilate P2 and P3.
C
         CALL DLARTG( P2, P3, C2, S2, E1 )
         CALL DLARTG( P1, E1, C1, S1, E2 )
      END IF
C
      RETURN
C *** Last line of MB03AH ***
      END


program energy
  real, parameter :: mass = 3.00, gravity = 9.81, height = 4.20
  real :: potential_energy

  real, parameter :: velocity = 4.00
  real :: kinetic_energy, total_energy

  potential_energy = mass * gravity * height
  kinetic_energy = 0.5 * mass * square(velocity)

  total_energy = potential_energy + kinetic_energy

  contains

  real function square(x)
    real x
    square = x * x
  end function square
end program energy


module vd1d_mms_mod

    implicit none

    real, parameter :: pi = 3.1415926535897932
    real, private :: rho0,rho1,k1,k2,w0,D,r01,k12,k21

    !functions
    PUBLIC :: vd1d_mms_init
    PUBLIC :: vd2d_mms_u
    PUBLIC :: vd2d_mms_rho
    PUBLIC :: vd2d_mms_z
    PUBLIC :: vd2d_mms_z_src
    PUBLIC :: vd2d_mms_v
    PUBLIC :: vd2d_mms_p
    PUBLIC :: vd2d_mms_u_src
    PUBLIC :: vd2d_mms_v_src
    PUBLIC :: sech
!————————————————————————————————————————————————————————————————————————————————————————
contains
    !————————————————————————————————————————————————————————————————————————————————————————
    subroutine vd1d_mms_init (rho_0,rho_1,k_1,k_2,width,diff)
        implicit none
        real,intent (in) :: rho_0,rho_1,k_1,k_2,width,diff
        k1 = k_1 ! inverse time-scale for amplitude decay
        k2 = k_2 ! inverse time-scale for width decay
        rho0 = rho_0! density at z = 0
        rho1 = rho_1! density at z = 1
        w0 = width ! initial width
        D = diff ! z-diffusivity
        r01 = rho0-rho1
        k12 = k1-k2
        k21 = k2-k1
    end subroutine vd1d_mms_init
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_u (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_u
        real :: u1

        u1 = exp(w0 * exp (-k2 * t) * abs(x)) + 0.0*y
        vd2d_mms_u=(2.0 * k2 * abs(x) * r01 * exp(- k1 * t) * u1/(u1**2 + 1.0)+r01 * k12 * &
            exp(-k12 * t)/w0 * (2.0 * atan(u1) - pi/2.0)) / vd2d_mms_rho(x,y,t)
    end function vd2d_mms_u
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_rho (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_rho
        real :: z

        z = vd2d_mms_z (abs(x),y,t) + 0.0*y
        vd2d_mms_rho = 1.0/(z/rho1+(1.0-z)/rho0)
    end function vd2d_mms_rho
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_z (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_z
        real :: z1,z2

        z1 = exp (-k1 * t) + 0.0*y
        z2 = cosh (w0 * exp (-k2 * t) * abs(x))
        vd2d_mms_z=(z1-z2)/(z1 * (1.0-rho0/rho1)-z2)
    end function vd2d_mms_z
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_z_src (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_z_src
        real :: s0,s1,s2,s3,s4,s5,s6,s7,s8,s9,s10, &
            s11,s12,s13,s14,s15,s16,s17,s18,s19,s20

        !x = abs(x) ! for double mesh

        s0 = exp (-k1 * t) + 0.0*y
        s1 = exp (-k2 * t)
        s2 = sech (w0 * s1 * abs(x))
        s3 = tanh (w0 * s1 * abs(x))
        s4 = exp (w0 * s1 * abs(x))
        s5 = s0/s1
        s6 = s0 * s2
        s7 = rho1 + r01 * s6! rho
        s8 = 2.0 * atan (s4)-1.0/2.0 * pi
        s9 = s6 * s3 * w0
        s10=-s9 * s1 * rho1 + s9 * s1 * rho0
        s11 = 1.0-s3 ** 2
        s12 = s4 ** 2 + 1.0
        s13=-1.0 + s6
        s14=-k1 * s6 + s9 * k2 * s1 * abs(x)
        s15 = s6 * s3 ** 2 * w0 ** 2 * s1 ** 2
        s16 = 2.0 * k2 * r01 * s0 * s4/s12
        s17 = s6 * s11 * w0 ** 2 * s1 ** 2
        s18 = s16 * abs(x) + r01 * k12 * s5/w0 * s8
        s19 = s9 * s1 * s7-s13 * s10
        s20 = 2.0 * s4 * s1 * (s16 * abs(x) * s4 * w0-r01 * k12 * s5)/s12
        vd2d_mms_z_src=&
            -s14 * rho1-(s16 + s16 * abs(x) * w0 * s1-s20) * s13 * rho1/s7 + s18 * rho1 * s19/s7 ** 2 &
            -D * rho1 * (-s15 * s7 ** 2 + s17 * s7 ** 2 + 2.0 * s9 * s1 * s10 * s7-2.0 * s13 * s10 ** 2 &
            -s13 * r01 * s7 * s17 + s13 * r01 * s7 * s15)/s7 ** 3

    end function vd2d_mms_z_src
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_v (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_v
        vd2d_mms_v=x*0.0 + y*0.0 + t*0.0
    end function vd2d_mms_v
    !!————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_p (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_p
        real :: xl, Pl !, xr, dummy
        integer          :: i, ReadWriteError
        character        :: PressureFile*250, CurrentPath*250,dummmy_str*10

        xl = 0.0
        Pl = 0.0 + 0.0*y
        vd2d_mms_p = 0.0

        PressureFile="pressure/xxxx.dat"
        PressureFile=trim(adjustl(PressureFile))
        write(PressureFile(10:13),'(I4.4)') int(t*800)
        !call getcwd(CurrentPath) !current path
        CurrentPath="/home/palkine/Fortran/Projects/combustion/T-FlowS-comb-Problem-3-CalcPS/Problem1/"

        open(unit=998, file=trim(adjustl(CurrentPath))//PressureFile, &
            iostat=ReadWriteError, access="sequential", status="old", action="read")
        !write(*,*) "Opened file is:", trim(adjustl(CurrentPath))//PressureFile
        if ( ReadWriteError /= 0 ) stop "Error opening pressure file"

        if ( abs(x) > 0.0 ) then
            read(998,*) dummmy_str ! dummy

            !if ( x > 0 ) then
            do i = 1, nint(abs(x)*4096/2)+1
                read(998,*) xl,Pl
            end do
            !else
            !  do i = 1, int(-x*4096/2)
            !    read(998,*) xl,Pl
            !  end do
            !end if
            !read(998,*) xr,dummy
            close(998)

            !write(*,*) xl-2/4096, x, xr-2/4096

            vd2d_mms_p = Pl
        end if

        if (isnan(Pl)) then
            write (*,*) "x=", x, "xl=", x, "P=", Pl
            stop '"Pl" is a NaN'
        end if

    end function vd2d_mms_p
    !————————————————————————————————————————————————————————————————————————————————————————

    function vd2d_mms_u_src (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_u_src
        vd2d_mms_u_src = x*0.0 + y*0.0 + t*0.0
    end function vd2d_mms_u_src
    !————————————————————————————————————————————————————————————————————————————————————————
    function vd2d_mms_v_src (x,y,t)
        implicit none
        real,intent (in) :: x,y,t
        real :: vd2d_mms_v_src
        vd2d_mms_v_src = x*0.0 + y*0.0 + t*0.0
    end function vd2d_mms_v_src
    !————————————————————————————————————————————————————————————————————————————————————————

    function sech (x)
        implicit none
        real,intent (in) :: x
        real :: sech

        sech = 1.0/cosh (abs(x))
    end function sech
end module vd1d_mms_mod

!function vd2d_mms_p_old (x,y,t)
!implicit none
!real,intent (in) :: x,y,t
!real :: vd2d_mms_p
!real,allocatable :: xf(:), Pf(:)
!integer          :: i, ReadWriteError
!character        :: PressureFile*250, CurrentPath*250,dummmy_str*10

!allocate(xf(4098)); xf = 0.0
!allocate(Pf(4098)); Pf = 0.0
!vd2d_mms_p = 0.0

!PressureFile="pressure/xxxx.dat"
!PressureFile=trim(adjustl(PressureFile))
!write(PressureFile(10:13),'(I4.4)') int(t*800)
!!call getcwd(CurrentPath) !current path
!CurrentPath="/home/palkine/Fortran/Projects/combustion/T-FlowS-comb-Problem-3-CalcPS/Problem1/"

!open(unit=998, file=trim(adjustl(CurrentPath))//PressureFile, iostat=ReadWriteError, status="old", action="read")
!!write(*,*) "Opened file is:", trim(adjustl(CurrentPath))//PressureFile
!if ( ReadWriteError /= 0 ) stop "Error opening pressure file"

!  read(998,*) dummmy_str ! dummy

!  do i = 1, 4098 !int(x*4096/2)+1
!    read(998,*) xf(i),Pf(i)
!  end do
!  close(998)

!if ( x == 0 ) then
!  vd2d_mms_p = 0.0*y !+ xr*0  + dummy*0
!  return
!else if ( x == 2 ) then
!  vd2d_mms_p = Pf(4098)
!  return
!end if

!do i = 2, 4097
!  if ( xf(i)-2/4096 <= x .and. x <= xf(i+1)-2/4096 ) then !caught the cell
!    vd2d_mms_p = Pf(i)
!    return
!  end if
!end do

!deallocate(xf)
!deallocate(Pf)

!end function vd2d_mms_p_old


subroutine diffuse_scalar2D (field,fluxb,fluxt,tkh,rho,rhow,flux)

use grid
use params, only: docolumn,dowallx,dosgs
use sgs,only: grdf_x,grdf_z
implicit none

! input
real field(dimx1_s:dimx2_s, dimy1_s:dimy2_s, nzm)	! scalar
real tkh(0:nxp1, 1-YES3D:nyp1, nzm)	! eddy conductivity
real fluxb(nx,ny)		! bottom flux
real fluxt(nx,ny)		! top flux
real rho(nzm)
real rhow(nz)
real flux(nz)

! local
real flx(0:nx,1,0:nzm)
real dfdt(nx,ny,nzm)
real rdx2,rdz2,rdz,rdx5,rdz5,tmp
real dxz,dzx,tkx,tkz,rhoi
integer i,j,k,ib,ic,kc,kb

if(.not.dosgs.and..not.docolumn) return

rdx2=1./(dx*dx)
rdz2=1./(dz*dz)
rdz=1./dz
dxz=dx/dz
dzx=dz/dx

j=1

dfdt(:,:,:)=0.

if(dowallx) then

  if(mod(rank,nsubdomains_x).eq.0) then
    do k=1,nzm
         field(0,j,k) = field(1,j,k)
    end do
  end if
  if(mod(rank,nsubdomains_x).eq.nsubdomains_x-1) then
    do k=1,nzm
         field(nx+1,j,k) = field(nx,j,k)
    end do
  end if

end if

if(.not.docolumn) then

do k=1,nzm

  rdx5=0.5*rdx2  *grdf_x(k)

  do i=0,nx
    ic=i+1
    tkx=rdx5*(tkh(i,j,k)+tkh(ic,j,k))
    flx(i,j,k)=-tkx*(field(ic,j,k)-field(i,j,k))
  end do
  do i=1,nx
    ib=i-1
    dfdt(i,j,k)=dfdt(i,j,k)-(flx(i,j,k)-flx(ib,j,k))
  end do

end do

end if

flux(1) = 0.
tmp=1./adzw(nz)
do i=1,nx
   flx(i,j,0)=fluxb(i,j)*rdz*rhow(1)
   flx(i,j,nzm)=fluxt(i,j)*rdz*tmp*rhow(nz)
   flux(1) = flux(1) + flx(i,j,0)
end do

do k=1,nzm-1
 kc=k+1
 flux(kc)=0.
 rhoi = rhow(kc)/adzw(kc)
 rdz5=0.5*rdz2 * grdf_z(k)
 do i=1,nx
    tkz=rdz5*(tkh(i,j,k)+tkh(i,j,kc))
    flx(i,j,k)=-tkz*(field(i,j,kc)-field(i,j,k))*rhoi
    flux(kc) = flux(kc) + flx(i,j,k)
 end do
end do

do k=1,nzm
 kb=k-1
 rhoi = 1./(adz(k)*rho(k))
 do i=1,nx
  dfdt(i,j,k)=dtn*(dfdt(i,j,k)-(flx(i,j,k)-flx(i,j,kb))*rhoi)
  field(i,j,k)=field(i,j,k) + dfdt(i,j,k)
 end do
end do

end subroutine diffuse_scalar2D


module update
  use functype
  use readingfile
  implicit none

contains

!updating the objects' data using algorithm
subroutine updatesub()
  implicit none
  INTEGER :: i
  TYPE(acceleration) :: vanha_acc     !to save the acceleration of an object from the current step

  !updating all the objects' data
  do i=1, lkm
    call accsub(objectarray(i), vanha_acc)
    call velsub(objectarray(i), vanha_acc)
    call possub(objectarray(i))
  end do

end subroutine updatesub

!updating one object's acceleration
subroutine accsub(kappale, vanha_acc)
  implicit none
  TYPE(object) , intent(inout) :: kappale         !object to be updated
  type(acceleration), INTENT(OUT) :: vanha_acc   !object's current acceleration
  type(acceleration) :: uusi_acc     ! adding acceleration contributions from other objects
  INTEGER :: i

  vanha_acc=kappale%a
  uusi_acc=acceleration(0,0,0)

  !lasketaan kappaleen uusi kiihtyvyys käymällä objectarray-vektorin läpi
  !calculating the contributions to the object's acceleration from the other objects.
  do i= 1, lkm
    if (i==kappale%nimi) cycle    !skipping the object kappale itself
    uusi_acc=uusi_acc+accfunc(kappale, objectarray(i))
  end do

  kappale%a=uusi_acc         !acceleration update

end subroutine accsub

!laskee kappaleen kiihtyvyys ollessaan kohde-objectin gravitaatiopotentiaalissa
!to calculate the acceleration of the object kappale when in the other one's gravitational potential field.
type(acceleration) function accfunc(kappale, kohde)
  implicit none
  type(object), INTENT(IN) :: kappale, kohde
  type(position) :: r1, r2    !saving the positions of objects kappale and kohde
  real(rk) :: skalaariosa

  r1= kappale%r
  r2=kohde%r

  !Scalar part of Newton's law of gravitation
  skalaariosa = G * kohde%m/(r2.etaisyys.r1)**3

  !calculating acceleration components
  accfunc%ax = skalaariosa* (r2%x-r1%x  )
  accfunc%ay = skalaariosa* (r2%y-r1%y  )
  accfunc%az = skalaariosa* (r2%z-r1%z  )

end function accfunc

!updating the velocity of one object.
subroutine velsub(kappale, vanha_acc)
  implicit none
  type(object) , INTENT(INOUT) :: kappale
  type(acceleration) , INTENT(IN) :: vanha_acc

  kappale%v%vx =kappale%v%vx + 1.0/2.0* (vanha_acc%ax+ kappale%a%ax)*deltaT
  kappale%v%vy =kappale%v%vy + 1.0/2.0* (vanha_acc%ay+ kappale%a%ay)*deltaT
  kappale%v%vz =kappale%v%vz + 1.0/2.0* (vanha_acc%az+ kappale%a%az)*deltaT
end subroutine velsub

  !updating the position of one object
  subroutine possub(kappale )
    implicit none
    TYPE(object), INTENT(inout) :: kappale
    !calculating new position coordinates
    kappale%r%x= kappale%r%x+ kappale%v%vx*deltaT+ (1.0/2.0)*kappale%a%ax*deltaT**2
    kappale%r%y= kappale%r%y+ kappale%v%vy*deltaT+ (1.0/2.0)*kappale%a%ay*deltaT**2
    kappale%r%z= kappale%r%z+ kappale%v%vz*deltaT+ (1.0/2.0)*kappale%a%az*deltaT**2

  end subroutine possub

end module update


      subroutine evintrp2d(xget,yget,x,nx,y,ny,jspline,
     >                   f,icoeff,ixdim,iydim,ict,fval,ier)
C
      implicit NONE
C
C  evaluate a 2d Hybrid interpolant on a rectilinear grid
C
C  this subroutine calls three subroutines:
C     vecin2d_argchk -- error check
C     herm2xy  -- find cell containing (xget,yget)
C     fvintrp2d  -- evaluate interpolant function and (optionally) derivatives
C
C  input arguments:
C  ================
C
      real xget,yget                    ! target of this interpolation
      integer nx,ny                     ! grid sizes
      real x(nx)                        ! ordered x grid
      real y(ny)                        ! ordered y grid
C
      integer :: jspline(2)             ! interpolation method for each
C           dimension: jspline(1) for x; jspline(2) for y
C           =-1: zonal step function; =0: pc lin; =1: Hermite; =2: Spline
C
      integer :: icoeff                 ! no. of coefficients per data point

      integer :: ixdim,iydim            ! dimensioning:
      !  ixdim=nx-1 for zonal step in x; otherwise nx
      !  iydim=ny-1 for zonal step in y; otherwise ny

      real f(icoeff,ixdim,iydim)        ! function data w/coefficients
C
      integer ict(6)                    ! code specifying output desired
C
C  Note on derivatives: for dimensions along which zonal step function
C    interpolation is done, ANY derivative request returns ZERO.
C    For dimensions along which piecewise linear or Hermite interpolation
C    are done, more than one differentiation returns ZERO!
C
C  Derivative controls are the same as for the compact 2d spline evaluation
C  routine (evbicub):
C
C  ict(1)=1 -- return f  (0, don't)
C  ict(2)=1 -- return df/dx  (0, don't)
C  ict(3)=1 -- return df/dy  (0, don't)
C  ict(4)=1 -- return d2f/dx2  (0, don't)
C  ict(5)=1 -- return d2f/dy2  (0, don't)
C  ict(6)=1 -- return d2f/dxdy (0, don't)
c                   the number of non zero values ict(1:6)
c                   determines the number of outputs...
c
c  new dmc December 2005 -- access to higher derivatives (even if not
c  continuous-- but can only go up to 3rd derivatives on any one coordinate.
c     if ict(1)=3 -- want 3rd derivatives
c          ict(2)=1 for d3f/dx3
c          ict(3)=1 for d3f/dx2dy
c          ict(4)=1 for d3f/dxdy2
c          ict(5)=1 for d3f/dy3
c               number of non-zero values ict(2:5) gives no. of outputs
c     if ict(1)=4 -- want 4th derivatives
c          ict(2)=1 for d4f/dx3dy
c          ict(3)=1 for d4f/dx2dy2
c          ict(4)=1 for d4f/dxdy3
c               number of non-zero values ict(2:4) gives no. of outputs
c     if ict(1)=5 -- want 5th derivatives
c          ict(2)=1 for d5f/dx3dy2
c          ict(3)=1 for d5f/dx2dy3
c               number of non-zero values ict(2:3) gives no. of outputs
c     if ict(1)=6 -- want 6th derivatives
c          d6f/dx3dy3 -- one value is returned.
C
C output arguments:
C =================
C
      real fval(6)                      ! output data
      integer ier                       ! error code =0 ==> no error
C
C  fval(1) receives the first output (depends on ict(...) spec)
C  fval(2) receives the second output (depends on ict(...) spec)
C  fval(3) receives the third output (depends on ict(...) spec)
C  fval(4) receives the fourth output (depends on ict(...) spec)
C  fval(5) receives the fifth output (depends on ict(...) spec)
C  fval(6) receives the sixth output (depends on ict(...) spec)
C
C  examples:
C    on input ict = [1,1,1,0,0,1]
C   on output fval= [f,df/dx,df/dy,d2f/dxdy], elements 5 & 6 not referenced.
C
C    on input ict = [1,0,0,0,0,0]
C   on output fval= [f] ... elements 2 -- 6 never referenced.
C
C    on input ict = [0,0,0,1,1,0]
C   on output fval= [d2f/dx2,d2f/dy2] ... elements 3 -- 6 never referenced.
C
C    on input ict = [0,0,1,0,0,0]
C   on output fval= [df/dy] ... elements 2 -- 6 never referenced.
C
C  ier -- completion code:  0 means OK
C-------------------
C  local:
C
      integer, dimension(1) :: i,j                       ! cell indices
C
C  normalized displacement from (x(i),y(j)) corner of cell.
C    xparam=0 @x(i)  xparam=1 @x(i+1)
C    yparam=0 @y(j)  yparam=1 @y(j+1)
C
      real, dimension(1) ::  xparam,yparam
C
C  cell dimensions and
C  inverse cell dimensions hxi = 1/(x(i+1)-x(i)), hyi = 1/(y(j+1)-y(j))
C
      real, dimension(1) ::   hx,hy
      real, dimension(1) ::   hxi,hyi
C
C  0 .le. xparam .le. 1
C  0 .le. yparam .le. 1
C
C  ** the interface is very similar to herm2ev.for; can use herm2xy **
C---------------------------------------------------------------------
C
      call vecin2d_argchk('evintrp2d',jspline,
     >     icoeff,nx,ny,ixdim,iydim,ier)
      if(ier.ne.0) return
c
      call herm2xy(xget,yget,x,nx,y,ny,0,0,
     >     i,j,xparam,yparam,hx,hxi,hy,hyi,ier)
      if(ier.ne.0) return
c
      call fvintrp2d(ict,1,1,
     >     fval,i,j,xparam,yparam,hx,hxi,hy,hyi,
     >     jspline,f,icoeff,ixdim,iydim)
C
      return
      end
C---------------------------------------------------------------------
C  evaluate Hybrid function interpolation -- 2d fcn
C   NO ERROR CHECKING
C
      subroutine fvintrp2d(ict,ivec,ivecd,
     >     fval,ii,jj,xparam,yparam,hx,hxi,hy,hyi,
     >     jspline,fin,icoeff,ixdim,iydim)
C
      implicit NONE
C
      integer ict(6)                    ! requested output control
      integer ivec                      ! vector length
      integer ivecd                     ! vector dimension (1st dim of fval)
C
      integer ii(ivec),jj(ivec)         ! target cells (i,j)
      real xparam(ivec),yparam(ivec)
                          ! normalized displacements from (i,j) corners
C
      real hx(ivec),hy(ivec)            ! grid spacing, and
      real hxi(ivec),hyi(ivec)          ! inverse grid spacing 1/(x(i+1)-x(i))
                                        ! & 1/(y(j+1)-y(j))
C
      integer :: jspline(2)             ! interpolation type control
      integer :: icoeff                 ! coefficient dimension
      integer :: ixdim,iydim            ! x & y dimensions
      real fin(icoeff,ixdim,iydim)      ! data on which to interpolate...
C
      real fval(ivecd,6)                ! output returned
C
C  for detailed description of fin, ict and fval see subroutine
C  evintrp2d comments.  Note ict is not vectorized; the same output
C  is expected to be returned for all input vector data points.
C
C  note that the index inputs ii,jj and parameter inputs
C     xparam,yparam,hx,hxi,hy,hyi are vectorized, and the
C     output array fval has a vector ** 1st dimension ** whose
C     size must be given as a separate argument
C
C  to use this routine in scalar mode, pass in ivec=ivecd=1
C
C---------------
C  local:
C
      integer :: i,j,i1,i2,zonrank,linrank,cubrank
      integer :: imaxx,imaxy,imaxdlin,imaxdcub
      logical :: splin_flag
C
      integer :: inum,indx
      integer :: iderivs(2,6),ict4(4),maxd_lin(4),maxd_cub(4)
C
      integer :: idcub,idlin
      real :: xp,dxlin,dxlini,h,hi,ans1,ans2
      real :: f22(2,2)
C
      real, parameter :: ONE = 1.0d0
C---------------
C
      linrank=0
      cubrank=0
      zonrank=0

      do i=1,2
         if(jspline(i).eq.-1) then
            zonrank = zonrank + 1

         else if(jspline(i).eq.0) then
            linrank = linrank + 1

         else
            splin_flag = jspline(i).eq.2
            cubrank = cubrank + 1
         endif
      enddo
C
      if(cubrank.eq.2) then
         if(splin_flag) then
            !  bicubic spline
            call fvbicub(ict,ivec,ivecd,
     >           fval,ii,jj,xparam,yparam,hx,hxi,hy,hyi,
     >           fin,ixdim,iydim)

         else
            !  bicubic Hermite -- translate derivative code
            call dtrans
            do i=1,inum
               if(maxval(iderivs(1:2,i)).le.1) then
                  ict4=0
                  indx=2*iderivs(2,i) + iderivs(1,i) + 1
                  ict4(indx)=1
                  call herm2fcn(ict4,ivec,ivecd,
     >                 fval(1,i),ii,jj,xparam,yparam,hx,hxi,hy,hyi,
     >                 fin,ixdim,iydim)

               else
                  fval(1:ivec,i) = 0
               endif
            enddo
         endif

         return
      endif

      call dtrans

      if(cubrank.eq.0) then
         if(jspline(1).eq.0) then
            imaxx=1
         else
            imaxx=0
         endif

         if(jspline(2).eq.0) then
            imaxy=1
         else
            imaxy=0
         endif

         do i=1,inum
            if((iderivs(1,i).le.imaxx).and.(iderivs(2,i).le.imaxy)) then
               if(linrank.eq.2) then
                  !  bilinear interpolation
                  ict4=0
                  indx=2*iderivs(2,i) + iderivs(1,i) + 1
                  ict4(indx)=1
                  call pc2fcn(ict4,ivec,ivecd,
     >                 fval(1,i),ii,jj,xparam,yparam,hx,hxi,hy,hyi,
     >                 fin,ixdim,iydim)
               else if(zonrank.eq.2) then
                  !  2d step function
                  do j=1,ivec
                     fval(j,i)=fin(1,ii(j),jj(j))
                  enddo
               else
                  !  linear-zone hybrid
                  if(jspline(1).eq.0) then
                     ! linear in x, zonal in y
                     do j=1,ivec
                        if(iderivs(1,i).eq.0) then
                           fval(j,i)=(ONE-xparam(j))*fin(1,ii(j),jj(j))
     >                          +xparam(j)*fin(1,ii(j)+1,jj(j))
                        else
                           fval(j,i)=(fin(1,ii(j)+1,jj(j)) -
     >                          fin(1,ii(j),jj(j)))*hxi(j)
                        endif
                     enddo
                  else
                     ! linear in y, zonal in x
                     do j=1,ivec
                        if(iderivs(2,i).eq.0) then
                           fval(j,i)=(ONE-yparam(j))*fin(1,ii(j),jj(j))
     >                          +yparam(j)*fin(1,ii(j),jj(j)+1)
                        else
                           fval(j,i)=(fin(1,ii(j),jj(j)+1) -
     >                          fin(1,ii(j),jj(j)))*hyi(j)
                        endif
                     enddo
                  endif
               endif
            else
               fval(1:ivec,i) = 0
            endif
         enddo

         return
      endif

      !  OK -- Hybrid: Hermite or Spline in one dimension
      !        pclin or step in the other dimension

      if(splin_flag) then
         imaxdcub=3
      else
         imaxdcub=1
      endif

      if(linrank.eq.1) then
         imaxdlin=1
      else
         imaxdlin=0
      endif

      do i=1,inum
         if(jspline(1).ge.1) then
            idcub=iderivs(1,i)
            idlin=iderivs(2,i)
         else
            idcub=iderivs(2,i)
            idlin=iderivs(1,i)
         endif

         if((maxd_lin(i).le.imaxdlin).and.
     >        (maxd_cub(i).le.imaxdcub)) then
            if(linrank.eq.1) then
               !  spline or Hermite in 1 dimension; pclin in other.
               if(jspline(1).ge.1) then
                  do j=1,ivec
                     xp=xparam(j)
                     dxlin=yparam(j)
                     dxlini=1.0-dxlin
                     h=hx(j)
                     hi=hxi(j)
                     i1=ii(j)
                     i2=jj(j)
                     f22(1:2,1)=fin(1:2,i1,i2)
                     f22(1:2,2)=fin(1:2,i1+1,i2)
                     if(splin_flag) then
                        call sp1d(ans1)
                     else
                        call hm1d(ans1)
                     endif
                     f22(1:2,1)=fin(1:2,i1,i2+1)
                     f22(1:2,2)=fin(1:2,i1+1,i2+1)
                     if(splin_flag) then
                        call sp1d(ans2)
                     else
                        call hm1d(ans2)
                     endif
                     if(idlin.eq.1) then
                        fval(j,i)=(ans2-ans1)*hyi(j)
                     else
                        fval(j,i)=ans2*dxlin + ans1*dxlini
                     endif
                  enddo

               else
                  do j=1,ivec
                     xp=yparam(j)
                     dxlin=xparam(j)
                     dxlini=1.0-dxlin
                     h=hy(j)
                     hi=hyi(j)
                     i1=ii(j)
                     i2=jj(j)
                     f22(1:2,1)=fin(1:2,i1,i2)
                     f22(1:2,2)=fin(1:2,i1,i2+1)
                     if(splin_flag) then
                        call sp1d(ans1)
                     else
                        call hm1d(ans1)
                     endif
                     f22(1:2,1)=fin(1:2,i1+1,i2)
                     f22(1:2,2)=fin(1:2,i1+1,i2+1)
                     if(splin_flag) then
                        call sp1d(ans2)
                     else
                        call hm1d(ans2)
                     endif
                     if(idlin.eq.1) then
                        fval(j,i)=(ans2-ans1)*hxi(j)
                     else
                        fval(j,i)=ans2*dxlin + ans1*dxlini
                     endif
                  enddo
               endif

            else
               !  spline or Hermite in 1 dimension; step fcn in other.
               if(jspline(1).ge.1) then
                  do j=1,ivec
                     xp=xparam(j)
                     h=hx(j)
                     hi=hxi(j)
                     i1=ii(j)
                     i2=jj(j)
                     f22(1:2,1)=fin(1:2,i1,i2)
                     f22(1:2,2)=fin(1:2,i1+1,i2)
                     if(splin_flag) then
                        call sp1d(ans1)
                     else
                        call hm1d(ans1)
                     endif
                     fval(j,i)=ans1
                  enddo
               else
                  do j=1,ivec
                     xp=yparam(j)
                     h=hy(j)
                     hi=hyi(j)
                     i1=ii(j)
                     i2=jj(j)
                     f22(1:2,1)=fin(1:2,i1,i2)
                     f22(1:2,2)=fin(1:2,i1,i2+1)
                     if(splin_flag) then
                        call sp1d(ans1)
                     else
                        call hm1d(ans1)
                     endif
                     fval(j,i)=ans1
                  enddo
               endif
            endif
         else
            fval(1:ivec,i) = 0
         endif
      enddo

      return

      contains
        subroutine dtrans

          ! convert ict(...) codes into intermediate coding:
          !   iderivs(i,j) = number of differentiations in coordinate i
          !                  of the j'th value evaluated

          !   maxd_lin(j) = max # of differentiations for a linear dimension
          !   maxd_cub(j) = max # of differentiations for a spline dimension

          integer :: i

          inum=0   ! actual number of vectors to be evaluated (summed here)

          if(ict(1).le.(2)) then
             if(ict(1).eq.1) then
                call add1(0,0)
             endif
             if(ict(2).eq.1) then
                call add1(1,0)
             endif
             if(ict(3).eq.1) then
                call add1(0,1)
             endif
             if(ict(4).eq.1) then
                call add1(2,0)
             endif
             if(ict(5).eq.1) then
                call add1(0,2)
             endif
             if(ict(6).eq.1) then
                call add1(1,1)
             endif

          else if(ict(1).eq.3) then
             if(ict(2).eq.1) then
                call add1(3,0)
             endif
             if(ict(3).eq.1) then
                call add1(2,1)
             endif
             if(ict(4).eq.1) then
                call add1(1,2)
             endif
             if(ict(5).eq.1) then
                call add1(0,3)
             endif

          else if(ict(1).eq.4) then
             if(ict(2).eq.1) then
                call add1(3,1)
             endif
             if(ict(3).eq.1) then
                call add1(2,2)
             endif
             if(ict(4).eq.1) then
                call add1(1,3)
             endif

          else if(ict(1).eq.5) then
             if(ict(2).eq.1) then
                call add1(3,2)
             endif
             if(ict(3).eq.1) then
                call add1(2,3)
             endif

          else if(ict(1).eq.6) then
             if(ict(2).eq.1) then
                call add1(3,3)
             endif
          endif

        end subroutine dtrans

        subroutine add1(idx,idy)
          ! insert record of derivative d[idx+idy]f/dx[idx]dy[idy]

          integer, intent(in) :: idx,idy

          inum=inum+1

          iderivs(1,inum)=idx
          if(jspline(1).le.0) then
             maxd_lin(inum)=idx
             maxd_cub(inum)=0
          else
             maxd_lin(inum)=0
             maxd_cub(inum)=idx
          endif

          iderivs(2,inum)=idy
          if(jspline(2).le.0) then
             maxd_lin(inum)=max(maxd_lin(inum),idy)
          else
             maxd_cub(inum)=max(maxd_cub(inum),idy)
          endif

        end subroutine add1

C===========================================================================
        subroutine sp1d(ans)
          real, intent(out) :: ans

C----------------------
C contained 1d spline evaluation
C----------------------

          real :: xpi,xp2,xpi2,cx,cxi,hx2,cxd,cxdi
C
          real sixth
          data sixth/0.166666666666666667/
C----------------------

          if(idcub.eq.0) then
             ! value

             xpi=1.0-xp
             xp2=xp*xp
             xpi2=xpi*xpi

             cx=xp*(xp2-1.0)
             cxi=xpi*(xpi2-1.0)
             hx2=h*h

             ans=xpi*f22(1,1) +xp*f22(1,2)
             ans=ans+sixth*hx2*(cxi*f22(2,1) +cx*f22(2,2))

          else if(idcub.eq.1) then
             ! 1st derivative

             xpi=1.0-xp
             xp2=xp*xp
             xpi2=xpi*xpi

             cxd=3.0*xp2-1.0
             cxdi=-3.0*xpi2+1.0

             ans=hi*(f22(1,2)-f22(1,1))
             ans=ans+sixth*h*(cxdi*f22(2,1) +cxd*f22(2,2))

          else if(idcub.eq.2) then
             ! 2nd derivative

             xpi=1.0-xp
             ans=xpi*f22(2,1) + xp*f22(2,2)

          else
             ! 3rd derivative

             ans = hi*(f22(2,2)-f22(2,1))

          endif

        end subroutine sp1d

C===========================================================================
        subroutine hm1d(ans)
          real, intent(out) :: ans

C----------------------
C contained 1d hermite evaluation
C----------------------

          real :: xpi,xp2,xpi2,ax,axbar,bx,bxbar
          real :: axp,axbarp,bxp,bxbarp

          xpi=1.0-xp
          xp2=xp*xp
          xpi2=xpi*xpi

          if(idcub.eq.0) then
             ! value

             ax=xp2*(3.0-2.0*xp)
             axbar=1.0-ax
             bx=-xp2*xpi
             bxbar=xpi2*xp

             ans=axbar*f22(1,1) + ax*f22(1,2)
             ans=ans + h*(bxbar*f22(2,1) + bx*f22(2,2))

          else
             ! 1st derivative

             axp=6.0*xp*xpi
             axbarp=-axp
             bxp=xp*(3.0*xp-2.0)
             bxbarp=xpi*(3.0*xpi-2.0)

             ans=hi*(axbarp*f22(1,1) +axp*f22(1,2))
             ans=ans + bxbarp*f22(2,1) + bxp*f22(2,2)

          endif

        end subroutine hm1d

      end