!!.. This module computes the illumination dip and scattering angle gathers
!!.. for a single frequency
module hes_gather_mod
  use  base_struct_mod
  use  base_fftwv3_mod
  use  base_fftsft_mod
  use  base_window_mod
  use  wem_wnbmap_mod
  use  quantile_mod
  implicit none

  type(hes_grid)                        , private :: hgrid
  type(axis)                            , private :: aphs
  type(axis)                            , private :: akss 
  type(axis)                            , private :: akx, aky, akz
  integer                               , private :: ntpx, ntpy, ntpz
  integer*8                             , private :: plan_clhs_3dxF
  real   , dimension(:)    , allocatable, private :: tapx, tapy, tapz
  complex, dimension(:,:,:), allocatable, private :: clhs
  real                                  , private :: pct 
  logical                               , private :: norm
contains

  integer function hes_gahter_init(hgrid_in, aphs_in, pct_in, ntpx_in, ntpy_in, ntpz_in, norm_in) result(stat)
    type(hes_grid) :: hgrid_in
    type(axis)     :: aphs_in
    real           :: pct_in
    integer        :: ntpx_in, ntpy_in, ntpz_in
    logical        :: norm_in
    real           :: PI
    integer        :: mkx, mky, mkz

    PI = 3.14159265359

    hgrid = hgrid_in
    aphs  = aphs_in
    pct   = pct_in
    ntpx  = ntpx_in
    ntpy  = ntpy_in
    ntpz  = ntpz_in
    norm  = norm_in

    akx%n = hgrid%ax_offd%n + hgrid%nx_offd_pad_beg + hgrid%nx_offd_pad_end
    aky%n = hgrid%ay_offd%n + hgrid%ny_offd_pad_beg + hgrid%ny_offd_pad_end
    akz%n = hgrid%az_offd%n + hgrid%nz_offd_pad_beg + hgrid%nz_offd_pad_end

    akx%d = 2*PI/akx%n/hgrid%ax_offd%d
    aky%d = 2*PI/aky%n/hgrid%ay_offd%d
    akz%d = 2*PI/akz%n/hgrid%az_offd%d

    mkx = akx%n/2; mky = aky%n/2; mkz = akz%n/2

    akx%o = (mkx-akx%n+1)*akx%d
    aky%o = (mky-aky%n+1)*aky%d
    akz%o = (mkz-akz%n+1)*akz%d
  
    allocate(clhs(akx%n,aky%n,akz%n))
    allocate(tapx(ntpx),tapy(ntpy),tapz(ntpz))

    call hammiwin1d(tapx, ntpx)
    call hammiwin1d(tapy, ntpy)
    call hammiwin1d(tapz, ntpz)

    call sfftw_dft_3dx_init(plan_clhs_3dxF, clhs, akx%n, aky%n, akz%n, +1)

write(0,*)"okx,oky,okz", akx%o, aky%o, akz%o

    stat = 0
  end function

  integer function hes_gather_clos() result(stat)
    deallocate(clhs)
    deallocate(tapx, tapy, tapz)
    stat = 0
  end function

  integer function hes_gather_oper(rlhs, ahes, dip_agl, sca_agl, sca_map, vref, ws, ix_modl, iy_modl, iz_modl) result(stat)
    real   , dimension(:,:,:) :: rlhs             !!.. off-diagonals of Hessian
    real   , dimension(:,:,:) :: ahes
    real   , dimension(:)     :: dip_agl          !!.. illumination dip angles for one image point
    real   , dimension(:)     :: sca_agl          !!.. illumination scattering angles for one image point
    real   , dimension(:,:,:) :: sca_map
    real                      :: vref, ws, clhs_max
    integer                   :: ix_modl, iy_modl, iz_modl
    integer                   :: ix_offd, iy_offd, iz_offd

    do iz_offd=1, hgrid%az_offd%n
    do iy_offd=1, hgrid%ay_offd%n
       rlhs(1:ntpx                                ,iy_offd,iz_offd) = rlhs(1:ntpx                                ,iy_offd,iz_offd)*tapx(1:ntpx:+1)
       rlhs(hgrid%ax_offd%n-ntpx+1:hgrid%ax_offd%n,iy_offd,iz_offd) = rlhs(hgrid%ax_offd%n-ntpx+1:hgrid%ax_offd%n,iy_offd,iz_offd)*tapx(ntpx:1:-1)
    end do
    end do

    do iz_offd=1, hgrid%az_offd%n
    do ix_offd=1, hgrid%ax_offd%n
       rlhs(ix_offd,1:ntpy                                ,iz_offd) = rlhs(ix_offd,1:ntpy                                ,iz_offd)*tapy(1:ntpy:+1)
       rlhs(ix_offd,hgrid%ay_offd%n-ntpy+1:hgrid%ay_offd%n,iz_offd) = rlhs(ix_offd,hgrid%ay_offd%n-ntpy+1:hgrid%ay_offd%n,iz_offd)*tapy(ntpy:1:-1)
    end do
    end do

    do iy_offd=1, hgrid%ay_offd%n
    do ix_offd=1, hgrid%ax_offd%n
       rlhs(ix_offd,iy_offd,1:ntpz                                ) = rlhs(ix_offd,iy_offd,1:ntpz                                )*tapz(1:ntpz:+1)
       rlhs(ix_offd,iy_offd,hgrid%az_offd%n-ntpz+1:hgrid%az_offd%n) = rlhs(ix_offd,iy_offd,hgrid%az_offd%n-ntpz+1:hgrid%az_offd%n)*tapz(ntpz:1:-1)
    end do
    end do

    clhs = 0.

    clhs(hgrid%nx_offd_pad_beg+1:hgrid%nx_offd_pad_beg+hgrid%ax_offd%n &
        ,hgrid%ny_offd_pad_beg+1:hgrid%ny_offd_pad_beg+hgrid%ay_offd%n &
        ,hgrid%nz_offd_pad_beg+1:hgrid%nz_offd_pad_beg+hgrid%az_offd%n ) = rlhs

    call sfftw_dft_3dx_oper(plan_clhs_3dxF, clhs, akx%n, aky%n, akz%n)

    call base_fftsft_3d(clhs, akx%n, aky%n, akz%n)

    clhs_max = pct*maxval(cabs(clhs))

    where(cabs(clhs)<clhs_max)
      clhs = 0.
    end where


    ahes = cabs(clhs)

    !call gather_sca_kernel(sca_agl, sca_map, vref, ws)
    call gather_sca_kernel2(sca_agl, sca_map, vref, ws)

    stat = 0
  end function


  subroutine gather_dip_kernel()
    
  end subroutine

  subroutine gather_sca_kernel(sca_agl, sca_map, vref, ws)
    real   , dimension(:)     :: sca_agl
    real   , dimension(:,:,:) :: sca_map
    real                      :: vref
    real                      :: ws

    integer                   :: index_kx, index_kz
    integer                   :: isca_agl, iphs
    integer                   :: counter
    real                      :: ssca_agl, kscatter
    real                      :: kscatter_x, kscatter_z, phs

    integer                   :: ikxm, ikxp, ikzm, ikzp
    real                      :: kxm, kxp, kzm, kzp
    real                      :: y1, y2, y3, y4, y, t, u

    sca_agl = 0.
    sca_map = 0.
    do isca_agl=1, hgrid%asca_agl%n
       ssca_agl = hgrid%asca_agl%o + (isca_agl-1)*hgrid%asca_agl%d
       kscatter = 2*ws/vref * cos(ssca_agl)
       counter  = 0
       do iphs=1, aphs%n
          phs = aphs%o + (iphs-1)*aphs%d
          kscatter_x = kscatter * cos(phs)
          kscatter_z = kscatter * sin(phs)

          index_kx   = 1.5 + (kscatter_x-akx%o)/akx%d
          index_kz   = 1.5 + (kscatter_z-akz%o)/akz%d

          ikxm = 1 + (kscatter_x-akx%o)/akx%d; ikxp = ikxm + 1
          ikzm = 1 + (kscatter_z-akz%o)/akz%d; ikzp = ikzm + 1
!write(0,*)"iii", ikxm, ikzm, index_kx, index_kz
!          if (ikxm>=1 .and. ikxm<akx%n .and. ikzm>=1 .and. ikzm<akz%n) then 
!             y1 = cabs(clhs(ikxm, 1, ikzm))
!             y2 = cabs(clhs(ikxp, 1, ikzm))
!             y3 = cabs(clhs(ikxp, 1, ikzp))
!             y4 = cabs(clhs(ikxm, 1, ikzp))
!             kxm = akx%o + (ikxm-1)*akx%d
!             kxp = akx%o + (ikxp-1)*akx%d
!             kzm = akz%o + (ikzm-1)*akz%d
!             kzp = akz%o + (ikzp-1)*akz%d
!             t   = (kscatter_x - kxm)/(kxp-kxm)
!             u   = (kscatter_z - kzm)/(kzp-kzm)
!             y   = (1-t)*(1-u)*y1+t*(1-u)*y2+t*u*y3+(1-t)*u*y4
!             sca_agl(isca_agl) = sca_agl(isca_agl)+y*abs(kscatter)
!             counter = counter + 1
!          end if
          if (index_kx>=1.and.index_kx<=akx%n.and.index_kz>=1.and.index_kz<=akz%n) then
             sca_agl(isca_agl) = sca_agl(isca_agl)+cabs(clhs(index_kx,1,index_kz))*abs(kscatter)
             sca_map(index_kx,1,index_kz) = 1
             counter = counter + 1
          end if
       end do
       sca_agl(isca_agl) = sca_agl(isca_agl)/counter
    end do

    if(norm) sca_agl = sca_agl / maxval(sca_agl)

    return
  end subroutine

  subroutine gather_sca_kernel2(sca_agl, sca_map, vref, ws)
    real   , dimension(:)     :: sca_agl
    real   , dimension(:,:,:) :: sca_map
    real                      :: vref
    real                      :: ws
    integer                   :: isca_agl, ikx, ikz, counter, i, im
    real                      :: ssca_agl, ks, kx, kz, kz2, f, fx, gx
    sca_agl = 0.
    do isca_agl=1, hgrid%asca_agl%n
       ssca_agl = hgrid%asca_agl%o + (isca_agl-1)*hgrid%asca_agl%d
       ks       = 2*ws/vref * cos(ssca_agl)
       counter  = 0
       do ikx=1, akx%n
          kx = akx%o + (ikx-1)*akx%d
          kz2 = ks*ks - kx*kx
          if (kz2 >= 0.) then
             kz = sqrt(kz2)
             !ikz = 1.5 + (kz-akz%o)/akz%d
             !sca_agl(isca_agl) = sca_agl(isca_agl)+cabs(clhs(ikx,1,ikz))
             !counter = counter + 1

             f = 1+(kz-akz%o)/akz%d
             i = f; im = i+1
             if (1<=im .and. im<akz%n) then
                ikz = 1.5 + (kz-akz%o)/akz%d
                fx = f - i; gx = 1. - fx
                sca_agl(isca_agl)  = sca_agl(isca_agl)+(gx*cabs(clhs(ikx,1,im))+fx*cabs(clhs(ikx,1,im+1)))*abs(ks)
                sca_map(ikx,1,ikz) = 1.
                counter = counter + 1
             end if
          end if
       end do
       sca_agl(isca_agl) = sca_agl(isca_agl)/counter
    end do

    if (norm) sca_agl = sca_agl / maxval(sca_agl)
    return
  end subroutine
end module



