!
!  Spreaded profiles
!
!
!  Copyright © 2016-7 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 :: zoom = 2**4

  type SprayType
     integer :: nbox, zoom
     real(dbl), dimension(:,:), allocatable :: psf
  end type SprayType

  real(dbl), dimension(:,:), allocatable, private :: psf
  integer, private :: nbox

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

contains

  subroutine spray_init(profile,spread,hwhm,airy,beta,spsf)

    character(len=*), intent(in) :: profile, spread
    real(dbl), intent(in) :: hwhm, airy, beta
    type(SprayType), intent(out) :: spsf
    integer :: n

    if( profile == 'GAUSS' ) then

       call gauss(hwhm)

    else if( profile == 'MOFFAT' ) then

       call moffat(hwhm,beta)

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

       if( spread == 'RANDOM' ) then

          call seeing(hwhm,airy)

       else if( spread == 'TEST' ) then

          call cseeing(hwhm,airy)

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

          call fseeing(hwhm,airy)

       end if
    end if

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

    ! keep results
    spsf%zoom = zoom
    spsf%nbox = nbox
    n = ubound(psf,1)
    allocate(spsf%psf(-n:n,-n:n))
    spsf%psf = psf
    deallocate(psf)

  end subroutine spray_init


  subroutine seeing(hwhm,airy)

    ! 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, because simmulates the real process
    ! of an image generation. On the other side, it is slow for many iterations
    ! It gives asymetric and randomly noised images.

    use noise

    integer, parameter :: maxiter = 5555

    real(dbl), intent(in) :: hwhm, airy

    real(dbl), dimension(:,:), allocatable :: bessel
    integer :: n,m,i,j,i0,j0,iter,l1,l2,k1,k2
    real(dbl) :: x,y,r,s

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

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

    ! fill-up rest of the PSF
    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)

    ! spread the profile by random spread with N(0,s)
    s = hwhm*zoom
    psf = 0
    do iter = 1, maxiter
       x = gnoise(0.0_dbl,s)
       y = gnoise(0.0_dbl,s)
       i0 = nint(x)
       j0 = nint(y)
       k1 = i0 - m
       k2 = i0 + m
       l1 = j0 - m
       l2 = j0 + m
       psf(k1:k2,l1:l2) = psf(k1:k2,l1:l2) + bessel
    end do

    deallocate(bessel)

  end subroutine seeing


  subroutine cseeing(hwhm,airy)

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

    real(dbl), intent(in) :: hwhm, airy

    real(dbl), dimension(:,:), allocatable :: bessel, kernel
    integer :: n,m,i,j,k,l
    real(dbl) :: x,y,r,s,s2,x2,y2

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

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

    ! fill-up rest of the PSF
    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)


    ! fill-up first quarter
    s2 = 2*(zoom*hwhm)**2
    do i = 0,n
       x2 = i**2
       do j = 0,n
          y2 = j**2
          kernel(i,j) = exp(-(x2 + y2)/s2)
       end do
    end do

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


    ! spread profile by direct convolution
    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
          psf(i,j) = s
       end do
    end do

    deallocate(bessel,kernel)

  end subroutine cseeing

  subroutine fseeing(hwhm,airy)

    ! 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

    real(dbl), intent(in) :: hwhm, airy

    complex(dbl), dimension(:,:), allocatable :: bessel, kernel, qpsf, &
         zbessel, zkernel, zpsf
    integer :: n,m,i,j,n2
    real(dbl) :: x,y,r,s,s2,x2,y2

    n = nint(13*max(hwhm,1.0)*max(airy,5.0))
    m = int(log(real(n*zoom))/log(2.0) + 1)  ! power o 2 for FFT
    n = 2**m
    nbox = n / zoom
    n = nbox*zoom
    s = zoom*airy
    n2 = n / 2
    allocate(psf(-n2+1:n2-1,-n2+1:n2-1),bessel(n,n),zbessel(n,n), &
         kernel(n,n),zkernel(n,n),qpsf(n,n),zpsf(n,n))

    ! fill-up Airy disk
    do i = 1,n
       x = (i - n2) / s
       do j = 1,n
          y = (j - n2) / s
          r = sqrt(x**2 + y**2)
          ! https://en.wikipedia.org/wiki/Airy_disk
          if( i == n2 .and. j == n2 ) then
             bessel(i,j) = 1
          else
             bessel(i,j) = (bessel_j1(r)/r)**2
          end if
       end do
    end do

    ! fill-up Gaussian spread
    s2 = 2*(zoom*hwhm)**2
    do i = 1,n
       x2 = (i - n2)**2
       do j = 1,n
          y2 = (j - n2)**2
          kernel(i,j) = exp(-(x2 + y2)/s2)
       end do
    end do

    ! spread profile by convolution by Fourier transformation
    call ftra(bessel,zbessel,m,1)
    call ftra(kernel,zkernel,m,1)

    zpsf = zbessel * zkernel
    call ftra(zpsf,qpsf,m,-1)

    ! re-arrange quarters
    psf(0:n2-1,0:n2-1) = real(qpsf(1:n2,1:n2))
    psf(-n2+1:0,0:n2-1) = real(qpsf(n2+1:n,1:n2))
    psf(0:n2-1,-n2+1:0) = real(qpsf(1:n2,n2+1:n))
    psf(-n2+1:0,-n2+1:0) = real(qpsf(n2+1:n,n2+1:n))

    deallocate(bessel,kernel,qpsf,zbessel,zkernel,zpsf)

  end subroutine fseeing

  subroutine gauss(hwhm)

    real(dbl), intent(in) :: hwhm
    real(dbl) :: x2,y2,s2
    integer :: i,j,n

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

    ! fill-up first quarter
    s2 = 2*(zoom*hwhm)**2
    do i = 0,n
       x2 = i**2
       do j = 0,n
          y2 = j**2
          psf(i,j) = exp(-(x2 + y2)/s2)
       end do
    end do

    ! fill-up rest of the PSF
    forall( i = -n:0, j = 0:n ) psf(i,j) = psf(-i,j)
    forall( i = -n:n, j = -n:-1 ) psf(i,j) = psf(i,-j)

  end subroutine gauss


  subroutine moffat(hwhm, beta)

    real(dbl), intent(in) :: hwhm, beta
    real(dbl) :: x2,y2,s2
    integer :: i,j,n

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

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

    ! fill-up first quarter
    s2 = (zoom*hwhm)**2
    do i = 0,n
       x2 = i**2
       do j = 0,n
          y2 = j**2
          psf(i,j) = (1 + (x2 + y2)/s2)**(-beta)
       end do
    end do

    ! fill-up rest of the PSF
    forall( i = -n:0, j = 0:n ) psf(i,j) = psf(-i,j)
    forall( i = -n:n, j = -n:-1 ) psf(i,j) = psf(i,-j)

  end subroutine moffat


  subroutine spray_pixelize(spsf,dx,dy,xpsf,n)

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

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

    nbox = spsf%nbox
    nn = ubound(spsf%psf,1)
    allocate(psf(-nn:nn,-nn:nn))
    psf = spsf%psf
    x = nint(zoom*dx)
    y = nint(zoom*dy)
    m = zoom
    m2 = m / 2

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

!    do i = -n,n
!       do j = -n,n
!          xpsf(i,j) = psf(i*m,j*m)
!       end do
!    end do

    deallocate(psf)

  end subroutine spray_pixelize


  subroutine spray_clear(spsf)

    type(SprayType), intent(in out) :: spsf
    deallocate(spsf%psf)

  end subroutine spray_clear


end module spray
