***********************************************************************
c                    ORBEL_FLON.F
***********************************************************************
*     PURPOSE:  Solves Kepler's eqn. for hyperbola using hybrid approach.
*
*             Input:
*                           e ==> eccentricity anomaly. (real scalar)
*                        capn ==> hyperbola mean anomaly. (real scalar)
*             Returns:
*                  orbel_flon ==>  eccentric anomaly. (real scalar)
*
*     ALGORITHM: Uses power series for N in terms of F and Newton,s method
*     REMARKS: ONLY GOOD FOR LOW VALUES OF N (N < 0.636*e -0.6)
*     AUTHOR: M. Duncan
*     DATE WRITTEN: May 26, 1992.
*     REVISIONS:
***********************************************************************

	real*8 function orbel_flon(e,capn)

      include '../swift.inc'

c...  Inputs Only:
	real*8 e,capn

c...  Internals:
	integer iflag,i,IMAX
	real*8 a,b,sq,biga,bigb
	real*8 x,x2
	real*8 f,fp,dx
	real*8 diff
	real*8 a0,a1,a3,a5,a7,a9,a11
	real*8 b1,b3,b5,b7,b9,b11
	PARAMETER (IMAX = 10)
	PARAMETER (a11 = 156.d0,a9 = 17160.d0,a7 = 1235520.d0)
	PARAMETER (a5 = 51891840.d0,a3 = 1037836800.d0)
	PARAMETER (b11 = 11.d0*a11,b9 = 9.d0*a9,b7 = 7.d0*a7)
	PARAMETER (b5 = 5.d0*a5, b3 = 3.d0*a3)

c----
c...  Executable code

c Function to solve "Kepler's eqn" for F (here called
c x) for given e and CAPN. Only good for smallish CAPN

	iflag = 0
	if( capn .lt. 0.d0) then
	   iflag = 1
	   capn = -capn
	endif

	a1 = 6227020800.d0 * (1.d0 - 1.d0/e)
	a0 = -6227020800.d0*capn/e
	b1 = a1

c  Set iflag nonzero if capn < 0., in which case solve for -capn
c  and change the sign of the final answer for F.
c  Begin with a reasonable guess based on solving the cubic for small F

	a = 6.d0*(e-1.d0)/e
	b = -6.d0*capn/e
	sq = sqrt(0.25*b*b +a*a*a/27.d0)
	biga = (-0.5*b + sq)**0.3333333333333333d0
	bigb = -(+0.5*b + sq)**0.3333333333333333d0
	x = biga + bigb
c	write(6,*) 'cubic = ',x**3 +a*x +b
	orbel_flon = x
c If capn is tiny (or zero) no need to go further than cubic even for
c e =1.
	if( capn .lt. TINY) go to 100

	do i = 1,IMAX
	  x2 = x*x
	  f = a0 +x*(a1+x2*(a3+x2*(a5+x2*(a7+x2*(a9+x2*(a11+x2))))))
	  fp = b1 +x2*(b3+x2*(b5+x2*(b7+x2*(b9+x2*(b11 + 13.d0*x2)))))
	  dx = -f/fp
c	  write(6,*) 'i,dx,x,f : '
c	  write(6,432) i,dx,x,f
432	  format(1x,i3,3(2x,1p1e22.15))
	  orbel_flon = x + dx
c   If we have converged here there's no point in going on
	  if(abs(dx) .le. TINY) go to 100
	  x = orbel_flon
	enddo

c Abnormal return here - we've gone thru the loop
c IMAX times without convergence
	if(iflag .eq. 1) then
	   orbel_flon = -orbel_flon
	   capn = -capn
	endif
	write(6,*) 'FLON : RETURNING WITHOUT COMPLETE CONVERGENCE'
	  diff = e*sinh(orbel_flon) - orbel_flon - capn
	  write(6,*) 'N, F, ecc*sinh(F) - F - N : '
	  write(6,*) capn,orbel_flon,diff
	return

c  Normal return here, but check if capn was originally negative
100	if(iflag .eq. 1) then
	   orbel_flon = -orbel_flon
	   capn = -capn
	endif

	return
	end     ! orbel_flon
c------------------------------------------------------------------


!---------------------------------------------------!
! Copyright (c) 2016 Shunsuke A. Sato               !
! Released under the MIT license                    !
! https://opensource.org/licenses/mit-license.php   !
!---------------------------------------------------!
!-------10--------20--------30--------40--------50--------60--------70--------80--------90
subroutine dt_evolve(it)
  use global_variables
  implicit none
  integer :: it
  real(8) :: lambda_v,lambda_c,theta_p,theta_m,eps_p,eps_m
  real(8) :: Etz0,Etz1,Act_t1,Act_t2,alpha,ss
  complex(8) :: zx,zy
  integer :: ikr,ikz
  complex(8) :: zeig_vec_v(2), zeig_vec_c(2)

  select case(npump_probe_type)
  case(N_COMBINED_PUMP_PROBE)
    Etz0 = -0.5d0*(Act(it+1)-Act(it-1))/dt
    Etz1 = -0.5d0*(Act(it+2)-Act(it))/dt
    Act_t1 = Act(it)
    Act_t2 = Act(it+1)
  case(N_DECOMPOSED_PUMP_PROBE)
    Etz0 = -0.5d0*(Act_probe(it+1)-Act_probe(it-1))/dt
    Etz1 = -0.5d0*(Act_probe(it+2)-Act_probe(it))/dt
    Act_t1 = Act_pump(it)
    Act_t2 = Act_pump(it+1)
  case default
    write(*,"(A,2x,A)")"Invalid npump_probe_type",npump_probe_type
  end select

  if(if_pure_intraband_fields_exist)then
    Act_t1 = Act_t1 + Act_intra(it)
    Act_t2 = Act_t2 + Act_intra(it+1)
  end if

!$omp parallel do private(ikz)
  do ikz = -NKz,NKz
    kz(ikz) = kz0(ikz) + Act_t1*fact_intra
  end do

  call set_deps

!$omp parallel do private(ikz,ikr,alpha,lambda_v,lambda_c,zx,zy,ss,zeig_vec_v,zeig_vec_c)
  do ikz = -NKz,NKz
  do ikr = 1,NKr

    alpha = piz_vc*Etz0/deps(ikr,ikz)
    lambda_v = 0.5d0*(deps(ikr,ikz)-sqrt(deps(ikr,ikz)**2+4d0*alpha**2))
    lambda_c = 0.5d0*(deps(ikr,ikz)+sqrt(deps(ikr,ikz)**2+4d0*alpha**2))
    zx = zi*alpha/(deps(ikr,ikz)-lambda_v)
    zy = zi*alpha/lambda_c

    ss = 1d0/sqrt(1d0+abs(zx)**2)
    zeig_vec_v(1) = ss; zeig_vec_v(2) = zx*ss

    ss = 1d0/sqrt(abs(zy)**2+1d0)
    zeig_vec_c(1) = zy*ss; zeig_vec_c(2) = ss

    zx = sum(conjg(zeig_vec_v(:))*zCt(:,ikr,ikz))
    zy = sum(conjg(zeig_vec_c(:))*zCt(:,ikr,ikz))

    zx = zx * exp(-zI*lambda_v*0.5d0*dt)
    zy = zy * exp(-zI*lambda_c*0.5d0*dt)

    zCt(:,ikr,ikz) = zx*zeig_vec_v(:) + zy*zeig_vec_c(:)

  end do
  end do

!$omp parallel do private(ikz)
  do ikz = -NKz,NKz
    kz(ikz) = kz0(ikz) + Act_t2*fact_intra
  end do

  call set_deps

!$omp parallel do private(ikz,ikr,alpha,lambda_v,lambda_c,zx,zy,ss,zeig_vec_v,zeig_vec_c)
  do ikz = -NKz,NKz
  do ikr = 1,NKr

    alpha = piz_vc*Etz1/deps(ikr,ikz)
    lambda_v = 0.5d0*(deps(ikr,ikz)-sqrt(deps(ikr,ikz)**2+4d0*alpha**2))
    lambda_c = 0.5d0*(deps(ikr,ikz)+sqrt(deps(ikr,ikz)**2+4d0*alpha**2))
    zx = zi*alpha/(deps(ikr,ikz)-lambda_v)
    zy = zi*alpha/lambda_c

    ss = 1d0/sqrt(1d0+abs(zx)**2)
    zeig_vec_v(1) = ss; zeig_vec_v(2) = zx*ss

    ss = 1d0/sqrt(abs(zy)**2+1d0)
    zeig_vec_c(1) = zy*ss; zeig_vec_c(2) = ss

    zx = sum(conjg(zeig_vec_v(:))*zCt(:,ikr,ikz))
    zy = sum(conjg(zeig_vec_c(:))*zCt(:,ikr,ikz))

    zx = zx * exp(-zI*lambda_v*0.5d0*dt)
    zy = zy * exp(-zI*lambda_c*0.5d0*dt)

    zCt(:,ikr,ikz) = zx*zeig_vec_v(:) + zy*zeig_vec_c(:)

  end do
  end do

  return
end subroutine dt_evolve


MODULE PS5PARA
    IMPLICIT NONE
    INTEGER:: I
    INTEGER, PARAMETER:: N=66, RETIRE_AGE=46 ! LIFE EXPECTANCY AND RETIREMENT
    REAL(KIND=8), PARAMETER:: A_MAX=15.0, A_MIN=0.1, STEP=0.1
    INTEGER, PARAMETER:: NZ=2, NA=(A_MAX-A_MIN)/STEP+1
    REAL(KIND=8):: ALPHA=0.36, BETA=0.97
    REAL(KIND=8), PARAMETER:: THETA=0.0, 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
    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 PS5RES
    use PS5PARA
    IMPLICIT NONE

    REAL(KIND=8):: WAGE, AGG_LABOR=1.0, AGG_CAP=2.0, BENEFIT, INTEREST, ERROR=100.
    REAL(KIND=8), DIMENSION(NA,NZ,N):: VFUNC, LFUNC, STAT_DIST, VFUNC_OLD
    INTEGER, DIMENSION(NA,NZ,N):: PFUNC=1
END MODULE

PROGRAM PS5
    USE PS5PARA
    USE PS5RES
    IMPLICIT NONE
    REAL(KIND=8):: CRIT_P=1e-3, WEL, CV
    INTEGER:: SROWIDX, AGEIDX

    CALL INIT_ARRAY()
    PRINT*, NA
    DO WHILE(ERROR>CRIT_P)
        ! UPDATE VARS
        WAGE = (1-ALPHA)*(AGG_CAP**(ALPHA))*(AGG_LABOR**(-ALPHA))! FOC WRT LABOR
        INTEREST = ALPHA*(AGG_CAP**(ALPHA-1))*(AGG_LABOR**(1-ALPHA))-DELTA! FOC WRT CAPITAL
        BENEFIT = THETA*WAGE*AGG_LABOR/SUM(COHORT_POP(RETIRE_AGE:N)) ! TAX REV/RETIRED_POP
        PRINT*, "=================================================="
        PRINT*, "R, B,W:", INTEREST, BENEFIT, WAGE
        PRINT*,  "L, K:",AGG_LABOR, AGG_CAP

        LFUNC = 0.
        VFUNC=0.
        PFUNC=1
        ! BACKWARD INDUCTION USING NEW VARS (BELLMAN)
        CALL BACKWARD_INDUCTION()
        ! UPDATE STAT_DIST
        CALL FIND_STAT_DIST()
        ! OBSERVE ERROR AND UPDATE PRICE
        CALL UPDATE_PRICE()
    ENDDO

    CALL WRITE_ALL()

    WAGE = (1-ALPHA)*(AGG_CAP**ALPHA)*(AGG_LABOR**(-ALPHA))! FOC WRT LABOR
    INTEREST = ALPHA*(AGG_CAP**(ALPHA-1))*(AGG_LABOR**(1-ALPHA))-DELTA! FOC WRT CAPITAL
    BENEFIT = THETA*WAGE*AGG_LABOR/SUM(COHORT_POP(RETIRE_AGE:N)) ! TAX REV/RETIRED_POP
    DO SROWIDX=1,NA
        DO AGEIDX=1,N
            WEL=WEL+VFUNC(SROWIDX,1,AGEIDX)*STAT_DIST(SROWIDX,1,AGEIDX)+VFUNC(SROWIDX,2,AGEIDX)*STAT_DIST(SROWIDX,2,AGEIDX)
            CV = CV + A_GRID(SROWIDX)*SUM(STAT_DIST(SROWIDX, :, AGEIDX))
        ENDDO
    ENDDO
    PRINT*, "=================================================="
    PRINT*, "R, B,W:", INTEREST, BENEFIT, WAGE
    PRINT*,  "L, K:",AGG_LABOR, AGG_CAP
    PRINT*, "WELFARE", WEL
    PRINT*, "WEALTH", CV

END PROGRAM PS5

! SUBROUTINES
SUBROUTINE FIND_LABOR(CHOICEIDX, SROWIDX, SCOLIDX, AGE, WORKING)
    USE PS5RES
    USE PS5PARA
    IMPLICIT NONE
    INTEGER, INTENT(IN):: CHOICEIDX, SROWIDX, SCOLIDX, AGE
    REAL(KIND=8), INTENT(OUT):: WORKING
    REAL(KIND=8):: LAST_CHOICE, CUR_CHOICE, OPT_LABOR, NOM

    CUR_CHOICE = A_GRID(CHOICEIDX)
    IF (AGE/=1) THEN
        LAST_CHOICE = (1+INTEREST)*A_GRID(SROWIDX) ! DISCOUNTED
    ELSE
        LAST_CHOICE = 0.
    ENDIF
    NOM = (GAMMA*(1-THETA)*AGE_EFF(AGE,SCOLIDX)*WAGE-(1-GAMMA)*(LAST_CHOICE-CUR_CHOICE))
    OPT_LABOR = NOM/((1-THETA)*WAGE*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 BACKWARD_INDUCTION()
    USE PS5RES
    USE PS5PARA
    IMPLICIT NONE
    INTEGER:: AGEIDX, AGE, SROWIDX
    REAL(KIND=8):: CONSUM, UTIL, ERROR_VFI, CRIT_VFI=1e-3

    ERROR_VFI=100.
    VFUNC=0.
    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)*A_GRID(SROWIDX) + BENEFIT ! SPEND ALL IN LAST PERIOD
                PFUNC(:,:,N) = 1 ! ACTUALLY DOES NOT MATTER
                UTIL = CONSUM**((1-SIGMA)*GAMMA)/(1-SIGMA) ! NO NEXT PERIOD VALUE
                VFUNC(SROWIDX,1,N) = UTIL
                VFUNC(SROWIDX,2,N) = UTIL
                ! ASSET MUST BE ZERO FOR AGE
        ENDDO

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

SUBROUTINE BELLMAN(AGE)
    USE PS5RES
    USE PS5PARA
    IMPLICIT NONE
    INTEGER, INTENT(IN):: AGE
    REAL(KIND=8):: COND_MAX_UTIL, CONSUM, UTIL, WORKING, PROD, NEXTU
    INTEGER:: SROWIDX, SCOLIDX, CHOICEIDX

    IF (AGE< RETIRE_AGE) THEN ! WOKRING BELLMAN
        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)
                    CONSUM = (1+INTEREST)*A_GRID(SROWIDX) + (1-THETA)*PROD*WORKING*WAGE - A_GRID(CHOICEIDX)
                    IF (CONSUM > 0.) THEN
                        NEXTU = BETA*SUM(PROD_MARKOV(SCOLIDX,:)*VFUNC(CHOICEIDX,:,AGE+1))
                        UTIL = (CONSUM**GAMMA*(1-WORKING)**(1-GAMMA))**(1-SIGMA)/(1-SIGMA)+NEXTU
                        IF (UTIL>COND_MAX_UTIL) THEN
                            PFUNC(SROWIDX, SCOLIDX, AGE) = CHOICEIDX
                            LFUNC(SROWIDX,SCOLIDX, AGE) = WORKING
                            COND_MAX_UTIL = UTIL
                        ENDIF
                    ENDIF
                ENDDO ! END LOOP CHOICE SPACE FOR ONE STATE
                    IF (COND_MAX_UTIL<-1E10) THEN
                    PRINT*, VFUNC(PFUNC(SROWIDX,1, AGE),1,AGE+1), VFUNC(PFUNC(SROWIDX,2, AGE),1,AGE+1), AGE, SROWIDX
                    PRINT*, COND_MAX_UTIL, CONSUM, UTIL, (CONSUM**GAMMA*(1-WORKING)**(1.-GAMMA)), WORKING
                    STOP
                    ENDIF
                VFUNC(SROWIDX, SCOLIDX, AGE) = COND_MAX_UTIL
            ENDDO
        ENDDO
    ELSE ! BELLMAN FOR RETIRED PEOPLE
        DO SROWIDX=1, NA
            COND_MAX_UTIL = -1e12
            DO CHOICEIDX=1, NA
                CONSUM = (1+INTEREST)*A_GRID(SROWIDX) + BENEFIT - A_GRID(CHOICEIDX)
                IF (CONSUM>0.) THEN
                    UTIL = CONSUM**((1-SIGMA)*GAMMA)/(1-SIGMA)+ BETA*(VFUNC(CHOICEIDX,1,AGE+1))
                    IF (UTIL>COND_MAX_UTIL) THEN
                        PFUNC(SROWIDX,:, AGE) = CHOICEIDX ! NO PRODUCTIVITY DIFFERENCE
                        COND_MAX_UTIL = UTIL
                    ENDIF
                ENDIF
            ENDDO
                IF (COND_MAX_UTIL<-1E10) THEN
                    PRINT*, VFUNC(PFUNC(SROWIDX,1, AGE),1,AGE+1), VFUNC(PFUNC(SROWIDX,2, AGE),1,AGE+1), AGE, SROWIDX
                    PRINT*, COND_MAX_UTIL
                    STOP
                ENDIF
            VFUNC(SROWIDX, 1, AGE) = COND_MAX_UTIL ! PROD_STATE DO NOT AFFECT RETIRED PEOPLE
            VFUNC(SROWIDX, 2, AGE) = COND_MAX_UTIL ! PROD_STATE DO NOT AFFECT RETIRED PEOPLE
        ENDDO
    ENDIF
END SUBROUTINE

SUBROUTINE FIND_STAT_DIST()
    USE PS5PARA
    USE PS5RES
    IMPLICIT NONE
    INTEGER:: SROWIDX, SCOLIDX, AGEIDX, LAST_AGE, LAST_CHOICE
    REAL(KIND=8), DIMENSION(NA,NZ,N):: STAT_DIST_NEW

    PRINT*, COUNT(STAT_DIST/=0), SUM(STAT_DIST), MAXLOC(STAT_DIST), SUM(COHORT_POP)
    ! EVERYONE HAVE ZERO ASSET AT STARTING AGE
    STAT_DIST_NEW = 0.
    STAT_DIST_NEW(1,1,1) = 0.2037 ! ALWAYS NO ASSET AT FIRST AGE,
    STAT_DIST_NEW(1,2,1) = 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) ! TWO POSSIBLE
                LAST_AGE = AGEIDX-1
                STAT_DIST_NEW(LAST_CHOICE, 1, AGEIDX) = &
                STAT_DIST_NEW(LAST_CHOICE, 1, AGEIDX)+ STAT_DIST_NEW(SROWIDX,SCOLIDX,LAST_AGE)*PROD_MARKOV(SCOLIDX,1)
                STAT_DIST_NEW(LAST_CHOICE, 2, AGEIDX) = &
                STAT_DIST_NEW(LAST_CHOICE, 2, AGEIDX) + STAT_DIST_NEW(SROWIDX,SCOLIDX,LAST_AGE)*PROD_MARKOV(SCOLIDX,2)
           ENDDO
        ENDDO
    ENDDO
    DO AGEIDX=1,N
        STAT_DIST(:,:,AGEIDX) = STAT_DIST_NEW(:,:,AGEIDX)*COHORT_POP(AGEIDX)
    ENDDO
END SUBROUTINE

SUBROUTINE UPDATE_PRICE()
    ! COMPUTE ERROR AND UPDATE PRICE
    USE PS5RES
    USE PS5PARA
    IMPLICIT NONE
    INTEGER:: SROWIDX, SCOLIDX, AGEIDX
    REAL(KIND=8):: AGG_LABOR_NEW=0., AGG_CAP_NEW=0., PORTION, ERROR_LAB, ERROR_CAP

    AGG_LABOR_NEW=0.
    AGG_CAP_NEW=0.
    DO AGEIDX=1, N ! CAL ERRORS
        DO SROWIDX=1,NA
            DO SCOLIDX=1,NZ
                PORTION = STAT_DIST(SROWIDX, SCOLIDX, AGEIDX)
                    AGG_CAP_NEW = AGG_CAP_NEW+ A_GRID(SROWIDX)*PORTION
                IF (AGEIDX<RETIRE_AGE) THEN
                    AGG_LABOR_NEW = AGG_LABOR_NEW+LFUNC(SROWIDX,SCOLIDX,AGEIDX)*PORTION*AGE_EFF(AGEIDX,SCOLIDX)
                ENDIF
            ENDDO
        ENDDO
    ENDDO
    ERROR_LAB = ABS(AGG_LABOR_NEW-AGG_LABOR)
    ERROR_CAP = ABS(AGG_CAP_NEW-AGG_CAP)
    PRINT*, "UPDATE_PRICE ERROR", ERROR_LAB, ERROR_CAP
    ERROR = MAX(ERROR_CAP, ERROR_LAB)
    AGG_LABOR = 0.01*AGG_LABOR_NEW + 0.99*AGG_LABOR
    AGG_CAP = 0.01*AGG_CAP_NEW + 0.99*AGG_CAP
    ! UPDATE INTEREST, WAGE, BENEFIT AT THE START OF NEXT PERIOD (IT IS THE SAME)

END SUBROUTINE

SUBROUTINE WRITE_ALL()
    USE PS5RES
    USE PS5PARA
    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)
        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)), A_GRID(PFUNC(SROWIDX,2, 20))
        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, :, 2)
        ENDDO
        CLOSE(UNIT=3)

        FILE_NAME = TRIM(PATH)//"LFUNC"
        OPEN(UNIT=4, FILE=FILE_NAME, STATUS='REPLACE') ! FOR HAVING THE X-AXIS OF PLOT
        DO SROWIDX=1, NA
            WRITE(UNIT=4,FMT=*) LFUNC(SROWIDX, :, 50)
        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)//"COHORTPOP"
        OPEN(UNIT=32, FILE=FILE_NAME, STATUS='REPLACE') ! START WITH THE TWO VALUE FUNCTIONS
        DO SROWIDX=1, N
            WRITE(UNIT=32,FMT=*) COHORT_POP(SROWIDX)
        ENDDO
        CLOSE(UNIT=32)

END SUBROUTINE


      FUNCTION zbrent(func,x1,x2,tol)
      INTEGER ITMAX
      REAL zbrent,tol,x1,x2,func,EPS
      EXTERNAL func
      PARAMETER (ITMAX=100,EPS=3.e-8)
      INTEGER iter
      REAL a,b,c,d,e,fa,fb,fc,p,q,r,s,tol1,xm
      a=x1
      b=x2
      fa=func(a)
      fb=func(b)
      if((fa.gt.0..and.fb.gt.0.).or.(fa.lt.0..and.fb.lt.0.))pause
     *'root must be bracketed for zbrent'
      c=b
      fc=fb
      do 11 iter=1,ITMAX
        if((fb.gt.0..and.fc.gt.0.).or.(fb.lt.0..and.fc.lt.0.))then
          c=a
          fc=fa
          d=b-a
          e=d
        endif
        if(abs(fc).lt.abs(fb)) then
          a=b
          b=c
          c=a
          fa=fb
          fb=fc
          fc=fa
        endif
        tol1=2.*EPS*abs(b)+0.5*tol
        xm=.5*(c-b)
        if(abs(xm).le.tol1 .or. fb.eq.0.)then
          zbrent=b
          return
        endif
        if(abs(e).ge.tol1 .and. abs(fa).gt.abs(fb)) then
          s=fb/fa
          if(a.eq.c) then
            p=2.*xm*s
            q=1.-s
          else
            q=fa/fc
            r=fb/fc
            p=s*(2.*xm*q*(q-r)-(b-a)*(r-1.))
            q=(q-1.)*(r-1.)*(s-1.)
          endif
          if(p.gt.0.) q=-q
          p=abs(p)
          if(2.*p .lt. min(3.*xm*q-abs(tol1*q),abs(e*q))) then
            e=d
            d=p/q
          else
            d=xm
            e=d
          endif
        else
          d=xm
          e=d
        endif
        a=b
        fa=fb
        if(abs(d) .gt. tol1) then
          b=b+d
        else
          b=b+sign(tol1,xm)
        endif
        fb=func(b)
11    continue
      pause 'zbrent exceeding maximum iterations'
      zbrent=b
      return
      END

NR-Functions/Numerical Recipes- Example & Functions/Functions/zbrent.for


!-------------------------------------------------------------------------------
! helper functions for arrays; these are utilized by neural network structures,
! but can be used on arrays in general
!-------------------------------------------------------------------------------

module net_helper_procedures
implicit none
contains

!===============================================================================
!===============================================================================
! procedures with 2D array input require variables-as-columns form;
! procedures with 4D array input require (rows, columns, channels, batches) form
!===============================================================================
!===============================================================================

!-------------------------------------------------------------------------------
! wrapper for DGEMM from BLAS; matrix multiplication: c = a * b,
! where either or both of a, b can be transposed
!-------------------------------------------------------------------------------
! a:        (real(:,:))
! b:        (real(:,:))
! c:        (real(:,:)) resulting array
!
! transa    (optional - logical) transpose a before multiplication
! transb    (optional - logical) transpose b before multiplication
!-------------------------------------------------------------------------------
! alters :: result of matmul(a, b) stored in c
!-------------------------------------------------------------------------------
subroutine dgemm_wrapper(a, b, c, transa, transb)
    real(kind=8), intent(in)      :: a(:,:), b(:,:)
    real(kind=8), allocatable     :: c(:,:)
    logical, intent(in), optional :: transa, transb
    integer                       :: m, n, k, lda, ldb, ldc
    character(len=1)              :: a_t, b_t

    lda = size(a, dim=1)
    ldb = size(b, dim=1)

    ! adjust variables for transpose a
    if (present(transa) .and. transa) then
        a_t = 't'
        m   = size(a, dim=2)
        k   = size(a, dim=1)
    else
        a_t = 'n'
        m   = size(a, dim=1)
        k   = size(a, dim=2)
    end if

    ! adjust variables for transpose b
    if (present(transb) .and. transb) then
        b_t = 't'
        n   = size(b, dim=1)
    else
        b_t = 'n'
        n   = size(b, dim=2)
    end if

    ! create new array if not correct size
    if (allocated(c)) then
        if (.not. all(shape(c) == [m,n])) then
            deallocate(c)
        end if
    end if

    if (.not. allocated(c)) then
        allocate(c(m, n))
    end if

    c = 0
    ldc = m

    ! matrix multiplication: c = 1*a * b + 0*c = a * b;
    ! (d in 1d0 and 0d0 means double precision)
    call DGEMM(a_t, b_t, m, n, k, 1d0, a, lda, b, ldb, 0d0, c, ldc)
end subroutine

!-------------------------------------------------------------------------------
! write 2D array as CSV-formatted row appended to given file, row-major order;
!
! appends to file at filepath if it exists, otherwise creates new file;
! filepath should be relative to overall project folder with executable
!-------------------------------------------------------------------------------
! a:        (real(:,:)) array to write
! filepath: (characters) desired full path with filename
! append:   (logical) .true. to append row to existing file, else new file
!-------------------------------------------------------------------------------
! alters :: filepath has values of a appended as CSV-formatted row
!-------------------------------------------------------------------------------
subroutine write_array_2D(a, filepath, append)
    real(kind=8), intent(in)      :: a(:,:)
    character(*), intent(in)      :: filepath
    logical, intent(in)           :: append
    integer                       :: ios, r, c, rows, cols
    character(len=30)             :: val
    character(len=:), allocatable :: cleanval, pos

    if (append) then
        pos = 'append'
    else
        pos = 'rewind'
    end if

    open(unit=1, file=filepath, form='formatted', position=pos, iostat=ios)

    if (ios /= 0) then
        print *, 'Error opening file.'
        stop -1
    endif

    rows = size(a, dim=1)
    cols = size(a, dim=2)

    ! write values in one row in file, row by row in matrix
    do r = 1, rows
        do c = 1, cols
            write(val, *) a(r,c)          ! real -> string
            cleanval = trim(adjustl(val)) ! move lead pad to end, then cut end

            write(unit=1, fmt='(A)', advance='no') cleanval

            if (r < rows .or. c < cols) then
                ! don't write ',' after final value
                write(unit=1, fmt='(A)', advance='no') ','
            else
                ! done row; flush to next line
                write(unit=1, fmt=*)
            endif
        end do
    end do
end subroutine

!===============================================================================
! activation functions and their derivatives
!   *** elemental functions can also be applied to arrays of the input type
!===============================================================================

!-------------------------------------------------------------------------------
! sigmoid activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function sigmoid(z)
    real(kind=8), intent(in) :: z
    sigmoid = exp(z) / (exp(z) + 1)
end function

!-------------------------------------------------------------------------------
! derivative of sigmoid activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function sigmoid_deriv(z)
    real(kind=8), intent(in) :: z
    sigmoid_deriv = exp(z) / (exp(z) + 1)**2
end function

!-------------------------------------------------------------------------------
! relu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function relu(z)
    real(kind=8), intent(in) :: z
    if (z >= 0) then
        relu = z
    else
        relu = 0
    end if
end function

!-------------------------------------------------------------------------------
! derivative of relu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function relu_deriv(z)
    real(kind=8), intent(in) :: z
    if (z > 0) then
        relu_deriv = 1
    else
        relu_deriv = 0
    end if
end function

!-------------------------------------------------------------------------------
! leaky relu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function leaky_relu(z)
    real(kind=8), intent(in) :: z
    if (z > 0) then
        leaky_relu = z
    else
        leaky_relu = 0.01 * z
    end if
end function

!-------------------------------------------------------------------------------
! derivative of leaky relu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function leaky_relu_deriv(z)
    real(kind=8), intent(in) :: z
    if (z > 0) then
        leaky_relu_deriv = 1
    else
        leaky_relu_deriv = 0.01
    end if
end function

!-------------------------------------------------------------------------------
! elu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function elu(z)
    real(kind=8), intent(in) :: z
    if (z > 0) then
        elu = z
    else
        elu = 0.01 * (exp(z) - 1)
    end if
end function

!-------------------------------------------------------------------------------
! derivative of elu activation function
!-------------------------------------------------------------------------------
! z:         (real)
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function elu_deriv(z)
    real(kind=8), intent(in) :: z
    if (z > 0) then
        elu_deriv = 1
    else
        elu_deriv = 0.01 * exp(z)
    end if
end function

!-------------------------------------------------------------------------------
! softmax activation function (applied to rows)
!-------------------------------------------------------------------------------
! z:        (real(:,:))
!-------------------------------------------------------------------------------
! alters :: softmax occurs in-place on input a
!-------------------------------------------------------------------------------
subroutine softmax(z)
    real(kind=8) :: z(:,:)
    integer      :: r
    z = exp(z)
    do r = 1, size(z, dim=1)
        ! vals in each row / sum of that row
        z(r,:) = z(r,:) / sum(z(r,:))
    end do
end subroutine

!-------------------------------------------------------------------------------
! wrapper for element-wise activation functions
! SUPPORTED: sigmoid, relu, leaky_relu, elu;
! must be checked by caller, so that this function can remain elemental
!-------------------------------------------------------------------------------
! z:         (real)
! activ:     (characters) activation function
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function activfunc(z, activ)
    real(kind=8), intent(in) :: z
    character(*), intent(in) :: activ

    select case(activ)
        case ('sigmoid')
            activfunc = sigmoid(z)
        case ('relu')
            activfunc = relu(z)
        case ('leaky_relu')
            activfunc = leaky_relu(z)
        case ('elu')
            activfunc = elu(z)
        case default
            activfunc = 0
    end select
end function

!-------------------------------------------------------------------------------
! wrapper for element-wise activation function derivatives;
! SUPPORTED: sigmoid, relu, leaky_relu, elu;
! must be checked by caller, so that this function can remain elemental
!-------------------------------------------------------------------------------
! z:         (real)
! activ:     (characters) activation function
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
elemental real(kind=8) function activfunc_deriv(z, activ)
    real(kind=8), intent(in) :: z
    character(*), intent(in) :: activ

    select case(activ)
        case ('sigmoid')
            activfunc_deriv = sigmoid_deriv(z)
        case ('relu')
            activfunc_deriv = relu_deriv(z)
        case ('leaky_relu')
            activfunc_deriv = leaky_relu_deriv(z)
        case ('elu')
            activfunc_deriv = elu_deriv(z)
        case default
            activfunc_deriv = 0
    end select
end function

!-------------------------------------------------------------------------------
! wrapper for output-layer activation functions;
! subroutine needed because some output activation functions (like softmax) are
! not element-wise, so we must manipulate an array, not return a value
!-------------------------------------------------------------------------------
! z:         (real(:,:))
! out_activ: (characters) output activation function
! res:       (real(:,:)) stores the output
!-------------------------------------------------------------------------------
! alters ::  res becomes activation applied to z
!-------------------------------------------------------------------------------
subroutine out_activfunc_2D(z, out_activ, res)
    real(kind=8)              :: z(:,:)
    character(*), intent(in)  :: out_activ
    real(kind=8), allocatable :: res(:,:)
    integer                   :: z_rows, z_cols

    z_rows = size(z, dim=1)
    z_cols = size(z, dim=2)

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == shape(z))) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(z_rows, z_cols))
    end if

    select case(out_activ)
        case ('softmax')
            res = z
            call softmax(res)
        case ('sigmoid')
            res = sigmoid_deriv(z)
        case ('relu')
            res = relu_deriv(z)
        case ('leaky_relu')
            res = leaky_relu_deriv(z)
        case ('elu')
            res = elu_deriv(z)
        case default
            print *, '--------------------------------------------------'
            print *, '(net_helper_functions :: out_activfunc)'
            print *, 'invalid output activation function.'
            print *, 'supported: softmax, sigmoid, relu, leaky_relu, elu'
            print *, '--------------------------------------------------'
            stop -1
    end select
end subroutine

!===============================================================================
! loss/metric functions
!===============================================================================

!-------------------------------------------------------------------------------
! mean square error between predictions and true targets
!-------------------------------------------------------------------------------
! preds:     (real(:,:)) predictions
! targets:   (real(:,:)) targets we want to predict
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function mse_func_2D(preds, targets)
    real(kind=8), intent(in) :: preds(:,:), targets(:,:)
    mse_func_2D = sum((preds - targets) ** 2) / (2 * size(preds, dim=1))
end function

!-------------------------------------------------------------------------------
! mean square error between predictions and true targets
!-------------------------------------------------------------------------------
! preds:     (real(:,:,:,:)) predictions
! targets:   (real(:,:,:,:)) targets we want to predict
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function mse_func_4D(preds, targets)
    real(kind=8), intent(in) :: preds(:,:,:,:), targets(:,:,:,:)
    mse_func_4D = sum((preds - targets) ** 2) / (2 * size(preds, dim=1))
end function

!-------------------------------------------------------------------------------
! categorical cross entropy between predictions and true targets
!-------------------------------------------------------------------------------
! preds:     (real(:,:)) predictions
! targets:   (real(:,:)) targets we want to predict
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function cross_entropy_func_2D(preds, targets)
    real(kind=8), intent(in) :: preds(:,:), targets(:,:)
    cross_entropy_func_2D = -sum(targets * log(preds)) / size(preds, dim=1)
end function

!-------------------------------------------------------------------------------
! calculate the accuracy between predictions and one-hot label rows
!-------------------------------------------------------------------------------
! preds:     (real(:,:)) predictions
! targets:   (real(:,:)) ONE-HOT ENCODED targets we want to predict
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function one_hot_accuracy_2D(preds, targets)
    real(kind=8), intent(in) :: preds(:,:), targets(:,:)
    real(kind=8)             :: correct

    ! correct where strongest predictions match one-hot targets
    correct = count(maxloc(preds, dim=2) == maxloc(targets, dim=2))
    one_hot_accuracy_2D = correct / size(preds, dim=1)
end function

!-------------------------------------------------------------------------------
! wrapper for loss functions
!-------------------------------------------------------------------------------
! preds:     (real(:,:)) predictions
! targets:   (real(:,:)) targets we want to predict
! loss:      (characters) loss function
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function lossfunc_2D(preds, targets, loss)
    real(kind=8), intent(in) :: preds(:,:), targets(:,:)
    character(*), intent(in) :: loss

    select case(loss)
        case ('mse')
            lossfunc_2D = mse_func_2D(preds, targets)
        case ('cross_entropy')
            lossfunc_2D = cross_entropy_func_2D(preds, targets)
        case default
            print *, '----------------------------------'
            print *, '(net_helper_functions :: lossfunc)'
            print *, 'invalid loss function.'
            print *, 'supported: mse, cross_entropy'
            print *, '----------------------------------'
            stop -1
    end select
end function

!-------------------------------------------------------------------------------
! wrapper for loss functions
!-------------------------------------------------------------------------------
! preds:     (real(:,:,:,:)) predictions
! targets:   (real(:,:,:,:)) targets we want to predict
! loss:      (characters) loss function
!-------------------------------------------------------------------------------
! returns :: (real)
!-------------------------------------------------------------------------------
real(kind=8) function lossfunc_4D(preds, targets, loss)
    real(kind=8), intent(in) :: preds(:,:,:,:), targets(:,:,:,:)
    character(*), intent(in) :: loss

    select case(loss)
        case ('mse')
            lossfunc_4D = mse_func_4D(preds, targets)
        case default
            print *, '----------------------------------'
            print *, '(net_helper_functions :: lossfunc)'
            print *, 'invalid loss function.'
            print *, 'supported: mse'
            print *, '----------------------------------'
            stop -1
    end select
end function

!===============================================================================
! more complex array-based helpers
!===============================================================================

!-------------------------------------------------------------------------------
! shuffle rows of two 2D arrays in corresponding order (Fisher-Yates shuffle)
!-------------------------------------------------------------------------------
! a:        (real(:,:))
! b:        (real(:,:))
!-------------------------------------------------------------------------------
! alters :: rows of a and b are shuffled (correspondingly) in-place
!-------------------------------------------------------------------------------
subroutine pair_shuffle_2D_2D(a, b)
    real(kind=8)              :: a(:,:), b(:,:), randn
    real(kind=8), allocatable :: row(:)
    integer                   :: i, j

    ! loop through rows from high to low indices
    do i = size(a, dim=1), 2, -1
        j = i + 1 ! arbitrary assignment to enter loop below

        ! generate random j in range [1,i]
        do while (j > i)
            call random_number(randn) ! range [0, 1)
            ! scale to [0, i) => floor to [0,i-1] => shift to [1,i]
            j = 1 + floor(randn * i)
        end do

        ! swap if random index in different place
        ! (if random j = i, don't need to swap i with i)
        if (i /= j) then
            ! array a
            row = a(i,:)
            a(i,:) = a(j,:)
            a(j,:) = row

            ! array b
            row = b(i,:)
            b(i,:) = b(j,:)
            b(j,:) = row
        end if
    end do
end subroutine

!-------------------------------------------------------------------------------
! shuffle 3D images in each 4D batch in corresponding order with rows of 2D
! batch (Fisher-Yates shuffle)
!-------------------------------------------------------------------------------
! a:        (real(:,:,:,:))
! b:        (real(:,:))
!-------------------------------------------------------------------------------
! alters :: a and b shuffled (correspondingly) in-place
!-------------------------------------------------------------------------------
subroutine pair_shuffle_4D_2D(a, b)
    real(kind=8)              :: a(:,:,:,:), b(:,:), randn
    real(kind=8), allocatable :: channel(:,:,:), row(:)
    integer                   :: i, j

    ! loop through channels from high to low indices
    do i = size(a, dim=4), 2, -1
        j = i + 1 ! arbitrary assignment to enter loop below

        ! generate random j in range [1,i]
        do while (j > i)
            call random_number(randn) ! range [0, 1)
            ! scale to [0, i) => floor to [0,i-1] => shift to [1,i]
            j = 1 + floor(randn * i)
        end do

        ! swap if random index in different place
        ! (if random j = i, don't need to swap i with i)
        if (i /= j) then
            channel = a(:,:,:,i)
            a(:,:,:,i) = a(:,:,:,j)
            a(:,:,:,j) = channel

            row = b(i,:)
            b(i,:) = b(j,:)
            b(j,:) = row
        end if
    end do
end subroutine

!-------------------------------------------------------------------------------
! shuffle 3D images in each 4D batch in corresponding order
! (Fisher-Yates shuffle)
!-------------------------------------------------------------------------------
! a:        (real(:,:,:,:))
! b:        (real(:,:,:,:))
!-------------------------------------------------------------------------------
! alters :: a and b images shuffled (correspondingly) in-place
!-------------------------------------------------------------------------------
subroutine pair_shuffle_4D_4D(a, b)
    real(kind=8)              :: a(:,:,:,:), b(:,:,:,:), randn
    real(kind=8), allocatable :: channel(:,:,:)
    integer                   :: i, j

    ! loop through channels from high to low indices
    do i = size(a, dim=4), 2, -1
        j = i + 1 ! arbitrary assignment to enter loop below

        ! generate random j in range [1,i]
        do while (j > i)
            call random_number(randn) ! range [0, 1)
            ! scale to [0, i) => floor to [0,i-1] => shift to [1,i]
            j = 1 + floor(randn * i)
        end do

        ! swap if random index in different place
        ! (if random j = i, don't need to swap i with i)
        if (i /= j) then
            channel = a(:,:,:,i)
            a(:,:,:,i) = a(:,:,:,j)
            a(:,:,:,j) = channel

            channel = b(:,:,:,i)
            b(:,:,:,i) = b(:,:,:,j)
            b(:,:,:,j) = channel
        end if
    end do
end subroutine

!-------------------------------------------------------------------------------
! one-hot encode a single variable;
!   *** special format: the type should be real, but values expected to have 0s
!       in decimal place (like an integer variable that was converted to reals);
!
!   *** values should start with: 0.0, 1.0, 2.0, 3.0, . . .
!-------------------------------------------------------------------------------
! a:        (real(:)) variable
! classes:  (integer) classes encoded in the variable
! res:      (real(:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes one-hot encoding of a
!-------------------------------------------------------------------------------
subroutine one_hot_encode_special(a, classes, res)
    real(kind=8), intent(in)  :: a(:)
    integer, intent(in)       :: classes
    real(kind=8), allocatable :: res(:,:)
    integer                   :: row, col
    integer(kind=4)           :: val ! help cast 8-bit double to 4-bit integer

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [size(a, dim=1), classes])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(size(a, dim=1), classes))
    end if

    res = 0

    do row = 1, size(a, dim=1)
        val = int(a(row), kind=4) ! cast 8-bit double to 4-bit integer
        col = val + 1             ! convert 0-indexed val to 1-indexed col
        res(row, col) = 1
    end do
end subroutine

!-------------------------------------------------------------------------------
! returns the amount to add (NOT INCLUDING base dims) needed to support a given
! padding scheme, with a given kernel passing over a given array
!-------------------------------------------------------------------------------
! a_dim:      (integer) base array dimension
! kernel_dim: (integer) kernel dimension
! stride_dim: (integer) size of kernel moves in y direction
! padding:    (characters) padding type
!-------------------------------------------------------------------------------
! returns ::  amount to pad to support given padding scheme
!-------------------------------------------------------------------------------
integer function pad_calc(a_dim, kernel_dim, stride_dim, padding)
    integer, intent(in)      :: a_dim, kernel_dim, stride_dim
    character(*), intent(in) :: padding

    select case(padding)
        case ('same')
            ! pads input so output has same size as input
            pad_calc = (stride_dim - 1) * a_dim - stride_dim + kernel_dim
        case ('full')
            ! pads input so kernel can overlap it by one unit on all sides
            pad_calc = 2 * kernel_dim - 2
        case ('valid')
            ! no padding
            pad_calc = 0
        case default
            print *, '----------------------------------'
            print *, '(net_helper_functions :: pad_calc)'
            print *, 'invalid padding type.'
            print *, 'supported: same, valid, full'
            print *, '----------------------------------'
            stop -1
    end select
end function

!-------------------------------------------------------------------------------
! returns the total size of a dimension after being expanded with stride;
! see expand_with_stride_3D subroutine below for description
!-------------------------------------------------------------------------------
! a_dim:      (integer) base array dimension
! stride_dim: (integer) size of kernel moves in y direction
!-------------------------------------------------------------------------------
! returns ::  resulting dim size
!-------------------------------------------------------------------------------
integer function expand_calc(a_dim, stride_dim)
    integer, intent(in) :: a_dim, stride_dim
    expand_calc = a_dim + (a_dim - 1) * (stride_dim - 1)
end function

!-------------------------------------------------------------------------------
! returns the size of a dimension after being passed over with 'valid'
! cross correlation; typically used by first calculating the pad with pad_calc
! and including that total amount of padding in the a_dim
!-------------------------------------------------------------------------------
! a_dim:      (integer) base array dimension
! kernel_dim: (integer) kernel dimension
! stride_dim: (integer) size of kernel moves in y direction
!-------------------------------------------------------------------------------
! returns ::  resulting dim size
!-------------------------------------------------------------------------------
integer function res_calc(a_dim, kernel_dim, stride_dim)
    integer, intent(in) :: a_dim, kernel_dim, stride_dim
    res_calc = (a_dim - kernel_dim) / stride_dim + 1
end function

!-------------------------------------------------------------------------------
! pads a 2D array based on a kernel (for cross correlation or convolution)
!-------------------------------------------------------------------------------
! a:           (real(:,:)) base array
! padding:     (characters) padding type
! kernel_dims: (integer(2)) (height, width) of kernel
! stride:      (integer(2)) size of kernel moves in (y, x) directions
! res:         (real(:,:)) stores the output
!-------------------------------------------------------------------------------
! alters ::    res becomes result of padding a around its (height, width)
!-------------------------------------------------------------------------------
subroutine pad_2D(a, padding, kernel_dims, stride, res)
    real(kind=8), intent(in)  :: a(:,:)
    integer, intent(in)       :: kernel_dims(2), stride(2)
    character(*), intent(in)  :: padding
    real(kind=8), allocatable :: res(:,:)
    integer                   :: a_rows, a_cols, row_pad, col_pad, top_pad, &
                                 left_pad

    a_rows = size(a, dim=1)
    a_cols = size(a, dim=2)

    row_pad = pad_calc(a_rows, kernel_dims(1), stride(1), padding)
    col_pad = pad_calc(a_cols, kernel_dims(2), stride(2), padding)

    ! if total pad number odd: top/left get rounded down amount,
    ! (bottom/right get remainder)
    top_pad = row_pad / 2
    left_pad = col_pad / 2

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [a_rows+row_pad, a_cols+col_pad])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(row_pad+a_rows, col_pad+a_cols))
    end if

    ! set a in res
    res = 0
    res(top_pad+1:top_pad+a_rows, left_pad+1:left_pad+a_cols) = a
end subroutine

!-------------------------------------------------------------------------------
! pads a 3D array based on a kernel (for cross correlation or convolution)
!-------------------------------------------------------------------------------
! a:           (real(:,:,:)) base array
! padding:     (characters) padding type
! kernel_dims: (integer(2)) (height, width) of kernel
! stride:      (integer(2)) size of kernel moves in (y, x) directions
! res:         (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters ::    res becomes result of padding a around its (height, width)
!-------------------------------------------------------------------------------
subroutine pad_3D(a, padding, kernel_dims, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:)
    integer, intent(in)       :: kernel_dims(2), stride(2)
    character(*), intent(in)  :: padding
    real(kind=8), allocatable :: res(:,:,:)
    integer                   :: a_rows, a_cols, a_channels, row_pad, col_pad, &
                                 top_pad, left_pad

    a_rows     = size(a, dim=1)
    a_cols     = size(a, dim=2)
    a_channels = size(a, dim=3)

    row_pad = pad_calc(a_rows, kernel_dims(1), stride(1), padding)
    col_pad = pad_calc(a_cols, kernel_dims(2), stride(2), padding)

    ! if total pad number odd: top/left get rounded down amount,
    ! (bottom/right get remainder)
    top_pad = row_pad / 2
    left_pad = col_pad / 2

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [a_rows+row_pad, a_cols+col_pad, &
                                     a_channels])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(a_rows+row_pad, a_cols+col_pad, a_channels))
    end if

    ! set a in res
    res = 0
    res(top_pad+1:top_pad+a_rows, left_pad+1:left_pad+a_cols, :) = a
end subroutine

!-------------------------------------------------------------------------------
! expands 3D array so that each row is stride(1) away from the next row,
! and so that each col is stride(2) away from the next col; gaps filled with 0s
!
! as a top-down 2D example, if a = [1 2], and stride = [3, 2],
!                                  [3 4]
!
! res = [1 0 2], becuase it takes 2 moves to go from column 1 to 2,
!       [0 0 0]  and it takes 3 moves to go from row 1 to 2
!       [0 0 0]
!       [3 0 4]
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes result of expanding a with stride
!-------------------------------------------------------------------------------
subroutine expand_with_stride_3D(a, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:)
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:,:)
    integer                   :: rows, cols, channels, res_rows, res_cols, r, c

    if (stride(1) == 1 .and. stride(2) == 1) then
        res = a ! no striding to do
    else
        rows     = size(a, dim=1)
        cols     = size(a, dim=2)
        channels = size(a, dim=3)

        res_rows = expand_calc(rows, stride(1))
        res_cols = expand_calc(cols, stride(2))

        ! create new array if not correct size
        if (allocated(res)) then
            if (.not. all(shape(res) == [res_rows, res_cols, channels])) then
                deallocate(res)
            end if
        end if

        if (.not. allocated(res)) then
            allocate(res(res_rows, res_cols, channels))
        end if

        res = 0

        do r = 1, rows
            do c = 1, cols
                ! fill in parts from a
                res((r-1)*stride(1)+1,(c-1)*stride(2)+1,:) = a(r,c,:)
            end do
        end do
    end if
end subroutine

!-------------------------------------------------------------------------------
! calculates the cross-correlation between a 2D array and a 2D kernel
!-------------------------------------------------------------------------------
! a:        (real(:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:)) kernal to pass over a
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes cross-correlation result
!-------------------------------------------------------------------------------
subroutine cross_correlate_2D(a, padding, kernel, stride, res)
    real(kind=8), intent(in)  :: a(:,:), kernel(:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:), padded(:,:)
    integer                   :: a_rows, a_cols, k_rows, k_cols, &
                                 padded_rows, padded_cols, &
                                 res_rows, res_cols, r, c, kr, kc

    a_rows = size(a, dim=1)
    a_cols = size(a, dim=2)
    k_rows = size(kernel, dim=1)
    k_cols = size(kernel, dim=2)

    if (padding /= 'valid') then
        call pad_2D(a, padding, [k_rows,k_cols], stride, padded)
    else
        padded = a ! no padding to account for
    end if

    padded_rows = size(padded, dim=1)
    padded_cols = size(padded, dim=2)

    ! dimensions of cross-correlation result
    res_rows = res_calc(padded_rows, k_rows, stride(1))
    res_cols = res_calc(padded_cols, k_cols, stride(2))

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [res_rows, res_cols])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(res_rows, res_cols))
    end if

    do r = 1, res_rows
        do c = 1, res_cols
            res(r,c) = 0

            do kr = 1, k_rows
                do kc = 1, k_cols
                    ! each index in result is dot product between kernel
                    ! and matching-sized section of the padded array
                    res(r,c) = res(r,c) + padded((r-1)*stride(1)+kr, &
                                                 (c-1)*stride(2)+kc) * &
                                          kernel(kr,kc)
                end do
            end do
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! calculates the cross-correlation between a 3D array and a 3D kernel;
! kernel channels match base array channels, so the output is 2D
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:,:)) kernal to pass over a
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:)) stores the output (only 1 channel output from dot-prod)
!-------------------------------------------------------------------------------
! alters :: res becomes cross-correlation result
!-------------------------------------------------------------------------------
subroutine cross_correlate_3D(a, padding, kernel, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernel(:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:), padded(:,:,:)
    integer                   :: a_rows, a_cols, a_channels, &
                                 k_rows, k_cols, k_channels, &
                                 padded_rows, padded_cols, &
                                 res_rows, res_cols, r, c, kr, kc, kd

    a_rows     = size(a, dim=1)
    a_cols     = size(a, dim=2)
    a_channels = size(a, dim=3)
    k_rows     = size(kernel, dim=1)
    k_cols     = size(kernel, dim=2)
    k_channels = size(kernel, dim=3)

    if (a_channels /= k_channels) then
        print *, '--------------------------------------------'
        print *, '(net_helper_functions :: cross_correlate_3D)'
        print *, 'input and kernel channels differ.'
        print *, '--------------------------------------------'
        stop -1
    end if

    if (padding /= 'valid') then
        call pad_3D(a, padding, [k_rows,k_cols], stride, padded)
    else
        padded = a ! no padding to account for
    end if

    padded_rows = size(padded, dim=1)
    padded_cols = size(padded, dim=2)

    ! dimensions of cross-correlation result
    res_rows = res_calc(padded_rows, k_rows, stride(1))
    res_cols = res_calc(padded_cols, k_cols, stride(2))

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [res_rows, res_cols])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(res_rows, res_cols))
    end if

    do r = 1, res_rows
        do c = 1, res_cols
            res(r,c) = 0

            do kr = 1, k_rows
                do kc = 1, k_cols
                    do kd = 1, k_channels
                        ! each index in result is dot product between kernel
                        ! and matching-sized section of the padded array
                        res(r,c) = res(r,c) + padded((r-1)*stride(1)+kr, &
                                                     (c-1)*stride(2)+kc, kd) * &
                                              kernel(kr,kc,kd)
                    end do
                end do
            end do
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! calculates the convolution between a 2D array and a 2D kernel
!-------------------------------------------------------------------------------
! a:        (real(:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:)) kernal to pass over a
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes convolution result
!-------------------------------------------------------------------------------
subroutine convolve_2D(a, padding, kernel, stride, res)
    real(kind=8), intent(in)  :: a(:,:), kernel(:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:), rot_kernel(:,:)
    integer                   :: rows, cols, c

    rows = size(kernel, dim=1)
    cols = size(kernel, dim=2)

    allocate(rot_kernel(rows, cols))

    ! rotate kernel 180 degrees
    do c = 1, cols
        ! res column from left is reverse of column of a from right
        rot_kernel(:,c) = kernel(rows:1:-1, cols+1-c)
    end do

    call cross_correlate_2D(a, padding, rot_kernel, stride, res)
end subroutine

!-------------------------------------------------------------------------------
! calculates the convolution between a 3D array and a 3D kernel;
! kernel channels must match base array channels, so the output is 2D
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:,:)) kernal to pass over a
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:)) stores the output (only 1 channel output from dot-prod)
!-------------------------------------------------------------------------------
! alters :: res becomes convolution result
!-------------------------------------------------------------------------------
subroutine convolve_3D(a, padding, kernel, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernel(:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:), rot_kernel(:,:,:)
    integer                   :: rows, cols, chans, c

    rows  = size(kernel, dim=1)
    cols  = size(kernel, dim=2)
    chans = size(kernel, dim=3)

    allocate(rot_kernel(rows, cols, chans))

    ! rotate kernel 180 degrees
    do c = 1, cols
        ! res column from left is reverse of column of a from right (all depth)
        rot_kernel(:,c,:) = kernel(rows:1:-1, cols+1-c, :)
    end do

    call cross_correlate_3D(a, padding, rot_kernel, stride, res)
end subroutine

!-------------------------------------------------------------------------------
! cross-correlate each 3D kernel (in 4D array) with a 3D input; return a 3D
! array where each layer corresponds to a cross-correlation result by kernel;
!
! used in forward prop for convolutional layers
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernels:  (real(:,:,:,:)) channels of 3D kernels
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes 3D stack of cross-correlations (kernel count depth)
!-------------------------------------------------------------------------------
subroutine cross_correlate_3D_kernels(a, padding, kernels, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernels(:,:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:,:), res_channel(:,:)
    integer                   :: k_count, k

    k_count = size(kernels, dim=4)

    ! cross correlate each kernel with input
    do k = 1, k_count
        call cross_correlate_3D(a, padding, kernels(:,:,:,k), stride, &
                                res_channel)

        ! create new array if not correct size
        if (allocated(res)) then
            if (.not. all(shape(res(:,:,1)) == shape(res_channel)) .or. &
                size(res, dim=3) /= k_count) then
                deallocate(res)
            end if
        end if

        if (.not. allocated(res)) then
            allocate(res(size(res_channel, dim=1), size(res_channel, dim=2), &
                         k_count))
        end if

        res(:,:,k) = res_channel
    end do
end subroutine

!-------------------------------------------------------------------------------
! transpose-convolve each 3D kernel (in 4D array) with a 3D input; return a 3D
! array where each layer corresponds to a transpose-convolution result by kernel
!
! used in forward prop for deconvolutional layers
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernels:  (real(:,:,:,:)) channels of 3D kernels
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes 3D stack of transpose-convolutions (kernel count depth)
!-------------------------------------------------------------------------------
subroutine transpose_convolve_3D_kernels(a, padding, kernels, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernels(:,:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: exp_a(:,:,:), res(:,:,:), res_channel(:,:)
    integer                   :: k_count, k

    if (padding /= 'full') then
        print *, '-------------------------------------------------------'
        print *, '(net_helper_functions :: transpose_convolve_3D_kernels)'
        print *, 'must use full padding.'
        print *, '-------------------------------------------------------'
        stop -1
    end if

    k_count = size(kernels, dim=4)

    ! account for stride
    call expand_with_stride_3D(a, stride, exp_a)

    ! convolve each kernel with input
    do k = 1, k_count
        call convolve_3D(exp_a, 'full', kernels(:,:,:,k), [1,1], res_channel)

        ! create new array if not correct size
        if (allocated(res)) then
            if (.not. all(shape(res(:,:,1)) == shape(res_channel)) .or. &
                size(res, dim=3) /= k_count) then
                deallocate(res)
            end if
        end if

        if (.not. allocated(res)) then
            allocate(res(size(res_channel, dim=1), size(res_channel, dim=2), &
                         k_count))
        end if

        res(:,:,k) = res_channel
    end do
end subroutine

!-------------------------------------------------------------------------------
! cross-correlate all permutations between channels in base array and kernel
! array; result is 4D array, where first 2D are cross-correlation results,
! 3D corresponds to base array channels, 4D corresponds to kernel channels;
! we group the cross-correlation results by kernel channel;
! user can choose to stride either (but not both) arrays with stride;
!
! used in backprop wrt kernels for convolutional layers
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:,:)) kernel array
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:,:)) stores the output
!
! exp_side: (optional - characters) stride expand 'left' (a) or 'right' (kernel)
!-------------------------------------------------------------------------------
! alters :: res becomes cross-correlation permutation result, grouped by kernel
!-------------------------------------------------------------------------------
subroutine cross_correlate_3D_perms_group_kernel(a, padding, kernel, stride, &
                                                 exp_side, res)
    real(kind=8), intent(in)           :: a(:,:,:), kernel(:,:,:)
    character(*), intent(in)           :: padding
    integer, intent(in)                :: stride(2)
    character(*), intent(in), optional :: exp_side
    real(kind=8), allocatable          :: res(:,:,:,:), res_channel(:,:), &
                                          exp_a(:,:,:), exp_k(:,:,:)
    integer                            :: a_channels, k_channels, a_i, k_i

    a_channels = size(a, dim=3)
    k_channels = size(kernel, dim=3)

    ! handle any stride expansions
    if (present(exp_side)) then
        if (exp_side == 'left') then
            call expand_with_stride_3D(a, stride, exp_a) ! a = left
            exp_k = kernel                               ! default
        else if (exp_side == 'right') then
            call expand_with_stride_3D(kernel, stride, exp_k) ! kernel = right
            exp_a = a                                         ! default
        end if
    else
        ! no expanding
        exp_a = a
        exp_k = kernel
    end if

    ! cross-correlate each corresponding channel pair, grouped by b
    do k_i = 1, k_channels
        do a_i = 1, a_channels
            call cross_correlate_2D(exp_a(:,:,a_i), padding, exp_k(:,:,k_i), &
                                    [1,1], res_channel)

            ! create new array if not correct size
            if (allocated(res)) then
                ! deallocate if wrong size
                if (.not. all(shape(res) == [size(res_channel, dim=1), &
                                             size(res_channel, dim=2), &
                                             a_channels, k_channels])) then
                    deallocate(res)
                end if
            end if

            if (.not. allocated(res)) then
                allocate(res(size(res_channel, dim=1), &
                             size(res_channel, dim=2), &
                             a_channels, k_channels))
            end if

            ! in each group of k_i, we store cross-correlate pairs of all a_i
            res(:,:,a_i,k_i) = res_channel
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! cross-correlate all permutations between channels in base array and kernel
! array; result is 4D array, where first 2D are cross-correlation results,
! 3D corresponds to kernel channels, 4D corresponds to base array channels;
! we group the cross-correlation results by base array channel;
! user can choose to stride either (but not both) arrays with stride;
!
! used in backprop wrt kernels for deconvolutional layers
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernel:   (real(:,:,:)) kernel array
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:,:)) stores the output
!
! exp_side: (optional - characters) stride expand 'left' (a) or 'right' (kernel)
!-------------------------------------------------------------------------------
! alters :: res becomes cross-correlation permutation result, grouped by base
!-------------------------------------------------------------------------------
subroutine cross_correlate_3D_perms_group_base(a, padding, kernel, stride, &
                                               exp_side, res)
    real(kind=8), intent(in)           :: a(:,:,:), kernel(:,:,:)
    character(*), intent(in)           :: padding
    integer, intent(in)                :: stride(2)
    character(*), intent(in), optional :: exp_side
    real(kind=8), allocatable          :: res(:,:,:,:), res_channel(:,:), &
                                          exp_a(:,:,:), exp_k(:,:,:)
    integer                            :: a_channels, k_channels, a_i, k_i

    if (padding /= 'full') then
        print *, '-------------------------------------------------------'
        print *, '(net_helper_functions :: transpose_convolve_3D_kernels)'
        print *, 'must use full padding.'
        print *, '-------------------------------------------------------'
        stop -1
    end if

    a_channels = size(a, dim=3)
    k_channels = size(kernel, dim=3)

    ! handle any stride expansions
    if (present(exp_side)) then
        if (exp_side == 'left') then
            call expand_with_stride_3D(a, stride, exp_a) ! a = left
            exp_k = kernel                               ! default
        else if (exp_side == 'right') then
            call expand_with_stride_3D(kernel, stride, exp_k) ! kernel = right
            exp_a = a                                         ! default
        end if
    else
        ! no expanding
        exp_a = a
        exp_k = kernel
    end if

    ! cross-correlate each corresponding channel pair, grouped by b
    do a_i = 1, a_channels
        do k_i = 1, k_channels
            ! assumed for deconvolution; use valid rather than full padding
            call cross_correlate_2D(exp_a(:,:,a_i), 'valid', exp_k(:,:,k_i), &
                                    [1,1], res_channel)

            ! create new array if not correct size
            if (allocated(res)) then
                ! deallocate if wrong size
                if (.not. all(shape(res) == [size(res_channel, dim=1), &
                                             size(res_channel, dim=2), &
                                             k_channels, a_channels])) then
                    deallocate(res)
                end if
            end if

            if (.not. allocated(res)) then
                allocate(res(size(res_channel, dim=1), &
                             size(res_channel, dim=2), &
                             k_channels, a_channels))
            end if

            ! in each group of a_i, we store cross-correlate pairs of all k_i
            res(:,:,k_i,a_i) = res_channel
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! transpose-convolve all permutations between channels in base array and kernels
! in kernel array; result is 3D array, where first 2D are transpose-convolution
! results, and 3D sums up the transpose-convolution by kernel channel (3D of
! kernels, where 3D for channels in each kernel, and 4D is different kernels);
!
! used in backprop wrt inputs for convolutional layers, so we must handle
! undoing padding to match input dimensions in this case
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernels:  (real(:,:,:,:)) array of 3D kernels
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes transpose-convolution permutation result, sum by kernel
!-------------------------------------------------------------------------------
subroutine transpose_convolve_3D_perms_sum_kernel(a, padding, kernels, &
                                                  stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernels(:,:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:,:), res_channel(:,:), exp_a(:,:,:), &
                                 unpad_channel(:,:)
    integer                   :: a_rows, a_cols, a_channels, k_channels, &
                                 res_row_pad, res_col_pad, top_pad, left_pad, &
                                 a_i, k_i
    logical                   :: zero_init

    zero_init  = .false.
    a_rows     = size(a, dim=1)
    a_cols     = size(a, dim=2)
    a_channels = size(a, dim=3)
    k_channels = size(kernels, dim=3)

    ! account for stride
    call expand_with_stride_3D(a, stride, exp_a)

    ! convolve each delta channel with corresponding kernel
    ! channel, then sum up the results by input channel
    do k_i = 1, k_channels
        do a_i = 1, a_channels
            call convolve_2D(exp_a(:,:,a_i), 'full', kernels(:,:,k_i,a_i), &
                             [1,1], res_channel)

            if (padding == 'same') then
                ! res layer results currently includes padding; must remove it
                res_row_pad = size(res_channel, dim=1) - a_rows
                res_col_pad = size(res_channel, dim=2) - a_cols

                ! if odd number pad: top/left get rounded down amount
                top_pad = res_row_pad / 2
                left_pad = res_col_pad / 2

                ! select array (without padding) from res_channel
                unpad_channel = res_channel(top_pad+1:top_pad+a_rows, &
                                            left_pad+1:left_pad+a_cols)
            else if (padding == 'valid') then
                unpad_channel = res_channel ! no padding to remove
            end if

            ! create new array if not correct size
            if (.not. zero_init) then
                if (allocated(res)) then
                    ! deallocate if wrong size
                    if (.not. all(shape(res(:,:,1)) == &
                                  shape(unpad_channel)) .or. &
                        size(res, dim=3) /= k_channels) then
                        deallocate(res)
                    else
                        res = 0 ! allocated, and correct size
                        zero_init = .true.
                    end if
                end if

                if (.not. allocated(res)) then
                    allocate(res(size(unpad_channel, dim=1), &
                                 size(unpad_channel, dim=2), &
                                 k_channels))
                    res = 0
                    zero_init = .true.
                end if
            end if

            ! sum up results by kernel channel
            res(:,:,k_i) = res(:,:,k_i) + unpad_channel
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! cross-correlate all permutations between channels in base array and kernels
! in kernel array; result is 3D array, where first 2D are cross-correlation
! results, and 3D sums up the cross-correlation by kernel channel (3D of
! kernels, where 3D for channels in each kernel, and 4D is different kernels);
!
! used in backprop wrt inputs for deconvolutional layers, so we do not need to
! undo any padding, because deconvolution is assumed have no extra padding
!-------------------------------------------------------------------------------
! a:        (real(:,:,:)) base array
! padding:  (characters) padding type
! kernels:  (real(:,:,:,:)) array of 3D kernels
! stride:   (integer(2)) size of kernel moves in (y, x) directions
! res:      (real(:,:,:)) stores the output
!-------------------------------------------------------------------------------
! alters :: res becomes cross-correlation permutation result, sum by kernel
!-------------------------------------------------------------------------------
subroutine cross_correlate_3D_perms_sum_kernel(a, padding, kernels, stride, res)
    real(kind=8), intent(in)  :: a(:,:,:), kernels(:,:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: stride(2)
    real(kind=8), allocatable :: res(:,:,:), res_channel(:,:)
    integer                   :: a_rows, a_cols, a_channels, k_channels, &
                                 a_i, k_i
    logical                   :: zero_init

    if (padding /= 'full') then
        print *, '-------------------------------------------------------'
        print *, '(net_helper_functions :: transpose_convolve_3D_kernels)'
        print *, 'must use full padding.'
        print *, '-------------------------------------------------------'
        stop -1
    end if

    zero_init  = .false.
    a_rows     = size(a, dim=1)
    a_cols     = size(a, dim=2)
    a_channels = size(a, dim=3)
    k_channels = size(kernels, dim=3)

    ! convolve each delta channel with corresponding kernel
    ! channel, then sum up the results by input channel
    do k_i = 1, k_channels
        do a_i = 1, a_channels
            ! assumed for deconvolution; use valid rather than full padding
            call cross_correlate_2D(a(:,:,a_i), 'valid', &
                                    kernels(:,:,k_i,a_i), stride, res_channel)

            if (.not. zero_init) then
                ! create new array if not correct size
                if (allocated(res)) then
                    ! deallocate if wrong size
                    if (.not. all(shape(res(:,:,1)) == &
                                  shape(res_channel)) .or. &
                        size(res, dim=3) /= k_channels) then
                        deallocate(res)
                    else
                        res = 0 ! allocated, and correct size
                        zero_init = .true.
                    end if
                end if

                if (.not. allocated(res)) then
                    allocate(res(size(res_channel, dim=1), &
                                 size(res_channel, dim=2), &
                                 k_channels))
                    res = 0
                    zero_init = .true.
                end if
            end if

            ! sum up results by kernel channel
            res(:,:,k_i) = res(:,:,k_i) + res_channel
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! calculates the max pooling between a 2D array and a 2D kernel
!-------------------------------------------------------------------------------
! a:           (real(:,:)) base array
! padding:     (characters) pool padding type
! kernel_dims: (integer(2)) (height, width) of kernel
! stride:      (real(2)) size of kernel moves in (y, x) directions
! res:         (real(:,:)) stores pooled output
! res_idxs:    (real(:,:,:)) stores chosen max indices
!-------------------------------------------------------------------------------
! alters :: * res becomes max pooling result
!           * res_idxs matches res (rows, cols), where third dimension of size 2
!             is the index (Y, X) of the max value in the padded original array;
!             -1 where max value is in padding, not original array
!-------------------------------------------------------------------------------
subroutine max_pool_2D(a, padding, kernel_dims, stride, res, res_idxs)
    real(kind=8), intent(in)  :: a(:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: kernel_dims(2), stride(2)
    real(kind=8), allocatable :: res(:,:), padded(:,:)
    integer, allocatable      :: res_idxs(:,:,:)
    integer                   :: max_idx(2), a_rows, a_cols, k_rows, k_cols, &
                                 padded_rows, padded_cols, top_pad, left_pad, &
                                 res_rows, res_cols, r, c

    a_rows = size(a, dim=1)
    a_cols = size(a, dim=2)
    k_rows = kernel_dims(1)
    k_cols = kernel_dims(2)

    if (padding /= 'valid') then
        call pad_2D(a, padding, [k_rows,k_cols], stride, padded)
    else
        padded = a ! no padding to account for
    end if

    padded_rows = size(padded, dim=1)
    padded_cols = size(padded, dim=2)

    top_pad = (padded_rows - a_rows) / 2
    left_pad = (padded_cols - a_cols) / 2

    ! dimensions of pooling result
    res_rows = res_calc(padded_rows, k_rows, stride(1))
    res_cols = res_calc(padded_cols, k_cols, stride(2))

    ! create new array if not correct size
    if (allocated(res)) then
        if (.not. all(shape(res) == [res_rows, res_cols])) then
            deallocate(res)
        end if
    end if

    if (.not. allocated(res)) then
        allocate(res(res_rows, res_cols))
    end if

    ! (rows, cols) of res_idxs matches up entries to res;
    ! depth of 2 corresponds to (row, col) of idx;
    ! create new array if not correct size
    if (allocated(res_idxs)) then
        if (.not. all(shape(res_idxs) == [res_rows, res_cols, 2])) then
            deallocate(res_idxs)
        end if
    end if

    if (.not. allocated(res_idxs)) then
        allocate(res_idxs(res_rows, res_cols, 2))
    end if

    do r = 1, res_rows
        do c = 1, res_cols
            ! each index in the result of the maximum between the
            ! in a section of the padded array matching the kernel;
            ! max_idx is index WITHIN THE KERNEL; NOT overall array
            max_idx = maxloc(padded((r-1)*stride(1)+1:(r-1)*stride(1)+k_rows, &
                                    (c-1)*stride(2)+1:(c-1)*stride(2)+k_cols))

            ! shift in-kernel max index to index in padded array
            max_idx(1) = max_idx(1) + (r-1)*stride(1)
            max_idx(2) = max_idx(2) + (c-1)*stride(2)

            ! set actual max value and index in padded array
            res_idxs(r,c,:) = max_idx
            res(r,c) = padded(max_idx(1), max_idx(2))

            ! if index is in the pad, not the original array, set (y,x) = -1
            if ((max_idx(1) <= top_pad .or. max_idx(1) > top_pad+a_rows) .or. &
                (max_idx(2) <= left_pad .or. max_idx(2) > left_pad+a_cols)) then
                res_idxs(r,c,:) = [-1,-1]
            end if
        end do
    end do
end subroutine

!-------------------------------------------------------------------------------
! calculates the max pooling between channels of arrays and a kernel
!-------------------------------------------------------------------------------
! a:           (real(:,:,:)) channels of base arrays
! padding:     (characters) pool padding type
! kernel_dims: (integer(2)) (height, width) of kernel
! stride:      (real(2)) size of kernel moves in (y, x) directions
! res:         (real(:,:,:)) stores pooled output; overwritten
! res_idxs:    (real(:,:,:,:)) stores chosen max indices; overwritten
!-------------------------------------------------------------------------------
! alters ::    * res becomes 3D max pooling result
!              * res_idxs becomes 4D array of max indices;
!                4th dimension channels,
!                first three dimensions match those returned by res_idxs
!                from max_pool_2D (see description above);
!                -1 where max value in padding, not original array
!-------------------------------------------------------------------------------
subroutine max_pool_3D(a, padding, kernel_dims, stride, res, res_idxs)
    real(kind=8), intent(in)  :: a(:,:,:)
    character(*), intent(in)  :: padding
    integer, intent(in)       :: kernel_dims(2), stride(2)
    real(kind=8), allocatable :: res_channel(:,:), res(:,:,:)
    integer, allocatable      :: channel_idxs(:,:,:), res_idxs(:,:,:,:)
    integer                   :: k, a_channels

    a_channels = size(a, dim=3)

    do k = 1, a_channels
        call max_pool_2D(a(:,:,k), padding, kernel_dims, stride, res_channel, &
                         channel_idxs)

        ! create new array if not correct size
        if (allocated(res)) then
            if (.not. all(shape(res) == [size(res_channel, dim=1), &
                                         size(res_channel, dim=2), &
                                         a_channels])) then
                deallocate(res)
            end if
        end if

        if (.not. allocated(res)) then
            allocate(res(size(res_channel, dim=1), size(res_channel, dim=2), &
                         a_channels))
        end if

        ! create new array if not correct size
        if (allocated(res_idxs)) then
            if (.not. all(shape(res_idxs) == [size(channel_idxs, dim=1), &
                                              size(channel_idxs, dim=2), &
                                              2, a_channels])) then
                deallocate(res_idxs)
            end if
        end if

        if (.not. allocated(res_idxs)) then
            allocate(res_idxs(size(channel_idxs, dim=1), &
                              size(channel_idxs, dim=2), &
                              2, a_channels))
        end if

        ! set max pool and max indices for channel
        res(:,:,k) = res_channel
        res_idxs(:,:,:,k) = channel_idxs
    end do
end subroutine
end module


! decimals.f90 --
!     Module for decimal arithmetic:
!     - the result of something like 0.1+0.2 is
!       exactly 0.3 and not 0.299999998
!     - the numbers as implemented are fixed-point,
!       not floating-point, so the range is
!       somewhat limited.
!     - the precision of the numbers is
!       around 9 digits (exactly)
!
!     NOTE:
!     - No provision yet to deal with overflows!
!     - No rounding modes taken care of yet
!
!     TODO:
!     - Rounding modes
!     - string to decimal
!
!     $Id$
!
module decimal_arithmetic
    implicit none

    !
    ! Use large integers ...
    ! (Perhaps double-precision reals would make it
    ! faster and simpler to implement?)
    !
    integer, parameter, private :: dec_kind = selected_int_kind( 10 )
    type DECIMAL
        integer(kind=dec_kind) :: int
        integer(kind=dec_kind) :: frac
    end type DECIMAL

    integer, parameter, private :: nodecimals = range( 1_dec_kind ) / 2
    integer, parameter, private :: supfrac = 10 ** nodecimals

    interface operator(-)
        module procedure dec_sub
        module procedure dec_umin
    end interface

    interface operator(+)
        module procedure dec_add
    end interface

    interface operator(*)
        module procedure dec_mult
    end interface

    interface operator(/)
        module procedure dec_div
    end interface

contains
subroutine print_decs
    write(*,*) 'Type:      ', dec_kind
    write(*,*) 'Precision: ', nodecimals
end subroutine

! real_to_dec
!     Convert a real to a decimal value
!
! Arguments:
!     value         Decimal number
!
type(DECIMAL) function real_to_dec( value )
    real, intent(in)          :: value

    real_to_dec%int  = int(value)
    real_to_dec%frac = int( (value-aint(value))*supfrac )
    if ( real_to_dec%frac >= supfrac ) then
        real_to_dec%frac = real_to_dec%frac - supfrac
        real_to_dec%int  = real_to_dec%int  + 1
    endif
    if ( real_to_dec%frac <= -supfrac ) then
        real_to_dec%frac = real_to_dec%frac + supfrac
        real_to_dec%int  = real_to_dec%int  - 1
    endif
end function real_to_dec

! dec_to_string
!     Convert a decimal value to a string
!
! Arguments:
!     value         Decimal number
!     decimals      Number of decimals (optional)
!
function dec_to_string( value, decimals ) result(string)
    type(DECIMAL), intent(in)     :: value
    integer, optional             :: decimals
    character(len=2*nodecimals+1) :: string

    integer                       :: dec
    character(len=20)             :: form

    if ( present(decimals) ) then
        dec = decimals
    else
        dec = nodecimals
    endif

    write( form, '(a,i2,a,i2,a,i2,a)' ) '(i', nodecimals, ',''.'',i', &
        nodecimals, '.', dec, ')'

    write( string, form ) value%int, value%frac
end function dec_to_string

! dec_abs
!     Return the absolute value of a decimal number
!
! Arguments:
!     value         Decimal number
!
type(DECIMAL) function dec_abs( value)
    type(DECIMAL), intent(in) :: value

    dec_abs%int  = abs(value%int)
    dec_abs%frac = abs(value%frac)
end function dec_abs

! dec_umin
!     Return the negative value of a decimal number
!     (unary minus)
!
! Arguments:
!     value         Decimal number
!
type(DECIMAL) function dec_umin( value)
    type(DECIMAL), intent(in) :: value

    dec_umin%int  = -value%int
    dec_umin%frac = -value%frac
end function dec_umin

! dec_add
!     Return the sum of two decimal numbers
!
! Arguments:
!     x             First decimal number
!     y             Second decimal number
!
type(DECIMAL) function dec_add( x, y )
    type(DECIMAL), intent(in) :: x
    type(DECIMAL), intent(in) :: y

    dec_add%int  = x%int  + y%int
    dec_add%frac = x%frac + y%frac

    !
    ! Normalise the number
    !
    if ( dec_add%frac > supfrac ) then
        dec_add%frac = dec_add%frac - supfrac
        dec_add%int  = dec_add%int  + 1
    elseif ( dec_add%frac < -supfrac ) then
        dec_add%frac = dec_add%frac + supfrac
        dec_add%int  = dec_add%int  - 1
    endif

    if ( dec_add%int > 0 .and. dec_add%frac < 0 ) then
        dec_add%frac = dec_add%frac + supfrac
        dec_add%int  = dec_add%int  + 1
    elseif ( dec_add%int < 0 .and. dec_add%frac > 0 ) then
        dec_add%frac = dec_add%frac - supfrac
        dec_add%int  = dec_add%int  - 1
    endif

end function dec_add

! dec_sub
!     Return the difference of two decimal numbers
!
! Arguments:
!     x             First decimal number
!     y             Second decimal number
!
type(DECIMAL) function dec_sub( x, y )
    type(DECIMAL), intent(in) :: x
    type(DECIMAL), intent(in) :: y

    dec_sub =  x + (-y)
end function dec_sub

! dec_mult
!     Return the product of two decimal numbers
!
! Arguments:
!     x             First decimal number
!     y             Second decimal number
!
type(DECIMAL) function dec_mult( x, y )
    type(DECIMAL), intent(in) :: x
    type(DECIMAL), intent(in) :: y

    integer(kind=dec_kind) :: i1, i2, f1, f2
    integer(kind=dec_kind) :: if1, if2, if3

    i1 = x%int
    f1 = x%frac
    i2 = y%int
    f2 = y%frac

    if1 = i1 * f2
    if2 = i2 * f1
    if1 = if1 + if2
    if3 = (if1 + if2) / supfrac
    if1 =  if1 - if3  * supfrac

    dec_mult%int  = i1*i2 + if3
    dec_mult%frac = if1   + (f1*f2)/supfrac

end function dec_mult

! dec_div
!     Return the quotient of two decimal numbers
!
! Arguments:
!     x             First decimal number
!     y             Second decimal number
!
type(DECIMAL) function dec_div( x, y )
    type(DECIMAL), intent(in) :: x
    type(DECIMAL), intent(in) :: y

    integer(kind=dec_kind) :: x2, y2
    integer(kind=dec_kind) :: i1, f1, f2, r
    integer                :: k

    x2 = x%int * supfrac + x%frac
    y2 = y%int * supfrac + y%frac

    i1 = x2 / y2
    r  = x2 - y2 * i1
    f2 = 0

    do k = 1,nodecimals
        f1 = r / y2
        f2 = f2 * 10 + f1
        r  = ( r - f1 * y2 ) * 10
    enddo

    dec_div%int  = i1
    dec_div%frac = f2

end function dec_div

end module decimal_arithmetic