! Copyright 2008-2021 Earth Sciences Department, BSC-CNS
!
! This file is part of MONARCH.
!
! MONARCH is free software: you can redistribute it and/or modify
! it under the terms of the GNU General Public License as published by
! the Free Software Foundation, either version 3 of the License, or
! (at your option) any later version.
!
! MONARCH is distributed in the hope that it will be useful,
! but WITHOUT ANY WARRANTY; without even the implied warranty of
! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
! GNU General Public License for more details.
!
! You should have received a copy of the GNU General Public License
! along with MONARCH. If not, see <http://www.gnu.org/licenses/>.

!> \file
!> MONARCH utility functions

!> MONARCH utility functions
module monarch_util
! MONARCH always has MPI
#define MONARCH_USE_MPI

  use monarch_constants,                    only : mint, mflt, mdbl

  implicit none

  !> Max length for string conversion
  integer(kind=mint), parameter, private :: MONARCH_UTIL_CONVERT_STRING_LEN = 100

  !> Unknown vertical interface
  integer(kind=mint), parameter, private :: VERTICAL_INTERFACE_UNKNOWN = 0
  !> Lower vertical interface
  integer(kind=mint), parameter, private :: VERTICAL_INTERFACE_LOWER   = 1
  !> Upper vertical interface
  integer(kind=mint), parameter, private :: VERTICAL_INTERFACE_UPPER   = 2

  !> String type for building arrays of string of various size
  type string_t
    !> String value
    character(len=:), allocatable :: string
  contains
    !> Equality comparison
    procedure :: equals => string_t_equals
    generic :: operator(==) => equals
  end type string_t

  !> String assignment interface
  interface assignment (=)
    module procedure string_t_assign_from_char_array
  end interface

  !> Temporal coordinate object
  type :: monarch_point_t
  private
    !> Number of vertical layers
    integer(kind=mint) :: n_layers = 0
    !> Latitudinal index
    integer(kind=mint) :: i_WE = 0
    !> Longitudinal index
    integer(kind=mint) :: i_SN = 0
    !> Vertical index( starting at the surface )
    integer(kind=mint) :: i_vert = 0
    !> Vertical index( starting at the top layer )
    integer(kind=mint) :: i_vert_flip = 0
    !> Upper or Lower interface
    integer(kind=mint) :: i_vert_interface = VERTICAL_INTERFACE_UNKNOWN
    !> Time step
    integer(kind=mint) :: i_time = 0
  contains
    !> Get the latitudinal index
    procedure :: WE_idx
    !> Set the latitudinal index
    procedure :: set_WE_idx
    !> Get the longitudinal index
    procedure :: SN_idx
    !> Set the longitudinal index
    procedure :: set_SN_idx
    !> Get the vertical index
    procedure :: vertical_idx
    !> Get the flipped vertical index
    procedure :: flipped_vertical_idx
    !> Set the vertical index
    procedure :: set_vertical_idx
    !> Get the vertical interface index
    procedure :: vertical_interface_idx
    !> Get the flipped vertical interface index
    procedure :: flipped_vertical_interface_idx
    !> Reset the vertical interface
    procedure :: reset_vertical_interface
    !> Set to the lower vertical interface
    procedure :: set_to_lower_interface
    !> Set to the upper vertical interface
    procedure :: set_to_upper_interface
    !> Get the time index
    procedure :: time_idx
    !> Set the time index
    procedure :: set_time_idx
  end type monarch_point_t

  !> Constructor for the monarch_point_t object
  interface monarch_point_t
    procedure :: monarch_point_constructor
  end interface monarch_point_t

  !> Dimensional object
  !!
  !! Upper and lower bounds of a spatiotemporal region
  type :: monarch_dimensions_t
    !> Lower bound
    type(monarch_point_t) :: lower
    !> Upper bound
    type(monarch_point_t) :: upper
  end type monarch_dimensions_t

  !> Interface for to_string() functions
  interface to_string
    module procedure integer_to_string, float_to_string, double_to_string,   &
                     logical_to_string, float_complex_to_string,             &
                     double_complex_to_string
  end interface to_string

contains

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Equality comparison for strings
  function string_t_equals( this, other ) result( is_equal )

    !> Flag indicating equality
    logical :: is_equal
    !> This string
    class(string_t), intent(in) :: this
    !> String to compare with
    type(string_t), intent(in) :: other

    is_equal = .false.

    if( allocated( this%string ) .and. allocated( other%string ) ) then
      is_equal = this%string == other%string
    end if

  end function string_t_equals

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Assignment of a character array to a string_t
  subroutine string_t_assign_from_char_array( this, char_array )

    !> The string object
    type(string_t), intent(out) :: this
    !> The character array used to set the string value
    character(len=*), intent(in) :: char_array

    this%string = char_array

  end subroutine string_t_assign_from_char_array

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Prints a warning message.
  subroutine warn_msg(code, warning_msg, already_warned)

    !> Status code to use.
    integer, intent(in) :: code
    !> Message to display.
    character(len=*), intent(in) :: warning_msg
    !> Flag to control warning only once (should be a save variable).
    logical, intent(inout), optional :: already_warned

    if (present(already_warned)) then
       if (already_warned) return
       ! set already_warned so next time we will immediately return
       already_warned = .true.
    end if
    write(0,'(a)') 'WARNING (MONARCH-' // trim(integer_to_string(code)) &
         // '): ' // trim(warning_msg)

  end subroutine warn_msg

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Prints a warning message if condition_ok is false.
  subroutine warn_assert_msg(code, condition_ok, warning_msg)

    !> Status code to use.
    integer, intent(in) :: code
    !> Whether the assertion is ok.
    logical, intent(in) :: condition_ok
    !> Message to display.
    character(len=*), intent(in) :: warning_msg

    if (.not. condition_ok) then
       call warn_msg(code, warning_msg)
    end if

  end subroutine warn_assert_msg

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Prints a test failure message if condition_ok is false.
  function test_assert(code, condition_ok, failure_msg)

    !> Returns 1 on failure, 0 on pass
    integer :: test_assert
    !> Status code to use.
    integer, intent(in) :: code
    !> Whether the assertion is ok.
    logical, intent(in) :: condition_ok
    !> Message to display.
    character(len=*), intent(in) :: failure_msg

    test_assert = 0
    if (.not. condition_ok) then
      write(0,'(a)') 'TEST FAILURE (MONARCH-' // trim(integer_to_string(code)) &
         // '): ' // trim(failure_msg)
      test_assert = 1
    end if

  end function test_assert

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Errors unless condition_ok is true.
  subroutine assert_msg(code, condition_ok, error_msg)
#ifdef MONARCH_USE_MPI
  use mpi
#endif

    !> Status code to use if assertion fails.
    integer, intent(in) :: code
    !> Whether the assertion is ok.
    logical, intent(in) :: condition_ok
    !> Msg if assertion fails.
    character(len=*), intent(in) :: error_msg

    integer :: ierr

    if (.not. condition_ok) then
       write(0,'(a)') 'ERROR (MONARCH-' // trim(integer_to_string(code)) &
            // '): ' // trim(error_msg)
#ifdef MONARCH_USE_MPI
       call mpi_abort(MPI_COMM_WORLD, code, ierr)
#else
       stop 1
#endif
    end if

  end subroutine assert_msg

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Errors unless condition_ok is true.
  subroutine assert(code, condition_ok)

    !> Status code to use if assertion fails.
    integer, intent(in) :: code
    !> Whether the assertion is ok.
    logical, intent(in) :: condition_ok

    if (.not. condition_ok) then
       call assert_msg(code, condition_ok, 'assertion failed')
    end if

  end subroutine assert

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Error immediately.
  subroutine die(code)

    !> Status code to use if assertion fails.
    integer, intent(in) :: code

    call assert(code, .false.)

  end subroutine die

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Error immediately.
  subroutine die_msg(code, error_msg)

    !> Status code to use if assertion fails.
    integer, intent(in) :: code
    !> Msg if assertion fails.
    character(len=*), intent(in) :: error_msg

    call assert_msg(code, .false., error_msg)

  end subroutine die_msg

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert an integer to a string format.
  function integer_to_string( val ) result( str_val )

    !> Integer as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    integer(kind=mint), intent(in) :: val

    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: ret_val

    ret_val = ""
    write(ret_val, '(i30)') val
    str_val = adjustl(ret_val)

  end function integer_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert a single precision floating-point number to a string format.
  function float_to_string( val ) result( str_val )

    !> Number as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    real(kind=mflt), intent(in) :: val

    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: ret_val

    ret_val = ""
    write(ret_val, '(g30.20)') val
    str_val = adjustl(ret_val)

  end function float_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert a double precision floating-point number to a string format.
  function double_to_string( val ) result( str_val )

    !> Number as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    real(kind=mdbl), intent(in) :: val

    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: ret_val

    ret_val = ""
    write(ret_val, '(g30.20)') val
    str_val = adjustl(ret_val)

  end function double_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert a logical value to a string format.
  function logical_to_string( val ) result( str_val )

    !> Logical as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    logical, intent(in) :: val

    str_val = ""
    if( val ) then
      str_val = "TRUE"
    else
      str_val = "FALSE"
    end if

  end function logical_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert a complex to a string format.
  function double_complex_to_string( val ) result( str_val )

    !> Complex real number as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    complex(kind=mdbl), intent(in) :: val

    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: ret_val

    ret_val = ""
    ret_val = "(" // trim(to_string(real(val))) &
         // ", " // trim(to_string(aimag(val))) // ")"
    str_val = adjustl(ret_val)

  end function double_complex_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Convert a complex to a string format.
  function float_complex_to_string( val ) result( str_val )

    !> Complex real number as a string
    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: str_val
    !> Value to convert.
    complex(kind=mflt), intent(in) :: val

    character(len=MONARCH_UTIL_CONVERT_STRING_LEN) :: ret_val

    ret_val = ""
    ret_val = "(" // trim(to_string(real(val))) &
         // ", " // trim(to_string(aimag(val))) // ")"
    str_val = adjustl(ret_val)

  end function float_complex_to_string

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Constructor for the monarch_point_t type
  function monarch_point_constructor( num_layers ) result( new_obj )

    !> New monarch_point_t object
    type(monarch_point_t) :: new_obj
    !> Total number of model layers
    integer(kind=mint), intent(in) :: num_layers

    new_obj%n_layers = num_layers

  end function monarch_point_constructor

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the latitudinal index
  pure function WE_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    idx = this%i_WE

  end function WE_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set the latitudinal index
  subroutine set_WE_idx( this, idx )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this
    !> Index
    integer(kind=mint), intent(in) :: idx

    this%i_WE = idx

  end subroutine set_WE_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the longitudinal index
  pure function SN_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    idx = this%i_SN

  end function SN_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set the longitudinal index
  subroutine set_SN_idx( this, idx )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this
    !> Index
    integer(kind=mint), intent(in) :: idx

    this%i_SN = idx

  end subroutine set_SN_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the vertical index
  pure function vertical_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    idx = this%i_vert

  end function vertical_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the flipped vertical index
  pure function flipped_vertical_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    idx = this%i_vert_flip

  end function flipped_vertical_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set the vertical index
  subroutine set_vertical_idx( this, idx )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this
    !> Index
    integer(kind=mint), intent(in) :: idx

    this%i_vert      = idx
    this%i_vert_flip = this%n_layers + 1 - idx

  end subroutine set_vertical_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the vertical interface index
  function vertical_interface_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    select case( this%i_vert_interface )
      case( VERTICAL_INTERFACE_LOWER )
        idx = this%i_vert
      case( VERTICAL_INTERFACE_UPPER )
        idx = this%i_vert + 1
      case default
        call die_msg( 996188522, "Trying to use un-set interface index" )
    end select

  end function vertical_interface_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the flipped vertical interface index
  function flipped_vertical_interface_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    select case( this%i_vert_interface )
      case( VERTICAL_INTERFACE_LOWER )
        idx = this%i_vert_flip + 1
      case( VERTICAL_INTERFACE_UPPER )
        idx = this%i_vert_flip
      case default
        call die_msg( 984206669, "Trying to use un-set interface index" )
    end select

  end function flipped_vertical_interface_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Reset the vertical interface
  subroutine reset_vertical_interface( this )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this

    this%i_vert_interface = VERTICAL_INTERFACE_UNKNOWN

  end subroutine reset_vertical_interface

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set to the lower vertical interface
  subroutine set_to_lower_interface( this )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this

    this%i_vert_interface = VERTICAL_INTERFACE_LOWER

  end subroutine set_to_lower_interface

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set to the upper vertical interface
  subroutine set_to_upper_interface( this )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this

    this%i_vert_interface = VERTICAL_INTERFACE_UPPER

  end subroutine set_to_upper_interface

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Get the time index
  pure function time_idx( this ) result( idx )

    !> Index
    integer(kind=mint) :: idx
    !> Temporal indices
    class(monarch_point_t), intent(in) :: this

    idx = this%i_time

  end function time_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !> Set the time index
  subroutine set_time_idx( this, idx )

    !> Temporal indices
    class(monarch_point_t), intent(inout) :: this
    !> Index
    integer(kind=mint), intent(in) :: idx

    this%i_time = idx

  end subroutine set_time_idx

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

end module monarch_util
