!
!  Spreaded profiles
!
!
!  Copyright © 2016-2022 F.Hroch (hroch@physics.muni.cz)
!
!  This file is part of Munipack.
!
!  Munipack 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, either version 3 of the License, or
!  (at your option) any later version.
!
!  Munipack 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 Munipack.  If not, see <http://www.gnu.org/licenses/>.
!


module spray

  implicit none

  integer, parameter, private :: dbl = selected_real_kind(15)
  integer, parameter, private :: defzoom = 2**4
  real(dbl), parameter, private :: rad = 57.295779513082322865_dbl
  real(dbl), parameter, private :: sqrt2 = sqrt(2.0)

  type SprayType
     integer :: zoom = defzoom
     integer :: nbox
     real(dbl), dimension(:,:), allocatable :: psf
   contains
     procedure :: Init, Pixelize, Load
     procedure :: gauss, moffat, seeing, cseeing, fseeing
  end type SprayType

  private :: gauss, moffat, seeing, cseeing, fseeing, fill_gaussian, fill_bessel

contains

  subroutine Init(this,profile,spread,hwhm,eccentricity,inclination, &
       airy,beta)

    class(SprayType) :: this
    character(len=*), intent(in) :: profile, spread
    real(dbl), intent(in) :: hwhm, airy, beta, eccentricity,inclination

    if( profile == 'GAUSS' ) then

       call this%gauss(hwhm,eccentricity,inclination)

    else if( profile == 'MOFFAT' ) then

       call this%moffat(hwhm,beta,eccentricity,inclination)

    else ! if( profile == 'SEEING' ) then

       if( spread == 'RANDOM' ) then

          call this%seeing(hwhm,airy,eccentricity,inclination)

       else if( spread == 'TEST' ) then

          call this%cseeing(hwhm,airy,eccentricity,inclination)

       else ! if (spread == 'FFT' ) then

          call this%fseeing(hwhm,airy,eccentricity,inclination)

       end if
    end if

    ! normalisation
    this%psf = this%psf * this%zoom**2 / sum(this%psf)

  end subroutine Init


  subroutine seeing(this,hwhm,airy,e,incl)

    ! The seeing is simulated by summing of Airy profiles with random
    ! offsets. The distribution of offsets is Gaussian N(0,hwhm).
    ! This method is very illustrative; it simulates the real process
    ! of an image generation. On the other side, it is slow for many
    ! iterations. It gives an asymmetric and randomly noised images.

    use noise

    integer, parameter :: maxiter = 5000

    class(SprayType) :: this
    real(dbl), intent(in) :: hwhm, airy, e,incl

    real(dbl), dimension(:,:), allocatable :: bessel
    integer :: n,m,i,j,iter,l1,l2,k1,k2
    real(dbl) :: x,y,u,v,a,b,c,s,sig

    this%nbox = nint(13*max(hwhm,1.0)*max(airy,5.0))
    n = this%nbox*this%zoom
    m = n / 2
    allocate(this%psf(-n:n,-n:n),bessel(-m:m,-m:m))

    ! fill-up by Bessel
    call fill_bessel(m,this%zoom*airy,bessel)

    ! spread the profile by random shifts with N(0,s)
    sig = hwhm*this%zoom
    a = sqrt2*sig
    b = a*sqrt(1 - e**2)
    c = cos(incl/rad)
    s = sin(incl/rad)

    this%psf = 0
    do iter = 1, maxiter
       u = gnoise(0.0_dbl,a)
       v = gnoise(0.0_dbl,b)
       x = u*c + v*s
       y =-u*s + v*c
       i = nint(x)
       j = nint(y)

       k1 = i - m
       k2 = i + m
       l1 = j - m
       l2 = j + m
       this%psf(k1:k2,l1:l2) = this%psf(k1:k2,l1:l2) + bessel
    end do

    deallocate(bessel)

  end subroutine seeing


  subroutine cseeing(this,hwhm,airy, e,incl)

    ! The seeing is simulated by convolution of Airy disk function and
    ! gaussian giving spread. The convolution is computed by direct
    ! sum, so by very slow way.

    class(SprayType) :: this
    real(dbl), intent(in) :: hwhm, airy, e,incl

    real(dbl), dimension(:,:), allocatable :: bessel, kernel
    integer :: n,m,i,j,k,l
    real(dbl) :: s

    this%nbox = nint(13*max(hwhm,1.0)*max(airy,5.0))
    n = this%nbox*this%zoom
    m = n / 2
    allocate(this%psf(-n:n,-n:n),bessel(-m:m,-m:m),kernel(-n:n,-n:n))

    ! fill-up the Bessel intensity
    call fill_bessel(m,this%zoom*airy,bessel)

    ! fill kernel
    call fill_gaussian(m,this%zoom*hwhm,e,incl,kernel)

    ! spread profile by direct convolution
    this%psf = 0
    do i = -m,m
       do j = -m,m
          s = 0
          do l = -m,m
             do k = -m,m
                s = s + bessel(k,l)*kernel(i-l,j-k)
             end do
          end do
          this%psf(i,j) = s
       end do
    end do

    deallocate(bessel,kernel)

  end subroutine cseeing

  subroutine fseeing(this,hwhm,airy,e,incl)

    ! The seeing is simulated by convolution of Airy disk function and
    ! gaussian giving spread. The convolution is computed by using
    ! os Fourier transformation. It is most fast and precise method.

    use ftransform

    class(SprayType) :: this
    real(dbl), intent(in) :: hwhm, airy,e,incl

    complex(dbl), dimension(:,:), allocatable :: bessel, kernel, qpsf, &
         zbessel, zkernel, zpsf
    real(dbl), dimension(:,:), allocatable :: rkernel, rbessel
    integer :: n,m,n2

    n = nint(13*max(hwhm,1.0)*max(airy,5.0))
    m = int(log(real(n*this%zoom))/log(2.0) + 1)  ! power o 2 for FFT
    n = 2**m
    this%nbox = n / this%zoom
    n = this%nbox * this%zoom
    n2 = n / 2

    ! fill-up Bessel
    allocate(bessel(n,n),rbessel(-n2:n2,-n2:n2))
    call fill_bessel(n2,this%zoom*airy,rbessel)
    bessel = rbessel(-n2+1:n2,-n2+1:n2)
    deallocate(rbessel)

    ! fill-up Gaussian spread
    allocate(kernel(n,n),rkernel(-n2:n2,-n2:n2))
    call fill_gaussian(n2,this%zoom*hwhm,e,incl,rkernel)
    ! conversion Complex to Real numbers
    kernel = rkernel(-n2+1:n2,-n2+1:n2)
    deallocate(rkernel)

    ! forward FFT
    allocate(zbessel(n,n),zkernel(n,n))
    call ftra(bessel,zbessel,m,1)
    call ftra(kernel,zkernel,m,1)
    deallocate(bessel,kernel)

    ! the convolution
    allocate(zpsf(n,n),qpsf(n,n))
    zpsf = zbessel * zkernel

    ! back FFT
    call ftra(zpsf,qpsf,m,-1)
    deallocate(zbessel,zkernel,zpsf)

    ! fill PSF by re-arranged quarters
    allocate(this%psf(-n2+1:n2-1,-n2+1:n2-1))
    this%psf = 0
    this%psf(    0:n2-1,    0:n2-1) = real(qpsf(1:n2,1:n2))
    this%psf(-n2+1:-1,      0:n2-1) = real(qpsf(n2+1:n-1,1:n2))
    this%psf(    0:n2-1,-n2+1:-1)   = real(qpsf(1:n2,n2+1:n-1))
    this%psf(-n2+1:-1,  -n2+1:-1)   = real(qpsf(n2+1:n-1,n2+1:n-1))

    deallocate(qpsf)

  end subroutine fseeing

  subroutine fill_bessel(m,sig,bessel)

    integer, intent(in) :: m
    real(dbl), intent(in) :: sig
    real(dbl), dimension(-m:,-m:), intent(out) :: bessel

    real(dbl) :: x,y,r
    integer :: i,j

    ! fill-up the top-right quarter of Airy disk
    do i = 0,m
       x = i / sig
       do j = 0,m
          y = j / sig
          ! https://en.wikipedia.org/wiki/Airy_disk
          if( i == 0 .and. j == 0 ) then
             bessel(i,j) = 1
          else
             r = sqrt(x**2 + y**2)
             bessel(i,j) = (bessel_j1(r)/r)**2
          end if
       end do
    end do

    ! fill-up rest of the kernel
    forall( i = -m:0, j = 0:m ) bessel(i,j) = bessel(-i,j)
    forall( i = -m:m, j = -m:-1 ) bessel(i,j) = bessel(i,-j)

  end subroutine fill_bessel


  subroutine fill_gaussian(m,sig,e,incl,kernel)

    integer, intent(in) :: m
    real(dbl), intent(in) :: sig, e, incl
    real(dbl), dimension(-m:,-m:), intent(out) :: kernel

    real(dbl) :: x,y,a,b,c,s
    integer :: i,j

    a = sqrt2*sig
    b = a*sqrt(1 - e**2)
    c = cos(incl/rad)
    s = sin(incl/rad)

    do i = -m, m
       do j = -m, m
          x = ( i*c + j*s) / a
          y = (-i*s + j*c) / b
          kernel(i,j) = exp(-(x**2 + y**2))
       end do
    end do

  end subroutine fill_gaussian

  subroutine gauss(this,hwhm,e,incl)

    class(SprayType) :: this
    real(dbl), intent(in) :: hwhm, e, incl
    integer :: n

    this%nbox = nint(7*max(hwhm,1.0))
    n = this%zoom * this%nbox
    allocate(this%psf(-n:n,-n:n))

    call fill_gaussian(n,this%zoom*hwhm,e,incl,this%psf)

  end subroutine gauss


  subroutine moffat(this, hwhm, beta, e,incl)

    class(SprayType) :: this
    real(dbl), intent(in) :: hwhm, beta, e, incl
    real(dbl) :: a, b, c, s, x, y
    integer :: i,j,n

    this%nbox = nint(30*max(hwhm,1.0))
    n = this%zoom * this%nbox

    allocate(this%psf(-n:n,-n:n))

    a = this%zoom * hwhm
    b = a*sqrt(1 - e**2)
    c = cos(incl/rad)
    s = sin(incl/rad)

    do i = -n, n
       do j = -n, n
          x = ( i*c + j*s) / a
          y = (-i*s + j*c) / b
          this%psf(i,j) = (1 + (x**2 + y**2))**(-beta)
       end do
    end do

  end subroutine moffat


  subroutine Pixelize(this,dx,dy,xpsf,n)

    class(SprayType) :: this
    real(dbl), intent(in) :: dx, dy
    real(dbl), dimension(-n:n,-n:n), intent(out) :: xpsf
    integer, intent(in) :: n
    ! The default lower dimension in subroutine is 1 (allocatable array?),
    ! we pass the dimensions which we needs.

    integer :: x,y,k,l,i,j,k1,k2,l1,l2,m,m2,nmax,kk,ll

    nmax = ubound(this%psf,1)
    x = nint(this%zoom*dx)
    y = nint(this%zoom*dy)
    m = this%zoom
    m2 = m / 2

    xpsf = 0
    do i = -n,n
       k = i*m + x
       k1 = max(k - m2,-nmax)
       k2 = min(k + m2, nmax)
       kk = k2 - k1 + 1
       do j = -n,n
          l = j*m + y
          l1 = max(l - m2,-nmax)
          l2 = min(l + m2, nmax)
          ll = l2 - l1 + 1
          if( kk > 0 .and. ll > 0 ) then
             xpsf(i,j) = sum(this%psf(k1:k2,l1:l2)) / (kk*ll)
          end if
       end do
    end do

  end subroutine Pixelize

  subroutine Load(this,filename)

    use titsio
    use iso_fortran_env

    class(SprayType) :: this
    character(len=*), intent(in) :: filename
    type(fitsfiles) :: fits
    integer :: naxis,status,n,n1,zoom
    integer, dimension(2) :: naxes
    character(len=FLEN_COMMENT) :: comment
    logical :: anynull

    status = 0
    call fits_open_file(fits,filename,FITS_READONLY,status)
    if( status == 0 ) then
       call fits_get_img_dim(fits,naxis,status)
       if( naxis /= 2 ) stop 'spray.Load(): PSF should by 2D image array.'
       call fits_get_image_size(fits,naxes,status)
       if( naxes(1) /= naxes(2) ) stop 'spray.Load(): PSF should be a square.'
       call fits_read_key(fits,'ZOOM',zoom,comment,status)
       if( status /= 0 .or. zoom < 1 ) stop 'spray.Load(): zoom undefined.'

       n = naxes(1) / 2
       n1 = naxes(1) - (n + 1)
       allocate(this%psf(-n:n1,-n:n1))
       this%nbox = n
       this%zoom = zoom

       call fits_read_image(fits,1,0.0_dbl,this%psf,anynull,status)
       call fits_close_file(fits,status)
    end if
    call fits_report_error(error_unit,status)
    if( status /= 0 ) stop 'spray.Load() failed to read a FITS file with PSF.'

  end subroutine Load


end module spray