! Module  : comp_range_m
! Purpose : define comp_range_t type, representing a composite of other
!           ranges
!
! Copyright 2022 Rich Townsend & The MSG Team
!
! This file is part of MSG. MSG 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, version 3.
!
! MSG 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 this program.  If not, see <http://www.gnu.org/licenses/>.

#:include 'forum.inc'

module comp_range_m

   ! Uses

   use forum_m
   
   use file_m
   use range_m
   use stat_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Derived-type definitions

   type, extends(range_t) :: comp_range_t
      private
      type(range_element_t), allocatable :: range_elements(:)
      integer, allocatable               :: i_bdy(:)
      integer                            :: n_ranges
   contains
      private
      procedure, public :: get_n
      procedure, public :: get_x_min
      procedure, public :: get_x_max
      procedure, public :: get_ssize
      procedure, public :: unpack
      procedure, public :: subset
      procedure, public :: read
      procedure, public :: write
   end type comp_range_t

   type :: range_element_t
      class(range_t), allocatable :: range
   end type range_element_t

   ! Interfaces

   interface comp_range_t
      module procedure comp_range_t_
   end interface comp_range_t

   ! Access specifiers

   private

   public :: comp_range_t

contains

   function comp_range_t_(ranges) result(range)

      class(range_t), intent(in) :: ranges(:)
      type(comp_range_t)         :: range

      type(range_element_t) :: range_elements(SIZE(ranges))
      integer               :: j

      ! Construct range from the ranges

      do j = 1, SIZE(range_elements)
         ! Straight assignment here seems to trigger a segfault with gfortran
         allocate(range_elements(j)%range, SOURCE=ranges(j))
      end do

      range = comp_range_t_elements_(range_elements)

      ! Finish

      return

   end function comp_range_t_

   !****

   function comp_range_t_elements_(range_elements) result(range)

      type(range_element_t), intent(in) :: range_elements(:)
      type(comp_range_t)                :: range

      integer  :: n_ranges
      integer  :: n
      real(RD) :: x_min
      real(RD) :: x_max
      real(RD) :: x_max_prev
      integer  :: j

      ! Construct range from the range elements

      n_ranges = SIZE(range_elements)

      @:ASSERT(n_ranges >= 1,'insufficient data')

      range%range_elements = range_elements

      allocate(range%i_bdy(n_ranges+1))

      range%i_bdy(1) = 1

      call range_elements(1)%range%get_n(n)
      call range_elements(1)%range%get_x_max(x_max)

      do j = 2, n_ranges

         range%i_bdy(j) = range%i_bdy(j-1) + n

         x_max_prev = x_max
         
         call range_elements(j)%range%get_n(n)
         call range_elements(j)%range%get_x_min(x_min)
         call range_elements(j)%range%get_x_max(x_max)

         @:ASSERT(x_min > x_max_prev,'overlapping ranges')

      end do

      range%i_bdy(n_ranges+1) = range%i_bdy(n_ranges) + n

      range%n_ranges = n_ranges

      ! Finish

      return

   end function comp_range_t_elements_

   !****

   subroutine get_n(self, n)

      class(comp_range_t), intent(in) :: self
      integer, intent(out)            :: n

      ! Get the size

      n = self%i_bdy(self%n_ranges+1)-1

      ! Finish

      return

   end subroutine get_n

   !****

   subroutine get_x_min(self, x_min)

      class(comp_range_t), intent(in) :: self
      real(RD), intent(out)           :: x_min

      ! Get the x-minimum

      call self%range_elements(1)%range%get_x_min(x_min)

      ! Finish

      return

   end subroutine get_x_min

   !****

   subroutine get_x_max(self, x_max)

      class(comp_range_t), intent(in) :: self
      real(RD), intent(out)           :: x_max

      ! Get the x-maximum

      call self%range_elements(self%n_ranges)%range%get_x_max(x_max)

      ! Finish

      return

   end subroutine get_x_max

   !****

   subroutine get_ssize(self, ssize)

      class(comp_range_t), intent(in) :: self
      integer(ID), intent(out)        :: ssize

      integer     :: j
      integer(ID) :: ssize_range

      ! Get the storage size

      ssize = STORAGE_SIZE(self)/8 +                                          &
              STORAGE_SIZE(self%range_elements)/8*SIZE(self%range_elements) + &
              STORAGE_SIZE(self%i_bdy)/8*SIZE(self%i_bdy)

      ranges_loop: do j = 1, self%n_ranges

         call self%range_elements(j)%range%get_ssize(ssize_range)

         ssize = ssize + ssize_range

      end do ranges_loop

      ! Finish

      return

   end subroutine get_ssize

   !****

   subroutine unpack(self, x)

      class(comp_range_t), intent(in) :: self
      real(RD), intent(out)           :: x(:)

      integer :: j

      @:CHECK_BOUNDS(SIZE(x), self%i_bdy(self%n_ranges+1)-1)

      ! Unpack the range into the array

      ranges_loop: do j = 1, self%n_ranges
         
         associate (i_a => self%i_bdy(j), &
                    i_b => self%i_bdy(j+1)-1)
           call self%range_elements(j)%range%unpack(x(i_a:i_b))
         end associate
      end do ranges_loop

      ! Finish

      return

   end subroutine unpack

   !****

   subroutine subset(self, i_a, i_b, range, stat)

      class(comp_range_t), intent(in)          :: self
      integer, intent(in)                      :: i_a
      integer, intent(in)                      :: i_b
      class(range_t), allocatable, intent(out) :: range
      integer, intent(out), optional           :: stat

      integer                            :: j_a
      integer                            :: j_b
      integer                            :: n_ranges
      type(range_element_t), allocatable :: range_elements(:)
      integer                             :: j
      integer                             :: k_a
      integer                             :: k_b

      ! Subset the range

      if (i_a < 1) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_RANGE_LO
            return
         else
            @:ABORT('out of bounds (lo) range')
         end if
      end if

      if (i_b > self%i_bdy(self%n_ranges+1)-1) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_RANGE_HI
            return
         else
            @:ABORT('out of bounds (hi) range')
         end if
      end if

      if (i_b < i_a) then
         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         end if
      end if

      ! First determine where in ranges i_a and i_b are located

      call locate(self%i_bdy, i_a, j_a)
      call locate(self%i_bdy, i_b, j_b)

      @:ASSERT(j_a >= 1,'out-of-bounds')
      @:ASSERT(j_b <= self%n_ranges,'out-of-bounds')

      ! Create the subset

      n_ranges = j_b - j_a + 1

      allocate(range_elements(n_ranges))

      ranges_loop : do j = j_a, j_b

         if (j == j_a) then
            k_a = i_a - self%i_bdy(j) + 1
         else
            k_a = 1
         end if

         if (j == j_b) then
            k_b = i_b - self%i_bdy(j) + 1
         else
            k_b = self%i_bdy(j+1) - self%i_bdy(j)
         end if

         call self%range_elements(j)%range%subset(k_a, k_b, range_elements(j-j_a+1)%range)

      end do ranges_loop

      range = comp_range_t_elements_(range_elements)

      ! Finish

      return

   end subroutine subset

   !****

   subroutine read(self, hdf5io, stat)

      class(comp_range_t), intent(out) :: self
      type(hdf5io_t), intent(inout)    :: hdf5io
      integer, intent(out), optional   :: stat

      integer :: revision

      ! Read the range

      call check_type(hdf5io, 'comp_range_t', stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      call hdf5io%read_attr('REVISION', revision)

      select case(revision)
      case(1)
         call read_rev1_(self, hdf5io, stat)
      case default
         if (PRESENT(stat)) then
            stat = STAT_INVALID_GROUP_REVISION
            return
         else
            @:ABORT('invalid group revision')
         end if
      end select

      ! Finish

      return

   contains

      subroutine read_rev1_(self, hdf5io, stat)

         class(comp_range_t), intent(out) :: self
         type(hdf5io_t), intent(inout)    :: hdf5io
         integer, intent(out), optional   :: stat

         integer                            :: n_ranges
         type(range_element_t), allocatable :: range_elements(:)
         type(hdf5io_t)                     :: hdf5io_range
         integer                            :: j
         character(TYPE_LEN)                :: type
      
         call hdf5io%read_attr('n_ranges', n_ranges)

         allocate(range_elements(n_ranges))

         ranges_loop : do j = 1, n_ranges

            hdf5io_range = hdf5io_t(hdf5io, range_group_name_(j))

            call hdf5io_range%read_attr('TYPE', type)
            allocate(range_elements(j)%range, MOLD=range_t(type))

            call range_elements(j)%range%read(hdf5io_range, stat)

            call hdf5io_range%final()
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

         end do ranges_loop

         select type(self)
         type is (comp_range_t)
            self = comp_range_t_elements_(range_elements)
         class default
            @:ABORT('invalid type')
         end select

         if (PRESENT(stat)) stat = STAT_OK

         return
      
      end subroutine read_rev1_

   end subroutine read

   !****

   subroutine write(self, hdf5io, stat)

      class(comp_range_t), intent(in) :: self
      type(hdf5io_t), intent(inout)   :: hdf5io
      integer, intent(out), optional  :: stat

      integer        :: j
      type(hdf5io_t) :: hdf5io_range

      ! Write the range

      call hdf5io%write_attr('TYPE', 'comp_range_t')
      call hdf5io%write_attr('REVISION', 1)

      call hdf5io%write_attr('n_ranges', self%n_ranges)

      ranges_loop : do j = 1, self%n_ranges

         hdf5io_range = hdf5io_t(hdf5io, range_group_name_(j))
         call self%range_elements(j)%range%write(hdf5io_range, stat)
         call hdf5io_range%final()
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         end if

      end do ranges_loop

      ! Finish

      return

   end subroutine write

   !****

   function range_group_name_(i) result(name)

      integer, intent(in) :: i
      character(256)      :: name

      write(name, 100) i
100   format('ranges[',I0,']')
      
   end function range_group_name_

end module comp_range_m
