! Module  : specgrid_m
! Purpose : Define specgrid_t type, for representing spectral intensity
!           grids
!
! Copyright 2021-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 specgrid_m

   ! Uses
   
   use forum_m

   use axis_m
   use file_m
   use passband_m
   use photgrid_m
   use speccache_m
   use specint_m
   use specsource_m
   use hdf5_specsource_m
   use spec_photsource_m
   use stat_m
   use vgrid_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Derived-type definitions

   type :: specgrid_t
      private
      class(specsource_t), allocatable :: specsource
      type(speccache_t)                :: speccache
      type(vgrid_t)                    :: vgrid
      real(RD)                         :: lam_min
      real(RD)                         :: lam_max
   contains
      private
      procedure, public :: get_rank
      procedure, public :: get_shape
      procedure, public :: get_axis
      procedure, public :: get_lam_min
      procedure, public :: get_lam_max
      procedure, public :: get_cache_lam_min
      procedure, public :: get_cache_lam_max
      procedure, public :: get_cache_limit
      procedure, public :: get_cache_usage
      procedure, public :: set_cache_lam_min
      procedure, public :: set_cache_lam_max
      procedure, public :: set_cache_limit
      procedure, public :: flush_cache
      #:for name in ('intensity', 'E_moment', 'D_moment', 'flux')
         procedure, public :: interp_${name}$
      #:endfor
      procedure, public :: adjust_x_vec
      procedure, public :: filter
      procedure, public :: read
      procedure, public :: write
   end type specgrid_t

   ! Interfaces

   interface specgrid_t
      module procedure specgrid_t_
   end interface specgrid_t

   ! Access specifiers

   private

   public :: specgrid_t

   ! Procedures

contains

   function specgrid_t_(specsource, vgrid) result (specgrid)

      class(specsource_t), intent(in) :: specsource
      type(vgrid_t), intent(in)       :: vgrid
      type(specgrid_t)                :: specgrid

      integer  :: n
      real(RD) :: lam_min
      real(RD) :: lam_max
      integer  :: n_seq

      ! Construct specgrid from the supplied data

      call specsource%get_n(n)
      call specsource%get_lam_min(lam_min)
      call specsource%get_lam_max(lam_max)
      
      call vgrid%get_n_seq(n_seq)

      @:CHECK_BOUNDS(n, n_seq)

      specgrid%specsource = specsource
      specgrid%speccache = speccache_t(n, lam_min, lam_max)

      specgrid%vgrid = vgrid

      specgrid%lam_min = lam_min
      specgrid%lam_max = lam_max

      ! Finish

      return

   end function specgrid_t_

   !****

   subroutine get_rank(self, rank)

      class(specgrid_t), intent(in) :: self
      integer, intent(out)          :: rank

      ! Get the rank

      call self%vgrid%get_rank(rank)

      ! Finish

      return

   end subroutine get_rank

   !****

   subroutine get_shape(self, shape)

      class(specgrid_t), intent(in) :: self
      integer, intent(out)          :: shape(:)

      ! Get the shape

      call self%vgrid%get_shape(shape)

      ! Finish

      return

   end subroutine get_shape

   !****

   subroutine get_axis(self, i, axis)

      class(specgrid_t), intent(in) :: self
      integer, intent(in)           :: i
      type(axis_t), intent(out)     :: axis

      ! Get the i'th axis

      call self%vgrid%get_axis(i, axis)

      ! Finish

      return

   end subroutine get_axis

   !****

   #:for name, type in (('lam_min', 'real(RD)'), &
                        ('lam_max', 'real(RD)'))
                        
      subroutine get_${name}$(self, ${name}$)

         class(specgrid_t), intent(in) :: self
         ${type}$, intent(out)         :: ${name}$

         ! Get ${name}$

         ${name}$ = self%${name}$

         ! Finish

         return

      end subroutine get_${name}$

   #:endfor

   !****

   #:for name, type in (('lam_min', 'real(RD)'), &
                        ('lam_max', 'real(RD)'), &
                        ('limit', 'integer'),    &
                        ('usage', 'integer'))

      subroutine get_cache_${name}$(self, cache_${name}$)

         class(specgrid_t), intent(in) :: self
         ${type}$, intent(out)         :: cache_${name}$

         ! Get cache_${name}$

         call self%speccache%get_${name}$(cache_${name}$)

         ! Finish

         return

      end subroutine get_cache_${name}$

   #:endfor

   !****

   subroutine set_cache_lam_min(self, cache_lam_min, stat)

      class(specgrid_t), intent(inout) :: self
      real(RD), intent(in)             :: cache_lam_min
      integer, intent(out), optional   :: stat

      ! Set the cache minimum wavelength

      if (cache_lam_min < self%lam_min) then
         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         end if
      end if

      call self%speccache%set_lam_min(cache_lam_min, stat)

      ! Finish

      return

   end subroutine set_cache_lam_min

   !****

   subroutine set_cache_lam_max(self, cache_lam_max, stat)

      class(specgrid_t), intent(inout) :: self
      real(RD), intent(in)             :: cache_lam_max
      integer, intent(out), optional   :: stat

      ! Set the cache maximum wavelength

      if (cache_lam_max > self%lam_max) then
         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         end if
      end if

      call self%speccache%set_lam_max(cache_lam_max, stat)

      ! Finish

      return

   end subroutine set_cache_lam_max

   !****

   subroutine set_cache_limit(self, cache_limit, stat)

      class(specgrid_t), intent(inout) :: self
      integer, intent(in)              :: cache_limit
      integer, intent(out), optional   :: stat

      ! Set the cache memory usage limit

      call self%speccache%set_limit(cache_limit, stat)

      ! Finish

      return

   end subroutine set_cache_limit

   !****

   subroutine flush_cache(self)

      class(specgrid_t), intent(inout) :: self

      ! Flush the cache

      call self%speccache%flush()

      ! Finish

      return

   end subroutine flush_cache

   !****

   #:for name, arg_var, arg_type in (('intensity', 'mu', 'real(RD)'), &
                                     ('E_moment', 'k', 'integer'), &
                                     ('D_moment', 'l', 'integer'))

      subroutine interp_${name}$(self, x_vec, ${arg_var}$, lam, res, stat, deriv_vec, order)

         class(specgrid_t), target, intent(inout) :: self
         real(RD), intent(in)                     :: x_vec(:)
         ${arg_type}$, intent(in)                 :: ${arg_var}$
         real(RD), intent(in)                     :: lam(:)
         real(RD), intent(out)                    :: res(:)
         integer, intent(out), optional           :: stat
         logical, intent(in), optional            :: deriv_vec(:)
         integer, intent(in), optional            :: order

         ! Interpolate ${name}$

         call self%vgrid%interp(data_proc_, x_vec, res, stat, deriv_vec, order)

         ! Finish

      contains

         subroutine data_proc_(v_seq, f, stat)

            integer, intent(in)            :: v_seq
            real(RD), intent(out)          :: f(:)
            integer, intent(out), optional :: stat

            class(specint_t), pointer :: specint

            call self%speccache%fetch(v_seq, fetch_proc_, specint, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

            call specint%interp_${name}$(${arg_var}$, lam, f, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

            call self%speccache%release(v_seq, stat)

            return

         end subroutine data_proc_

         subroutine fetch_proc_(i, specint, stat)

            integer, intent(in)                        :: i
            class(specint_t), allocatable, intent(out) :: specint
            integer, intent(out), optional             :: stat

            call self%specsource%fetch(i, specint, stat)

            return

         end subroutine fetch_proc_

      end subroutine interp_${name}$

   #:endfor

   !****

   #:for name in ('flux',)

      subroutine interp_${name}$(self, x_vec, lam, res, stat, deriv_vec, order)

         class(specgrid_t), intent(inout) :: self
         real(RD), intent(in)             :: x_vec(:)
         real(RD), intent(in)             :: lam(:)
         real(RD), intent(out)            :: res(:)
         integer, intent(out), optional   :: stat
         logical, intent(in), optional    :: deriv_vec(:)
         integer, intent(in), optional    :: order

         ! Interpolate ${name}$

         call self%vgrid%interp(data_proc_, x_vec, res, stat, deriv_vec, order)

         ! Finish

      contains

         subroutine data_proc_(v_seq, f, stat)

            integer, intent(in)            :: v_seq
            real(RD), intent(out)          :: f(:)
            integer, intent(out), optional :: stat

            class(specint_t), pointer :: specint

            call self%speccache%fetch(v_seq, fetch_proc_, specint, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

            call specint%interp_${name}$(lam, f, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

            call self%speccache%release(v_seq, stat)

            return

         end subroutine data_proc_

         subroutine fetch_proc_(i, specint, stat)

            integer, intent(in)                        :: i
            class(specint_t), allocatable, intent(out) :: specint
            integer, intent(out), optional             :: stat

            call self%specsource%fetch(i, specint, stat)

            return

         end subroutine fetch_proc_

      end subroutine interp_${name}$

   #:endfor

   !****

   subroutine adjust_x_vec(self, x_vec, dx_vec, x_adj, stat)

      class(specgrid_t), intent(in)  :: self
      real(RD), intent(in)           :: x_vec(:)
      real(RD), intent(in)           :: dx_vec(:)
      real(RD), intent(out)          :: x_adj(:)
      integer, intent(out), optional :: stat

      ! Adjust x_vec in the direction dx_vec, until it falls within a
      ! valid cell

      call self%vgrid%adjust_x_vec(x_vec, dx_vec, x_adj, stat)

      ! Finish

      return

   end subroutine adjust_x_vec

   !****   

   subroutine filter(self, passband, photgrid, stat)

      class(specgrid_t), intent(in)  :: self
      type(passband_t), intent(in)   :: passband
      type(photgrid_t), intent(out)  :: photgrid
      integer, intent(out), optional :: stat

      ! Apply the passband to create a photgrid

      photgrid = photgrid_t(spec_photsource_t(self%specsource, passband), self%vgrid)

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine filter
      
   !****

   subroutine read(self, hdf5io, stat)

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

      integer :: revision

      ! Read the specgrid

      call check_type(hdf5io, 'specgrid_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(specgrid_t), intent(out) :: self
         type(hdf5io_t), intent(inout)  :: hdf5io
         integer, intent(out), optional :: stat

         type(hdf5io_t)            :: hdf5io_specsource
         integer                   :: n
         real(RD)                  :: lam_min
         real(RD)                  :: lam_max
         character(:), allocatable :: file_name
         character(:), allocatable :: group_name
         integer                   :: group_names_len
         character(:), allocatable :: file_names(:)
         character(:), allocatable :: group_names(:)
         integer                   :: i
         type(hdf5_specsource_t)   :: specsource
         type(hdf5io_t)            :: hdf5io_vgrid
         type(vgrid_t)             :: vgrid
         
         hdf5io_specsource = hdf5io_t(hdf5io, 'specsource')

         call hdf5io_specsource%inquire(file_name=file_name, group_name=group_name)

         call hdf5io_specsource%read_attr('n', n)
         call hdf5io_specsource%read_attr('lam_min', lam_min)
         call hdf5io_specsource%read_attr('lam_max', lam_max)

         call hdf5io_specsource%final()

         group_names_len = LEN_TRIM(group_name) + 1 + LEN_TRIM(specint_group_name_(HUGE(0)))

         allocate(character(LEN_TRIM(file_name))::file_names(n))
         allocate(character(group_names_len)::group_names(n))

         name_loop : do i = 1, n

            file_names(i) = TRIM(file_name)
            group_names(i) = TRIM(group_name)//'/'//TRIM(specint_group_name_(i))

         end do name_loop

         specsource = hdf5_specsource_t(file_names, group_names, lam_min, lam_max)

         hdf5io_vgrid = hdf5io_t(hdf5io, 'vgrid')
         call vgrid%read(hdf5io_vgrid, stat)
         call hdf5io_vgrid%final()
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         end if

         select type(self)
         type is(specgrid_t)
            self = specgrid_t(specsource, vgrid)
         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(specgrid_t), intent(inout) :: self
      type(hdf5io_t), intent(inout)    :: hdf5io
      integer, intent(out), optional   :: stat

      integer                   :: n
      real(RD)                  :: lam_min
      real(RD)                  :: lam_max
      type(hdf5io_t)            :: hdf5io_specsource
      integer                   :: i
      class(specint_t), pointer :: specint
      type(hdf5io_t)            :: hdf5io_specint
      type(hdf5io_t)            :: hdf5io_vgrid

      ! Write the specgrid

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

      call self%vgrid%get_n_seq(n)
      call self%speccache%get_lam_min(lam_min)
      call self%speccache%get_lam_max(lam_max)

      hdf5io_specsource = hdf5io_t(hdf5io, 'specsource')

      call hdf5io_specsource%write_attr('n', n)
      call hdf5io_specsource%write_attr('lam_min', lam_min)
      call hdf5io_specsource%write_attr('lam_max', lam_max)

      write_loop: do i = 1, n

         call self%speccache%fetch(i, fetch_proc_, specint, stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) exit write_loop
         end if

         hdf5io_specint = hdf5io_t(hdf5io_specsource, TRIM(specint_group_name_(i)))

         call specint%write(hdf5io_specint, stat)

         call hdf5io_specint%final()
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) exit write_loop
         end if

         call self%speccache%release(i, stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) exit write_loop
         end if

      end do write_loop

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

      hdf5io_vgrid = hdf5io_t(hdf5io, 'vgrid')
      call self%vgrid%write(hdf5io_vgrid, stat)
      call hdf5io_vgrid%final()
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   contains

      subroutine fetch_proc_(i, specint, stat)

         integer, intent(in)                        :: i
         class(specint_t), allocatable, intent(out) :: specint
         integer, intent(out), optional             :: stat

         call self%specsource%fetch(i, specint, stat)

         return

      end subroutine fetch_proc_

   end subroutine write

   !****

   function specint_group_name_(i) result(name)

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

      write(name, 100) i
100   format('specints[',I0,']')
      
   end function specint_group_name_

end module specgrid_m
