! Module  : limb_specint_m
! Purpose : define limb_specint_t type, for representing spectral
!           intensity parameterized by a limb-darkening law
!
! 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 limb_specint_m

   ! Uses

   use forum_m

   use file_m
   use limb_m
   use math_m
   use passband_m
   use photint_m
   use limb_photint_m
   use range_m
   use specint_m
   use stat_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Derived-type definitions

   type, extends(specint_t) :: limb_specint_t
      private
      real(RD), allocatable       :: c(:,:)
      real(RD), allocatable       :: lam(:)
      class(range_t), allocatable :: range
      type(limb_t)                :: limb
      integer                     :: n_b
      integer                     :: n_lam
      logical                     :: precise
   contains
      private
      procedure, public :: get_lam_min
      procedure, public :: get_lam_max
      procedure, public :: get_precise
      procedure, public :: get_ssize
      #:for name in ('intensity', 'E_moment', 'D_moment', 'flux')
         procedure, public :: interp_${name}$
      #:endfor
      procedure         :: interp_f_
      procedure, public :: subset
      procedure, public :: rebin
      procedure, public :: filter
      procedure, public :: read
      procedure, public :: write
   end type limb_specint_t

   ! Interfaces

   interface limb_specint_t
      module procedure limb_specint_t_
      module procedure limb_specint_t_fit_
   end interface limb_specint_t

   ! Access specifiers

   private

   public :: limb_specint_t

   ! Procedures

contains

   function limb_specint_t_(c, range, limb, precise) result(specint)

      real(RD), intent(in)          :: c(:,:)
      class(range_t), intent(in)    :: range
      type(limb_t), intent(in)      :: limb
      logical, intent(in), optional :: precise
      type(limb_specint_t)          :: specint

      logical :: precise_
      integer :: n_b
      integer :: n_lam 

      if (PRESENT(precise)) then
         precise_ = precise
      else
         precise_ = .TRUE.
      endif

      ! Construct specint from the supplied data

      call limb%get_n(n_b)
      call range%get_n(n_lam)

      @:CHECK_BOUNDS(n_b, SIZE(c, 1))
      @:CHECK_BOUNDS(n_lam, SIZE(c, 2)+1)

      specint%c = c

      allocate(specint%lam(n_lam))
      call range%unpack(specint%lam)

      specint%range = range
      specint%limb = limb

      specint%n_b = n_b
      specint%n_lam = n_lam

      specint%precise = precise_
      
      ! Finish

      return

   end function limb_specint_t_

   !****

   function limb_specint_t_fit_(I, mu, range, limb, precise) result(specint)

      real(RD), intent(in)          :: I(:,:)
      real(RD), intent(in)          :: mu(:)
      class(range_t), intent(in)    :: range
      type(limb_t), intent(in)      :: limb
      logical, intent(in), optional :: precise
      type(limb_specint_t)          :: specint

      real(RD), allocatable :: c(:,:)

      ! Construct specint by fitting the supplied I(mu,lambda)
      ! intensity data

      call limb%fit_intensity(I, mu, c)

      specint = limb_specint_t(c, range, limb, precise)

      ! Finish

      return

   end function limb_specint_t_fit_

   !****

   subroutine get_lam_min(self, lam_min)

      class(limb_specint_t), intent(in) :: self
      real(RD), intent(out)             :: lam_min

      ! Get the minimum wavelength

      lam_min = self%lam(1)

      ! Finish

      return

   end subroutine get_lam_min

   !****

   subroutine get_lam_max(self, lam_max)

      class(limb_specint_t), intent(in) :: self
      real(RD), intent(out)             :: lam_max

      ! Get the maximum wavelength

      lam_max = self%lam(self%n_lam)

      ! Finish

      return

   end subroutine get_lam_max

   !****

   subroutine get_precise(self, precise)

      class(limb_specint_t), intent(in) :: self
      logical, intent(out)              :: precise

      ! Get the precision flag

      precise = self%precise

      ! Finish

      return

   end subroutine get_precise

   !****

   subroutine get_ssize(self, ssize)

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

      integer(ID) :: ssize_range
      integer(ID) :: ssize_limb

      ! Get the storage size

      call self%range%get_ssize(ssize_range)
      call self%limb%get_ssize(ssize_limb)

      ssize = STORAGE_SIZE(self)/8 +                    &
              STORAGE_SIZE(self%c)/8*SIZE(self%c) +     &
              STORAGE_SIZE(self%lam)/8*SIZE(self%lam) + &
              ssize_range +                             &
              ssize_limb

      ! Finish

      return

   end subroutine get_ssize

   !****

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

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

         class(limb_specint_t), intent(in) :: self
         ${arg_type}$, intent(in)          :: ${arg_var}$
         real(RD), intent(in)              :: lam(:)
         real(RD), intent(out)             :: res(:)
         integer, intent(out), optional    :: stat

         real(RD) :: b(self%n_b)

         ! Set up the basis functions

         call self%limb%eval_${name}$_basis(${arg_var}$, b, stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         endif

         ! Do the interpolation

         call self%interp_f_(b, lam, res, stat)

         ! Finish

         return

      end subroutine interp_${name}$

   #:endfor
      
   !****

   #:for name, var in (('flux', 'F'),)

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

         class(limb_specint_t), intent(in) :: self
         real(RD), intent(in)              :: lam(:)
         real(RD), intent(out)             :: res(:)
         integer, intent(out), optional    :: stat

         real(RD) :: b(self%n_b)

         ! Set up the basis functions

         call self%limb%eval_${name}$_basis(b, stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         endif

         ! Do the interpolation

         call self%interp_f_(b, lam, res, stat)

         ! Finish

         return

      end subroutine interp_${name}$

   #:endfor

   !****
      
   subroutine interp_f_(self, b, lam, f, stat)

      class(limb_specint_t), intent(in) :: self
      real(RD), intent(in)              :: b(:)
      real(RD), intent(in)              :: lam(:)
      real(RD), intent(out)             :: f(:)
      integer, intent(out), optional    :: stat

      integer  :: n_lam
      integer  :: i
      integer  :: j
      real(RD) :: lam_a
      real(RD) :: lam_b
      real(RD) :: dlam

      @:CHECK_BOUNDS(SIZE(b), self%n_b)
      @:CHECK_BOUNDS(SIZE(f), SIZE(lam)-1)

      ! Check lam is valid

      n_lam = SIZE(lam)

      if (lam(1) < self%lam(1)) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_LAM_LO
            return
         else
            @:ABORT('out-of-bounds (lo) lam')
         endif
      endif

      if (lam(n_lam) > self%lam(self%n_lam)) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_LAM_HI
            return
         else
            @:ABORT('out-of-bounds (hi) lam')
         endif
      endif

      ! Perform the interpolation (conservative piecewise-constant
      ! rebin)

      i = 0

      !$OMP PARALLEL DO PRIVATE(lam_a, lam_b, dlam) FIRSTPRIVATE(i) SCHEDULE(static)
      out_loop: do j = 1, n_lam-1

         if (i == 0) then
            call locate(self%lam, lam(j), i)
            if (i == self%n_lam) i = i - 1
         end if

         f(j) = 0._RD

         dlam = lam(j+1) - lam(j)

         in_loop: do

            lam_a = MAX(self%lam(i), lam(j))
            lam_b = MIN(self%lam(i+1), lam(j+1))

            f(j) = f(j) + (lam_b - lam_a)/dlam*DOT_PRODUCT(b, self%c(:,i))

            if (lam_b == self%lam(i+1)) i = i + 1
            if (lam_b == lam(j+1)) exit in_loop

         end do in_loop

      end do out_loop

      if (PRESENT(stat)) stat = STAT_OK

      ! Finish

      return

   end subroutine interp_f_

   !****

   subroutine subset(self, lam_min, lam_max, specint, stat)

      class(limb_specint_t), intent(in)          :: self
      real(RD), intent(in)                       :: lam_min
      real(RD), intent(in)                       :: lam_max
      class(specint_t), allocatable, intent(out) :: specint
      integer, intent(out), optional             :: stat

      integer                     :: i_a
      integer                     :: i_b
      class(range_t), allocatable :: range

      ! Subset specint

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

      call locate(self%lam, lam_min, i_a, right=.FALSE.)
      call locate(self%lam, lam_max, i_b, right=.TRUE.)

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

      if (i_a > self%n_lam) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_LAM_LO
            return
         else
            @:ABORT('out-of-bounds (lo) lam')
         end if
      end if

      call self%range%subset(i_a, i_b, range, stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      specint = limb_specint_t(self%c(:,i_a:i_b-1), range, self%limb, self%precise)

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine subset

   !****

   subroutine rebin(self, range, specint, stat)

      class(limb_specint_t), intent(in)          :: self
      class(range_t), intent(in)                 :: range
      class(specint_t), allocatable, intent(out) :: specint
      integer, intent(out), optional             :: stat

      integer               :: n_lam
      real(RD), allocatable :: lam(:)
      real(RD), allocatable :: c(:,:)
      integer               :: p
      real(RD)              :: b(self%n_b)

      ! Rebin specint

      call range%get_n(n_lam)

      allocate(lam(n_lam))
      allocate(c(self%n_b,n_lam-1))

      call range%unpack(lam)

      do p = 1, self%n_b

         b = 0._RD
         b(p) = 1._RD

         call self%interp_f_(b, lam, c(p,:), stat)
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         end if

      end do

      specint = limb_specint_t(c, range, self%limb, self%precise)

      ! Finish

      return

   end subroutine rebin

   !****

   subroutine filter(self, passband, photint, stat)

      class(limb_specint_t), intent(in)          :: self
      type(passband_t), intent(in)               :: passband
      class(photint_t), allocatable, intent(out) :: photint
      integer, intent(out), optional             :: stat

      real(RD)              :: lam_min
      real(RD)              :: lam_max
      real(RD)              :: lam_min_pb
      real(RD)              :: lam_max_pb
      integer               :: i_min
      integer               :: i_max
      integer               :: n_lam
      real(RD), allocatable :: lam(:)
      real(RD), allocatable :: P(:)
      real(RD)              :: c(self%n_b)

      ! Filter the specint through the passband to yield photint

      ! First check that the spectrum is defined across the passband support

      call self%get_lam_min(lam_min)
      call self%get_lam_max(lam_max)

      call passband%get_lam_min(lam_min_pb)
      call passband%get_lam_max(lam_max_pb)

      if (lam_min_pb < lam_min) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_LAM_LO
            return
         else
            @:ABORT('out-of-bounds (lo) lam')
         endif
      endif

      if (lam_max_pb > lam_max) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_LAM_HI
            return
         else
            @:ABORT('out-of-bounds (hi) lam')
         endif
      endif

      ! Evaluate the coefficients

      ! First, determine where the passband is located, and set up a
      ! wavelength axis for the integration

      call locate(self%lam, lam_min_pb, i_min)
      call locate(self%lam, lam_max_pb, i_max, right=.TRUE.)

      lam = self%lam(i_min:i_max)

      n_lam = SIZE(lam)

      lam(1) = lam_min_pb
      lam(n_lam) = lam_max_pb

      ! Construct the passthru function

      P = passband%passthru(lam)

      ! Integrate

      c = MATMUL(self%c(:,i_min:i_max-1), P)

      ! Construct the photint

      select type(self)
      type is(limb_specint_t)
         photint = limb_photint_t(c, self%limb)
      class default
         @:ABORT('invalid type')
      end select

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine filter

   !****

   subroutine read(self, hdf5io, stat)

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

      integer :: revision

      ! Read the limb_specint_t

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

         logical                     :: precise
         real(RD), allocatable       :: c(:,:)
         type(hdf5io_t)              :: hdf5io_range
         character(TYPE_LEN)         :: type
         class(range_t), allocatable :: range
         type(hdf5io_t)              :: hdf5io_limb
         type(limb_t)                :: limb

         call hdf5io%read_attr('precise', precise)

         call hdf5io%alloc_read_dset('c', c)

         hdf5io_range = hdf5io_t(hdf5io, 'range')

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

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

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

         select type(self)
         type is(limb_specint_t)
            self = limb_specint_t(c, range, limb, precise)
         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(limb_specint_t), intent(in) :: self
      type(hdf5io_t), intent(inout)     :: hdf5io
      integer, intent(out), optional    :: stat

      real(RD)       :: lam_min
      real(RD)       :: lam_max
      type(hdf5io_t) :: hdf5io_range
      type(hdf5io_t) :: hdf5io_limb

      ! Write the limb_specint_t

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

      call self%get_lam_min(lam_min)
      call self%get_lam_max(lam_max)

      call hdf5io%write_attr('lam_min', lam_min)
      call hdf5io%write_attr('lam_max', lam_max)

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

      if (self%precise) then
         call hdf5io%write_dset('c', REAL(self%c, RD), comp_level=6)
      else
         call hdf5io%write_dset('c', REAL(self%c, RS), comp_level=6)
      endif

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

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

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine write

end module limb_specint_m
