! 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
!> The monarch_state_var_t object and related functions

!> Structure for storing variables on the SOLVER_INTERNAL_STATE
module monarch_state_var

  use monarch_constants,                    only : mint, mflt, mdbl

  implicit none
  private

  public :: monarch_state_var_t

  !> State variable types
  integer(kind=mint), parameter :: STATE_VAR_TYPE_UNALLOCATED = 0
  integer(kind=mint), parameter :: STATE_VAR_TYPE_2D_REAL     = 1
  integer(kind=mint), parameter :: STATE_VAR_TYPE_3D_REAL     = 2
  integer(kind=mint), parameter :: STATE_VAR_TYPE_4D_REAL     = 3

  !> State variable in the SOLVER_INTERNAL_STATE
  !!
  !! Contains the set of variables allocated on the SOLVER_INTERNAL_STATE
  !! from monarch_state_changer_t
  type monarch_state_var_t
    private
    !> Variable type
    integer(kind=mint) :: var_type = STATE_VAR_TYPE_UNALLOCATED
    !> 2D real array
    real(kind=mflt), pointer :: var_data_2D_real( :,: ) => null( )
    !> 3D real array
    real(kind=mflt), pointer :: var_data_3D_real( :,:,: ) => null( )
    !> 4D real array
    real(kind=mflt), pointer :: var_data_4D_real( :,:,:,: ) => null( )
  contains
    !> Set variable
    generic :: set_variable => set_variable_2D_real,                         &
                               set_variable_3D_real,                         &
                               set_variable_4D_real
    !> Set a 2D real variable
    procedure :: set_variable_2D_real
    !> Set a 3D real variable
    procedure :: set_variable_3D_real
    !> Set a 4D real variable
    procedure :: set_variable_4D_real
    !> Set the variable pointer in the solver internal state VARS array
    procedure :: set_variable_ptr
    !> Attach to an already-allocated solver intergnal state VARS element
    procedure :: attach_variable_ptr
    !> Reset an array to a given value
    procedure :: reset_array
    !> Determine whether the state array has been properly allocated
    procedure :: is_allocated
  end type monarch_state_var_t

contains

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

  !> Set a 2D real variable
  elemental subroutine set_variable_2D_real( this, dim1, dim2, new_value )

    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Dimension 1
    integer(kind=mint), intent(in) :: dim1
    !> Dimension 2
    integer(kind=mint), intent(in) :: dim2
    !> New value
    real(kind=mflt), intent(in) :: new_value

    this%var_data_2D_real( dim1, dim2 ) = new_value

  end subroutine set_variable_2D_real

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

  !> Set a 3D real variable
  elemental subroutine set_variable_3D_real( this, dim1, dim2, dim3,         &
      new_value )

    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Dimension 1
    integer(kind=mint), intent(in) :: dim1
    !> Dimension 2
    integer(kind=mint), intent(in) :: dim2
    !> Dimension 3
    integer(kind=mint), intent(in) :: dim3
    !> New value
    real(kind=mflt), intent(in) :: new_value

    this%var_data_3D_real( dim1, dim2, dim3 ) = new_value

  end subroutine set_variable_3D_real

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

  !> Set a 4D real variable
  elemental subroutine set_variable_4D_real( this, dim1, dim2, dim3, dim4,   &
      new_value )

    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Dimension 1
    integer(kind=mint), intent(in) :: dim1
    !> Dimension 2
    integer(kind=mint), intent(in) :: dim2
    !> Dimension 3
    integer(kind=mint), intent(in) :: dim3
    !> Dimension 4
    integer(kind=mint), intent(in) :: dim4
    !> New value
    real(kind=mflt), intent(in) :: new_value

    this%var_data_4D_real( dim1, dim2, dim3, dim4 ) = new_value

  end subroutine set_variable_4D_real

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

  !> Set a pointer for this array in the solver internal state VARS array
  subroutine set_variable_ptr( this, vars, num_vars, unique_name,            &
      lower_bounds, upper_bounds )

    use module_VARS,                        only : VAR, SET_VAR_PTR
    use monarch_util,                       only : assert, die_msg, to_string

    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Variable array from interal state
    type(VAR), dimension( : ), intent(inout) :: vars
    !> Number of state variables
    integer(kind=mint), intent(in) :: num_vars
    !> Unique array name
    character(len=*), intent(in) :: unique_name
    !> Lower array bounds
    integer(kind=mint), dimension( : ), intent(in) :: lower_bounds
    !> Upper array bounds
    integer(kind=mint), dimension( : ), intent(in) :: upper_bounds

    ! Figure out the shape of the variable
    call assert( 255691772, size( lower_bounds ) .eq. size( upper_bounds ) )
    select case( size( lower_bounds ) )
      case( 2 )
        this%var_type = STATE_VAR_TYPE_2D_REAL
      case( 3 )
        this%var_type = STATE_VAR_TYPE_3D_REAL
      case( 4 )
        this%var_type = STATE_VAR_TYPE_4D_REAL
      case default
        call die_msg( 747463085, "Invalid shape for state variable: "//      &
                      to_string( size( lower_bounds ) ) )
    end select

    ! Set a pointer to the appropriate array by type
    select case( this%var_type )
      case( STATE_VAR_TYPE_2D_REAL )
        call SET_VAR_PTR( vars, num_vars, unique_name,                       &
                          this%var_data_2D_real, lower_bounds, upper_bounds )
        call assert( 995585589, associated( this%var_data_2D_real ) )
      case( STATE_VAR_TYPE_3D_REAL )
        call SET_VAR_PTR( vars, num_vars, unique_name,                       &
                          this%var_data_3D_real, lower_bounds, upper_bounds )
        call assert( 265742494, associated( this%var_data_3D_real ) )
      case( STATE_VAR_TYPE_4D_REAL )
        call SET_VAR_PTR( vars, num_vars, unique_name,                       &
                          this%var_data_4D_real, lower_bounds, upper_bounds )
        call assert( 770948899, associated( this%var_data_4D_real ) )
      case default
        call die_msg( 357812155, "Unknown state array type: "//              &
                      to_string( this%var_type ) )
    end select

  end subroutine set_variable_ptr

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

  !> Attach to an already-allocated solver internal state VARS element
  !!
  !! This function finds an array in the internal state VARS array, checks
  !! its shape, and returns the index for the array in VARS
  function attach_variable_ptr( this, vars, num_vars, unique_name,           &
      lower_bounds, upper_bounds ) result( indx )

    use module_VARS,                        only : VAR, SET_VAR_PTR,         &
                                                   FIND_VAR_INDX
    use monarch_util,                       only : assert, die_msg, to_string

    !> Index of array in VARS
    integer(kind=mint) :: indx
    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Variable array from interal state
    type(VAR), dimension( : ), intent(inout) :: vars
    !> Number of state variables
    integer(kind=mint), intent(in) :: num_vars
    !> Unique array name
    character(len=*), intent(in) :: unique_name
    !> Lower array bounds
    integer(kind=mint), dimension( : ), intent(in) :: lower_bounds
    !> Upper array bounds
    integer(kind=mint), dimension( : ), intent(in) :: upper_bounds

    integer(kind=mint) :: dim1, dim2, dim3, dim4

    ! Figure out the shape of the variable
    call assert( 255691772, size( lower_bounds ) .eq. size( upper_bounds ) )
    select case( size( lower_bounds ) )
      case( 2 )
        this%var_type = STATE_VAR_TYPE_2D_REAL
        dim1 = upper_bounds( 1 ) - lower_bounds( 1 ) + 1
        dim2 = upper_bounds( 2 ) - lower_bounds( 2 ) + 1
      case( 3 )
        this%var_type = STATE_VAR_TYPE_3D_REAL
        dim1 = upper_bounds( 1 ) - lower_bounds( 1 ) + 1
        dim2 = upper_bounds( 2 ) - lower_bounds( 2 ) + 1
        dim3 = upper_bounds( 3 ) - lower_bounds( 3 ) + 1
      case( 4 )
        this%var_type = STATE_VAR_TYPE_4D_REAL
        dim1 = upper_bounds( 1 ) - lower_bounds( 1 ) + 1
        dim2 = upper_bounds( 2 ) - lower_bounds( 2 ) + 1
        dim3 = upper_bounds( 3 ) - lower_bounds( 3 ) + 1
        dim4 = upper_bounds( 4 ) - lower_bounds( 4 ) + 1
      case default
        call die_msg( 654179030, "Invalid shape for state variable: "//      &
                      to_string( size( lower_bounds ) ) )
    end select

    ! Set a pointer to the appropriate array by type
    ! Ensure that the array is at least as big as the specified dimensions
    select case( this%var_type )
      case( STATE_VAR_TYPE_2D_REAL )
        call FIND_VAR_INDX( unique_name, vars, num_vars, indx )
        call assert( 619851882, associated( vars( indx )%r2d ) )
        call assert( 560560084, size( vars( indx )%r2d, 1 ) .ge. dim1 )
        call assert( 422824794, size( vars( indx )%r2d, 2 ) .ge. dim2 )
        this%var_data_2D_real => vars( indx )%r2d
      case( STATE_VAR_TYPE_3D_REAL )
        call FIND_VAR_INDX( unique_name, vars, num_vars, indx )
        call assert( 293712584, associated( vars( indx )%r3d ) )
        call assert( 406030929, size( vars( indx )%r3d, 1 ) .ge. dim1 )
        call assert( 235874025, size( vars( indx )%r3d, 2 ) .ge. dim2 )
        call assert( 130725521, size( vars( indx )%r3d, 3 ) .ge. dim3 )
        this%var_data_3D_real => vars( indx )%r3d
      case( STATE_VAR_TYPE_4D_REAL )
        call FIND_VAR_INDX( unique_name, vars, num_vars, indx )
        call assert( 243043866, associated( vars( indx )%r4d ) )
        call assert( 407936463, size( vars( indx )%r4d, 1 ) .ge. dim1 )
        call assert( 520254808, size( vars( indx )%r4d, 2 ) .ge. dim2 )
        call assert( 685147405, size( vars( indx )%r4d, 3 ) .ge. dim3 )
        call assert( 232515252, size( vars( indx )%r4d, 4 ) .ge. dim4 )
        this%var_data_4D_real => vars( indx )%r4d
      case default
        indx = 0
        call die_msg( 816165149, "Unknown state array type: "//              &
                      to_string( this%var_type ) )
    end select

  end function attach_variable_ptr

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

  !> Reset an array to a given value
  subroutine reset_array( this, new_value )

    use monarch_util,                            only : assert_msg, die_msg, &
                                                        to_string

    !> State variable
    class(monarch_state_var_t), intent(inout) :: this
    !> Value to set array with
    real(kind=mflt), intent(in) :: new_value

    call assert_msg( 577490939, this%is_allocated( ),                        &
                     "Trying to intialize an unallocated array" )

    select case( this%var_type )
      case( STATE_VAR_TYPE_2D_REAL )
        this%var_data_2D_real = new_value
      case( STATE_VAR_TYPE_3D_REAL )
        this%var_data_3D_real = new_value
      case( STATE_VAR_TYPE_4D_REAL )
        this%var_data_4D_real = new_value
      case default
        call die_msg( 623799940, "Unknown state array type: "//              &
                      to_string( this%var_type ) )
    end select

  end subroutine reset_array

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

  !> Determine whether the state variable has been properly allocated
  function is_allocated( this )

    use monarch_util,                       only : assert, die_msg, to_string

    !> Flag indicating whether the state variable is allocated
    logical :: is_allocated
    !> State variable
    class(monarch_state_var_t), intent(in) :: this

    is_allocated = .false.

    select case( this%var_type )
      case( STATE_VAR_TYPE_UNALLOCATED )
        call assert( 640357853,                                              &
                     .not. associated( this%var_data_2D_real ) .and.         &
                     .not. associated( this%var_data_3D_real ) .and.         &
                     .not. associated( this%var_data_4D_real ) )
      case( STATE_VAR_TYPE_2D_REAL )
        call assert( 409455912,                                              &
                           associated( this%var_data_2D_real ) .and.         &
                     .not. associated( this%var_data_3D_real ) .and.         &
                     .not. associated( this%var_data_4D_real ) )
        is_allocated = .true.
      case( STATE_VAR_TYPE_3D_REAL )
        call assert( 737787867,                                              &
                     .not. associated( this%var_data_2D_real ) .and.         &
                           associated( this%var_data_3D_real ) .and.         &
                     .not. associated( this%var_data_4D_real ) )
        is_allocated = .true.
      case( STATE_VAR_TYPE_4D_REAL )
        call assert( 957160250,                                              &
                     .not. associated( this%var_data_2D_real ) .and.         &
                     .not. associated( this%var_data_3D_real ) .and.         &
                           associated( this%var_data_4D_real ) )
        is_allocated = .true.
      case default
        call die_msg( 953801477, "Unknown state array type: "//              &
                      to_string( this%var_type ) )
    end select

  end function is_allocated

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

end module monarch_state_var
