subroutine Kz_func(z_pbl,L,u_star0_in,z,K_min,K_z)
implicit none
real, intent(in) :: z_pbl,L,u_star0_in,z,K_min
real, intent(out) :: K_z
real kappa
parameter (kappa=0.4)
real phih,phih_i
real phih_hs,phih_i_hs
real phih_hs_p1,phih_i_hs_p1
real K_zpbl,K_zhs,h_s,delta_K_zhs,u_star0
h_s=0.04*z_pbl
!ustar0 cannot be 0. Set to a low value
!u_star0=max(u_star0_in,0.01)
u_star0=u_star0_in
call phih_func(z,L,phih,phih_i)
K_zpbl=K_min
K_z=K_min
if (L.ge.0) then
!K_z=0.39*u_star0*z*exp(-0.5*(z/0.21/z_pbl)*(z/0.21/z_pbl)) !As in EMEP
!K_z=0.39*u_star0*z*exp(-0.5*(z/0.32/z_pbl)*(z/0.32/z_pbl)) !Adjusted to match under neutral conditions
K_z=0.39*u_star0*z/phih*exp(-0.5*(z/0.32/z_pbl)*(z/0.32/z_pbl)) !Adjusted but with stability added
else
call phih_func(h_s,L,phih_hs,phih_i_hs)
call phih_func(h_s+1.,L,phih_hs_p1,phih_i_hs_p1)
K_zhs=u_star0*kappa*h_s/phih_hs
delta_K_zhs=u_star0*kappa*((h_s+1.)/phih_hs_p1-(h_s)/phih_hs)
if (z.le.h_s) K_z=u_star0*kappa*z/phih
if (z.gt.h_s.and.z.le.z_pbl) K_z=K_zpbl+((z_pbl-z)/(z_pbl-h_s))*((z_pbl-z)/(z_pbl-h_s))*(K_zhs-K_zpbl+(z-h_s)*(delta_K_zhs+2.*(K_zhs-K_zpbl)/(z_pbl-h_s)))
endif
K_z=max(K_z,K_min)
end subroutine Kz_func