!WRF:MODEL_LAYER:PHYSICS
!FROM:ALEJANDRO_SOTO
!UPDATED: AUDREY_CHATAIN 'AC' (2021)
!
!-----------------------------------------------------------------------
!
! This module contains the idealized two stream code from the Tapio    
! Schneider FMS model. This includes elements of the Earth, Titan,     
! Jupiter and Saturn models.                                           
!
! This scheme is based on the radiation scheme from Frierson et al. (2006)
! and Frierson et al. (2007), with updates and modifications from
! O'Gorman and Schneider (2008). Much of the actual code is modified from
! the O'Gorman and Schneider (2008) implementation in FMS. Hearafter,
! O'Gorman and Schneider (2008) is referred to as OGS2008.
! Updates by AC: complete computaion of insolation (new subroutine)
! and scattering of the solar radiations (important on Titan)
! from Schneider (2012) and Thomas & Stamnes (1999). 
!
! Converted to a WRF module by Alejandro Soto and Mark Richardson
! Adapted to mtWRF by AC
!
! References:
!
! Frierson, D. M. W., Held, I. M., and Zurita-Gotor, P. (2006). 
!     A gray-radiation aquaplanet moist GCM. part i: Static stability and 
!     eddy scale. Journal of the Atmospheric Sciences, 63(10):2548–2566.
! 
! Frierson, D. M. W., Held, I. M., and Zurita-Gotor, P. (2007). 
!     A gray-radiation aqua- planet moist GCM. part II: Energy transports 
!     in altered climates. Journal of the Atmospheric Sciences, 64(5):1680–1693.
!
! O'Gorman, P. A. and Schneider, T. (2008). The hydrological cycle over a 
!     wide range of climates simulated with an idealized gcm. Journal of 
!     Climate, 21(15):3815–3832.
!
! Schneider T., Graves S., Schaller E. et al (2012). Polar methane 
!      accumulation and rainstorms on Titan from simulations of the 
!      methane cycle. Nature. 481(7379). + supporting information
!
! Thomas G.E. and Stamnes K. (1999). Radiative Transfer in the Atmosphere
!      and Ocean. Cambridge Atmospheric and Space Science Series
!
!-----------------------------------------------------------------------

MODULE module_ra_gray


USE MODULE_MODEL_CONSTANTS, ONLY : cp, g, stbolt, degrad, obliquity, eccentricity, solar_cst, Ls_perih
USE MODULE_WRF_ERROR
      
implicit none

!private
!public :: gray_radiation


!=====================================================================
! PARAMETER VALUES
!=====================================================================

!#if ( WRF_EARTH)
!real, parameter :: solar_constant = 1360.  ! values from O'Gorman and Schneider [2008]
!real, parameter :: del_sol = 1.2           ! fudge factor to make annual mean work out
!#endif

!! constants (defined in share/MODULE_MODEL_CONSTANTS)
!real, parameter ::  cp = 1044.0, &      ! in the air [J.Kg-1.K-1] 
!                    stbolt = 5.67e-8, & ! [J.s-1.m-2.K-4]
!                    g = 1.352, &     	 ! [m.s-2]
!                    degrad = 0.017453, &   ! pi/180
!                    eccentricity = 0.0565, &    ! for the orbit of Saturn
!                    obliquity = 26.7, &    ! (deg)
!					 solar_cst = 15.04, &  ! (W/m^2)
!                    Ls_perih = 278.		  ! solar longitude at perihelion (deg)
        
        
contains


!=====================================================================
! 'MAIN' SUBROUTINE: CALL compute_insolation AND radiation_down_up
! AND ADAPT TO WRF FORMALISM
!=====================================================================

subroutine gray_radiation(P8W, T3D, PI3D, XLAT, XLON, TSK, ALBEDO, EMISS,             &
                          Lsol, ltst_init, xtime, solar_cst, eccentricity,            &
                          Ls_perih, obliquity, ra_F_lw_d_ztop,                        &
                          SWDOWN, GSW, GLW, RTHRATEN, RTHRATENSW, RTHRATENLW,         &
                          ltst, FLUX_SW_D, FLUX_SW_U, FLUX_LW_D, FLUX_LW_U,           &
                          ids,ide, jds,jde, kds,kde, ims,ime, jms,jme,                &
                          kms,kme, its,ite, jts,jte, kts,kte)
  

!---------------------------------------------------------------------
! DECLARATIONS
!---------------------------------------------------------------------

IMPLICIT NONE

! in variables: ------
!-- TSK           "skin" temperature of the surface, i.e, the ground temp. (K)
!-- XLAT          latitude, south is negative (degree)
!-- XLON          longitude
!-- ALBEDO        albedo (between 0 and 1)
!-- EMISS         surface emissivity (between 0 and 1)

INTEGER, INTENT(IN) ::            ids,ide, jds,jde, kds,kde,      &
                                  ims,ime, jms,jme, kms,kme,      &
                                  its,ite, jts,jte, kts,kte

REAL, DIMENSION(ims:ime, kms:kme, jms:jme), INTENT(IN) ::         &
                                   P8W,                           &
                                   T3D,                           &
                                   PI3D       
                                
REAL, DIMENSION(ims:ime, jms:jme), INTENT(IN) ::                  &
                                   XLAT,                          &
                                   XLON,                          &
                                   TSK,                           &
                                   ALBEDO,                        &
                                   EMISS

REAL, INTENT(IN) :: Lsol, ltst_init, xtime,                       &
                solar_cst, eccentricity, Ls_perih, obliquity,     &
                ra_F_lw_d_ztop
                         
                         
! out variables: ------
!-- RTHRATEN      THETA TENDENCY DUE TO RADIATION
!-- RTHRATENSW    THETA TENDENCY DUE TO LONG WAVE RADIATION
!-- RTHRATENLW    THETA TENDENCY DUE TO SHORT WAVE RADIATION
!-- GLW           downward long wave flux at ground surface (W/m^2)
!-- GSW           net short wave flux at ground surface (W/m^2)

REAL, INTENT(OUT) :: ltst

REAL, DIMENSION(ims:ime, kms:kme, jms:jme), INTENT(INOUT) ::      & ! true tendencies
                                    RTHRATEN,                     & ! in pot'l K/s
                                    RTHRATENSW,                   & 
                                    RTHRATENLW

REAL, DIMENSION(ims:ime, kms:kme, jms:jme), INTENT(OUT) ::      & 
                                    FLUX_SW_D, FLUX_SW_U,         &
                                    FLUX_LW_D, FLUX_LW_U

REAL, DIMENSION( ims:ime, jms:jme ), INTENT(INOUT) :: SWDOWN,     &
                                    GSW,                          &
                                    GLW


! --- local variables 
!-- TOA_LW_U      long wave flux leaving at the top of the atmosphere (W/m^2)
!-- TOA_SW_D      short wave flux entering at the top of the atmosphere (W/m^2)


! definition of the input parameters for the grey scheme model
! most values are from Schneider et al (2012) SI
! some are adjusted to fit fluxes in Tomasko et al (2008)
REAL, PARAMETER ::  ra_gray_odp = 1, &  ! should be =1.5? but it seems already included in lw_tau0
                    ra_gray_lw_tau0_eqtr = 10, & 
                    ra_gray_lw_tau0_pole = 10, &
                    ra_gray_atm_abs = 8, &
                    ra_gray_sw_diff = 0, &
                    ra_gray_lw_linear_frac = 0.15, & ! 0.15
                    ra_gray_lw_tau_exponent = 2.,  &
                    ra_gray_sw_tau_exponent = 0.21, &
                    ra_gray_single_albedo = 0.9, & ! single scattering albedo 0.95
                    ra_gray_asymmetry_factor = 0.85 ! for scattering 0.65

REAL, DIMENSION( ims:ime, jms:jme ) :: mu0, insolation 

REAL, DIMENSION(ims:ime, jms:jme, kms:kme ) ::  p8w_fms,        &
                                                t3d_fms,        &
                                                hr_fms,         &
                                                hr_sw_fms,      &
                                                hr_lw_fms,      &
                                                sw_up_fms,      &
                                                lw_up_fms,      &
                                                sw_down_fms,    &
                                                lw_down_fms,    &
                                                tau_lw_fms,     &
                                                tau_sw_fms 

REAL, DIMENSION(ims:ime, kms:kme, jms:jme) ::                   & 
                                    HR_VIS, HR_IR,              &
                                    tau_gray_sw, tau_gray_lw 

REAL, DIMENSION( ims:ime, jms:jme ) ::  TOA_SW_D, TOA_SW_U,     &
                                        TOA_LW_D, TOA_LW_U,     &
                                        TOA_NET

integer :: i,j,k, is, ie, js, je


! COMMENT ON THE DEFINITION OF VARIABLES --------------------------------
!
! p_full = p3d    Pressure at FMS "full" model levels - layer centre pressures
! p_half = p8w    Pressure at FMS "half" (interface) model levels - layer edge pressures
!
! NOTE that radiation_up_down assumes the surface is at the largest index but WRF physics has the surface at kms
!
! NOTE that because of the inherent meaninglessness of the "half" and "full" level terminology, WRF and FMS 
! arbitrarily pick opposite definitions (might as well call them "pumpkin" and "twinkleberry" levels):
!
! In WRF, kms (smallest number) is the bottom level and kme (largest        In FMS k = kmax is surface and we flip
! number) is the top level.  In your scheme, if 1 is at the top level,      the "meaning" of half and full:
! then you have to reverse the order in the k direction.
!
!                WRF WORLD:                               English Translation:            FMS WORLD:            Whimsy:
!         kme      -   half level (no data at this level)                    
!         kme    ----- full level                         top layer top edge           k=1   ----- half level   pumpkin level
!         kme-1    -   half level                          top layer centre            k=1     -   full level   twinkleberry level
!         kme-1  ----- full level                        top layer bottom edge         k=2   ----- half level   pumpkin level
!         .                                                       .                            .
!         .
!         .
!         kms+2    -   half level
!         kms+2  ----- full level
!         kms+1    -   half level                                 .                            .
!         kms+1  ----- full level                                 .                   kmax   ----- half level   pumpkin level
!         kms      -   half level                        bottom layer centre          kmax     -   full level   twinkleberry level
!         kms    ----- full level                      bottom layer bottom edge      kmax+1  ----- half level   pumkin level
!
!------------------------------------------------------------------------


!---------------------------------------------------------------------
! INITIALIZATIONS
!---------------------------------------------------------------------

ltst = modulo(ltst_init + xtime/60.,24.) ! time at the center of the domain in hours

! zero everything we're going to output over whole tile
do i=its,ite   
    do j=jts,jte
        do k=kts,kte+1
            RTHRATEN(i,k,j) = 0.
            RTHRATENSW(i,k,j) = 0.
            RTHRATENLW(i,k,j) = 0.
            FLUX_SW_U(i,k,j) = 0.
            FLUX_SW_D(i,k,j) = 0.
            FLUX_LW_U(i,k,j) = 0.
            FLUX_LW_D(i,k,j) = 0.
            HR_VIS(i,k,j) = 0.
            HR_IR(i,k,j) = 0.
        enddo
        SWDOWN(i,j) = 0.
        GSW(i,j) = 0.
        GLW(i,j) = 0.
        TOA_SW_D(i,j) = 0.
        TOA_SW_U(i,j) = 0.
        TOA_LW_D(i,j) = 0.
        TOA_LW_U(i,j) = 0.
        TOA_NET(i,j) = 0.
        insolation(i,j) = 0.
        mu0(i,j) = 0.
    enddo
enddo

! define adapted indices for x and y
is = its
ie = min(ite,ide-1)
js = jts
je = min(jte,jde-1)

! WRF orders 3D variables i,k,j - FMS orders them i,j,k
do i=is,ie
    do j=js,je
        do k=kts,kte   
            t3d_fms(i,j,k) = t3d(i,k,j)
            p8w_fms(i,j,k) = p8w(i,k,j)
        enddo
        p8w_fms(i,j,kte+1) = p8w(i,kte+1,j)
    enddo
enddo


!---------------------------------------------------------------------
! COMPUTE INSOLATION (TOA_SW_D)
!---------------------------------------------------------------------

call compute_insolation(solar_cst,                    &
                        eccentricity,                 &
                        Ls_perih,                     &
                        obliquity,                    &
                        ltst,                         &
                        Lsol,                         &
                        XLAT(is:ie,js:je),            &
                        XLON(is:ie,js:je),            &
                        insolation(is:ie,js:je),      &
                        mu0(is:ie,js:je))

! OLD VERSIONS (for information):
!do i=is,ie
!do j=js,je
!#if ( WRF_EARTH == 1 )
!  if ( ra_gray_perpetual_equinox ) then  ! O'Gorman and Schneider [2008] Eqn 2.
!    TOA_SW_D(i,j) = 0.25*solar_constant*(1.0 + del_sol* (1. - 3.*sin(degrad*xlat(i,j))**2)/4. )
!  else
!#endif
!    TOA_SW_D(i,j) = real(dble(solcon)*dble(sunfrac(i,j))/(dble(coszen(i,j))+tiny(1.d0)))  ! tiny gives smallest number, prevent divide by 0
!#if ( WRF_EARTH == 1 )
!  endif
!#endif
!enddo
!enddo


!---------------------------------------------------------------------
! COMPUTATION: SUBROUTINE RADIATION_DOWN_UP
!---------------------------------------------------------------------

call radiation_down_up (odp              = ra_gray_odp,                            &  ! config_flags parameters
                        lw_tau0_eqtr     = ra_gray_lw_tau0_eqtr,                   &
                        lw_tau0_pole     = ra_gray_lw_tau0_pole,                   &
                        atm_abs          = ra_gray_atm_abs,                        &
                        sw_diff          = ra_gray_sw_diff,                        &
                        lw_linear_frac   = ra_gray_lw_linear_frac,                 &
                        lw_tau_exponent  = ra_gray_lw_tau_exponent,                &
                        sw_tau_exponent  = ra_gray_sw_tau_exponent,                &
                        lw_down_ztop     = ra_F_lw_d_ztop,                         & 
                        single_albedo    = ra_gray_single_albedo,                  &
                        asymmetry_factor = ra_gray_asymmetry_factor,               &
                        lat              =        XLAT(is:ie,js:je),               &  ! variables
                        surf_sw_down     =      SWDOWN(is:ie,js:je),               &
                        net_surf_sw_down =         GSW(is:ie,js:je),               &
                        surf_lw_down     =         GLW(is:ie,js:je),               &
                        insolation       =  insolation(is:ie,js:je),               &
                        mu0              =         mu0(is:ie,js:je),               &
                        albedo           =      ALBEDO(is:ie,js:je),               &
                        emis             =       EMISS(is:ie,js:je),               &
                        t_surf           =         TSK(is:ie,js:je),               &
                        p_half           =     p8w_fms(is:ie,js:je,kte+1:kts:-1),  & ! remember to put your thing down, flip it and reverse it
                        t                =     t3d_fms(is:ie,js:je,kte:kts:-1),    & ! (i.e., WRF and FMS are on reverse grids)
                        TOA_LW_D         =    toa_lw_d(is:ie,js:je),               &
                        TOA_SW_U         =    toa_sw_u(is:ie,js:je),               &
                        net_lw_toa       =    toa_lw_u(is:ie,js:je),               &
                        sw_up            =   sw_up_fms(is:ie,js:je,kte+1:kts:-1),  &
                        lw_up            =   lw_up_fms(is:ie,js:je,kte+1:kts:-1),  &
                        sw_down          = sw_down_fms(is:ie,js:je,kte+1:kts:-1),  &
                        lw_down          = lw_down_fms(is:ie,js:je,kte+1:kts:-1),  &
                        tdt              =      hr_fms(is:ie,js:je,kte:kts:-1),    &
                        tdt_ir           =   hr_lw_fms(is:ie,js:je,kte:kts:-1),    &
                        tdt_solar        =   hr_sw_fms(is:ie,js:je,kte:kts:-1),    &
                        lw_tau           =  tau_lw_fms(is:ie,js:je,kte+1:kts:-1),  &
                        sw_tau           =  tau_sw_fms(is:ie,js:je,kte+1:kts:-1)  )


!---------------------------------------------------------------------
! COMPUTATION: FLUXES, HR AND RTHRATEN
!---------------------------------------------------------------------

do i=is,ie
    do j=js,je

        TOA_SW_D(i,j) = insolation(i,j)*mu0(i,j)
        TOA_NET(i,j) = TOA_SW_D(i,j) + TOA_LW_D(i,j) - ( TOA_SW_U(i,j) + TOA_LW_U(i,j) )

        ! switch indices
        do k=kts,kte+1 
            flux_sw_u(i,k,j) = sw_up_fms(i,j,k)
            flux_lw_u(i,k,j) = lw_up_fms(i,j,k) 
            flux_sw_d(i,k,j) = sw_down_fms(i,j,k) 
            flux_lw_d(i,k,j) = lw_down_fms(i,j,k) 

            tau_gray_sw(i,k,j) = tau_sw_fms(i,j,k)
            tau_gray_lw(i,k,j) = tau_lw_fms(i,j,k)
        enddo

        do k=kts,kte
            HR_VIS(i,k,j) = hr_sw_fms(i,j,k)
            HR_IR(i,k,j)  = hr_lw_fms(i,j,k)

            ! divide by exner function (dimensionless)
            RTHRATEN(i,k,j)   = hr_fms(i,j,k)/PI3D(i,k,j)
            RTHRATENSW(i,k,j) = hr_sw_fms(i,j,k)/PI3D(i,k,j)
            RTHRATENLW(i,k,j) = hr_lw_fms(i,j,k)/PI3D(i,k,j)
        enddo
    enddo
enddo

return
end subroutine gray_radiation



!=====================================================================
! SUBROUTINE TO COMPUTE INSOLATION
!=====================================================================

subroutine compute_insolation(solar_cst,eccentricity,Ls_perih,obliquity,&
            ltst_avg,Lsol,lat,lon,insolation,cosSZA)

!---------------------------------------------------------------------
! DECLARATIONS
!---------------------------------------------------------------------

implicit none 

! input parameters and variables
real, intent(in) :: solar_cst, &        ! (W/m^2)
                    eccentricity, &     ! for the orbit of Saturn
                    Ls_perih, &         ! solar longitude at perihelion (deg)
                    obliquity           ! (deg)

real, intent(in) :: ltst_avg, & ! local true solar time at the center of the domain (h)
                    Lsol        ! solar longitute (deg)
      
real, intent(in), dimension(:,:) :: lat, lon    ! latitude and longitude (deg)


! output insolation
real, intent(inout), dimension(:,:) :: insolation,  &  ! (W/m^2)
                                       cosSZA          ! cosine of Solar Zenith Angle
       

! local variables
real :: delta           ! latitude of the subsolar point (rad)
real :: coef
      
real, dimension(size(lat,1), size(lat,2)) ::    &
            ltst, &     ! local true solar time on the grid (h)
            dlon        ! delta longitude between the subsolar point and grid data points (deg)
      
real :: degrad = 0.017453   ! pi/180
integer :: l,m, i,j

! find out the dimensions 
l = size(lat,1)
m = size(lat,2)


!---------------------------------------------------------------------
! COMPUTATION
!---------------------------------------------------------------------

! local true solar time at the grid data points
do i=1,l 
    do j=1,m
        ltst(i,j) = ltst_avg !+ (lon(i,j)-lon(i,m/2+1)) * 24./360.
    enddo
enddo

! delta longitude between the subsolar point and grid data points
dlon(:,:) = (ltst(:,:)-12.)/24.*360.   ! (deg)

! latitude of the subsolar point
delta = obliquity*sin(Lsol*degrad) ! (deg)

! cosine of the Solar Zenith Angle
! (spherical law of cosine)
cosSZA(:,:) = sin(lat(:,:)*degrad)*sin(delta*degrad)+cos(lat(:,:)*degrad)*cos(delta*degrad)*cos(dlon(:,:)*degrad)

! computation of the insolation
coef = solar_cst * ((1.+ eccentricity*cos((Lsol-Ls_perih)*degrad))/(1.-eccentricity**2))**2
do i=1,l 
    do j=1,m
        if (cosSZA(i,j)<=tiny(1.d0)) then
            insolation(i,j) = 0
        else
            insolation(i,j) = coef !* cosSZA(i,j)
        endif
    enddo
enddo

return
end subroutine compute_insolation



!=====================================================================
! SUBROUTINE TO COMPUTE DOWN/UP RADIATIONS (GRAY SCHEME)
!=====================================================================

subroutine radiation_down_up(  odp, lw_tau0_eqtr, lw_tau0_pole, atm_abs, sw_diff, &
                               lw_linear_frac, lw_tau_exponent, sw_tau_exponent,  &
                               lw_down_ztop, single_albedo, asymmetry_factor,     &
                               lat, surf_sw_down, net_surf_sw_down, surf_lw_down, &
                               insolation, mu0, albedo, emis, t_surf,             &
                               p_half, t, toa_lw_d, toa_sw_u, net_lw_toa, sw_up,  &
                               lw_up, sw_down, lw_down, tdt, tdt_ir, tdt_solar,   &
                               lw_tau, sw_tau ) 
                  
! /!\ lat is in deg
                  
! Begin the radiation calculation by computing downward fluxes.
! This part of the calculation does not depend on the surface temperature.
! Finish by calculating the upward fluxes. The downward and upward
! calculations are combined into one subroutine.

!---------------------------------------------------------------------
! DECLARATIONS
!---------------------------------------------------------------------

implicit none

! in variables:
real, intent(in) :: odp, lw_tau0_eqtr, lw_tau0_pole, atm_abs, sw_diff, &
                    lw_linear_frac, lw_tau_exponent,                   &
                    sw_tau_exponent, lw_down_ztop,                     &
                    single_albedo, asymmetry_factor  ! added for scattering

real, intent(in) , dimension(:,:)   :: lat, t_surf, albedo, emis, insolation, mu0

real, intent(in) , dimension(:,:,:) :: t, p_half ! size(p,3) = size(t,3)+1

! out variables:
real, intent(inout) , dimension(:,:)  :: surf_sw_down, net_surf_sw_down,  &
                                         surf_lw_down, &
                                         toa_lw_d, toa_sw_u, net_lw_toa

real, intent(inout), dimension(:,:,:) :: sw_up, lw_up, sw_down,           &
                                         lw_down, tdt, tdt_ir, tdt_solar, &
                                         lw_tau, sw_tau
                                       
! local variables:
real :: degrad = 0.017453   ! pi/180
real :: pi = 3.141593
real, dimension(size(t,1), size(t,2), size(t,3)) :: lw_dtrans, &
                                                    tdt_rad, b
real, dimension(size(t,1), size(t,2), (size(t,3)+1)) :: rad_flux, &
                                                        lw_flux, &
                                                        sw_flux
real, dimension(size(t,1), size(t,2)) :: b_surf, lw_tau0, sw_tau0
! for scattering
real :: albedo_inf, Ga   
real :: mu_avg, b_avg
!real, dimension(size(t,1), size(t,2)) :: intensity_inflation, albedo_mod
real, dimension(size(t,1), size(t,2)) :: X_up, X_down, &
                    Z_up, Z_down, rho_L, denom, K1, K2, b_mu0

integer :: k, l, m, n


! find out the vertical dimension and allocate arrays we will need:
l = size(t,1)
m = size(t,2)
n = size(t,3)


! initialize arrays
surf_sw_down(:,:) = 0.
net_surf_sw_down(:,:) = 0.
surf_lw_down(:,:) = 0.
toa_lw_d(:,:) = 0.
toa_sw_u(:,:) = 0.
net_lw_toa(:,:) = 0.

sw_up(:,:,:) = 0.
lw_up(:,:,:)=0.
sw_down(:,:,:)=0.
lw_down(:,:,:)=0.
tdt(:,:,:)=0.
tdt_ir(:,:,:)=0.
tdt_solar(:,:,:)=0.
lw_tau(:,:,:)=0.
sw_tau(:,:,:)=0.

X_up(:,:) = 0.
X_down(:,:) = 0.
Z_up(:,:) = 0.
Z_down(:,:) = 0.
rho_L(:,:) = 0.
denom(:,:) = 0.
K1(:,:) = 0.
K2(:,:) = 0.
b_mu0(:,:) = 0.


!---------------------------------------------------------------------
! DOWNWARD RADIATION CALCULATIONS
!---------------------------------------------------------------------

! Below, we calculate the right hand bracket on the right hand side of
! Equation 1 in OGS2008.
lw_tau0 = lw_tau0_eqtr + (lw_tau0_pole - lw_tau0_eqtr)*sin(degrad*lat)**2

! We can scale the lw_tau0 by the optical depth parameter, odp
lw_tau0 = lw_tau0*odp

! Calculate the shortwave optical depth. Using the default values, the
! SW optical depth is 0.2, similar to the 0.22 value in OGS2008. This 
! will need to be tuned for Mars and other planets.
sw_tau0 = (1.0 - sw_diff*sin(degrad*lat)**2)*atm_abs

! The first calculation in the loop below solves for the entire Equation 1
! in OGS2008, integrating in the vertical. The second line calculates
! part of the exponential term in Equation 3 in OGS. 
do k = 1, n+1
   lw_tau(:,:,k)  = lw_tau0(:,:) * ((lw_linear_frac * p_half(:,:,k)/p_half(:,:,n+1)) &
       + ((1.0 - lw_linear_frac) * (p_half(:,:,k)/p_half(:,:,n+1) )**lw_tau_exponent))
   sw_tau(:,:,k)  = sw_tau0(:,:) * (p_half(:,:,k)/p_half(:,:,n+1) )**sw_tau_exponent
end do

! Calculate the blackbody (for lw).
b = stbolt*t**4

! Calculate the longwave transmissivity.
do k = 1, n
  lw_dtrans(:,:,k) = exp(-(lw_tau(:,:,k+1)-lw_tau(:,:,k)))
end do

! Calculate the net, vertically-integrated downward longwave radiative flux.
lw_down(:,:,1) = lw_down_ztop ! depends on ztop
toa_lw_d(:,:) = lw_down(:,:,1)
do k = 1,n
    lw_down(:,:,k+1) = lw_down(:,:,k)*lw_dtrans(:,:,k) + b(:,:,k)*(1.0 - lw_dtrans(:,:,k))
end do


! Calculate parameters for the sw scattering computation:
! (albedo_inf = albedo of a semi-inifinite scattering atmosphere)
! (Ga = rescaling factor for optical depth due to scattering)

!! Version 2: with scattering included (Petty 2006 or Adamson 1975)
!albedo_inf  = ( sqrt(1. - asymmetry_factor*single_albedo) - sqrt(1. - single_albedo) )  &
!           / ( sqrt(1. - asymmetry_factor*single_albedo) + sqrt(1. - single_albedo) )
!albedo_mod(:,:)  = (albedo(:,:) - albedo_inf) / (1. - albedo(:,:)*albedo_inf)
!intensity_inflation(:,:)  = 1. / (1. + albedo_mod(:,:) * albedo_inf * exp(-2.*Ga*sw_tau0(:,:)))
!! x2 if Petty 2006 (isotropic); only x1 if Adamson 1975 (1D)
!Ga = 2.*sqrt( (1. - single_albedo) * (1. - asymmetry_factor*single_albedo) ) 

! Version 3: anisotropic solution, Thomas & Stamnes 1999 (TS99)
if (insolation(1,1)>0) then
    mu_avg = 1./sqrt(3.) ! guess (for beams) from TS99 p.269
    b_mu0(:,:) = 0.5*(1.-3.*asymmetry_factor*mu0(:,:)*mu_avg)
    b_avg = 0.5*(1.-3.*asymmetry_factor*mu_avg**2)
    albedo_inf = (sqrt(1.-single_albedo+2.*single_albedo*b_avg)-sqrt(1.-single_albedo)) &
             /(sqrt(1.-single_albedo+2.*single_albedo*b_avg)+sqrt(1.-single_albedo))
    Ga = 1./mu_avg*sqrt((1.-single_albedo)*(1.-single_albedo+2.*single_albedo*b_avg))

    X_up(:,:) = single_albedo/(2.*pi) *insolation(:,:) *b_mu0(:,:)
    X_down(:,:) = single_albedo/(2.*pi) *insolation(:,:) *(1.-b_mu0(:,:))
    Z_up(:,:) = (single_albedo*b_avg*X_down(:,:) &
           + (1.-single_albedo+single_albedo*b_avg-mu_avg/(mu0(:,:)+tiny(1.d0)))*X_up(:,:)) &
           /((1.-single_albedo)*(1.-single_albedo+2.*single_albedo*b_avg) &
           - (mu_avg/(mu0(:,:)+tiny(1.d0)))**2)
    Z_down(:,:) = (single_albedo*b_avg*X_up(:,:) &
             + (1.-single_albedo+single_albedo*b_avg+mu_avg/(mu0(:,:)+tiny(1.d0)))*X_down(:,:)) &
             /((1.-single_albedo)*(1.-single_albedo+2.*single_albedo*b_avg) &
             - (mu_avg/(mu0(:,:)+tiny(1.d0)))**2)
    rho_L(:,:) = albedo(:,:)/(2.*pi*mu_avg)
    denom(:,:)  = (1.-albedo(:,:)*albedo_inf)*exp(Ga*sw_tau0(:,:)) &
              + albedo_inf*(albedo(:,:)-albedo_inf)*exp(-Ga*sw_tau0(:,:))
    K1(:,:) = ((rho_L(:,:)*mu0(:,:)*insolation(:,:)-Z_up(:,:)+albedo(:,:)*Z_down(:,:)) &
          *exp(-sw_tau0(:,:)/(mu0(:,:)+tiny(1.d0))) + (albedo_inf-albedo(:,:))*Z_down(:,:) &
          *exp(-Ga*sw_tau0(:,:))) / denom(:,:)
    K2(:,:) = (-albedo_inf*(rho_L(:,:)*mu0(:,:)*insolation(:,:)-Z_up(:,:)+albedo(:,:)*Z_down(:,:)) &
          *exp(-sw_tau0(:,:)/(mu0(:,:)+tiny(1.d0))) - (1.-albedo(:,:)*albedo_inf)*Z_down(:,:) &
          *exp(Ga*sw_tau0(:,:))) / denom(:,:)
endif

! Calculate the net, vertically-integrated downward shortwave radiative flux.
!! Version 1: no scattering (this is the entirety of Equation 3 in OGS2008).
!do k = 1,n+1
!   sw_down(:,:,k) = insolation(:,:)*exp(-sw_tau(:,:,k))
!end do
!! Version 2: with scattering included (see Petty 2006 or Adamson 1975)
!do k = 1,n+1
!   sw_down(:,:,k) = insolation(:,:) * intensity_inflation(:,:) &
!                    * (albedo_inf * albedo_mod(:,:) &
!                       * exp(Ga * (sw_tau(:,:,k) - 2.*sw_tau0(:,:))) &
!                       + exp(-Ga * sw_tau(:,:,k)))
!end do
! Version 3: Thomas & Stamnes 1999 (TS99)
do k = 1,n+1
   if (insolation(1,1)>0) then
      sw_down(:,:,k) = 2.*pi*mu_avg*(albedo_inf*K1(:,:)*exp(Ga*sw_tau(:,:,k)) &
                                 + K2(:,:)*exp(-Ga*sw_tau(:,:,k)) &
                                 + Z_down(:,:)*exp(-sw_tau(:,:,k)/(mu0(:,:)+tiny(1.d0)))) &
                   + mu0(:,:)*insolation(:,:)*exp(-sw_tau(:,:,k)/(mu0(:,:)+tiny(1.d0)))
   else
      sw_down(:,:,k) = 0.
   end if
end do


! Save the results.
!------- downward lw flux surface -------
! surf_lw_down
!------- downward sw flux surface -------
! net_surf_sw_down
surf_lw_down     = lw_down(:,:,n+1)                      ! GLW in main code.
surf_sw_down     = sw_down(:,:,n+1)                      ! SWDOWN in main code.
net_surf_sw_down = sw_down(:,:,n+1)*(1. - albedo(:,:))   ! GSW in the main code.
    

!---------------------------------------------------------------------
! UPWARD RADIATION CALCULATIONS
!---------------------------------------------------------------------

!------- upward lw flux surface -------
! b_surf
!------- outgoing lw flux toa (olr) -------
! net_lw_toa
!------- net upward lw flux surface ------
! net_lw_surf (not used)
!------- outgoing sw flux toa -------
! toa_sw_u

! Calculate the blackbody radition of the surface.
b_surf = emis*stbolt*t_surf**4

! compute upward longwave flux by integrating upward
lw_up(:,:,n+1) = b_surf
do k = n,1,-1
   lw_up(:,:,k) = lw_up(:,:,k+1)*lw_dtrans(:,:,k) + b(:,:,k)*(1.0 - lw_dtrans(:,:,k))
end do

! compute shortwave flux
!! Version 1: no scattering
!do k = 1,n+1
!   sw_up(:,:,k) = albedo(:,:)*sw_down(:,:,n+1)
!end do
!! Version 2: with scattering included (see Petty 2006 or Adamson 1975)
!do k = 1,n+1
!   sw_up(:,:,k) = insolation(:,:) * intensity_inflation(:,:)  &
!                  * (albedo_mod(:,:) * exp(Ga * (sw_tau(:,:,k) - 2.*sw_tau0(:,:))) & 
!                     + albedo_inf * exp(-Ga * sw_tau(:,:,k)))
!end do
! Version 3: anisotropic solution, Thomas & Stamnes 1999 (TS99)
do k = 1,n+1
   if (insolation(1,1)>0) then
       sw_up(:,:,k) = 2.*pi*mu_avg*(K1(:,:)*exp(Ga*sw_tau(:,:,k)) &
                               + albedo_inf*K2(:,:)*exp(-Ga*sw_tau(:,:,k)) &
                               + Z_up(:,:)*exp(-sw_tau(:,:,k)/(mu0(:,:)+tiny(1.d0))))
   else
       sw_up(:,:,k) = 0.
   end if
end do

! Save the results.
net_lw_toa  = lw_up(:,:,1)
toa_sw_u    = sw_up(:,:,1)


!---------------------------------------------------------------------
! TOTAL FLUX AND TEMPERATURE TENDENCY CALCULATIONS
!---------------------------------------------------------------------

!------- longwave radiative flux (at half levels) --------
! lw_flux
!------- shortwave radiative flux (at half levels) --------
! sw_flux
!------- total radiative flux (at half levels) -----------
! rad_flux
!------- temperature tendency due to radiation ------------
! tdt_rad
!------- temperature tendency due to solar radiation ------------
! tdt_solar

! net fluxes (positive up)
lw_flux  = lw_up - lw_down
sw_flux  = sw_up - sw_down
rad_flux = lw_flux + sw_flux

! temperature tendency
do k = 1,n   ! tdt and tdt_... are all in kinetic K/s
  tdt_rad(:,:,k) = ( rad_flux(:,:,k+1) - rad_flux(:,:,k) )  &
             *g/(cp*(p_half(:,:,k+1)-p_half(:,:,k)))
  tdt_solar(:,:,k) = ( sw_flux(:,:,k+1) - sw_flux(:,:,k))  &
             *g/(cp*(p_half(:,:,k+1)-p_half(:,:,k)))
enddo

do k = 1,n
  tdt(:,:,k) = tdt(:,:,k) + tdt_rad(:,:,k)
  tdt_ir(:,:,k) = tdt_rad(:,:,k) - tdt_solar(:,:,k)
end do

return
end subroutine radiation_down_up


! ==================================================================================
! ==================================================================================

end module module_ra_gray
