module rad_flux

  use vtype_module

  implicit none

  private

  public :: RadFluxCalcFlux
  public :: RadFluxInit
  public :: RadFluxFinalize
  public :: RadFluxOutputInit
  public :: RadFluxOutputFinalize


  integer, save      :: IDMethod
  integer, parameter :: IDMethodKD         = 0
  integer, parameter :: IDMethodKDDevelop  = 1
  integer, parameter :: IDMethodLBL        = 2
  integer, parameter :: IDMethodLBLConstPF = 3

  real(8), save :: Grav
  real(8), save :: H2OMolWt

  integer              , save :: NBand
  real(DP), allocatable, save :: aa_BandWNBnds      (:,:)
  integer , allocatable, save :: aa_BandBinIndexBnds(:,:)


  logical         , save :: FlagOutputInited = .false.
  integer         , save :: NcIDFlux


  character(128), save :: ModuleName = 'flux'

  logical, save :: FlagInited       = .false.


  !----------------------------------------------------------------------------

contains

  !----------------------------------------------------------------------------

  subroutine RadFluxCalcFlux( &
    & kmax, NMolNum, &
    & r_Press, &
    & SurfTemp, r_Temp, rm_VMR, &
    & r_MeanMolWt, &
    & r_RadUwFlux, r_RadDwFlux, &
    & WNIntegStart, WNIntegEnd &
    & )

    use ni3_module

    use constants0, only : &
      & PI, &
      & AtomMassConst, &
      & AvogConst

    use rad_kd_utils, only : InqPFIntGQNum

    use rad_rte_two_stream_app_wrapper, only : RadRTETwoStreamAppLWWrapper

    use opt_prop, only : &
      & OptPropGetNMol, &
      & OptPropGetNWN, &
      & OptPropGetNBand, &
      & OptPropGetDelWN, &
      & OptPropGetWN, &
      & OptPropGetAbsCoefProf, &
      & OptPropGetPFInted, &
      & OptPropGetPFIntedRatio

    use planck_func_wrapper

    integer , intent(in ) :: kmax
    integer , intent(in ) :: NMolNum
    real(DP), intent(in ) :: r_Press     (0:kmax)
    real(DP), intent(in ) :: SurfTemp
    real(DP), intent(in ) :: r_Temp      (0:kmax)
    real(DP), intent(in ) :: rm_VMR      (0:kmax,1:NMolNum)
    real(DP), intent(in ) :: r_MeanMolWt (0:kmax)
    real(DP), intent(out) :: r_RadUwFlux (0:kmax)
    real(DP), intent(out) :: r_RadDwFlux (0:kmax)
    real(DP), intent(in ) :: WNIntegStart
    real(DP), intent(in ) :: WNIntegEnd


    ! local variables
    !
    real(DP), allocatable :: a_BandSurfAvePF   (:)
    real(DP), allocatable :: ra_BandAvePF      (:,:)
    real(DP), allocatable :: a_BandSurfAveDPFDT(:)

    integer :: NMol
    integer :: NWaveNum

    integer :: iBand
    integer :: iBandStart
    integer :: iBandEnd

    integer :: iWaveNum
    integer :: iWaveNumStart
    integer :: iWaveNumEnd
    real(8) :: WaveNum
    real(8) :: DelWaveNum

    real(8), allocatable :: pm_AbsCoefPerAtmMol(:,:)

    real(8) :: z_DelAtmMols   (1:kmax)
    real(8) :: DelOptDep

    real(DP) :: z_SSA          (1:kmax)
    real(DP) :: z_AF           (1:kmax)
    real(DP) :: r_OptDep       (0:kmax)
    real(DP) :: SurfAlbedo
    real(DP) :: r_PFInted      (0:kmax)
    real(DP) :: SurfPFInted
    real(DP) :: SurfDPFDTInted
    real(DP) :: r_RadUwFluxEach    (0:kmax)
    real(DP) :: r_RadDwFluxEach    (0:kmax)
    real(DP) :: ra_DelRadUwFluxEach(0:kmax,0:1)
    real(DP) :: ra_DelRadDwFluxEach(0:kmax,0:1)

    real(DP) :: r_PFIntedRatio (0:kmax)
    real(DP) :: SurfPFIntedRatio

    real(DP) :: WNs
    real(DP) :: WNe
    integer  :: GQNum

    integer :: k
    integer :: m
    integer :: n


    if ( .not. FlagInited ) then
      write( 6, * ) trim( ModuleName ), " is not initialized."
      stop
    end if


    select case ( IDMethod )
    case ( IDMethodLBLConstPF, IDMethodKDDevelop, IDMethodKD )

      allocate( a_BandSurfAvePF   (NBand) )
      allocate( ra_BandAvePF      (0:kmax,NBand) )
      allocate( a_BandSurfAveDPFDT(NBand) )
!!$      allocate( a_SolarFluxTOA(NBand) )

      do n = 1, NBand
        WNs = aa_BandWNBnds(1,n)
        WNe = aa_BandWNBnds(2,n)
        GQNum = InqPFIntGQNum( WNe - WNs )
        call Integ_PF_GQ_Array0D( &
          & WNs, WNe, GQNum, &
          & SurfTemp, &
          & a_BandSurfAvePF(n) &
          & )
        a_BandSurfAvePF(n) = PI * a_BandSurfAvePF(n) / ( WNe - WNs )
        call Integ_PF_GQ_Array1D( &
          & WNs, WNe, GQNum, &
          & 0, kmax, &
          & r_Temp, &
          & ra_BandAvePF(:,n) &
          & )
        ra_BandAvePF(:,n) = PI * ra_BandAvePF(:,n) / ( WNe - WNs )
        call Integ_DPFDT_GQ_Array0D( &
          & WNs, WNe, GQNum,                 & ! (in )
          & SurfTemp,                & ! (in )
          & a_BandSurfAveDPFDT(n) &
          & )
        a_BandSurfAveDPFDT(n) = PI * a_BandSurfAveDPFDT(n) / ( WNe - WNs )
      end do

    end select



    do k = 1, kmax
!!$      z_DelAtmMols(k) = &
!!$        &   ( r_Press(k-1) - r_Press(k) ) / Grav
      z_DelAtmMols(k) = &
        &   ( r_Press(k-1) - r_Press(k) ) / Grav &
        & / ( ( r_MeanMolWt(k-1) + r_MeanMolWt(k) ) / 2.0d0 &
        &     / AvogConst )
    end do


    NMol       = OptPropGetNMol()
    NWaveNum   = OptPropGetNWN()

    allocate( pm_AbsCoefPerAtmMol(0:kmax, NMol) )


    r_RadUwFlux = 0.0d0
    r_RadDwFlux = 0.0d0


    select case ( IDMethod )
    case ( IDMethodKDDevelop, IDMethodKD )
      if ( ( WNIntegStart < 0.0d0 ) .or. ( WNIntegEnd < 0.0d0 ) ) then
        iWaveNumStart = 1
        iWaveNumEnd   = NWaveNum
      else
        do iBand = 1, NBand
          if ( aa_BandWNBnds(1,iBand) >= WNIntegStart ) exit
        end do
        iBandStart = iBand
        iBandStart = min( iBandStart, NBand )
        do iBand = iBandStart, NBand
          if ( aa_BandWNBnds(2,iBand) >= WNIntegEnd ) exit
        end do
        iBandEnd = iBand
        iWaveNumStart = aa_BandBinIndexBnds(1,iBandStart)
        iWaveNumEnd   = aa_BandBinIndexBnds(2,iBandEnd  )
      end if
    case ( IDMethodLBL, IDMethodLBLConstPF )
      if ( ( WNIntegStart < 0.0d0 ) .or. ( WNIntegEnd < 0.0d0 ) ) then
        iWaveNumStart = 1
        iWaveNumEnd   = NWaveNum
      else
        loop_search_wnstart: do iWaveNum = 1, NWaveNum
          if ( OptPropGetWN( iWaveNum ) >= WNIntegStart ) exit
        end do loop_search_wnstart
        iWaveNumStart = iWaveNum
        iWaveNumStart = min( iWaveNumStart, NWaveNum )
        loop_search_wnend: do iWaveNum = iWaveNumStart, NWaveNum
          if ( OptPropGetWN( iWaveNum ) > WNIntegEnd ) exit
        end do loop_search_wnend
        iWaveNumEnd = iWaveNum - 1
      end if
    case default
      write( 6, * ) 'In ', trim( ModuleName )
      write( 6, * ) '  Unexpected case default.'
      stop
    end select


!!$    do iWaveNum = 1, NWaveNum
    do iWaveNum = iWaveNumStart, iWaveNumEnd

      DelWaveNum = OptPropGetDelWN( iWaveNum )

      z_SSA = 0.0d0
      z_AF  = 0.0d0
      !
      call OptPropGetAbsCoefProf( &
        & iWaveNum, kmax+1, NMol, r_Press, r_Temp, rm_VMR, &
        & pm_AbsCoefPerAtmMol &
        & )


      r_OptDep(kmax) = 0.0d0
      do k = kmax, 1, -1
        DelOptDep = 0.0d0
        do m = 1, NMol
!!$          DelOptDep = DelOptDep &
!!$            & + (   rm_AbsCoef(k-1,m) * rm_VMR(k-1,m)           &
!!$            &     + rm_AbsCoef(k  ,m) * rm_VMR(k  ,m) ) / 2.0d0 &
!!$            & * z_DelAtmMols(k)
          DelOptDep = DelOptDep                          &
            & + (   pm_AbsCoefPerAtmMol(k-1,m)           &
            &     + pm_AbsCoefPerAtmMol(k  ,m) ) / 2.0d0 &
            & * z_DelAtmMols(k)
        end do
        r_OptDep(k-1) = r_OptDep(k) + DelOptDep
      end do
      !
      SurfAlbedo = 0.0d0



      select case ( IDMethod )
      case ( IDMethodLBL )
        call OptPropGetPFInted( &
          & iWaveNum, &
          & kmax+1, &
          & r_Temp, SurfTemp, &
          & r_PFInted, SurfPFInted, SurfDPFDTInted &
          & )
      case ( IDMethodLBLConstPF, IDMethodKDDevelop, IDMethodKD )

        select case ( IDMethod )
        case ( IDMethodLBLConstPF )
          WaveNum = OptPropGetWN( iWaveNum )
          search_Band : do n = 1, NBand
            if ( WaveNum <= aa_BandWNBnds(2,n) ) exit search_Band
          end do search_Band
        case ( IDMethodKDDevelop, IDMethodKD )
          do n = 1, NBand
            if ( ( aa_BandBinIndexBnds(1,n) <= iWaveNum ) .and. &
              &  ( iWaveNum <= aa_BandBinIndexBnds(2,n) ) ) &
              exit
          end do
        case default
          write( 6, * ) 'In ', trim( ModuleName )
          write( 6, * ) '  Unexpected case default'
          stop
        end select
        if ( n >= NBand + 1 ) then
          write( 6, * ) 'In ', trim( ModuleName )
          write( 6, * ) '  Unable to find appropriate band number: ', n
          stop
        end if

        r_PFInted      = ra_BandAvePF(:,n)
        SurfPFInted    = a_BandSurfAvePF(n)
        SurfDPFDTInted = a_BandSurfAveDPFDT(n)
        !
        call OptPropGetPFIntedRatio( &
          & iWaveNum, kmax+1, r_Temp, SurfTemp, &
          & r_PFIntedRatio, SurfPFIntedRatio &
          & )
        !
        r_PFInted      = r_PFInted * r_PFIntedRatio
        SurfPFInted    = SurfPFInted * SurfPFIntedRatio
        SurfDPFDTInted = SurfDPFDTInted * SurfPFIntedRatio

        r_PFInted      = r_PFInted      * DelWaveNum
        SurfPFInted    = SurfPFInted    * DelWaveNum
        SurfDPFDTInted = SurfDPFDTInted * DelWaveNum

      case default
        write( 6, * ) 'In ', trim( ModuleName )
        write( 6, * ) '  Unexpected case default'
        stop
      end select

      call RadRTETwoStreamAppLWWrapper(            &
        & z_SSA, z_AF,                             & ! (in)
        & r_OptDep,                                & ! (in)
        & SurfAlbedo,                              & ! (in)
        & r_PFInted, SurfPFInted, SurfDPFDTInted,  & ! (in)
        & r_RadUwFluxEach, r_RadDwFluxEach,        & ! (out)
        & ra_DelRadUwFluxEach, ra_DelRadDwFluxEach & ! (out)
        & )

      if ( FlagOutputInited ) then
        call ni3_put_varss( NcIDFlux, "RadUwFlux", iWaveNum, r_RadUwFluxEach/DelWaveNum )
        call ni3_put_varss( NcIDFlux, "RadDwFlux", iWaveNum, r_RadDwFluxEach/DelWaveNum )
      end if

      r_RadUwFlux = r_RadUwFlux + r_RadUwFluxEach
      r_RadDwFlux = r_RadDwFlux + r_RadDwFluxEach

    end do

    deallocate( pm_AbsCoefPerAtmMol )


    select case ( IDMethod )
    case ( IDMethodLBLConstPF, IDMethodKDDevelop, IDMethodKD )
      deallocate( a_BandSurfAvePF    )
      deallocate( ra_BandAvePF       )
      deallocate( a_BandSurfAveDPFDT )
    end select


  end subroutine RadFluxCalcFlux

  !----------------------------------------------------------------------------

  subroutine RadFluxInit( &
    & ArgGrav, ArgH2OMolWt, &
    & OptPropNcFN, &
    & Method, &
    & a_ArgBandWNBnds &
    & )

    use opt_prop, only : &
      & OptPropInit, &
      & OptPropGetNBand, &
      & OptPropGetBandWNBnds, &
      & OptPropGetBandBinIndexBnds

    use rad_rte_two_stream_app_wrapper, only : RadRTETwoStreamAppLWWrapperInit

    real(8)     , intent(in )           :: ArgGrav
    real(8)     , intent(in )           :: ArgH2OMolWt
    character(*), intent(in )           :: OptPropNcFN
    character(*), intent(in ), optional :: Method
    real(DP)    , intent(in ), optional :: a_ArgBandWNBnds(:)


    ! local variables
    !
    integer :: n


    if ( FlagInited ) return


    if ( present( Method ) ) then
      select case ( Method )
      case ( 'KD' )
        IDMethod = IDMethodKD
      case ( 'KDDevelop' )
        IDMethod = IDMethodKDDevelop
      case ( 'LBL' )
        IDMethod = IDMethodLBL
      case ( 'LBLConstPF' )
        IDMethod = IDMethodLBLConstPF
      case default
        write( 6, * ) "In module, ", trim( ModuleName )
        write( 6, * ) "  Unexpected method: ", trim( Method )
        stop
      end select
    else
      IDMethod = IDMethodKD
    end if


    ! Check arguments
    !
    select case ( IDMethod )
    case ( IDMethodLBLConstPF )
      if ( present( a_ArgBandWNBnds ) ) then
      else
        write( 6, * ) "In module, ", trim( ModuleName )
        write( 6, * ) 'a_BandWNBnds has to be present when IDMethod is IDMethodLBLConstPF.'
        stop
      end if
    end select


    select case ( IDMethod )
    case ( IDMethodKD )
      call OptPropInit( &
        & 'KD', &
        & OptPropNcFN &
        & )
    case ( IDMethodKDDevelop )
      call OptPropInit( &
        & 'KDDevelop', &
        & OptPropNcFN &
        & )
    case ( IDMethodLBL )
      call OptPropInit( &
        & 'LBL', &
        & OptPropNcFN &
        & )
    case ( IDMethodLBLConstPF )
      call OptPropInit( &
        & 'LBLConstPF', &
        & OptPropNcFN &
        & )
    end select


    call RadRTETwoStreamAppLWWrapperInit


    Grav      = ArgGrav
    H2OMolWt  = ArgH2OMolWt


    select case ( IDMethod )
    case ( IDMethodKDDevelop, IDMethodKD )

      NBand = OptPropGetNBand()
      allocate( aa_BandWNBnds      (2,NBand) )
      allocate( aa_BandBinIndexBnds(2,NBand) )
      call OptPropGetBandWNBnds( &
        & NBand, &
        & aa_BandWNBnds &
        & )
      call OptPropGetBandBinIndexBnds( &
        & NBand, &
        & aa_BandBinIndexBnds &
        & )

    case ( IDMethodLBLConstPF )

      NBand = size( a_ArgBandWNBnds ) - 1
      allocate( aa_BandWNBnds(2,NBand) )

      do n = 1, NBand
        aa_BandWNBnds(1,n) = a_ArgBandWNBnds(n  )
        aa_BandWNBnds(2,n) = a_ArgBandWNBnds(n+1)
      end do

    end select



    FlagInited = .true.


  end subroutine RadFluxInit

  !----------------------------------------------------------------------------

  subroutine RadFluxFinalize

    use opt_prop, only : OptPropFinalize


    ! local variables
    !


    call OptPropFinalize

    if ( FlagOutputInited ) then
      call RadFluxOutputFinalize
    end if


    select case ( IDMethod )
    case ( IDMethodKDDevelop )
      deallocate( aa_BandWNBnds       )
      deallocate( aa_BandBinIndexBnds )
    case ( IDMethodLBLConstPF )
      deallocate( aa_BandWNBnds )
    end select


    FlagInited = .false.

  end subroutine RadFluxFinalize

  !----------------------------------------------------------------------------

  subroutine RadFluxOutputInit( &
    & OptPropNcFN, &
    & FluxOutNcFn, &
    & kmax, r_Press &
    & )

    use ni3_module
    use netcdf

    character(*), intent(in) :: OptPropNcFN
    character(*), intent(in) :: FluxOutNcFn
    integer     , intent(in) :: kmax
    real(8)     , intent(in) :: r_Press    (0:kmax)


    ! local variables
    !
    integer           :: NcID
    character(stdstr) :: Mode
    character(stdstr) :: Name
    character(stdstr) :: StdName
    character(stdstr) :: LongName
    character(stdstr) :: Units
    character(stdstr) :: aa_DimNames(2)

    integer :: LBLNWN
    real(DP), allocatable :: w_WaveNum(:)


    if ( .not. FlagInited ) then
      write( 6, * ) trim( ModuleName ), " is not initialized."
      stop
    end if


    Mode = "read"
    call ni3_open( OptPropNcFN, Mode, NcID )
    call ni3_inq_dimlen( NcID, 'WaveNum', LBLNWN )
    allocate( w_WaveNum( 1:LBLNWN ) )
    call ni3_get_var( NcID, 'WaveNum', w_WaveNum )
    call ni3_close( NcID )


    ! Output

    Mode = "new"
    call ni3_open( FluxOutNcFN, Mode, NcIDFlux )

    Name    = "Press"
    StdName = "plev"
    Units   = "Pa"
    call ni3_set_dim( NcIDFlux, Name, NF90_DOUBLE, r_Press  , &
      & stdname = stdname, units  = units )
    Name    = "WaveNum"
    StdName = "wavenumber"
    Units   = "m-1"
    call ni3_set_dim( NcIDFlux, Name, NF90_DOUBLE, w_WaveNum , &
      & stdname = stdname, units  = units )

    aa_DimNames(1) = "Press"
    aa_DimNames(2) = "WaveNum"

    Name     = "RadUwFlux"
    StdName  = Name
    LongName = "radiative upward flux"
    Units   = "W m-2 (m-1)-1"
    call ni3_def_var( NcIDFlux, Name, NF90_DOUBLE, 2, aa_DimNames, &
      & stdname = StdName, longname = LongName, Units = Units )

    Name     = "RadDwFlux"
    StdName  = Name
    LongName = "radiative downward flux"
    Units   = "W m-2 (m-1)-1"
    call ni3_def_var( NcIDFlux, Name, NF90_DOUBLE, 2, aa_DimNames, &
      & stdname = StdName, longname = LongName, Units = Units )


    deallocate( w_WaveNum )


    FlagOutputInited = .true.

  end subroutine RadFluxOutputInit

  !----------------------------------------------------------------------------

  subroutine RadFluxOutputFinalize

    use ni3_module

    ! local variables
    !


    if ( .not. FlagInited ) then
      write( 6, * ) trim( ModuleName ), " is not initialized."
      stop
    end if


    if ( FlagOutputInited ) then
      call ni3_close( NcIDFlux )
      FlagOutputInited = .false.
    end if


  end subroutine RadFluxOutputFinalize

  !--------------------------------------------------------------------------

end module rad_flux