! Module  : limb_m
! Purpose : define limb_t type, providing basis functions and fitting
!           for limb-darkening laws
!
! 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_m

   ! Uses

   use forum_m

   use file_m
   use fit_m
   use stat_m

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Parameter definitions

   integer, parameter :: CONST_LAW = 1
   integer, parameter :: LINEAR_LAW = 2
   integer, parameter :: SQRT_LAW = 3
   integer, parameter :: QUAD_LAW = 4
   integer, parameter :: CLARET_LAW = 5

   integer, parameter :: STR_LEN = 16

   ! Derived-type definitions

   type limb_t
      private
      integer :: law_id = -1
      integer :: n
   contains
      private
      procedure, public :: get_n
      procedure, public :: get_ssize
      procedure, public :: eval_intensity_basis
      procedure, public :: eval_flux_basis
      procedure, public :: eval_E_moment_basis
      procedure, public :: eval_D_moment_basis
      procedure         :: fit_intensity_0_
      procedure         :: fit_intensity_1_
      generic, public   :: fit_intensity => fit_intensity_0_, fit_intensity_1_
      procedure, public :: read
      procedure, public :: write
   end type limb_t

   ! Interfaces

   interface limb_t
      module procedure limb_t_
   end interface limb_t

   ! Access specifiers

   private

   public :: limb_t

   ! Procedures

contains

   function limb_t_(law_str) result (limb)

      character(*), intent(in) :: law_str
      type(limb_t)             :: limb

      ! Construct limb from the supplied data

      limb%law_id = str_to_id_(law_str)

      select case (limb%law_id)
      case (CONST_LAW)
         limb%n = 1
      case (LINEAR_LAW)
         limb%n = 2
      case (SQRT_LAW)
         limb%n = 3
      case (QUAD_LAW)
         limb%n = 3
      case (CLARET_LAW)
         limb%n = 5
      case default
         @:ABORT('invalid law_id')
      end select

      ! Finish

      return

   end function limb_t_

   !****

   function id_to_str_(law_id) result(law_str)

      integer, intent(in) :: law_id
      character(STR_LEN)  :: law_str

      ! Convert the law id into a string

      select case (law_id)
      case (CONST_LAW)
         law_str = 'CONST'
      case (LINEAR_LAW)
         law_str = 'LINEAR'
      case (SQRT_LAW)
         law_str = 'SQRT'
      case (QUAD_LAW)
         law_str = 'QUAD'
      case (CLARET_LAW)
         law_str = 'CLARET'
      case default
         @:ABORT('invalid law_id')
      end select

      ! Finish

      return

   end function id_to_str_

   !****

   function str_to_id_(law_str) result(law_id)

      character(*), intent(in) :: law_str
      integer                  :: law_id

      ! Convert the string into a law id

      select case (law_str)
      case ('CONST')
         law_id = CONST_LAW
      case ('LINEAR')
         law_id = LINEAR_LAW
      case ('SQRT')
         law_id = SQRT_LAW
      case ('QUAD')
         law_id = QUAD_LAW
      case ('CLARET')
         law_id = CLARET_LAW
      case default
         @:ABORT('invalid law_str')
      end select

      ! Finish

      return

   end function str_to_id_

   !****

   subroutine get_n(self, n)
   
      class(limb_t), intent(in) :: self
      integer, intent(out)      :: n

      ! Get the number of limb-darkening coefficients

      n = self%n

      ! Finish

      return

   end subroutine get_n

   !****

   subroutine get_ssize(self, ssize)
   
      class(limb_t), intent(in) :: self
      integer(ID), intent(out)  :: ssize

      ! Get the storage size

      ssize = STORAGE_SIZE(self)/8

      ! Finish

      return

   end subroutine get_ssize

   !****

   subroutine eval_intensity_basis(self, mu, b, stat)

      class(limb_t), intent(in)      :: self
      real(RD), intent(in)           :: mu
      real(RD)                       :: b(:)
      integer, intent(out), optional :: stat

      @:CHECK_BOUNDS(SIZE(b), self%n)

      ! Evaluate the intensity basis functions

      if (mu < 0._RD) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_MU_LO
            return
         else
            @:ABORT('out-of-bounds (lo) mu')
         endif
      endif

      if (mu > 1._RD) then
         if (PRESENT(stat)) then
            stat = STAT_OUT_OF_BOUNDS_MU_HI
            return
         else
            @:ABORT('out-of-bounds (hi) mu')
         endif
      endif

      select case (self%law_id)
      case (CONST_LAW)
         b = [      &
              1._RD &
              ]
      case (LINEAR_LAW)
         b = [           &
              1._RD,     &
              mu - 1._RD &
              ]
      case (SQRT_LAW)
         b = [                  &
              1._RD,            &
              sqrt(mu) - 1._RD, &
              mu - 1._RD        &
              ]
      case (QUAD_LAW)
         b = [              &
              1._RD,        &
              mu - 1._RD,   &
              mu**2 - 1._RD &
              ]
      case (CLARET_LAW)
         b = [                     &
              1._RD,               &
              sqrt(mu) - 1._RD,    &
              mu - 1._RD,          &
              sqrt(mu)*mu - 1._RD, &
              mu**2 - 1._RD        &
              ]
      case default
         @:ABORT('invalid law')
      end select

      if (PRESENT(stat)) stat = STAT_OK

      ! Finish

      return

   end subroutine eval_intensity_basis

   !****

   subroutine eval_E_moment_basis(self, k, b, stat)

      class(limb_t), intent(in)      :: self
      integer, intent(in)            :: k
      real(RD)                       :: b(:)
      integer, intent(out), optional :: stat

      @:CHECK_BOUNDS(SIZE(b), self%n)

      ! Evaluate the E (Eddington) moment basis functions

      if (k < 0) then
         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         endif
      endif

      select case (self%law_id)
      case (CONST_LAW)
         b = [                &
              E_ks_(k, 0._RD) &
              ]
      case (LINEAR_LAW)
         b = [                                  &
              E_ks_(k, 0._RD),                  &
              E_ks_(k, 1._RD) - E_ks_(k, 0._RD) &
              ]
      case (SQRT_LAW)
         b = [                                    &
              E_ks_(k, 0._RD),                    &
              E_ks_(k, 0.5_RD) - E_ks_(k, 0._RD), &
              E_ks_(k, 1._RD) - E_ks_(k, 0._RD)   &
              ]
      case (QUAD_LAW)
         b = [                                   &
              E_ks_(k, 0._RD),                   &
              E_ks_(k, 1._RD) - E_ks_(k, 0._RD), &
              E_ks_(k, 2._RD) - E_ks_(k, 0._RD)  &
              ]
      case (CLARET_LAW)
         b = [                                    &
              E_ks_(k, 0._RD),                    &
              E_ks_(k, 0.5_RD) - E_ks_(k, 0._RD), &
              E_ks_(k, 1._RD) - E_ks_(k, 0._RD),  &
              E_ks_(k, 1.5_RD) - E_ks_(k, 0._RD), &
              E_ks_(k, 2._RD) - E_ks_(k, 0._RD)   &
              ]
      case default
         @:ABORT('invalid law')
      end select

      if (PRESENT(stat)) stat = STAT_OK

      ! Finish

      return

   contains

      function E_ks_(k, s) result(E)

         integer, intent(in)  :: k
         real(RD), intent(in) :: s
         real(RD)             :: E

         ! Evaluate the integral 1/2 mu^(s+k) from 0 to 1

         E = 0.5_RD/(s + k + 1)

         ! Finish

         return

      end function E_ks_

   end subroutine eval_E_moment_basis

   !****

   subroutine eval_D_moment_basis(self, l, b, stat)

      class(limb_t), intent(in)      :: self
      integer, intent(in)            :: l
      real(RD)                       :: b(:)
      integer, intent(out), optional :: stat

      @:CHECK_BOUNDS(SIZE(b), self%n)

      ! Evaluate the D (differential flux function) moment basis
      ! functions

      if (l < 0) then
         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('invalid argument')
         endif
      endif

      select case (self%law_id)
      case (CONST_LAW)
         b = [                &
              D_ls_(l, 1._RD) &
              ]
      case (LINEAR_LAW)
         b = [                                  &
              D_ls_(l, 1._RD),                  &
              D_ls_(l, 2._RD) - D_ls_(l, 1._RD) &
              ]
      case (SQRT_LAW)
         b = [                                    &
              D_ls_(l, 1._RD),                    &
              D_ls_(l, 1.5_RD) - D_ls_(l, 1._RD), &
              D_ls_(l, 2._RD) - D_ls_(l, 1._RD)   &
              ]
      case (QUAD_LAW)
         b = [                                   &
              D_ls_(l, 1._RD),                   &
              D_ls_(l, 2._RD) - D_ls_(l, 1._RD), &
              D_ls_(l, 3._RD) - D_ls_(l, 1._RD)  &
              ]
      case (CLARET_LAW)
         b = [                                    &
              D_ls_(l, 1._RD),                    &
              D_ls_(l, 1.5_RD) - D_ls_(l, 1._RD), &
              D_ls_(l, 2._RD) - D_ls_(l, 1._RD),  &
              D_ls_(l, 2.5_RD) - D_ls_(l, 1._RD), &
              D_ls_(l, 3._RD) - D_ls_(l, 1._RD)   &
              ]
      case default
         @:ABORT('invalid law')
      end select

      if (PRESENT(stat)) stat = STAT_OK

      ! Finish

      return

   contains

      function D_ls_(l, s) result(D)

         integer, intent(in)  :: l
         real(RD), intent(in) :: s
         real(RD)             :: D

         integer :: l_0
         integer :: m

         ! Evaluate the integral mu^s Pl(mu) from 0 to 1, using the
         ! expressions (42-44) from Townsend (2003)

         l_0 = MOD(l, 2)

         D = 1._RD/(1._RD + s + l_0)

         do m = l_0, l-2, 2
            D = (s - m)/(s + m + 3._RD)*D
         end do

         ! Finish

         return

      end function D_ls_

   end subroutine eval_D_moment_basis

   !****

   subroutine eval_flux_basis(self, b, stat)

      class(limb_t), intent(in)      :: self
      real(RD)                       :: b(:)
      integer, intent(out), optional :: stat

      @:CHECK_BOUNDS(SIZE(b), self%n)

      ! Evaluate the flux basis functions

      select case (self%law_id)
      case (CONST_LAW)
         b = TWOPI*[      &
              1._RD/2._RD &
              ]
      case (LINEAR_LAW)
         b = TWOPI*[       &
              1._RD/2._RD, &
              -1._RD/6._RD &
              ]
      case (SQRT_LAW)
         b = TWOPI*[         &
              1._RD/2._RD,   &
              -1._RD/10._RD, &
              -1._RD/6._RD   &
              ]
      case (QUAD_LAW)
         b = TWOPI*[        &
              1._RD/2._RD,  &
              -1._RD/6._RD, &
              -1._RD/4._RD  &
              ]
      case (CLARET_LAW)
         b = TWOPI*[         &
              1._RD/2._RD,   &
              -1._RD/10._RD, &
              -1._RD/6._RD,  &
              -3._RD/14._RD, &
              -1._RD/4._RD   &
              ]
      case default
         @:ABORT('invalid law')
      end select

      if (PRESENT(stat)) stat = STAT_OK

      ! Finish

      return

   end subroutine eval_flux_basis

   !****

   subroutine fit_intensity_0_ (self, I, mu, c)

      class(limb_t), intent(in)          :: self
      real(RD), intent(in)               :: I(:)
      real(RD), intent(in)               :: mu(:)
      real(RD), allocatable, intent(out) :: c(:)

      real(RD), allocatable :: I_(:,:)
      real(RD), allocatable :: c_(:,:)

      ! Fit the limb-darkening law to the intensity data

      allocate(I_(SIZE(I), 1))

      I_(:,1) = I

      call self%fit_intensity(I_, mu, c_)

      c = c_(:,1)

      ! Finish

      return

   end subroutine fit_intensity_0_

   !****

   subroutine fit_intensity_1_ (self, I, mu, c)

      class(limb_t), intent(in)          :: self
      real(RD), intent(in)               :: I(:,:)
      real(RD), intent(in)               :: mu(:)
      real(RD), allocatable, intent(out) :: c(:,:)

      integer               :: n_mu
      integer               :: n_k
      integer               :: k
      real(RD)              :: y(SIZE(I, 1),SIZE(I, 2))
      real(RD), allocatable :: a(:,:)

      @:CHECK_BOUNDS(SIZE(I, 1), SIZE(mu))

      ! Fit the limb-darkening law to the intensity data. The fit is
      ! constrained to pass through I(mu=1)

      n_mu = SIZE(I, 1)
      n_k = SIZE(I, 2)

      if (self%n > 1) then

         ! Set up the fit data

         y_loop : do k = 1, n_k
            y(:,k) = 1._RD - I(:,k)/I(n_mu,k)
         end do y_loop

         ! Fit the a-coefficients

         allocate(a(self%n-1,n_k))

         call gls_fit(mu, y, f_, a)

         ! Tranform to c-coefficients

         allocate(c(self%n,n_k))

         c_loop : do k = 1, n_k
            c(1,k) = I(n_mu,k)
            c(2:,k) = -I(n_mu,k)*a(:,k)
         end do c_loop

      else

         ! When self%n == 1, no fitting is (or can be) done; just copy
         ! over the normal intensity

         allocate(c(1,n_k))

         c(1,:) = I(n_mu,:)

      endif

      ! Finish

      return

   contains

      function f_ (mu, n)

         real(RD), intent(in) :: mu
         integer, intent(in)  :: n
         real(RD)             :: f_(n)

         real(RD) :: b(n+1)

         call self%eval_intensity_basis(mu, b)

         f_ = b(2:)

      end function f_

   end subroutine fit_intensity_1_

   !****
      
   subroutine read(self, hdf5io, stat)

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

      integer :: revision

      ! Read the limb_t

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

         character(STR_LEN) :: law_str

         call hdf5io%read_attr('law', law_str)

         select type(self)
         type is(limb_t)
            self = limb_t(law_str)
         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_t), intent(in)      :: self
      type(hdf5io_t), intent(inout)  :: hdf5io
      integer, intent(out), optional :: stat

      ! Write the limb_t

      call hdf5io%write_attr('TYPE', 'limb_t')
      call hdf5io%write_attr('REVISION', 1)
      
      call hdf5io%write_attr('law', id_to_str_(self%law_id))

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine write

end module limb_m
