! This file has been modified for the use in ICON

! radiation_single_level.F90 - Derived type for single-level fields
!
! (C) Copyright 2014- ECMWF.
!
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
!
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!
! Author:  Robin Hogan
! Email:   r.j.hogan@ecmwf.int
!
! Modifications
!   2019-01-14  R. Hogan  Added out_of_physical_bounds routine
!   2019-01-18  R. Hogan  Added weighted albedo mapping

module radiation_single_level

  use parkind1, only : jprb
  use radiation_io, only : nulerr, radiation_abort

  implicit none
  public

  !---------------------------------------------------------------------
  ! Derived type to contain variables that don't vary with height;
  ! mostly they are surface properties
  type single_level_type
    ! Note that skin temperature is only defined if
    ! is_simple_surface==.true.
    real(jprb), allocatable, dimension(:) :: &
         &   cos_sza, &       ! (ncol) Cosine of solar zenith angle
         &   skin_temperature ! (ncol) Skin temperature (K)

    ! Shortwave albedo: if sw_albedo_direct is not allocated then
    ! sw_albedo will be used for both direct and diffuse solar
    ! radiation; otherwise, sw_albedo will be used for diffuse
    ! radiation and sw_albedo_direct for direct radiation.
    real(jprb), allocatable, dimension(:,:) :: &
         &   sw_albedo, &        ! (ncol,nalbedobands)
         &   sw_albedo_direct    ! (ncol,nalbedobands) 
    
    ! If is_simple_surface==.true., we provide longwave emissivity
    ! coarser spectral intervals, while if
    ! is_simple_surface==.false. it will be generated by the surface
    ! radiation scheme and may optionally be on the full spectral grid
    ! of the atmospheric model.
    real(jprb), allocatable, dimension(:,:) :: &
         &   lw_emissivity       ! (ncol,nemissbands) If
    ! is_simple_surface==.false. then we specify the emission instead
    ! of the skin temperature, and again this may be on the emissivity
    ! bands or the full spectral grid
    real(jprb), allocatable, dimension(:,:) :: &
         &   lw_emission         ! (ncol,nemissbands)

    ! Incoming solar irradiance at the Earth is specified as the total
    ! across the shortwave spectrum. Note that this needs to be
    ! adjusted for Earth-Sun distance, solar cycle etc. To get
    ! normalized fluxes out, simply set this to 1.0.
    real(jprb) :: solar_irradiance = 1366.0_jprb ! W m-2

    ! In addition to the effect of the solar cycle on total solar
    ! irradiance (which the user is expected to input via
    ! solar_irradiance), it can change the balance between different
    ! parts of the solar spectrum, e.g. more UV at solar maximum. This
    ! variable provides the scaling to be applied to the spectral
    ! amplitude of the solar cycle (assuming it has been provided),
    ! e.g. 1.0 for solar maximum, -1.0 for solar maximum and 0.0 for
    ! a mean solar spectrum.
    real(jprb) :: spectral_solar_cycle_multiplier = 0.0_jprb
    
    ! If config%use_spectral_solar_irradiance==true then this will be
    ! scaled by spectral_solar_scaling
    real(jprb), allocatable, dimension(:) :: &
         &   spectral_solar_scaling ! (n_bands_sw)
 
    ! Seed for random number generator in McICA; it is expected that
    ! the host model generates this from the model time, longitude
    ! and latitude, in order that the model is deterministic
    integer, allocatable, dimension(:) :: iseed ! (ncol)

    ! Is the underlying surface a simple single "flat" tile? If so we
    ! describe it with skin_temperature and lw_emissivity. Otherwise
    ! we describe it with lw_emission and lw_emissivity coming out of
    ! the surface radiation scheme.
    logical :: is_simple_surface = .true.

    ! If is_simple_surface==.false., do we use the full number of g
    ! points for dimensioning sw_albedo, sw_albedo_direct and
    ! lw_emission?
    !logical :: use_full_spectrum_sw = .false.
    !logical :: use_full_spectrum_lw = .true.

  contains
    procedure :: allocate   => allocate_single_level
    procedure :: deallocate => deallocate_single_level
    procedure :: init_seed_simple
    procedure :: get_albedos
    procedure :: out_of_physical_bounds
#ifdef _OPENACC
    procedure :: update_host
    procedure :: update_device
#endif

  end type single_level_type

contains
  
  !---------------------------------------------------------------------
  ! Allocate the arrays of a single-level type
  subroutine allocate_single_level(this, ncol, nalbedobands, nemisbands, &
       &                           use_sw_albedo_direct, is_simple_surface)

    use yomhook, only : lhook, dr_hook, jphook

    class(single_level_type), intent(inout) :: this
    integer,                  intent(in)    :: ncol, nalbedobands, nemisbands
    logical,        optional, intent(in)    :: use_sw_albedo_direct
    logical,        optional, intent(in)    :: is_simple_surface

    real(jphook) :: hook_handle

    if (lhook) call dr_hook('radiation_single_level:allocate',0,hook_handle)

    call this%deallocate()

    if (present(is_simple_surface)) then
      this%is_simple_surface = is_simple_surface
    else
      this%is_simple_surface = .true.
    end if

    allocate(this%cos_sza(ncol))
    !$ACC ENTER DATA CREATE(this%cos_sza) ASYNC(1)

    if (this%is_simple_surface) then
      allocate(this%skin_temperature(ncol))
      !$ACC ENTER DATA CREATE(this%skin_temperature) ASYNC(1)
    else
      allocate(this%lw_emission(ncol, nemisbands))
      !$ACC ENTER DATA CREATE(this%lw_emission) ASYNC(1)
    end if
    allocate(this%lw_emissivity(ncol, nemisbands))
    !$ACC ENTER DATA CREATE(this%lw_emissivity) ASYNC(1)

    allocate(this%sw_albedo(ncol, nalbedobands))
    !$ACC ENTER DATA CREATE(this%sw_albedo) ASYNC(1)

    if (present(use_sw_albedo_direct)) then
      if (use_sw_albedo_direct) then
        allocate(this%sw_albedo_direct(ncol, nalbedobands))
        !$ACC ENTER DATA CREATE(this%sw_albedo_direct) ASYNC(1)
      end if
    end if

    allocate(this%iseed(ncol))
    !$ACC ENTER DATA CREATE(this%iseed) ASYNC(1)

    if (lhook) call dr_hook('radiation_single_level:allocate',1,hook_handle)

  end subroutine allocate_single_level


  !---------------------------------------------------------------------
  ! Deallocate the arrays of a single-level type
  subroutine deallocate_single_level(this)

    use yomhook, only : lhook, dr_hook, jphook

    class(single_level_type), intent(inout) :: this

    real(jphook) :: hook_handle

    if (lhook) call dr_hook('radiation_single_level:deallocate',0,hook_handle)

    if (allocated(this%cos_sza)) then
      !$ACC EXIT DATA DELETE(this%cos_sza) WAIT(1)
      deallocate(this%cos_sza)
    end if
    if (allocated(this%skin_temperature)) then
      !$ACC EXIT DATA DELETE(this%skin_temperature) WAIT(1)
      deallocate(this%skin_temperature)
    end if
    if (allocated(this%sw_albedo)) then
      !$ACC EXIT DATA DELETE(this%sw_albedo) WAIT(1)
      deallocate(this%sw_albedo)
    end if
    if (allocated(this%sw_albedo_direct)) then
      !$ACC EXIT DATA DELETE(this%sw_albedo_direct) WAIT(1)
      deallocate(this%sw_albedo_direct)
    end if
    if (allocated(this%lw_emissivity)) then
      !$ACC EXIT DATA DELETE(this%lw_emissivity) WAIT(1)
      deallocate(this%lw_emissivity)
    end if
    if (allocated(this%lw_emission)) then
      !$ACC EXIT DATA DELETE(this%lw_emission) WAIT(1)
      deallocate(this%lw_emission)
    end if
    if (allocated(this%spectral_solar_scaling)) then
      !$ACC EXIT DATA DELETE(this%spectral_solar_scaling) WAIT(1)
      deallocate(this%spectral_solar_scaling)
    end if
    if (allocated(this%iseed)) then
      !$ACC EXIT DATA DELETE(this%iseed) WAIT(1)
      deallocate(this%iseed)
    end if

    if (lhook) call dr_hook('radiation_single_level:deallocate',1,hook_handle)

  end subroutine deallocate_single_level


  !---------------------------------------------------------------------
  ! Unimaginative initialization of random-number seeds
  subroutine init_seed_simple(this, istartcol, iendcol)
    class(single_level_type), intent(inout) :: this
    integer, intent(in)                     :: istartcol, iendcol

    integer :: jcol

    if (.not. allocated(this%iseed)) then
      allocate(this%iseed(istartcol:iendcol))
      !$ACC ENTER DATA CREATE(this%iseed)
    end if

    !$ACC PARALLEL DEFAULT(PRESENT)
    !$ACC LOOP GANG VECTOR
    do jcol = istartcol,iendcol
      this%iseed(jcol) = jcol
    end do
    !$ACC END PARALLEL

    !$ACC UPDATE HOST(this%iseed)

  end subroutine init_seed_simple


  !---------------------------------------------------------------------
  ! Extract the shortwave and longwave surface albedos in each g-point
  subroutine get_albedos(this, istartcol, iendcol, config, &
       &                 sw_albedo_direct, sw_albedo_diffuse, lw_albedo)

    use radiation_config, only : config_type
    use radiation_io,     only : nulerr, radiation_abort
    use yomhook,          only : lhook, dr_hook, jphook

    class(single_level_type), intent(in) :: this
    type(config_type),        intent(in) :: config
    integer,                  intent(in) :: istartcol, iendcol

    ! The longwave albedo of the surface in each longwave g-point;
    ! note that these are weighted averages of the values from
    ! individual tiles
    real(jprb), intent(out), optional &
         &  :: lw_albedo(config%n_g_lw, istartcol:iendcol)

    ! Direct and diffuse shortwave surface albedo in each shortwave
    ! g-point; note that these are weighted averages of the values
    ! from individual tiles
    real(jprb), intent(out), dimension(config%n_g_sw, istartcol:iendcol) &
         &  :: sw_albedo_direct, sw_albedo_diffuse

    ! Temporary storage of albedo in ecRad bands
    real(jprb) :: sw_albedo_band(istartcol:iendcol, config%n_bands_sw)
    real(jprb) :: lw_albedo_band(istartcol:iendcol, config%n_bands_lw)

    ! Number of albedo bands
    integer :: nalbedoband

    ! Loop indices for ecRad bands and albedo bands
    integer :: jband, jalbedoband, jg, jcol

    real(jphook) :: hook_handle

    if (lhook) call dr_hook('radiation_single_level:get_albedos',0,hook_handle)

    !$ACC DATA CREATE(sw_albedo_band, lw_albedo_band) &
    !$ACC     PRESENT(this, config, sw_albedo_direct, sw_albedo_diffuse, &
    !$ACC             lw_albedo)

    if (config%do_sw) then
      ! Albedos/emissivities are stored in single_level in their own
      ! spectral intervals and with column as the first dimension
      if (config%use_canopy_full_spectrum_sw) then
        ! Albedos provided in each g point
        if (size(this%sw_albedo,2) /= config%n_g_sw) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%sw_albedo does not have the expected ', &
               &  config%n_g_sw, ' spectral intervals'
          call radiation_abort()
        end if
        sw_albedo_diffuse = transpose(this%sw_albedo(istartcol:iendcol,:))
        if (allocated(this%sw_albedo_direct)) then
          sw_albedo_direct = transpose(this%sw_albedo_direct(istartcol:iendcol,:))
        end if
      else if (.not. config%do_nearest_spectral_sw_albedo) then
        ! Albedos averaged accurately to ecRad spectral bands
        nalbedoband = size(config%sw_albedo_weights,1)
        if (size(this%sw_albedo,2) /= nalbedoband) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%sw_albedo does not have the expected ', &
               &  nalbedoband, ' bands'
          call radiation_abort()
        end if

        !$ACC PARALLEL DEFAULT(NONE) ASYNC(1)
        !$ACC LOOP SEQ
        do jband = 1,config%n_bands_sw
        !$ACC LOOP GANG(STATIC:1) VECTOR
          do jcol = istartcol,iendcol
            sw_albedo_band(jcol,jband) = 0.0_jprb
          end do
        end do

        !$ACC LOOP SEQ
        do jband = 1,config%n_bands_sw
          !$ACC LOOP SEQ
          do jalbedoband = 1,nalbedoband
#ifndef _OPENACC
            if (config%sw_albedo_weights(jalbedoband,jband) /= 0.0_jprb) then
#endif
              !$ACC LOOP GANG(STATIC:1) VECTOR
              do jcol = istartcol,iendcol
                sw_albedo_band(jcol,jband) &
                    &  = sw_albedo_band(jcol,jband) & 
                    &  + config%sw_albedo_weights(jalbedoband,jband) &
                    &    * this%sw_albedo(jcol, jalbedoband)
              end do
#ifndef _OPENACC
            end if
#endif
          end do
        end do

#ifndef _OPENACC
        sw_albedo_diffuse = transpose(sw_albedo_band(istartcol:iendcol, &
             &                              config%i_band_from_reordered_g_sw))
#else
        !$ACC LOOP GANG(STATIC:1) VECTOR
        do jcol = istartcol,iendcol
          !$ACC LOOP SEQ
          do jg = 1,config%n_g_sw
            sw_albedo_diffuse(jg,jcol) = sw_albedo_band(jcol, &
                &                             config%i_band_from_reordered_g_sw(jg))
          end do
        end do
#endif
        if (allocated(this%sw_albedo_direct)) then
          !$ACC LOOP SEQ
          do jband = 1,config%n_bands_sw
            !$ACC LOOP GANG(STATIC:1) VECTOR
            do jcol = istartcol,iendcol
              sw_albedo_band(jcol,jband) = 0.0_jprb
            end do
          end do

          !$ACC LOOP SEQ
          do jband = 1,config%n_bands_sw
            !$ACC LOOP SEQ
            do jalbedoband = 1,nalbedoband
#ifndef _OPENACC
              if (config%sw_albedo_weights(jalbedoband,jband) /= 0.0_jprb) then
#endif
                !$ACC LOOP GANG(STATIC:1) VECTOR
                do jcol = istartcol,iendcol
                  sw_albedo_band(jcol,jband) &
                      &  = sw_albedo_band(jcol,jband) & 
                      &  + config%sw_albedo_weights(jalbedoband,jband) &
                      &    * this%sw_albedo_direct(jcol, jalbedoband)
                end do
#ifndef _OPENACC
              end if
#endif
            end do
          end do
#ifndef _OPENACC
          sw_albedo_direct = transpose(sw_albedo_band(istartcol:iendcol, &
               &                             config%i_band_from_reordered_g_sw))
#else
        !$ACC LOOP GANG(STATIC:1) VECTOR
        do jcol = istartcol,iendcol
          !$ACC LOOP SEQ
          do jg = 1,config%n_g_sw 
            sw_albedo_direct(jg,jcol) = sw_albedo_band(jcol, &
                &                         config%i_band_from_reordered_g_sw(jg))
          end do
        end do
#endif
        else
          !$ACC LOOP GANG(STATIC:1) VECTOR
          do jcol = istartcol,iendcol
            !$ACC LOOP SEQ
            do jg = 1,config%n_g_sw
              sw_albedo_direct(jg,jcol) = sw_albedo_diffuse(jg,jcol)
            end do
          end do
        end if
      !$ACC END PARALLEL
      else
        ! Albedos mapped less accurately to ecRad spectral bands
        if (maxval(config%i_albedo_from_band_sw) > size(this%sw_albedo,2)) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%sw_albedo has fewer than required ', &
               &  maxval(config%i_albedo_from_band_sw), ' bands'
          call radiation_abort()
        end if      
        sw_albedo_diffuse = transpose(this%sw_albedo(istartcol:iendcol, &
             &  config%i_albedo_from_band_sw(config%i_band_from_reordered_g_sw)))
        if (allocated(this%sw_albedo_direct)) then
          sw_albedo_direct = transpose(this%sw_albedo_direct(istartcol:iendcol, &
               &  config%i_albedo_from_band_sw(config%i_band_from_reordered_g_sw)))
        else
          sw_albedo_direct = sw_albedo_diffuse
        end if
      end if
    end if

    if (config%do_lw .and. present(lw_albedo)) then
      if (config%use_canopy_full_spectrum_lw) then
        if (config%n_g_lw /= size(this%lw_emissivity,2)) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%lw_emissivity does not have the expected ', &
               &  config%n_g_lw, ' spectral intervals'
          call radiation_abort()
        end if
        lw_albedo = 1.0_jprb - transpose(this%lw_emissivity(istartcol:iendcol,:))
      else if (.not. config%do_nearest_spectral_lw_emiss) then
        ! Albedos averaged accurately to ecRad spectral bands
        nalbedoband = size(config%lw_emiss_weights,1)
        if (nalbedoband /= size(this%lw_emissivity,2)) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%lw_emissivity does not have the expected ', &
               &  nalbedoband, ' bands'
          call radiation_abort()
        end if

        do jband = 1,config%n_bands_lw
          do jcol = istartcol,iendcol
            lw_albedo_band(jcol,jband) = 0.0_jprb
          end do
        end do

        do jband = 1,config%n_bands_lw
          do jalbedoband = 1,nalbedoband
            if (config%lw_emiss_weights(jalbedoband,jband) /= 0.0_jprb) then
              do jcol = istartcol,iendcol
                lw_albedo_band(jcol,jband) &
                    &  = lw_albedo_band(jcol,jband) & 
                    &  + config%lw_emiss_weights(jalbedoband,jband) &
                    &    * (1.0_jprb-this%lw_emissivity(jcol, jalbedoband))
              end do
            end if
          end do
        end do

        lw_albedo = transpose(lw_albedo_band(istartcol:iendcol, &
             &                config%i_band_from_reordered_g_lw))
      else
        if (maxval(config%i_emiss_from_band_lw) > size(this%lw_emissivity,2)) then
          write(nulerr,'(a,i0,a)') '*** Error: single_level%lw_emissivity has fewer than required ', &
               &  maxval(config%i_emiss_from_band_lw), ' bands'
          call radiation_abort()
        end if
#ifndef _OPENACC
        lw_albedo = 1.0_jprb - transpose(this%lw_emissivity(istartcol:iendcol, &
             &  config%i_emiss_from_band_lw(config%i_band_from_reordered_g_lw)))
#else
        !$ACC PARALLEL DEFAULT(NONE) ASYNC(1)
        !$ACC LOOP GANG VECTOR COLLAPSE(2)
        do jcol = istartcol,iendcol
          do jg = 1,config%n_g_lw
            lw_albedo(jg,jcol) = 1.0_jprb - this%lw_emissivity(jcol, &
                &  config%i_emiss_from_band_lw(config%i_band_from_reordered_g_lw(jg)))
          end do
        end do
        !$ACC END PARALLEL
#endif
      end if
    end if

    !$ACC WAIT
    !$ACC END DATA

    if (lhook) call dr_hook('radiation_single_level:get_albedos',1,hook_handle)

  end subroutine get_albedos


  !---------------------------------------------------------------------
  ! Return .true. if the contents of a single_level structure are out
  ! of a physically sensible range, optionally only considering only
  ! columns between istartcol and iendcol
  function out_of_physical_bounds(this, istartcol, iendcol, do_fix) result(is_bad)

    use yomhook,          only : lhook, dr_hook, jphook
    use radiation_check,  only : out_of_bounds_1d, out_of_bounds_2d

    class(single_level_type), intent(inout) :: this
    integer,         optional,intent(in) :: istartcol, iendcol
    logical,         optional,intent(in) :: do_fix
    logical                              :: is_bad

    logical    :: do_fix_local

    real(jphook) :: hook_handle

    if (lhook) call dr_hook('radiation_single_level:out_of_physical_bounds',0,hook_handle)

    if (present(do_fix)) then
      do_fix_local = do_fix
    else
      do_fix_local = .false.
    end if

    is_bad =    out_of_bounds_1d(this%cos_sza, "cos_sza", -1.0_jprb, 1.0_jprb, &
         &                       do_fix_local, i1=istartcol, i2=iendcol) &
         & .or. out_of_bounds_1d(this%skin_temperature, "skin_temperature", 173.0_jprb, 373.0_jprb, &
         &                       do_fix_local, i1=istartcol, i2=iendcol) &
         & .or. out_of_bounds_2d(this%sw_albedo, "sw_albedo", 0.0_jprb, 1.0_jprb, &
         &                       do_fix_local, i1=istartcol, i2=iendcol) &
         & .or. out_of_bounds_2d(this%sw_albedo_direct, "sw_albedo", 0.0_jprb, 1.0_jprb, &
         &                       do_fix_local, i1=istartcol, i2=iendcol) &
         & .or. out_of_bounds_2d(this%lw_emissivity, "lw_emissivity", 0.0_jprb, 1.0_jprb, &
         &                       do_fix_local, i1=istartcol, i2=iendcol)

    if (lhook) call dr_hook('radiation_single_level:out_of_physical_bounds',1,hook_handle)

  end function out_of_physical_bounds

#ifdef _OPENACC
  !---------------------------------------------------------------------
  ! updates fields on host
  subroutine update_host(this)

    class(single_level_type), intent(inout) :: this

    !$ACC UPDATE HOST(this%cos_sza) &
    !$ACC   IF(allocated(this%cos_sza))

    !$ACC UPDATE HOST(this%skin_temperature) &
    !$ACC   IF(allocated(this%skin_temperature))

    !$ACC UPDATE HOST(this%sw_albedo) &
    !$ACC   IF(allocated(this%sw_albedo))

    !$ACC UPDATE HOST(this%sw_albedo_direct) &
    !$ACC   IF(allocated(this%sw_albedo_direct))

    !$ACC UPDATE HOST(this%lw_emissivity) &
    !$ACC   IF(allocated(this%lw_emissivity))

    !$ACC UPDATE HOST(this%lw_emission) &
    !$ACC   IF(allocated(this%lw_emission))

    !$ACC UPDATE HOST(this%spectral_solar_scaling) &
    !$ACC   IF(allocated(this%spectral_solar_scaling))

    !$ACC UPDATE HOST(this%iseed) &
    !$ACC   IF(allocated(this%iseed))

  end subroutine update_host

  !---------------------------------------------------------------------
  ! updates fields on device
  subroutine update_device(this)

    class(single_level_type), intent(inout) :: this

    !$ACC UPDATE DEVICE(this%cos_sza) &
    !$ACC   IF(allocated(this%cos_sza))

    !$ACC UPDATE DEVICE(this%skin_temperature) &
    !$ACC   IF(allocated(this%skin_temperature))

    !$ACC UPDATE DEVICE(this%sw_albedo) &
    !$ACC   IF(allocated(this%sw_albedo))

    !$ACC UPDATE DEVICE(this%sw_albedo_direct) &
    !$ACC   IF(allocated(this%sw_albedo_direct))

    !$ACC UPDATE DEVICE(this%lw_emissivity) &
    !$ACC   IF(allocated(this%lw_emissivity))

    !$ACC UPDATE DEVICE(this%lw_emission) &
    !$ACC   IF(allocated(this%lw_emission))

    !$ACC UPDATE DEVICE(this%spectral_solar_scaling) &
    !$ACC   IF( allocated(this%spectral_solar_scaling))

    !$ACC UPDATE DEVICE(this%iseed) &
    !$ACC   IF(allocated(this%iseed))

  end subroutine update_device
#endif 

end module radiation_single_level
