module off2ang_mod
  !! Stolt Residual Migration
  !! Paul Sava (paul@sep.stanford.edu)
  !! 12/1/00
  use sep
  use sep3d_struct_mod
  use utilities_mod
  use pctrist_mod
  use fftw_flags_mod
  implicit none

  logical, private , save :: verb_in
  type(paxis) :: az,ahx,ahy,amx,amy,ar,aax,aay
  complex, pointer :: sft(:,:)
  logical, private, save :: push
  integer*8, private,save :: planF,planI
  real, private, save :: eps,amax,scaleF,scaleI
  complex, pointer, private, save :: temp(:,:,:)
  complex, dimension(:,:,:),       pointer, private :: sftF !! shared
  complex, dimension(:),           pointer, private :: sftI !! shared
  real,    dimension(:),   allocatable, private :: khi,kai,kzi
  complex, dimension(:,:), allocatable, private :: model, data
  real,    dimension(:,:), allocatable, private :: str
  real,    dimension(:,:), allocatable, private :: stretch




contains 

logical function init_off2ang_params(input,nang,oang,dang)
type(sep3d) :: input
real :: oang(2),dang(2)
integer :: nang(2)
;
init_off2ang_params=.false.
  call from_param("amax",amax,60.)
  if(input%o(1)/=0.  .or. input%o(2)/=0.) then
    write(0,*) "shifts questionalbe in off2ang"
    return
  end if
  oang=0.
  dang(1)=amax/(input%n(1)-1)
  if(input%n(2)>1) dang(2)=amax/(input%n(2)-1)
  nang=input%n
  
  call from_param(   "eps",   eps,    1.0)
  call from_param(  "push",  push,.false.)

init_off2ang_params=.true.

end function

logical function init_off2ang(input)
  type(sep3d) :: input
  integer :: stat,ihx,ihy
  real    :: ang,pi
  init_off2ang=.false.
  call axis_setup_s(input,1,az)
  call axis_setup_s(input,2,ahx)
  call axis_setup_s(input,3,ahy)
  call axis_setup_s(input,4,amx)
  call axis_setup_s(input,5,amy)

 aax%n = ahx%n
 if(ahx%n /=1) then
   if(ahx%o ==0.) then
      aax%o = 0.
      aax%d =   amax/aax%n
   else
      aax%o =  -amax
      aax%d = 2*amax/aax%n
   end if
 end if
 aay%n = ahy%n
 if(ahy%n /=1) then
   if(ahy%o ==0.) then
     aay%o = 0.
     aay%o = 0.
     aay%d =   amax/aay%n
  else
     aay%o =  -amax
     aay%d = 2*amax/aay%n
  end if
end if
allocate(temp(az%n,aax%n,aay%n))
call adcig_fft_init((/az%n,aax%n,aay%n/))
call adcig_shift_init(az,ahx,ahy)
init_off2ang=.true.

end function

subroutine adcig_fft_init(ni)
    integer, dimension(3), intent(in) :: ni
    integer :: stat,i,j
    integer, external :: npfao
    scaleF = 1./sqrt(1.*product(ni))
    scaleI = 1./sqrt(1.*ni(1))

    call sfftw_plan_dft(planF,3,ni,temp(1,1,1),temp(1,1,1),fFTW_FORWARD  &
    ,    fFTW_ESTIMATE )
                                                                                
    call sfftw_plan_dft(planI,1,(/ni(1)/),temp(1,1,1),temp(1,1,1),fFTW_BACKWARD &
    ,    fFTW_ESTIMATE )
end subroutine



subroutine adcig_shift_init(a1,a2,a3)
    type(paxis), intent(in):: a1,a2,a3
                                                                                  
    integer             :: i1,i2,i3
    real                :: pi, argF, argI
    integer             :: stat
                                                                                  
    allocate(sftF(a1%n,a2%n,a3%n), stat=stat)
    if(stat /=0) call seperr("cant allocate sftF")
    allocate(sftI(a1%n), stat=stat)
    if(stat /=0) call seperr("cant allocate sftI")
                                                                                  
    sftF=cmplx(1.,0.)
    sftI=cmplx(1.,0.)
                                                                                  
    if(a1%o/=0. .or. a2%o/=0. .or. a3%o/=0 ) then
                                                                                  
       pi = acos(-1.)
       do i1=1,a1%n
          do i2=1,a2%n
             do i3=1,a3%n
                                                                                  
                argF = 2*pi &
                * ( (i1-1.)/a1%n * a1%o/a1%d &
                +   (i2-1.)/a2%n * a2%o/a2%d &
                +   (i3-1.)/a3%n * a3%o/a3%d )
                                                                                  
                argI = 2*pi &
                * ( (i1-1.)/a1%n * a1%o/a1%d )
                                                                                  
                sftF(i1,i2,i3) = cmplx(cos(argF),-(+1)*sin(argF))
                sftI(i1      ) = cmplx(cos(argI),-(-1)*sin(argI))
                                                                                  
             end do
          end do
       end do
                                                                                  
    end if
                                                                                  
  end subroutine adcig_shift_init



 subroutine adcig_ang_init(nth,az,ah,aa,eps_in,push_in)
    integer,     intent(in) :: nth
    type(paxis), intent(in) :: az,ah,aa
    real,        intent(in) :: eps_in
    logical,     intent(in) :: push_in
    type(paxis)             :: akh
    integer :: iz,ih,ia
    real    :: kz,kh,ka
    real    :: pi
    integer :: stat
    integer,external :: sep_thread_num
                                                                                  
    push  = push_in
    eps   = eps_in
                                                                                  
    pi = acos(-1.)
                                                                                  
    allocate( data(nth,ah%n), stat=stat)
    if(stat /=0) call seperr("cant allocate data1")
    allocate(model(nth,aa%n), stat=stat)
    if(stat /=0) call seperr("cant allocate model")
    allocate(  str(nth,ah%n), stat=stat)
    if(stat /=0) call seperr("cant allocate str")
                                                                                  
    allocate(khi(ah%n), stat=stat)
    if(stat /=0) call seperr("cant allocate khi")
    allocate(kai(aa%n), stat=stat)
    if(stat /=0) call seperr("cant allocate kai")
    allocate(kzi(az%n), stat=stat)
    if(stat /=0) call seperr("cant allocate kzi")

                                                                                  
    !! kz
    do iz = 1,az%n/2
       kzi(iz       ) = (        iz-1)
       kzi(iz+az%n/2) = (-az%n/2+iz-1)
    end do
    kzi = 2*pi* kzi / (az%n*az%d)
                                                                                  
    !! ka
    do ia = 1,aa%n
       kai(ia) =  aa%o + (ia-1) * aa%d
    end do
                                                                                  
    !! kh
    do ih = 1,ah%n
       khi(ih) = (-ah%n/2+ih-1)
    end do
    khi = 2*pi* khi / (ah%n*ah%d)
    !! akh
    akh%o = khi(1)
    akh%d = 2*pi / (ah%n*ah%d)
    akh%n = ah%n
                                                                                  
    if(     push) then !! push
       allocate(stretch(az%n,ah%n), stat=stat)
       if(stat /=0) call seperr("cant allocate stretch")
                                                                                  
       stretch=0.
       do iz=1,az%n
          kz = kzi(iz)
          if (kz/=0.) then
             do ih = 1,ah%n
                kh = khi(ih)
                ka = - 180* atan(kh/kz) /pi
                stretch(iz,ih) = ka
!if(sep_thread_num()==0) writE(0,*) "STRETCH",ih,iz,ka, stretch(iz,ih),sep_thread_num()
             end do
          else
             stretch(iz,:) = 90.
          end if
       end do
    else               !! pull
       allocate(stretch(az%n,aa%n), stat=stat)
       if(stat /=0) call seperr("cant allocate stretch")
                                                                                  
       stretch=0.
       do iz=1,az%n
          kz = kzi(iz)
          do ia = 1,aa%n
             ka = pi* kai(ia) /180
             kh = - kz * tan(ka)
             stretch(iz,ia) = kh
          end do
       end do
    end if
                                                                                  
    call pctrist_init(1,akh,aa,push,eps)
  end subroutine adcig_ang_init
  !----------------------------------------------------------------
  subroutine adcig_ang_x(f,az,ahx,aax,ahy)
    type(paxis), intent(in) :: az,ahx,aax,ahy
    complex, dimension(:,:,:,:,:) :: f
                                                                                  
    integer :: ihx,iax,ihy,iz,i1,i2,ix
    integer, external :: sep_thread_num
    real    :: khx,kax,    kz
    integer :: jhx,i
    jhx = ahx%n/2
                                                                                  
    do iz=1,az%n
      do i=1,size(stretch,2)
        str(1,i) = stretch(iz,i)
      end do
       call pctrist_define(1,str)

       do i2=1,size(f,5)
         do i1=1,size(f,4)
                                                                                  
           do ihy = 1,ahy%n
             do i=1,jhx
              data(1,i    )=f(iz,i+jhx,ihy,i1,i2)
              data(1,i+jhx)=f(iz,i    ,ihy,i1,i2)
             end do
                                                                                  
              call pctrist_apply(1,model,data)
                                                                                  
             do i=1,jhx
               do ix=1,size(model,2)
                 f(iz,ix,ihy,i1,i2) = model(1,ix)
               end do
             end do
                                                                                  
                                                                              

           end do
        end do
       end do
    end do
                                                                                  
  end subroutine adcig_ang_x
  !----------------------------------------------------------------
  subroutine adcig_ang_y(f,az,ahy,aay,ahx)
    type(paxis), intent(in) :: az,ahy,aay,ahx
    complex, dimension(:,:,:,:,:) :: f
                                                                                  
    integer :: ihy,iay,ihx,iz,i2,i1
    real    :: khy,kay,    kz
    integer :: jhy
                                                                                  
  jhy = ahy%n/2
  do iz=1,az%n
    str(1,:) = stretch(iz,:)
    call pctrist_define(1,str)
     do i2=1,size(f,5)
       do i1=1,size(f,4)
         do ihx = 1,ahx%n
           data(1,1    :jhy    )=f(iz,ihx,1+jhy:jhy+jhy,i1,i2)
           data(1,1+jhy:jhy+jhy)=f(iz,ihx,1    :jhy    ,i1,i2)
           call pctrist_apply(1,model,data)
           f(iz,ihx,:,i1,i2) = model(1,:)
         end do
       end do
     end do
  end do
 end subroutine adcig_ang_y



  !----------------------------------------------------------------
  subroutine adcig_ang_close()
    integer :: stat
    call pctrist_close()
                                                                                  
    deallocate(khi, stat=stat)
    if(stat /=0) call seperr("cant deallocate khi")
    deallocate(kzi, stat=stat)
    if(stat /=0) call seperr("cant deallocate kzi")
    deallocate(kai, stat=stat)
    if(stat /=0) call seperr("cant deallocate kai")
                                                                                  
    deallocate(data, stat=stat)
    if(stat /=0) call seperr("cant deallocate data")
    deallocate(model, stat=stat)
    if(stat /=0) call seperr("cant deallocate model")
    deallocate(str, stat=stat)
    if(stat /=0) call seperr("cant deallocate str")
    deallocate(stretch, stat=stat)
    if(stat /=0) call seperr("cant deallocate stretch")
                                                                                  
  end subroutine adcig_ang_close
  !----------------------------------------------------------------




logical function off2ang(buffer)
complex :: buffer(:,:,:,:,:)
integer, external :: sep_thread_num
integer :: i

off2ang=.false.
 call adcig_F_fft(buffer)
 call adcig_F_shift(buffer)
 if(ahx%n >1) then
   call adcig_ang_init(1,az,ahx,aax,eps,push)
   call adcig_ang_x(buffer,az,ahx,aax,ahy)
   call adcig_ang_close()
 end if
 if(ahy%n >1) then
   call adcig_ang_init(1,az,ahy,aay,eps,push)
   call adcig_ang_y(buffer,az,ahy,aay,ahy)
   call adcig_ang_close()
 end if
 call adcig_I_shift(buffer)
 call adcig_I_fft(buffer)
off2ang=.true.
end function 

  subroutine adcig_I_shift(f)
    complex, dimension(:,:,:,:,:) :: f
    integer :: n2,n3
    integer :: i2,i3,i4,i5,n4,n5
                                                                                  
    n2=size(f,2)
    n3=size(f,3)
    n4=size(f,4)
    n5=size(f,5)
    do i5=1,n5
    do i4=1,n4
    do i2=1,n2
       do i3=1,n3
          f(:,i2,i3,i4,i5) = f(:,i2,i3,i4,i5) * sftI
       end do
    end do
    end do
    end do
  end subroutine adcig_I_shift

    subroutine adcig_I_fft(f)
    complex, dimension(:,:,:,:,:) :: f
    integer :: n1,n2,n3,n4,n5
    integer ::    i2,i3,i4,i5
                                                                                  
    n2=size(f,2)
    n3=size(f,3)
    n4=size(f,4)
    n5=size(f,5)
                                                                                  
    do i5=1,n5
    do i4=1,n4
    do i2=1,n2
       do i3=1,n3
          temp(:,1,1) = f(:,i2,i3,i4,i5)
          call sfftw_execute(planI) !! Xform in place
          f(:,i2,i3,i4,i5) = temp(:,1,1) * scaleI
       end do
    end do
    end do
    end do
  end subroutine adcig_I_fft



  subroutine adcig_F_fft(f)
    complex, dimension(:,:,:,:,:) :: F
    integer :: i1,i2
                                                                                  
    do i2=1,size(f,5)
      do i1=1,size(f,4)
        temp = f(:,:,:,i1,i2)
        call sfftw_execute(planF)
        f(:,:,:,i1,i2) = temp * scaleF
      end do
   end do
  end subroutine adcig_F_fft



logical function  off2ang_clean()
  deallocate(sftF,sftI,temp)
 call sfftw_destroy_plan(planF)
 call sfftw_destroy_plan(planI)
off2ang_clean=.true.


end function





subroutine adcig_F_shift(f)
 integer :: i1,i2
  complex, dimension(:,:,:,:,:) :: f
  do i2=1,size(f,5)
  do i1=1,size(f,4)
      f(:,:,:,i1,i2) = f(:,:,:,i1,i2) * sftF
  end do
  end do
end subroutine adcig_F_shift

subroutine axis_setup_s(struct,iax,axis)
 type(sep3d) :: struct
 integer :: iax
 type(paxis) :: axis
   if(iax <= sep3d_ndims(struct)) then
   axis%n=struct%n(iax)
   axis%o=struct%o(iax)
   axis%d=struct%d(iax)
   else
     axis%o=0
     axis%d=1.
     axis%n=1
  end if
end subroutine

end module
