! Version 1.0 released by Yi-Chuan Lu on May 20, 2022.
! 
! When using this code, please cite:
! 
! @article{20heatindex,
!   Title   = {Extending the Heat Index},
!   Author  = {Yi-Chuan Lu and David M. Romps},
!   Journal = {Journal of Applied Meteorology and Climatology},
!   Year    = {2022},
!   Volume  = {61},
!   Number  = {10},
!   Pages   = {1367--1383},
!   Year    = {2022},
! }
!
! This headindex function returns the Heat Index in Kelvin. The inputs are:
! - T, the temperature in Kelvin
! - RH, the relative humidity, which is a value from 0 to 1

module heatindex_mod
    implicit none
    double precision, private, parameter::     &
    sigma       = 5.67d-8                     ,&! W/m^2/K^4 , Boltzmann constant
    Ttrip       = 273.16                      ,&! K         , vapor temperature at triple point
    ptrip       = 611.65                      ,&! Pa        , vapor pressure at triple point
    E0v         = 2.3740d6                    ,&! J/kg
    E0s         = 0.3337d6                    ,&! J/kg
    rgasa       = 287.04                      ,&! J/kg/K 
    rgasv       = 461.                        ,&! J/kg/K 
    cva         = 719.                        ,&! J/kg/K
    cvv         = 1418.                       ,&! J/kg/K 
    cvl         = 4119.                       ,&! J/kg/K
    cvs         = 1861                        ,&! J/kg/K
    cpa         = cva + rgasa                 ,&
    cpv         = cvv + rgasv                 ,&
    epsilon     = 0.97                        ,&!           , emissivity of surface
    M           = 83.6                        ,&! kg        , mass of average US adults, fryar2018
    H           = 1.69                        ,&! m         , height of average US adults, fryar2018
    A           = 0.202*(M**0.425)*(H**0.725) ,&! m^2       , DuBois formula, parson2014
    cpc         = 3492.                       ,&! J/kg/K    , specific heat capacity of core, gagge1972
    C           = M*cpc/A                     ,&!           , heat capacity of core
    r           = 124.                        ,&! Pa/K      , Zf/Rf
    Q           = 180.                        ,&! W/m^2     , metabolic rate per skin area
    phi_salt    = 0.9                         ,&!           , vapor saturation pressure level of saline solution
    Tc          = 310.                        ,&! K         , core temeprature
    Pc          = 5612.299124343398           ,&!           , core vapor pressure, phi_salt*pvstar(Tc)
    p           = 1.013d5                     ,&! Pa        , atmospheric pressure
    eta         = 1.43d-6                     ,&! kg/J      , "inhaled mass" / "metabolic rate"
    Pa0         = 1.6d3                       ,&! Pa        , reference air vapor pressure in regions III, IV, V, VI, chosen by Steadman
    Za          = 60.6/17.4                   ,&! Pa m^2/W  , mass transfer resistance through air, exposed part of skin
    Za_bar      = 60.6/11.6                   ,&! Pa m^2/W  , mass transfer resistance through air, clothed part of skin
    Za_un       = 60.6/12.3                   ,&! Pa m^2/W  , mass transfer resistance through air, when being naked
    errT        = 1.d-8                       ,&!           , tolorence of the root solver for heatindex
    err         = 1.d-8                         !           , tolorence of the root solver
    integer, parameter, private:: maxIter = 100 !           , maximum number of iteration of the root solver
    
    private :: pvstar, Le, Qv,Ra, Ra_bar,Ra_un,find_eqvar,find_T,solve1,solve2,solve3,solve4,solve5,solveI,solveII,solveIII,solveIV
    public  :: heatindex
    
contains

pure double precision function pvstar (T)
    implicit none
    double precision, intent(in) :: T
    if (T<=0.) then
        pvstar = 0.
    else if (T<Ttrip) then
        pvstar = ptrip * (T/Ttrip)**((cpv-cvs)/rgasv) * exp( (E0v + E0s -(cvv-cvs)*Ttrip)/rgasv * (1./Ttrip - 1./T) )
    else
        pvstar = ptrip * (T/Ttrip)**((cpv-cvl)/rgasv) * exp( (E0v       -(cvv-cvl)*Ttrip)/rgasv * (1./Ttrip - 1./T) ) 
    end if
end function pvstar

pure double precision function Le (T)
    implicit none
    double precision, intent(in):: T
    Le = (E0v + (cvv-cvl)*(T-Ttrip) + rgasv*T)
end function Le

pure double precision function Qv(Ta,Pa) ! respiratory heat loss, W/m^2
    implicit none
    double precision, intent(in) :: Ta, Pa
    Qv = eta * Q *(cpa*(Tc-Ta) + Le(Tc)*rgasa/(p*rgasv) * ( Pc-Pa ) )
end function Qv
    
pure double precision function Zs(Rs) ! mass transfer resistance through skin, Pa m^2/W
    implicit none
    double precision, intent(in) :: Rs
    if (Rs == 0.0387) then
        Zs = 52.1
    else
        Zs = 6.0e8 * Rs**5
    end if
end function Zs

pure double precision function Ra(Ts,Ta) ! heat transfer resistance through air, exposed part of skin, K m^2/W
    implicit none
    double precision, intent(in) :: Ts,Ta
    double precision             :: hc,phi_rad,hr
    hc      = 17.4
    phi_rad = 0.85
    hr      = epsilon * phi_rad * sigma* (Ts**2 + Ta**2)*(Ts + Ta)
    Ra      = 1./(hc+hr)
end function Ra

pure double precision function Ra_bar(Tf,Ta) ! heat transfer resistance through air, clothed part of skin, K m^2/W
    implicit none
    double precision, intent(in) :: Tf,Ta
    double precision             :: hc,phi_rad,hr
    hc      = 11.6
    phi_rad = 0.79
    hr      = epsilon * phi_rad * sigma* (Tf**2 + Ta**2)*(Tf + Ta)
    Ra_bar  = 1./(hc+hr)
end function Ra_bar

pure double precision function Ra_un(Ts,Ta) ! heat transfer resistance through air, when being naked, K m^2/W
    implicit none
    double precision, intent(in) :: Ts,Ta
    double precision             :: hc,phi_rad,hr
    hc      = 12.3
    phi_rad = 0.80
    hr      = epsilon * phi_rad * sigma* (Ts**2 + Ta**2)*(Ts + Ta)
    Ra_un   = 1./(hc+hr)
end function Ra_un

function find_eqvar(Ta,RH, eqvar_name)
    implicit none
    double precision, intent(in)      :: Ta, RH
    character(len=5),intent(out)      :: eqvar_name
    double precision, dimension(4)    :: find_eqvar
    double precision :: Pa,phi,Rf,Rs,dTcdt,Ts,Ts_bar,Tf,ps, flux1, flux2, flux3, m, m_bar
    
    Pa    = RH*pvstar(Ta)
    phi   = 0.84
    Rs    = 0.0387
    dTcdt = 0
    
    m     = (Pc-Pa)/(Zs(Rs)+Za)
    m_bar = (Pc-Pa)/(Zs(Rs)+Za_bar)
    Ts = solve1(Ta,Pa,Rs,max(0.d0,min(Tc,Ta)-Rs*abs(m)),     max(Tc,Ta)+Rs*abs(m),    err,maxIter)
    Tf = solve2(Ta,Pa,Rs,max(0.d0,min(Tc,Ta)-Rs*abs(m_bar)), max(Tc,Ta)+Rs*abs(m_bar),err,maxIter)
    flux1 = Q-Qv(Ta,Pa)-(1.-phi)*(Tc-Ts)/Rs                   ! C*dTc/dt when Rf=Zf=\inf
    flux2 = Q-Qv(Ta,Pa)-(1.-phi)*(Tc-Ts)/Rs - phi*(Tc-Tf)/Rs  ! C*dTc/dt when Rf=Zf=0
    if (flux1 <= 0.) then ! region I
        eqvar_name = "phi"
        phi = 1.-(Q-Qv(Ta,Pa))*Rs/(Tc-Ts)
        Rf  = huge(0d0)
    else if (flux2 <=0.) then ! region II&III
        eqvar_name = "Rf"
        Ts_bar = Tc - (Q-Qv(Ta,Pa))*Rs/phi + (1./phi -1.)*(Tc-Ts)
        Tf = solve3(Ta,Pa,Rs,Ts_bar,Ta,Ts_bar,err,maxIter)
        Rf = Ra_bar(Tf,Ta)*(Ts_bar-Tf)/(Tf-Ta)
    else ! region IV,V,VI
        Rf = 0.
        flux3 =  Q-Qv(Ta,Pa)-(Tc-Ta)/Ra_un(Tc,Ta)-(phi_salt*pvstar(Tc)-Pa)/Za_un
        if (flux3 <= 0.) then ! region IV,V
            Ts = solve4(Ta,Pa,0d0,Tc,err,maxIter)
            Rs = (Tc-Ts)/(Q-Qv(Ta,Pa))
            eqvar_name = "Rs"
            Ps = Pc - (Pc-Pa)* Zs(Rs)/( Zs(Rs)+Za_un)
            if (Ps > phi_salt * pvstar(Ts)) then  ! region V
                Ts = solve5(Ta,Pa,0d0,Tc,err,maxIter)
                Rs = (Tc-Ts)/(Q-Qv(Ta,Pa))
                eqvar_name = "Rs*"
            end if
        else ! region VI
            Rs = 0.
            eqvar_name = "dTcdt"
            dTcdt = (1./C)* flux3
        end if
    end if
    find_eqvar = (/ phi,Rf,Rs,dTcdt /)
end function find_eqvar    

double precision function find_T( eqvar_name, eqvar, region)
    implicit none
    character(len=5), intent(in)   :: eqvar_name
    double precision, intent(in)   :: eqvar
    character(len=5), intent(out)  :: region
    double precision               :: T

    if ( eqvar_name == "phi") then
        T = solveI( eqvar)
        region = "I"
    else if ( eqvar_name == "Rf") then
        T = solveII( eqvar)
        if (Pa0>pvstar(T)) then
            region = "II"
        else
            region = "III"
        end if
    else if ( eqvar_name == "Rs" .or. eqvar_name == "Rs*") then
        T = solveIII( eqvar)
        if ( eqvar_name == "Rs") then
            region = "IV"
        else
            region = "V"
        end if
    else
        T = solveIV( eqvar)
        region = "VI"
    end if
    find_T = T
end function find_T

double precision function heatindex(Ta,RH)
    implicit none
    double precision, intent(in) :: Ta,RH
    double precision             :: eqvars(4)
    double precision             :: T
    character(len=5)             :: eqvar_name, region
    integer                      :: eqvar_index
    
    eqvars = find_eqvar(Ta,RH, eqvar_name)
    if ( eqvar_name == "phi")                         eqvar_index = 1
    if ( eqvar_name == "Rf")                          eqvar_index = 2
    if ( eqvar_name == "Rs" .or. eqvar_name == "Rs*") eqvar_index = 3
    if ( eqvar_name == "dTcdt")                       eqvar_index = 4
    T = find_T( eqvar_name, eqvars( eqvar_index), region)
    heatindex = T
    if (Ta==0.) heatindex=0.
end function heatindex

!!!!!!!!!!!!!!!!! root solvers !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

double precision function solve1(Ta,Pa,Rs,x1,x2,err,maxIter)
    implicit none
    double precision, intent(in)   :: Ta,Pa,Rs,x1,x2,err
    integer, intent(in):: maxIter
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    
    a = x1
    b = x2
    fa = (a-Ta)/Ra(a,Ta) + (Pc-Pa)/(Zs(Rs)+Za) - (Tc-a)/Rs
    fb = (b-Ta)/Ra(b,Ta) + (Pc-Pa)/(Zs(Rs)+Za) - (Tc-b)/Rs
    if (fa*fb >0.) then
        print*, 'error interval, solve1', Ta,Pa,Rs
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        fc = (c-Ta)/Ra(c,Ta) + (Pc-Pa)/(Zs(Rs)+Za) - (Tc-c)/Rs
        if (abs(b-a) <err) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solve1,"
    solve1 = c
end function solve1

double precision function solve2(Ta,Pa,Rs,x1,x2,err,maxIter)
    implicit none
    double precision, intent(in)   :: Ta,Pa,Rs,x1,x2,err
    integer, intent(in):: maxIter
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    
    a = x1
    b = x2
    fa = (a-Ta)/Ra_bar(a,Ta) + (Pc-Pa)/(Zs(Rs)+Za_bar) - (Tc-a)/Rs
    fb = (b-Ta)/Ra_bar(b,Ta) + (Pc-Pa)/(Zs(Rs)+Za_bar) - (Tc-b)/Rs
    if (fa*fb >0.) then
        print*, 'error interval, solve2'
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        fc = (c-Ta)/Ra_bar(c,Ta) + (Pc-Pa)/(Zs(Rs)+Za_bar) - (Tc-c)/Rs
        if (abs(b-a) <err) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solve2"
    solve2 = c
end function solve2

double precision function solve3(Ta,Pa,Rs,Ts_bar,x1,x2,err,maxIter)
    implicit none
    double precision, intent(in)   :: Ta,Pa,Rs,Ts_bar,x1,x2,err
    integer, intent(in):: maxIter
    integer            :: iter
    double precision   :: a,b,c,fa,fb,fc
    
    a = x1
    b = x2

    fa = (a-Ta)/Ra_bar(a,Ta) + (Pc-Pa)*(a-Ta)/((a-Ta)*(Zs(Rs)+Za_bar)+r*Ra_bar(a,Ta)*(Ts_bar-a)) - (Tc-Ts_bar)/Rs
    fb = (b-Ta)/Ra_bar(b,Ta) + (Pc-Pa)*(b-Ta)/((b-Ta)*(Zs(Rs)+Za_bar)+r*Ra_bar(b,Ta)*(Ts_bar-b)) - (Tc-Ts_bar)/Rs
    if (fa*fb >0.) then
        print*, 'error interval, solve3', Ta, Pa, Rs, Ts_bar, a,b,fa,fb
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        fc = (c-Ta)/Ra_bar(c,Ta) + (Pc-Pa)*(c-Ta)/((c-Ta)*(Zs(Rs)+Za_bar)+r*Ra_bar(c,Ta)*(Ts_bar-c)) - (Tc-Ts_bar)/Rs
        if (abs(b-a) <err) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solve3"
    solve3 = c
end function solve3

double precision function solve4(Ta,Pa,x1,x2,err,maxIter)
    implicit none
    double precision, intent(in)   :: Ta,Pa,x1,x2,err
    integer, intent(in):: maxIter
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    
    a = x1
    b = x2
    fa = (a-Ta)/Ra_un(a,Ta)+(Pc-Pa)/(Zs((Tc-a)/(Q-Qv(Ta,Pa)))+Za_un)-(Q-Qv(Ta,Pa))
    fb = (b-Ta)/Ra_un(b,Ta)+(Pc-Pa)/(Zs((Tc-b)/(Q-Qv(Ta,Pa)))+Za_un)-(Q-Qv(Ta,Pa))
    if (fa*fb >0.) then
        print*, 'error interval, solve4'
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        fc = (c-Ta)/Ra_un(c,Ta)+(Pc-Pa)/(Zs((Tc-c)/(Q-Qv(Ta,Pa)))+Za_un)-(Q-Qv(Ta,Pa))
        if (abs(b-a) <err) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solve4"
    solve4 = c
end function solve4

double precision function solve5(Ta,Pa,x1,x2,err,maxIter)
    implicit none
    double precision, intent(in)   :: Ta,Pa,x1,x2,err
    integer, intent(in):: maxIter
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    
    a = x1
    b = x2
    fa = (a-Ta)/Ra_un(a,Ta) + (phi_salt*pvstar(a)-Pa)/Za_un -(Q-Qv(Ta,Pa))
    fb = (b-Ta)/Ra_un(b,Ta) + (phi_salt*pvstar(b)-Pa)/Za_un -(Q-Qv(Ta,Pa))
    if (fa*fb >0.) then
        print*, 'error interval, solve5'
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        fc = (c-Ta)/Ra_un(c,Ta) + (phi_salt*pvstar(c)-Pa)/Za_un -(Q-Qv(Ta,Pa))
        if (abs(b-a) <err) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solve5"
    solve5 = c
end function solve5

double precision function solveI( eqvar)
    implicit none
    double precision, intent(in)   :: eqvar
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    double precision, dimension(4) :: tmp
    character(len=5)   :: region
    
    a = 0.
    b = 240.
    tmp= find_eqvar(a,1d0,region)
    fa = tmp(1)- eqvar
    tmp= find_eqvar(b,1d0,region)
    fb = tmp(1)- eqvar
    if (fa*fb >0.) then
        print*, 'error interval, solveI'
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        tmp= find_eqvar(c,1d0,region)
        fc = tmp(1)- eqvar
        if (abs(b-a) <errT) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solveI"
    solveI = c
end function solveI

double precision function solveII( eqvar)
    implicit none
    double precision, intent(in)   :: eqvar
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    double precision, dimension(4) :: tmp
    character(len=5)   :: region
    
    a = 230.
    b = 300.
    tmp= find_eqvar(a,min(1d0,Pa0/pvstar(a)),region)
    fa = tmp(2)- eqvar
    tmp= find_eqvar(b,min(1d0,Pa0/pvstar(b)),region)
    fb = tmp(2)- eqvar
    if (fa*fb >0.) then
        print*, 'error interval, solveII', a, b, fa, fb, eqvar
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        tmp= find_eqvar(c,min(1d0,Pa0/pvstar(c)),region)
        fc = tmp(2)- eqvar
        if (abs(b-a) <errT) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solveII"
    solveII = c
end function solveII

double precision function solveIII( eqvar)
    implicit none
    double precision, intent(in)   :: eqvar
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    double precision, dimension(4) :: tmp
    character(len=5)   :: region
    
    a = 295.
    b = 350.
    tmp=find_eqvar(a,Pa0/pvstar(a),region)
    fa = tmp(3)- eqvar
    tmp=find_eqvar(b,Pa0/pvstar(b),region)
    fb = tmp(3)- eqvar
    if (fa*fb >0.) then
        print*, 'error interval, solveIII'
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        tmp=find_eqvar(c,Pa0/pvstar(c),region)
        fc = tmp(3)- eqvar
        if (abs(b-a) <errT) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solveIII"
    solveIII = c
end function solveIII

double precision function solveIV( eqvar)
    implicit none
    double precision, intent(in)   :: eqvar
    integer            :: iter
    double precision               :: a,b,c,fa,fb,fc
    double precision, dimension(4) :: tmp
    character(len=5)   :: region
    
    a = 340.
    b = 800.
    tmp=find_eqvar(a,Pa0/pvstar(a),region)
    fa = tmp(4)- eqvar
    tmp=find_eqvar(b,Pa0/pvstar(b),region)
    fb = tmp(4)- eqvar
    if (fa*fb >0.) then
        print*, 'error interval, solveIV', eqvar
        call exit(-1)
    end if
    do iter = 1, maxIter
        c = (a+b)/2.
        tmp=find_eqvar(c,Pa0/pvstar(c),region)
        fc = tmp(4)- eqvar
        if (abs(b-a) <errT) exit
        if (fb*fc > 0.) then
            b = c
            fb = fc
        else
            a = c
            fa = fc
        end if
    end do
    if (iter == maxIter+1) print*, "maxIter, solveIV"
    solveIV = c
end function solveIV


end module heatindex_mod