! Module  : photgrid_m
! Purpose : Define photgrid_t type, for representing photometric intensity
!           grids
!
! Copyright 2021-2024 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 photgrid_m

   ! Uses
   
   use forum_m

   use axis_m
   use file_m
   use photcache_m
   use photint_m
   use photsource_m
   use hdf5_photsource_m
   use stat_m
   use vgrid_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Derived-type definitions

   type :: photgrid_t
      private
      class(photsource_t), allocatable :: photsource
      type(photcache_t)                :: photcache
      type(vgrid_t)                    :: vgrid
   contains
      private
      procedure, public :: get_rank
      procedure, public :: get_shape
      procedure, public :: get_axis
      procedure, public :: get_cache_limit
      procedure, public :: get_cache_usage 
      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 :: interp_irradiance
      procedure, public :: adjust_x_vec
      procedure, public :: read
      procedure, public :: write
   end type photgrid_t

   ! Interfaces

   interface photgrid_t
      module procedure photgrid_t_
   end interface photgrid_t

   ! Access specifiers

   private

   public :: photgrid_t

   ! Procedures
   
contains

   function photgrid_t_(photsource, vgrid) result (photgrid)

      class(photsource_t), intent(in) :: photsource
      type(vgrid_t), intent(in)       :: vgrid
      type(photgrid_t)                :: photgrid

      integer :: n
      integer :: n_seq

      ! Construct photgrid from the supplied data

      call photsource%get_n(n)
      call vgrid%get_n_seq(n_seq)

      @:CHECK_BOUNDS(n, n_seq)

      photgrid%photsource = photsource
      photgrid%photcache = photcache_t(n)

      photgrid%vgrid = vgrid

      ! Finish

      return

   end function photgrid_t_

   !****

   subroutine get_rank(self, rank)

      class(photgrid_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(photgrid_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(photgrid_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 (('limit', 'integer'), &
                        ('usage', 'integer'))
   
      subroutine get_cache_${name}$(self, cache_${name}$)

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

         ! Get cache_${name}$

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

         ! Finish

         return

      end subroutine get_cache_${name}$

   #:endfor

   !****

   subroutine set_cache_limit(self, cache_limit, stat)

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

      ! Set the cache memory usage limit

      call self%photcache%set_limit(cache_limit, stat)

      ! Finish

      return

   end subroutine set_cache_limit
 
   !****

   #: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}$, res, stat, deriv_vec, order)

         class(photgrid_t), target, intent(inout) :: self
         real(RD), intent(in)                     :: x_vec(:)
         ${arg_type}$, intent(in)                 :: ${arg_var}$
         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

         return

      contains

         subroutine data_proc_(v_seq, f, stat)

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

            class(photint_t), pointer :: photint

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

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

            call self%photcache%release(v_seq, stat)

            return

         end subroutine data_proc_

         subroutine fetch_proc_(i, photint, stat)

            integer, intent(in)                        :: i
            class(photint_t), allocatable, intent(out) :: photint
            integer, intent(out), optional             :: stat

            call self%photsource%fetch(i, photint, stat)

            return

         end subroutine fetch_proc_

      end subroutine interp_${name}$

   #:endfor

   !****

   subroutine interp_irradiance(self, x_vec, mu, dA, d, res, stat, deriv_vec, order)

      class(photgrid_t), target, intent(inout) :: self
      real(RD), intent(in)                     :: x_vec(:,:)
      real(RD), intent(in)                     :: mu(:)
      real(RD), intent(in)                     :: dA(:)
      real(RD), intent(in)                     :: d(:)
      real(RD), intent(out)                    :: res
      integer, intent(out), optional           :: stat
      logical, intent(in), optional            :: deriv_vec(:)
      integer, intent(in), optional            :: order

      integer  :: i
      real(RD) :: res_i

      @:CHECK_BOUNDS(SIZE(mu), SIZE(x_vec, 2))
      @:CHECK_BOUNDS(SIZE(dA), SIZE(x_vec, 2))
      @:CHECK_BOUNDS(SIZE(d), SIZE(x_vec, 2))

      ! Interpolate the irradiance

      res = 0._RD

      !$OMP PARALLEL DO REDUCTION(+:res)
      do i = 1, SIZE(x_vec, 2)
         call self%vgrid%interp(data_proc_, x_vec(:,i), res_i, stat, deriv_vec, order)
         res = res + res_i*dA(i)/d(i)**2
      end do
      !$OMP END PARALLEL DO

      ! Finish

      return

   contains

      subroutine data_proc_(v_seq, f, stat)

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

         class(photint_t), pointer :: photint

         !$OMP CRITICAL
         call self%photcache%fetch(v_seq, fetch_proc_, photint, stat)
         !$OMP END CRITICAL
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         end if

         call photint%interp_intensity(mu(i), f, stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         end if

         !$OMP CRITICAL
         call self%photcache%release(v_seq, stat)
         !$OMP END CRITICAL

         return

      end subroutine data_proc_

      subroutine fetch_proc_(i, photint, stat)

         integer, intent(in)                        :: i
         class(photint_t), allocatable, intent(out) :: photint
         integer, intent(out), optional             :: stat

         !$OMP CRITICAL
         call self%photsource%fetch(i, photint, stat)
         !$OMP END CRITICAL

         return

      end subroutine fetch_proc_

   end subroutine interp_irradiance

   !****

   subroutine flush_cache(self)

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

      ! Flush the cache

      call self%photcache%flush()

      ! Finish

      return

   end subroutine flush_cache

   !****

   #:for name in ('flux',)

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

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

         ! Interpolate the ${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(photint_t), pointer :: photint

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

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

            call self%photcache%release(v_seq, stat)

            return

         end subroutine data_proc_
 
        subroutine fetch_proc_(i, photint, stat)

            integer, intent(in)                         :: i
            class(photint_t), allocatable, intent(out)  :: photint
            integer, intent(out), optional              :: stat

            call self%photsource%fetch(i, photint, stat)

            return

         end subroutine fetch_proc_

      end subroutine interp_${name}$

   #:endfor

   !****

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

      class(photgrid_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 read(self, hdf5io, stat)

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

      integer :: revision

      ! Read the photgrid

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

         type(hdf5io_t)            :: hdf5io_photsource
         integer                   :: n
         character(:), allocatable :: file_name
         character(:), allocatable :: group_name
         integer                   :: group_names_len
         character(:), allocatable :: file_names(:)
         character(:), allocatable :: group_names(:)
         integer                   :: i
         type(hdf5_photsource_t)   :: photsource
         type(hdf5io_t)            :: hdf5io_vgrid
         type(vgrid_t)             :: vgrid
         
         hdf5io_photsource = hdf5io_t(hdf5io, 'photsource')

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

         call hdf5io_photsource%read_attr('n', n)

         call hdf5io_photsource%final()

         group_names_len = LEN_TRIM(group_name) + 1 + LEN_TRIM(photint_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(photint_group_name_(i))

         end do name_loop
      
         photsource = hdf5_photsource_t(file_names, group_names)

         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(photgrid_t)
            self = photgrid_t(photsource, 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(photgrid_t), intent(inout) :: self
      type(hdf5io_t), intent(inout)    :: hdf5io
      integer, intent(out), optional   :: stat

      integer                   :: n
      type(hdf5io_t)            :: hdf5io_photsource
      integer                   :: i
      class(photint_t), pointer :: photint
      type(hdf5io_t)            :: hdf5io_photint
      type(hdf5io_t)            :: hdf5io_vgrid

      ! Write the photgrid

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

      call self%vgrid%get_n_seq(n)

      hdf5io_photsource = hdf5io_t(hdf5io, 'photsource')

      call hdf5io_photsource%write_attr('n', n)

      write_loop: do i = 1, n

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

         hdf5io_photint = hdf5io_t(hdf5io_photsource, photint_group_name_(i))
         call photint%write(hdf5io_photint, stat)
         call hdf5io_photint%final()
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) exit write_loop
         end if

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

      end do write_loop

      call hdf5io_photsource%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, photint, stat)

         integer, intent(in)                        :: i
         class(photint_t), allocatable, intent(out) :: photint
         integer, intent(out), optional             :: stat

         call self%photsource%fetch(i, photint, stat)

         return

      end subroutine fetch_proc_

   end subroutine write

   !****

   function photint_group_name_(i) result(name)

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

      write(name, 100) i
100   format('photints[',I0,']')
      
   end function photint_group_name_

end module photgrid_m
