module shr_assert_mod

! Assert subroutines for common debugging operations.

use shr_kind_mod, only: &
     r4 => shr_kind_r4, &
     r8 => shr_kind_r8, &
     i4 => shr_kind_i4, &
     i8 => shr_kind_i8

use shr_sys_mod, only: &
     shr_sys_abort

use shr_log_mod, only: &
     shr_log_Unit

use shr_infnan_mod, only: shr_infnan_isnan

use shr_strconvert_mod, only: toString

implicit none
private
save

! Assert that a logical is true.
public :: shr_assert
public :: shr_assert_all
public :: shr_assert_any

! Assert that a numerical value satisfies certain constraints.
public :: shr_assert_in_domain

interface shr_assert_all
   module procedure shr_assert
   ! DIMS 1,2,3,4,5,6,7
   module procedure shr_assert_all_{DIMS}d
end interface

interface shr_assert_any
   module procedure shr_assert
   ! DIMS 1,2,3,4,5,6,7
   module procedure shr_assert_any_{DIMS}d
end interface

interface shr_assert_in_domain
   ! TYPE double,real,int,long
   ! DIMS 0,1,2,3,4,5,6,7
   module procedure shr_assert_in_domain_{DIMS}d_{TYPE}
end interface

! Private utilities.

interface print_bad_loc
   ! TYPE double,real,int,long
   ! DIMS 0,1,2,3,4,5,6,7
   module procedure print_bad_loc_{DIMS}d_{TYPE}
end interface

interface find_first_loc
   ! DIMS 0,1,2,3,4,5,6,7
   module procedure find_first_loc_{DIMS}d
end interface

interface within_tolerance
   ! TYPE double,real,int,long
   module procedure within_tolerance_{TYPE}
end interface

contains

subroutine shr_assert(var, msg, file, line)

  ! Logical being asserted
  logical, intent(in) :: var
  ! Optional error message if assert fails
  character(len=*), intent(in), optional :: msg
  ! Optional file and line of the caller, written out if given
  ! (line is ignored if file is absent)
  character(len=*), intent(in), optional :: file
  integer         , intent(in), optional :: line

  character(len=:), allocatable :: full_msg

  if (.not. var) then
     full_msg = 'ERROR'
     if (present(file)) then
        full_msg = full_msg // ' in ' // trim(file)
        if (present(line)) then
           full_msg = full_msg // ' at line ' // toString(line)
        end if
     end if
     if (present(msg)) then
        full_msg = full_msg // ': ' // msg
     end if
     call shr_sys_abort(full_msg)
  end if

end subroutine shr_assert

! DIMS 1,2,3,4,5,6,7
subroutine shr_assert_all_{DIMS}d(var, msg, file, line)

  ! Logical being asserted
  logical, intent(in) :: var{DIMSTR}
  ! Optional error message if assert fails
  character(len=*), intent(in), optional :: msg
  ! Optional file and line of the caller, written out if given
  ! (line is ignored if file is absent)
  character(len=*), intent(in), optional :: file
  integer         , intent(in), optional :: line

  call shr_assert(all(var), msg=msg, file=file, line=line)

end subroutine shr_assert_all_{DIMS}d

! DIMS 1,2,3,4,5,6,7
subroutine shr_assert_any_{DIMS}d(var, msg, file, line)

  ! Logical being asserted
  logical, intent(in) :: var{DIMSTR}
  ! Optional error message if assert fails
  character(len=*), intent(in), optional :: msg
  ! Optional file and line of the caller, written out if given
  ! (line is ignored if file is absent)
  character(len=*), intent(in), optional :: file
  integer         , intent(in), optional :: line

  call shr_assert(any(var), msg=msg, file=file, line=line)

end subroutine shr_assert_any_{DIMS}d

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

! TYPE double,real,int,long
! DIMS 0,1,2,3,4,5,6,7
subroutine shr_assert_in_domain_{DIMS}d_{TYPE}(var, varname, msg, &
     is_nan, lt, gt, le, ge, eq, ne, abs_tol)

!-----------------------------
! BEGIN defining local macros
!-----------------------------

! Flag for floating point types.

#if ({ITYPE} == TYPEREAL) || ({ITYPE} == TYPEDOUBLE)
#define TYPEFP
#else
#undef TYPEFP
#endif

! "Generalized" macro functions allow transformational intrinsic functions
! to handle both scalars and arrays.

#if ({DIMS} != 0)
! When given an array, use the intrinsics.
#define GEN_SIZE(x) size(x)
#define GEN_ALL(x) all(x)
#else

! Scalar extensions:
!   GEN_SIZE always returns 1 for a scalar.
!   GEN_ALL (logical reduction) is a no-op for a scalar.
!   GEN_[MAX,MIN]LOC should return a 1D, size 0 (empty), integer array.
#define GEN_SIZE(x) 1
#define GEN_ALL(x) x

#endif

!-----------------------------
! END macro section
!-----------------------------

  ! Variable being checked.
  {VTYPE}, intent(in) :: var{DIMSTR}
  ! Variable name to be used in error messages.
  character(len=*), intent(in), optional :: varname
  ! Optional error message if assert fails.
  character(len=*), intent(in), optional :: msg
  ! Assert that the variable is not (or is) NaN.
  logical, intent(in), optional :: is_nan
  ! Limits for (in)equalities.
  {VTYPE}, intent(in), optional :: lt
  {VTYPE}, intent(in), optional :: gt
  {VTYPE}, intent(in), optional :: le
  {VTYPE}, intent(in), optional :: ge
  {VTYPE}, intent(in), optional :: eq
  {VTYPE}, intent(in), optional :: ne
  {VTYPE}, intent(in), optional :: abs_tol

  ! Note that the following array is size 0 for scalars.
  integer :: loc_vec({DIMS})

  logical :: is_nan_passed
  logical :: lt_passed
  logical :: gt_passed
  logical :: le_passed
  logical :: ge_passed
  logical :: eq_passed
  logical :: ne_passed

  {VTYPE} :: abs_tol_loc

  ! Handling of abs_tol makes a couple of fairly safe assumptions.
  !  1. It is not the most negative integer.
  !  2. It is finite (not a floating point infinity or NaN).
  if (present(abs_tol)) then
     abs_tol_loc = abs(abs_tol)
  else
     abs_tol_loc = 0_i4
  end if

  is_nan_passed = .true.
  lt_passed = .true.
  gt_passed = .true.
  le_passed = .true.
  ge_passed = .true.
  eq_passed = .true.
  ne_passed = .true.

  ! Do one pass just to find out if we can return with no problem.

#ifdef TYPEFP
  ! Only floating-point values can actually be Inf/NaN.
  if (present(is_nan)) &
     is_nan_passed = GEN_ALL(shr_infnan_isnan(var) .eqv. is_nan)
#else
  if (present(is_nan)) &
     is_nan_passed = .not. is_nan .or. GEN_SIZE(var) == 0
#endif

  if (present(lt)) &
     lt_passed = GEN_ALL(var < lt)

  if (present(gt)) &
     gt_passed = GEN_ALL(var > gt)

  if (present(le)) &
     le_passed = GEN_ALL(var <= le)

  if (present(ge)) &
     ge_passed = GEN_ALL(var >= ge)

  if (present(eq)) then
     eq_passed = GEN_ALL(within_tolerance(eq, var, abs_tol_loc))
  end if

  if (present(ne)) then
     ne_passed = GEN_ALL(.not. within_tolerance(ne, var, abs_tol_loc))
  end if

  if ( is_nan_passed .and. &
       lt_passed .and. &
       gt_passed .and. &
       le_passed .and. &
       ge_passed .and. &
       eq_passed .and. &
       ne_passed) &
       return

  ! If we got here, assert will fail, so find out where so that we
  ! can try to print something useful.

  if (.not. is_nan_passed) then
#ifdef TYPEFP
     loc_vec = find_first_loc(shr_infnan_isnan(var) .neqv. is_nan)
     call print_bad_loc(var, loc_vec, varname)
     if (is_nan) then
        write(shr_log_Unit,*) "Expected value to be NaN."
     else
        write(shr_log_Unit,*) "Expected value to be a number."
     end if
#else
     loc_vec = spread(1,1,{DIMS})
     call print_bad_loc(var, loc_vec, varname)
     if (is_nan) then
        write(shr_log_Unit,*) &
             "Asserted NaN, but the variable is not floating-point!"
     end if
#endif
  end if

  if (.not. lt_passed) then
     loc_vec = find_first_loc(var >= lt)
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to be less than ",lt
  end if

  if (.not. gt_passed) then
     loc_vec = find_first_loc(var <= gt)
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to be greater than ",gt
  end if

  if (.not. le_passed) then
     loc_vec = find_first_loc(var > le)
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to be less than or &
          &equal to ",le
  end if

  if (.not. ge_passed) then
     loc_vec = find_first_loc(var < ge)
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to be greater than or &
          &equal to ",ge
  end if

  if (.not. eq_passed) then
     loc_vec = find_first_loc(.not. within_tolerance(eq, var, abs_tol_loc))
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to be equal to ",eq
     if (abs_tol_loc > 0) &
          write(shr_log_Unit,*) "Asserted with tolerance ", abs_tol_loc
  end if

  if (.not. ne_passed) then
     loc_vec = find_first_loc(within_tolerance(ne, var, abs_tol_loc))
     call print_bad_loc(var, loc_vec, varname)
     write(shr_log_Unit,*) "Expected value to never be equal to ",ne
     if (abs_tol_loc > 0) &
          write(shr_log_Unit,*) "Asserted with tolerance ", abs_tol_loc
  end if

  call shr_sys_abort(msg)

! Undefine local macros.
#undef TYPEFP
#undef GEN_SIZE
#undef GEN_ALL

end subroutine shr_assert_in_domain_{DIMS}d_{TYPE}

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

! TYPE double,real,int,long
! DIMS 0,1,2,3,4,5,6,7
subroutine print_bad_loc_{DIMS}d_{TYPE}(var, loc_vec, varname)
  ! Print information about a bad location in an variable.
  ! For scalars, just print value.

  {VTYPE}, intent(in) :: var{DIMSTR}
  integer, intent(in) :: loc_vec({DIMS})

  character(len=*), intent(in), optional :: varname

  character(len=:), allocatable :: varname_to_write

  if (present(varname)) then
     allocate(varname_to_write, source=varname)
  else
     allocate(varname_to_write, source="input variable")
  end if

  write(shr_log_Unit,*) &
       "ERROR: shr_assert_in_domain: ",trim(varname_to_write), &
       " has invalid value ", &
#if ({DIMS} != 0)
       var({REPEAT:loc_vec(#)}), &
       " at location: ",loc_vec
#else
       var

  ! Kill compiler spam for unused loc_vec.
  if (.false.) write(*,*) loc_vec
#endif

end subroutine print_bad_loc_{DIMS}d_{TYPE}

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

! DIMS 0,1,2,3,4,5,6,7
pure function find_first_loc_{DIMS}d(mask) result (loc_vec)
  ! Inefficient but simple subroutine for finding the location of
  ! the first .true. value in an array.
  ! If no true values, returns first value.

  logical, intent(in) :: mask{DIMSTR}
  integer :: loc_vec({DIMS})

#if ({DIMS} != 0)
  integer :: flags({REPEAT:size(mask,#)})

  where (mask)
     flags = 1
  elsewhere
     flags = 0
  end where

  loc_vec = maxloc(flags)
#else

! Remove compiler warnings (statement will be optimized out).

#if (! defined CPRPGI && ! defined CPRCRAY)
  if (.false. .and. mask) loc_vec = loc_vec
#endif

#endif

end function find_first_loc_{DIMS}d

! TYPE double,real,int,long
elemental function within_tolerance_{TYPE}(expected, actual, tolerance) &
     result(is_in_tol)
  ! Precondition: tolerance must be >= 0.
  {VTYPE}, intent(in) :: expected
  {VTYPE}, intent(in) :: actual
  {VTYPE}, intent(in) :: tolerance
  logical :: is_in_tol

  ! The following conditionals are to ensure that we don't overflow.

  ! This takes care of two identical infinities.
  if (actual == expected) then
     is_in_tol = .true.
  else if (actual > expected) then
     if (expected >= 0) then
        is_in_tol = (actual - expected) <= tolerance
     else
        is_in_tol = actual <= (expected + tolerance)
     end if
  else
     if (expected < 0) then
        is_in_tol = (expected - actual) <= tolerance
     else
        is_in_tol = actual >= (expected - tolerance)
     end if
  end if

end function within_tolerance_{TYPE}

end module shr_assert_mod
