module wemva_bornsc_mod
  use  base_fftwv3_mod
  implicit none

  real, dimension(:), pointer, private :: kx2, ky2
  real                       , private :: dx, dy, dz
  real                       , private :: cc(4), etabsc
  integer                    , private :: nx, ny, nopbsc, typbsc

contains
  
  integer function wemva_bornsc_init(kx2_in, ky2_in, dx_in, dy_in, dz_in, nx_in, ny_in, &
                                     etabsc_in, nopbsc_in, typbsc_in) result(stat)
    real, dimension(:), target :: kx2_in, ky2_in
    real                       :: dx_in, dy_in, dz_in, etabsc_in
    integer                    :: nx_in, ny_in, nopbsc_in, typbsc_in

    kx2 => kx2_in
    ky2 => ky2_in
    dx  =  dx_in
    dy  =  dy_in
    dz  =  dz_in
    nx  =  nx_in
    ny  =  ny_in

    etabsc = etabsc_in
    nopbsc = nopbsc_in
    typbsc = typbsc_in

    call wemva_phsscr_coef(cc)

    select case(typbsc) 
      case (1)
           write(0,*)"Born Scattering Operator: FOURIER DOMAIN"
      case (2)
           write(0,*)"Born Scattering Operator: SPACE DOMAIN"
      case (3) 
           write(0,*)"Born Scattering Operator: FINITE DIFFERENCE"
    end select

    write(0,*)"etabsc = ", etabsc
    write(0,*)"nopbsc = ", nopbsc 

    stat = 0
  end function

  integer function wemva_bornsc_oper(gwfd, bwfd, wtmp, bvel, ws, opsign, plan_2dxF, plan_2dxI) result(stat)
    complex, dimension(:,:) :: gwfd
    complex, dimension(:,:) :: bwfd
    complex, dimension(:,:) :: wtmp
    real   , dimension(:,:) :: bvel
    real                    :: ws
    integer                 :: opsign
    integer*8               :: plan_2dxF, plan_2dxI

    select case(typbsc)
       case (1)
            call phsscr_fouri(gwfd, bwfd, wtmp, bvel, ws, opsign, plan_2dxF, plan_2dxI)
       case (2)
            call phsscr_space(gwfd, bwfd, wtmp, bvel, ws, opsign, plan_2dxF, plan_2dxI)
       case (3)
            call phsscr_fidif(gwfd, bwfd, bvel, ws, opsign)
    end select

    stat = 0
  end function

  integer function wemva_bornsc_clos() result(stat)
    stat = 1
    kx2 => null()
    ky2 => null()
    dx = 0; dy = 0; dz = 0
    nx = 0; ny = 0
    etabsc = 0; nopbsc = 0; typbsc = 0
    stat = 0
  end function


! input background wfld, output scattered wfld 
  subroutine phsscr_space(gwfd, bwfd, wtmp, bvel, ws, opsign, plan_2dxF, plan_2dxI)
    complex, dimension(:,:) :: gwfd
    complex, dimension(:,:) :: bwfd
    complex, dimension(:,:) :: wtmp
    real   , dimension(:,:) :: bvel
    real                    :: ws
    integer                 :: opsign
    integer*8               :: plan_2dxF, plan_2dxI

    integer                 :: iop, ikx, iky
    complex                 :: iwdz
    real                    :: kk2

    iwdz = cmplx(0., 1.)*ws*dz*opsign

    gwfd = bwfd

    do iop=1, nopbsc
       wtmp = bwfd 
       call sfftw_dft_2dx_oper(plan_2dxF, wtmp, nx, ny)
       do iky=1, ny
       do ikx=1, nx
          kk2 = kx2(ikx)+ky2(iky)
          wtmp(ikx,iky) = cc(iop)*(kk2**iop)*wtmp(ikx,iky)
       end do
       end do
       call sfftw_dft_2dx_oper(plan_2dxI, wtmp, nx, ny)
       gwfd = gwfd + wtmp/ws/ws*bvel*bvel 
    end do

    gwfd = gwfd*iwdz

    return
  end subroutine

  subroutine phsscr_fouri(gwfd, bwfd, wtmp, bvel, ws, opsign, plan_2dxF, plan_2dxI)
    complex, dimension(:,:) :: gwfd
    complex, dimension(:,:) :: bwfd
    complex, dimension(:,:) :: wtmp
    real   , dimension(:,:) :: bvel
    real                    :: ws
    integer                 :: opsign
    integer*8               :: plan_2dxF, plan_2dxI

    integer                 :: iop, ikx, iky
    complex                 :: iwdz, damp
    real                    :: kk2, vavr, w2, s2, kz2

    iwdz = cmplx(0., 1.)*ws*dz*opsign
    damp = (1 - cmplx(0.,etabsc))**2
    vavr = sum(bvel)/nx/ny
    w2   = ws*ws
    s2   = 1./vavr/vavr

    wtmp = bwfd

    call sfftw_dft_2dx_oper(plan_2dxF, wtmp, nx, ny)

    do iky=1, ny
    do ikx=1, nx
       kk2 = kx2(ikx)+ky2(iky)
       kz2 = 1 - kk2/w2/s2

       if (kz2 >= etabsc) then
          wtmp(ikx,iky) = wtmp(ikx,iky)/sqrt(kz2)
       end if

    end do
    end do

    call sfftw_dft_2dx_oper(plan_2dxI, wtmp, nx, ny)

    gwfd = wtmp * iwdz

    return
  end subroutine

  subroutine phsscr_fidif(gwfd, bwfd, bvel, ws, opsign)
    complex, dimension(:,:) :: gwfd
    complex, dimension(:,:) :: bwfd
    real   , dimension(:,:) :: bvel
    real                    :: ws
    integer                 :: opsign

    real                    :: dx2, w2, s2
    complex                 :: iwdz, lapxo
    integer                 :: ix

    iwdz = cmplx(0., 1.)*ws*dz*opsign
    dx2  = dx*dx
    w2   = ws*ws
    do ix=1, nx
       s2 = 1./bvel(ix,1)/bvel(ix,1)
       if (ix == 1) then
          lapxo = (bwfd(ix+1,1)-bwfd(ix,1))/dx
       else if (ix == nx) then
          lapxo = (bwfd(ix,1)-bwfd(ix-1,1))/dx
       else
          lapxo = (bwfd(ix-1,1)-2.*bwfd(ix,1)+bwfd(ix+1,1))/dx2
       end if
       gwfd(ix,1) = iwdz*(bwfd(ix,1)+0.5/w2/s2*lapxo)
    end do
    return
  end subroutine

  subroutine wemva_phsscr_coef(cc)
    real, dimension(4) :: cc

    cc(1) = 1./2.
    cc(2) = 3./8.
    cc(3) = 5./16.
    cc(4) = 35./128.
    return
  end subroutine




end module
