module odcig2D
  !! offset to angle conversion

  use util_
  use pctrist_
  use mymath_

  implicit none

  !! FFT
  real,              private :: scaleF, scaleI
  integer*8, pointer,private ::  planF(:), planI(:)

  complex, allocatable, private :: cf2d(:,:),cf1d(:)
  complex, pointer, private :: sftI(:,:) 
  complex, pointer, private :: sftF(:)

  !! angle conversion
  real,  allocatable, private :: khi(:),kai(:),kzi(:),str(:,:),stretch(:,:)
  complex,  allocatable, private :: model(:,:), data(:,:)
  real,    private :: eps
  logical, private :: push
  integer, parameter   ::  FFTW_R2HC=0,FTW_HC2R=1,FFTW_DHT=2,FFTW_REDFT00=3,FFTW_REDFT01=4
  integer, parameter   ::  FFTW_REDFT10=5,FFTW_REDFT11=6,FFTW_RODFT00=7,FFTW_RODFT01=8,FFTW_RODFT10=9
  integer, parameter   ::  FFTW_RODFT11=10,FFTW_FORWARD=-1,FFTW_BACKWARD=+1,FFTW_DESTROY_INPUT=1
  integer, parameter   ::  FFTW_UNALIGNED=2,FFTW_CONSERVE_MEMORY=4,FFTW_EXHAUSTIVE=8,FTW_PRESERVE_INPUT=16
  integer, parameter   ::  FFTW_PATIENT=32,FFTW_ESTIMATE=64,FFTW_ESTIMATE_PATIENT=128
  integer, parameter   ::  FFTW_BELIEVE_PCOST=256,FFTW_DFT_R2HC_ICKY=512,FFTW_NONTHREADED_ICKY=1024
  integer, parameter   ::  FFTW_NO_BUFFERING=2048,FFTW_NO_INDIRECT_OP=4096,FFTW_ALLOW_LARGE_GENERIC=8192
  integer, parameter   ::  FFTW_NO_RANK_SPLITS=16384,FFTW_NO_VRANK_SPLITS=32768,FFTW_NO_VRECURSE=65536

contains
  !----------------------------------------------------------------
  subroutine odcig2D_fft_init(nth,ni)
    integer, intent(in) :: nth
    integer, dimension(2), intent(in) :: ni
    integer :: stat
    scaleI = 1./sqrt(1.*product(ni))
    scaleF = 1./sqrt(1.*ni(1))

    allocate( cf2d(ni(1),ni(2)) , cf1d(ni(1)) )
    allocate( planF(1),planI(1) )

    call sfftw_plan_dft(planI(1),2,ni(1:2),cf2d(1,1),cf2d(1,1),&
    FFTW_BACKWARD,FFTW_PATIENT)

    call sfftw_plan_dft(planF(1),1,ni(1),cf1d(1),cf1d(1),&
    FFTW_FORWARD,FFTW_PATIENT)

  end subroutine odcig2D_fft_init
  !----------------------------------------------------------------
  subroutine odcig2D_I_fft(ith,isl,f)
    integer, intent(in) :: ith,isl
    complex :: f(:,:,:)

    cf2d = f(:,:,isl)
    call sfftw_execute(planI(1))
    f(:,:,isl) = cf2d * scaleI

  end subroutine odcig2D_I_fft
  !----------------------------------------------------------------
  subroutine odcig2D_F_fft(ith,isl,f)
    integer, intent(in) :: ith,isl
    complex :: f(:,:,:)
    integer :: i2

    do i2=1,size(f,2)
       cf1d=f(:,i2,isl)
       call sfftw_execute(planF(1))
       f(:,i2,isl) = cf1d * scaleF
    end do

  end subroutine odcig2D_F_fft
  !----------------------------------------------------------------
  subroutine odcig2D_fft_close()

    call sfftw_destroy_plan(planF)
    call sfftw_destroy_plan(planI)
    deallocate(cf1d,cf2d)

  end subroutine odcig2D_fft_close
  !----------------------------------------------------------------
  subroutine odcig2D_shift_init(a1,a2)
    type(myaxis), intent(in):: a1,a2
    integer  :: i1,i2,stat
    real       :: pi, argF, argI

    allocate(sftI(a1%n,a2%n), stat=stat)
    if(stat /=0) call seperr("cant allocate sftI")
    allocate(sftF(a1%n), stat=stat)
    if(stat /=0) call seperr("cant allocate sftF")

    sftF=cmplx(1.,0.)
    sftI=cmplx(1.,0.)

    if(a1%o/=0. .or. a2%o/=0. ) then

       pi = acos(-1.)
       do i1=1,a1%n
          do i2=1,a2%n
             argI = 2*pi &
             * ( (i1-1.)/a1%n * a1%o/a1%d &
             +   (i2-1.)/a2%n * a2%o/a2%d)
             argF = 2*pi &
             * ( (i1-1.)/a1%n * a1%o/a1%d )
             sftI(i1,i2) = cmplx(cos(argI),-sin(argI))
             sftF(i1  ) = cmplx(cos(argF), sin(argF))
          end do
       end do
    end if

  end subroutine odcig2D_shift_init
  !----------------------------------------------------------------
  subroutine odcig2D_I_shift(isl,f)
    integer, intent(in) :: isl
    complex :: f(:,:,:)
    f(:,:,isl) = f(:,:,isl) * sftI
  end subroutine odcig2D_I_shift
  !----------------------------------------------------------------
  subroutine odcig2D_F_shift(isl,f)
    integer, intent(in) :: isl
    complex :: f(:,:,:)
    integer :: i2
    do i2=1,size(f,2)
       f(:,i2,isl) = f(:,i2,isl) * sftF
    end do
  end subroutine odcig2D_F_shift
  !----------------------------------------------------------------
  subroutine odcig2D_shift_close()
    integer :: stat
    deallocate(sftI, stat=stat)
    if(stat /=0) call seperr("cant deallocate sftI")
    deallocate(sftF, stat=stat)
    if(stat /=0) call seperr("cant deallocate sftF")
  end subroutine odcig2D_shift_close
  !----------------------------------------------------------------
  subroutine odcig2D_off_init(nth,az,ah,aa,eps_in,push_in)
    integer,     intent(in) :: nth
    type(myaxis), intent(in) :: az,ah,aa
    real,        intent(in) :: eps_in
    logical,     intent(in) :: push_in
    type(myaxis)             :: akh
    integer :: iz,ih,ia,stat
    real    :: kz,kh,ka,pi
    push  = push_in
    eps   = eps_in
    pi = acos(-1.)
    allocate( data(nth,aa%n), stat=stat)
    if(stat /=0) call seperr("cant allocate data")
    allocate(model(nth,ah%n), stat=stat)
    if(stat /=0) call seperr("cant allocate model")
    allocate(  str(nth,aa%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(.not. push) then !! pull
       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 * sign(kh/kz)
                stretch(iz,ih) = ka
             end do
          else
             stretch(iz,:) = 90.
          end if
       end do
    else                !! push
       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(nth,aa,akh,push,eps)
  end subroutine odcig2D_off_init
  !----------------------------------------------------------------
  subroutine odcig2D_off_x(ith,isl,f,az,ahx,aax)
    type(myaxis), intent(in) :: az,ahx,aax
    integer, intent(in) :: ith,isl
    integer :: ihx,iax,ihy,iz,jhx
    real    :: khx,kax,kz
    complex :: f(:,:,:)
    jhx = ahx%n/2
    do iz=1,az%n
       str(ith,:) = stretch(iz,:)
       call pctrist_define(ith,str)
       data(ith,:) = f(iz,:,isl)
       call pctrist_apply(ith,model,data)
       f(iz,1+jhx:jhx+jhx,isl) = model(ith,1:jhx)
       f(iz,1:jhx,isl) = model(ith,1+jhx:jhx+jhx)
    end do
  end subroutine odcig2D_off_x
  !----------------------------------------------------------------  
  subroutine odcig2D_off_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 odcig2D_off_close
  !----------------------------------------------------------------
end module odcig2D
