module fkb_utils

!---------------------------------------------------------------------------------
! Purpose:
! Provides utilties to implement FKB neural network inference calls
!
! Author:
! modified by Po-Lun Ma July 2024)
! Sungduk Yu (Fri 16 Feb 2024 11:40:41 AM PST)
!
!---------------------------------------------------------------------------------

! Note:
! The parameter "do_test" (inside subroutine  ml_activate_modal_adjust in ndrop.F90)
! shoul be set .false. to run a continuous simulation. Otherwise, a model will
! generate a test log (fort.101) and halt after the 1st timestep.

use shr_kind_mod,     only : shr_kind_cl, r8 => shr_kind_r8        

implicit none
save 
private

!!! user_nl_eam namelist variables
! aerosol activation  (ndrop.F90)
logical :: do_fkb_activation = .false. ! on/off switch for aerosol activation emulator correction
logical :: do_fkb_activation_test = .false. ! on/off switch for aerosol activation emulator test (pouplates fort.101)
character(len=shr_kind_cl) :: fkb_act_model_txt, & ! fkb weight txt file path for an aerosol activation emulator
                              fkb_smax_model_txt   ! fkb weight txt file path for a smax emulator

! aerosol optics (modal_aer_opt.F90)
!logical :: do_fkb_aerosol_optics = .false. ! on/off switch for aerosol activation emulator correction
integer :: do_fkb_aerosol_optics = 0
character(len=shr_kind_cl) :: fkb_aer_opt_model1_txt ! fkb weight txt file path for aerosol optics
character(len=shr_kind_cl) :: fkb_aer_opt_model2_txt ! fkb weight txt file path for aerosol optics

public fkb_readnl, &
       do_fkb_activation, do_fkb_activation_test, fkb_act_model_txt, fkb_smax_model_txt, &
       do_fkb_aerosol_optics, fkb_aer_opt_model1_txt,fkb_aer_opt_model2_txt, &
       erf_inv

contains

subroutine fkb_readnl(nlfile)
!-------------------------------------------------------------------------------
! Purpose:
! Read user_nl_eam namelist options for neural network (FKB) configurations
!-------------------------------------------------------------------------------

      use namelist_utils,  only: find_group_name
      use units,           only: getunit, freeunit
      use cam_abortutils,  only: endrun
      use spmd_utils,      only: masterproc
#ifdef SPMD
      use mpishorthand,    only: mpichar, mpicom, mpilog
#endif

      character(len=*), intent(in) :: nlfile  ! filepath for file containing namelist input

      ! Local variables
      integer :: unitn, ierr, f
      character(len=*), parameter :: subname = 'fkb_readnl'
      
      namelist /fkb_nl/ do_fkb_activation, do_fkb_activation_test, &
                        fkb_act_model_txt, fkb_smax_model_txt, &
                        do_fkb_aerosol_optics, fkb_aer_opt_model1_txt, &
                        fkb_aer_opt_model2_txt

      ierr = 0
      if (masterproc) then
         unitn = getunit()
         open( unitn, file=trim(nlfile), status='old' )
         call find_group_name(unitn, 'fkb_nl', status=ierr)
         if (ierr == 0) then
            read(unitn, fkb_nl, iostat=ierr)
            if (ierr /= 0) then
               call endrun(subname // ':: ERROR reading namelist')
            end if
         end if
         close(unitn)
         call freeunit(unitn)
      end if

#ifdef SPMD
      ! Broadcast namelist variables to all MPI tasks
      call mpibcast(fkb_act_model_txt,  len(fkb_act_model_txt),  mpichar, 0, mpicom)
      call mpibcast(fkb_smax_model_txt, len(fkb_smax_model_txt), mpichar, 0, mpicom)
      call mpibcast(do_fkb_activation,                        1,  mpilog, 0, mpicom)
      ! call mpibcast(do_fkb_activation_test,                   1,  mpilog, 0, mpicom) ! Not broadcasting to keep the test log (fort.101) short
      call mpibcast(fkb_aer_opt_model1_txt, len(fkb_aer_opt_model1_txt), mpichar, 0, mpicom)
      call mpibcast(fkb_aer_opt_model2_txt, len(fkb_aer_opt_model2_txt), mpichar, 0, mpicom)
      call mpibcast(do_fkb_aerosol_optics,                          1,  mpilog, 0, mpicom)
#endif

end subroutine fkb_readnl

pure function erf_inv(x) result(p)
!-------------------------------------------------------------------------------
! Inverse error function
! Adapted from https://github.com/catchorg/Catch2/blob/1078e7e95b3a06d4dadc75188de48bc4afffb955/src/catch2/benchmark/detail/catch_stats.cpp#L93-L167
! Original source: M. Giles, “Approximating the erfinv function.” In GPU Computing Gems Jade Edition, pp. 109-116. 2011.
! Solutions are verified against INTEL MKL's VML function vderfinv.
!-------------------------------------------------------------------------------

    implicit none
    real(r8), intent(in) :: x
    real(r8) :: p, w

    w = -log((1.0d0 - x) * (1.0d0 + x))

    if (w < 6.250000d0) then
        w = w - 3.125000d0
        p = -3.6444120640178196996d-21
        p = -1.685059138182016589d-19 + p * w
        p = 1.2858480715256400167d-18 + p * w
        p = 1.115787767802518096d-17 + p * w
        p = -1.333171662854620906d-16 + p * w
        p = 2.0972767875968561637d-17 + p * w
        p = 6.6376381343583238325d-15 + p * w
        p = -4.0545662729752068639d-14 + p * w
        p = -8.1519341976054721522d-14 + p * w
        p = 2.6335093153082322977d-12 + p * w
        p = -1.2975133253453532498d-11 + p * w
        p = -5.4154120542946279317d-11 + p * w
        p = 1.051212273321532285d-09 + p * w
        p = -4.1126339803469836976d-09 + p * w
        p = -2.9070369957882005086d-08 + p * w
        p = 4.2347877827932403518d-07 + p * w
        p = -1.3654692000834678645d-06 + p * w
        p = -1.3882523362786468719d-05 + p * w
        p = 0.0001867342080340571352d0 + p * w
        p = -0.00074070253416626697512d0 + p * w
        p = -0.0060336708714301490533d0 + p * w
        p = 0.24015818242558961693d0 + p * w
        p = 1.6536545626831027356d0 + p * w
    else if (w < 16.000000d0) then
        w = sqrt(w) - 3.250000d0
        p = 2.2137376921775787049d-09
        p = 9.0756561938885390979d-08 + p * w
        p = -2.7517406297064545428d-07 + p * w
        p = 1.8239629214389227755d-08 + p * w
        p = 1.5027403968909827627d-06 + p * w
        p = -4.013867526981545969d-06 + p * w
        p = 2.9234449089955446044d-06 + p * w
        p = 1.2475304481671778723d-05 + p * w
        p = -4.7318229009055733981d-05 + p * w
        p = 6.8284851459573175448d-05 + p * w
        p = 2.4031110387097893999d-05 + p * w
        p = -0.0003550375203628474796d0 + p * w
        p = 0.00095328937973738049703d0 + p * w
        p = -0.0016882755560235047313d0 + p * w
        p = 0.0024914420961078508066d0 + p * w
        p = -0.0037512085075692412107d0 + p * w
        p = 0.005370914553590063617d0 + p * w
        p = 1.0052589676941592334d0 + p * w
        p = 3.0838856104922207635d0 + p * w
    else
        w = sqrt(w) - 5.000000d0
        p = -2.7109920616438573243d-11
        p = -2.5556418169965252055d-10 + p * w
        p = 1.5076572693500548083d-09 + p * w
        p = -3.7894654401267369937d-09 + p * w
        p = 7.6157012080783393804d-09 + p * w
        p = -1.4960026627149240478d-08 + p * w
        p = 2.9147953450901080826d-08 + p * w
        p = -6.7711997758452339498d-08 + p * w
        p = 2.2900482228026654717d-07 + p * w
        p = -9.9298272942317002539d-07 + p * w
        p = 4.5260625972231537039d-06 + p * w
        p = -1.9681778105531670567d-05 + p * w
        p = 7.5995277030017761139d-05 + p * w
        p = -0.00021503011930044477347d0 + p * w
        p = -0.00013871931833623122026d0 + p * w
        p = 1.0103004648645343977d0 + p * w
        p = 4.8499064014085844221d0 + p * w
    end if

    p = p * x

    return

end function erf_inv

end module fkb_utils
