module wei_pcig
  !! compute CIGs
  !!        offset domain (weihcig)
  !! ray parameter domain (weipcig)

#ifndef USE_FFTW
use fftw_sep_fake_mod
#endif
 use wei_basics_mod


  !---------------
  implicit none

  type(axis) :: akhx, akhy
  type(axis) :: aphx, aphy

  integer(kind=dbl),pointer,  private :: planF(:) !! forward FFT
  integer(kind=dbl),pointer,  private :: planI(:) !! inverse FFT
  real,               private :: scale

  real,               private :: pcigeps
  integer,            private :: jhx,jhy !! half nhx,nhy - even # of samples

  !! temp FFT arrays
  complex, dimension(:,:,:), allocatable, private :: hht
  complex, dimension(:,:),       pointer, private :: sftF, sftI !! shared

  !! mapping arrays
  integer, dimension(:,:),       pointer, private :: xilk, yilk !! index to lower kh
  logical, dimension(:,:),       pointer, private :: xobf, yobf !! off-bounds flag
  real,    dimension(:,:),       pointer, private :: xwui, ywui !! weight to upper kh

  !! tmp arrays
  complex, dimension(:,:),       pointer, private :: hxtmp, hytmp
  complex, dimension(:,:),       pointer, private :: pxtmp, pytmp

  !! tridiagonal matrices
  complex, dimension(:,:),       pointer, private :: xdiag, xoffd
  complex, dimension(:,:),       pointer, private :: ydiag, yoffd
contains
  !----------------------------------------------------------------  
!!$=head1 NAME
!!$
!!$weipcig_init -initialize offset ray parameter imaging
!!$
!!$=head1 SYNOPSIS
!!$
!!$ierr= weihcig_init()
!!$
!!$=head1 DESCRIPTION
!!$
!!$ Initializes migration operator whose image (R) domain is p_h
!!$
!!$ Grabs pcig_eps (default=1.0)
!!$
!!$
!!$=head1 LIBRARY
!!$
!!$sepweif90
!!$
!!$=head1 SEE ALSO
!!$
!!$ L<weihcig_init> L<weipcig>
!!$
!!$=cut
!!$

  !----------------------------------------------------------------  
  !! p-domain CIG
  integer function weipcig_init() result(st)
    integer, dimension(3) :: ni
    logical               :: fftmeasure
    integer               :: fftsetup
    integer               :: er,i
    integer               :: ipx,ipy, jpx,jpy, ihx,ihy
    real                  :: arg
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_init'
    if(run%debug) call in(unit)
#endif


    !! Note: the offset axes should have an even number of samples
    jhx = down%wsep%n(3)/2
    jhy = down%wsep%n(4)/2

    call from_param("pcigeps",pcigeps,1.0)


    ni = (/ down%wsep%n(3), down%wsep%n(4), 1 /)
    allocate(planF(run%o_nth),planI(ruN%o_nth))
    call add_mem(size(planF))
    scale = 1./sqrt(1.*product(ni))
    !! ----- FFT tmp array -----
    allocate(hht(ni(1),ni(2),run%o_nth),stat=er)
    call add_mem(size(hht))
    if(er/=0) call seperr('cant allocate hht')
    !! ----- init FFT plans -----
    do i=1,run%o_nth
      call sfftw_plan_dft(planF(i),2,ni,hht(1,1,i),hht(1,1,i),FFTW_FORWARD,FFTW_PATIENT)
      call sfftw_plan_dft(planI(i),2,ni,hht(1,1,i),hht(1,1,i),FFTW_BACKWARD,FFTW_PATIENT)
    end do

    !! ----- init FFT shift -----
    allocate(sftF(down%wsep%n(3),down%wsep%n(4)), stat=er)
    call add_mem(size(sftF))
    if(er /=0) call seperr("cant allocate sftF")
    allocate(sftI(down%wsep%n(3),down%wsep%n(4)), stat=er)
    call add_mem(size(sftI))
    if(er /=0) call seperr("cant allocate sftI")

    do ihx=1,down%wsep%n(3)
       do ihy=1,down%wsep%n(4)
          arg = 2*pi &
          * ( (ihx-1.)/down%wsep%n(3) * down%wsep%o(3)/down%wsep%d(3) &
          +   (ihy-1.)/down%wsep%n(4) * down%wsep%o(4)/down%wsep%d(4) )

          sftF(ihx,ihy) = cmplx(cos(arg),-(+1)*sin(arg))
          sftI(ihx,ihy) = cmplx(cos(arg),-(-1)*sin(arg))
       end do
    end do

    !! ----- tmp arrays ------
    allocate(hxtmp(down%wsep%n(3),run%o_nth), stat=er)
    call add_mem(size(hxtmp))
    if(er /=0) call seperr("cant allocate hxtmp") !! hxtmp(nhx,nth)
    allocate(hytmp(down%wsep%n(4),run%o_nth), stat=er)
    call add_mem(size(hytmp))
    if(er /=0) call seperr("cant allocate hytmp") !! hytmp(nhy,nth)
    allocate(pxtmp(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(pxtmp))
    if(er /=0) call seperr("cant allocate pxtmp") !! pxtmp(npx,nth)
    allocate(pytmp(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(pytmp))
    if(er /=0) call seperr("cant allocate pytmp") !! pytmp(npy,nth)

    !! ------ stretch arrays ------
    allocate(xilk(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(xilk))
    if(er /=0) call seperr("cant allocate xilk") !! xilk(npx,nth)
    allocate(xobf(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(xobf))
    if(er /=0) call seperr("cant allocate xobf") !! xobf(npx,nth)
    allocate(xwui(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(xwui))
    if(er /=0) call seperr("cant allocate xwui") !! xwui(npx,nth)

    allocate(yilk(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(yilk))
    if(er /=0) call seperr("cant allocate yilk") !! yilk(npy,nth)
    allocate(yobf(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(yobf))
    if(er /=0) call seperr("cant allocate yobf") !! yobf(npy,nth)
    allocate(ywui(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(ywui))
    if(er /=0) call seperr("cant allocate ywui") !! ywui(npy,nth)

    !! ----- tridiagonal solver -----
    allocate(xdiag(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(xdiag))

    if(er /=0) call seperr("cant allocate xdiag") !! xdiag(npx,nth)
    allocate(xoffd(down%rsep%n(3),run%o_nth), stat=er)
    call add_mem(size(xoffd))
    if(er /=0) call seperr("cant allocate xoffd") !! xoffd(npx,nth)

    allocate(ydiag(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(ydiag))
    if(er /=0) call seperr("cant allocate ydiag") !! ydiag(npy,nth)
    allocate(yoffd(down%rsep%n(4),run%o_nth), stat=er)
    call add_mem(size(yoffd))
    if(er /=0) call seperr("cant allocate yoffd") !! yoffd(npy,nth)

    !! ----- init tridiagonal solver -----
    call weitds_init(down%rsep%n(3),down%rsep%n(4))

    !! ----- init kh and ph axes -----
    !! offset x axes (khx,phx)
    akhx%n = down%wsep%n(3)
    akhx%o = -   pi /  down%wsep%d(3)
    akhx%d = 2 *pi / (down%wsep%d(3) * down%wsep%n(3))
    if(down%wsep%n(3) ==1) then
       akhx%o=0.
       akhx%d=1.
    end if
    aphx%n =down%rsep%n(3)
    aphx%o =down%rsep%o(3)
    aphx%d =down%rsep%d(3)

    !! offset y axes (khy,phy)
    akhy%n = down%wsep%n(4)
    akhy%o = -   pi /  down%wsep%d(4)
    akhy%d = 2 * pi / (down%wsep%d(4) * down%wsep%n(4))
    if(down%wsep%n(4) ==1) then
       akhy%o=0.
       akhy%d=1.
    end if
    aphy%n = down%rsep%n(4)
    aphy%o = down%rsep%o(4)
    aphy%d = down%rsep%d(4)

#ifdef DEBUG
    if(run%debug) call out(unit)
#endif
    st=OK
  end function weipcig_init
!!$=head1 NAME
!!$
!!$weipcig - Perform offset ray parameter imaging 
!!$
!!$=head1 SYNOPSIS
!!$
!!$ ierr=weipcig(wfld,iws,ith)
!!$
!!$=head1 PARAMETERS
!!$
!!$=over 4
!!$
!!$=item iws - integer
!!$
!!$  current frequency slice
!!$
!!$=item ith - integer
!!$
!!$  current thread
!!$
!!$=item wfld - complex(:,:,:,:,:)
!!$
!!$  wavefield
!!$
!!$=back
!!$
!!$=head1 DESCRIPTION
!!$
!!$ Imaging operation when in the offset ray parameter (p_h) domain
!!$
!!$
!!$=head1 LIBRARY
!!$
!!$sepweif90
!!$
!!$=head1 SEE ALSO
!!$
!!$ L<weipcig_init> L<weihcig>
!!$
!!$=cut
  !----------------------------------------------------------------
  integer function weipcig(wfld,iws,ith) result(st)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in) ::iws,ith
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig'
    if(run%depar) call in(unit)
#endif

    !! interpolator map (A)
    call weipcig_map(iws,ith, aphx,akhx, xobf,xilk,xwui)
    call weipcig_map(iws,ith, aphy,akhy, yobf,yilk,ywui)

    if(down%adj) then !! MIGRATION
       call weipcig_fft(wfld,ith) !!            kh = Ft h
       call weipcig_h2p(wfld,ith) !! ph = At Mt kh
    else             !! MODELING
       call weipcig_p2h(wfld,ith) !!         kh = A  M  ph
       call weipcig_ift(wfld,ith) !!  h = F  kh
    end if
 
#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
    st=OK
  end function weipcig
  !----------------------------------------------------------------  
  !! ph to kh conversion
  subroutine weipcig_p2h(wfld,ith)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in) :: ith
    integer             :: imx,imy, ihx,ihy, ipx,ipy
    integer             :: il,iu
    real                :: wl,wu 
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_p2h'
    if(run%depar) call in(unit)
#endif

    xp2h: if(down%rsep%n(3) >1) then
       !!setup tridiagonal solver
       call weipcig_i3d(ith,down%rsep%n(3),xdiag,xoffd,xobf,xilk,xwui)

       do imx=1,down%wsep%n(1)
          do imy=1,down%wsep%n(2)
             do ihy=1,down%wsep%n(4)
                hxtmp(:,ith) = 0.
                pxtmp(:,ith) = wfld(imx,imy,:,ihy,ith)

                !! solve (At A + e^2 Rt R ) ph = ph 
                call weitdsx(ith,xdiag,xoffd,pxtmp)
                do ipx=1,down%rsep%n(3)
                   if(xobf(ipx,ith)) cycle
                   il=xilk(ipx,ith); iu= il +  1
                   wu=xwui(ipx,ith); wl= 1. - wu 

                   !! kh = A ph
                   hxtmp(il,ith) = hxtmp(il,ith) + pxtmp(ipx,ith) * wl
                   hxtmp(iu,ith) = hxtmp(iu,ith) + pxtmp(ipx,ith) * wu
                end do

                wfld(imx,imy,:,ihy,ith) = hxtmp(:,ith)
             end do
          end do
       end do
    end if xp2h

    yh2p: if(down%rsep%n(4) >1) then
       !!setup tridiagonal solver
       call weipcig_i3d(ith,down%rsep%n(4),ydiag,yoffd,yobf,yilk,ywui)

       do imx=1,down%wsep%n(1)
          do imy=1,down%wsep%n(2)
             do ihx=1,down%wsep%n(3)
                hytmp(:,ith) = 0.
                pytmp(:,ith) = wfld(imx,imy,ihx,:,ith)

                !! solve (At A + e^2 Rt R ) ph = ph
                call weitdsy(ith,ydiag,yoffd,pytmp)
                do ipy=1,down%rsep%n(4)
                   if(yobf(ipy,ith)) cycle
                   il=yilk(ipy,ith); iu= il +  1
                   wu=ywui(ipy,ith); wl= 1. - wu 

                   !! kh = A ph
                   hytmp(il,ith) = hytmp(il,ith) + pytmp(ipy,ith) * wl
                   hytmp(iu,ith) = hytmp(iu,ith) + pytmp(ipy,ith) * wu
                end do

                wfld(imx,imy,ihx,:,ith) = pytmp(:,ith)
             end do
          end do
       end do

    end if yh2p

#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_p2h
  !----------------------------------------------------------------
  !! kh to ph conversion
  subroutine weipcig_h2p(wfld,ith)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in) :: ith
    integer             :: imx,imy, ihx,ihy, ipx,ipy
    integer             :: il,iu
    real                :: wl,wu
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_h2p'
    if(run%depar) call in(unit)
#endif

    xh2p: if(down%rsep%n(3) >1) then
       !!setup tridiagonal solver
       call weipcig_i3d(ith,down%rsep%n(3),xdiag,xoffd,xobf,xilk,xwui)

       do imx=1,down%wsep%n(1)
          do imy=1,down%wsep%n(2)
             do ihy=1,down%wsep%n(4)
                hxtmp(:,ith) = wfld(imx,imy,:,ihy,ith)
                pxtmp(:,ith) = 0.

                do ipx=1,down%rsep%n(3)
                   if(xobf(ipx,ith)) cycle
                   il=xilk(ipx,ith); iu= il +  1
                   wu=xwui(ipx,ith); wl= 1. - wu 

                   !! ph = At kh
                   pxtmp(ipx,ith) = pxtmp(ipx,ith) + hxtmp(il,ith) * wl
                   pxtmp(ipx,ith) = pxtmp(ipx,ith) + hxtmp(iu,ith) * wu
                end do
                !! solve (At A + e^2 Rt R ) ph = At kh 
                call weitdsx(ith,xdiag,xoffd,pxtmp)

                wfld(imx,imy,:,ihy,ith) = pxtmp(:,ith)
             end do
          end do
       end do
    end if xh2p

    yh2p: if(down%rsep%n(4) >1) then
       !!setup tridiagonal solver
       call weipcig_i3d(ith,down%rsep%n(4),ydiag,yoffd,yobf,yilk,ywui)

       do imx=1,down%wsep%n(1)
          do imy=1,down%wsep%n(2)
             do ihx=1,down%wsep%n(3)
                hytmp(:,ith) = wfld(imx,imy,ihx,:,ith)
                pytmp(:,ith) = 0.

                do ipy=1,down%rsep%n(4)
                   if(yobf(ipy,ith)) cycle
                   il=yilk(ipy,ith); iu= il +  1
                   wu=ywui(ipy,ith); wl= 1. - wu 

                   !! ph = At kh
                   pytmp(ipy,ith) = pytmp(ipy,ith) + hytmp(il,ith) * wl
                   pytmp(ipy,ith) = pytmp(ipy,ith) + hytmp(iu,ith) * wu
                end do
                !! solve (At A + e^2 Rt R ) ph = At kh 
                call weitdsy(ith,ydiag,yoffd,pytmp)

                wfld(imx,imy,ihx,:,ith) = pytmp(:,ith)
             end do
          end do
       end do
    end if yh2p
#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_h2p
  !----------------------------------------------------------------
  !! PULL stretch mapping - definition of A and At
  !! (emphasize migration over modeling) 
  !! loop over ph; map every ph into kh
  subroutine weipcig_map( iws,ith, aph,akh, obf,ilk,wui)
    integer,    intent(in) ::iws,ith
    type(axis), intent(in) :: aph, akh
    logical, dimension(:,:), pointer :: obf
    integer, dimension(:,:), pointer :: ilk
    real,    dimension(:,:), pointer :: wui
    integer :: ip
    real    :: p2k,w,ph
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_map'
    if(run%depar) call in(unit)
#endif

!   w  = down%bw%o + (iws-1)*down%bw%d
    w=down%w_local(iws)

    p2k=0.
    obf(:,ith)=.false.
    do ip=1,aph%n; ph = aph%o + (ip-1) * aph%d
       p2k = ph * w  !! kh = f(ph) = ph * w

       !! ILK = Index to Lower kh
       ilk(ip,ith) = floor(1 + (p2k - akh%o)/akh%d)

       !! WUI = Weight to Upper Index in the kh array
       wui(ip,ith) = p2k - (akh%o + (ilk(ip,ith)-1)*akh%d)
       wui(ip,ith) = wui(ip,ith)/akh%d

    end do
    where(ilk(:,ith)<1 .or. ilk(:,ith)>akh%n-1)
       !! OBF = Off-Bounds Flag
       obf(:,ith) = .true.
    end where

#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_map
  !----------------------------------------------------------------
  !! init tridiagonal matrix
  subroutine weipcig_i3d(ith,n,diag,offd,obf,ilk,wui)
    integer, intent(in) :: n, ith
    complex, dimension(:,:), pointer :: diag, offd
    logical, dimension(:,:), pointer :: obf
    integer, dimension(:,:), pointer :: ilk
    real,    dimension(:,:), pointer :: wui
    integer :: i, il,iu
    real    ::    wl,wu
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_i3d'
    if(run%depar) call in(unit)
#endif

    diag(:,ith) = 0.
    offd(:,ith) = 0.
    do i=1,n
       if(obf(i,ith)) cycle
       il=ilk(i,ith); iu= il +  1
       wu=wui(i,ith); wl= 1. - wu
       
       !! (At A)
       diag(i,ith) = wl*wl + wu*wu
    end do
    !! (At A + e^2 Rt R)
    diag(:,ith) = diag(:,ith) + 2 * pcigeps*pcigeps
    offd(:,ith) =             -     pcigeps*pcigeps
#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_i3d
  !----------------------------------------------------------------
  subroutine weipcig_fft(wfld,ith)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in) :: ith
    integer :: imx,imy!!, ihx,ihy
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_fft'
    if(run%depar) call in(unit)
#endif
    !! FFT on hx and hy
    do imx=1,down%wsep%n(1)
       do imy=1,down%wsep%n(2)
          hht(:,:,ith) = wfld(imx,imy,:,:,ith)
          call sfftw_execute(planF(ith))
          wfld(imx,imy,:,:,ith) = hht(:,:,ith) * scale
       end do
    end do

    call weipcig_flp(wfld,ith)
    call weipcig_sft(wfld,ith,sftF)

#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_fft
  !----------------------------------------------------------------
  subroutine weipcig_ift(wfld,ith)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in) ::ith
    integer :: imx,imy!!, ihx,ihy
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_ift'
    if(run%depar) call in(unit)
#endif

    call weipcig_sft(wfld,ith,sftI)
    call weipcig_flp(wfld,ith)

    !! IFT on hx and hy
    do imx=1,down%wsep%n(1)
       do imy=1,down%wsep%n(2)
          hht(:,:,ith) = wfld(imx,imy,:,:,ith)
          call sfftw_execute(planI(ith))
          wfld(imx,imy,:,:,ith) = hht(:,:,ith) * scale
       end do
    end do
#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_ift
  !----------------------------------------------------------------
  !! origin shift
  subroutine weipcig_sft(wfld,ith,sft)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    complex, dimension(:,:),       pointer :: sft
    integer,                    intent(in) ::ith
    integer :: imx,imy
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_sft'
    if(run%depar) call in(unit)
#endif
    do imx=1,down%wsep%n(1)
       do imy=1,down%wsep%n(2)
          wfld(imx,imy,:,:,ith) = &
          wfld(imx,imy,:,:,ith) * sft
       end do
    end do
#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_sft
  !----------------------------------------------------------------
  subroutine weipcig_flp(wfld,ith)
    complex, dimension(:,:,:,:,:), pointer :: wfld
    integer, intent(in)::ith
    integer :: imx,imy, ihx,ihy
#ifdef DEBUG
    character(len=128), parameter :: unit='weipcig_flp'
    if(run%depar) call in(unit)
#endif

    do imx=1,down%wsep%n(1)
       do imy=1,down%wsep%n(2)

          if(down%wsep%n(3)>1) then
             do ihy=1,down%wsep%n(4)
                hxtmp(    1:    jhx,ith) = wfld(imx,imy,jhx+1:jhx+jhx,ihy,ith)
                hxtmp(jhx+1:jhx+jhx,ith) = wfld(imx,imy,    1:    jhx,ihy,ith)
                wfld(imx,imy,:,ihy,ith)  = hxtmp(:,ith)
             end do
          end if

          if(down%wsep%n(4)>1) then
             do ihx=1,down%wsep%n(3)
                hytmp(    1:    jhy,ith) = wfld(imx,imy,ihx,jhy+1:jhy+jhy,ith)
                hytmp(jhy+1:jhy+jhy,ith) = wfld(imx,imy,ihx,    1:    jhy,ith)
                wfld(imx,imy,ihx,:,ith)  = hytmp(:,ith)
             end do
          end if

       end do
    end do

#ifdef DEBUG
    if(run%depar) call out(unit)
#endif
  end subroutine weipcig_flp
  !----------------------------------------------------------------
end module wei_pcig

!! ---------------------------------------------------
!!
!!        A        ph  =  kh
!!  . . . . . . .   .      .
!!  . . . . . . .   .      .
!!  . . . . . . .   .      .
!!  . . . . . . .   .      .
!!                  .
!!                  .
!!                  .
!!           A             ph =    kh
!! LSQ:  (At A + e^2 Rt R) ph = At kh
!! 
!!
!! MIGRATION: forward pCIG transformation
!!       ph = [(At A + e^2 Rt R)^-1]  At  kh
!!          = Mt  At  kh
!!
!!  MODELING: adjoint pCIG transformation
!!       kh = A   M   ph
!!          = A   [(At A + e^2 Rt R)^-1]  ph  
!!
!! ---------------------------------------------------
!!
!! weipcig_map - defines the pull mapping (stretch) A
!!
!!
!! weipcig_h2p - pCIG forward transformation
!!               * compute       ph = At kh
!!               * apply weight  ph = Mt ph
!!
!! weipcig_p2h - pCIG adjoint transformation
!!               * apply weight  ph = M  ph
!!               * compute       kh = A  ph 
!!
!! ---------------------------------------------------

