! This is Sergey Fomels dip finder and Jon Claerbout Puck
module find_xz_dip_mod 

  use sep3d_struct_mod
  use tri_smooth_mod
  use grad3_dip_mod !good
  use allpass_3d    !good
  use allder_3d
  use divide

  implicit none
  integer, private       :: pn1,pn2,pn3, pniter
  real, private          :: pd1,pd2,pd3, peps1,peps2
  real, private, pointer :: paux1(:),pdip1(:)
  integer, private       :: pnwind(3)


contains


  ! --------------------------------------------------------------
  ! ======================  XSERGEY_DIP  =========================
  ! --------------------------------------------------------------

  subroutine get_dip_aux(dip1)
    real :: dip1(:)

    dip1=pdip1
    deallocate(pdip1)
  end subroutine 

  ! --------------------------------------------------------------        
  subroutine xsergey_dip_init(aux_s,aux1_in,niter,eps1,eps2)
   ! Initialize Xpuck function
    ! aux_s types 3D argument passes the dimensions 
    ! aux_in passes the soft data
    type(sep3d)                :: aux_s
    real, dimension(:), target :: aux1_in
    integer                    :: smooth(3),niter
    real                       :: eps1,eps2

    ! Set Sergeys iter and eps (30,0.00002,0.00002)
    pniter=niter
    peps1=eps1
    peps2=eps2

    ! Set size of aux
    pn1=aux_s%n(1)
    pn2=aux_s%n(2)
    pd1=aux_s%d(1)
    pd2=aux_s%d(2)

    if (sep3d_ndims(aux_s) >2) then
       pn3=aux_s%n(3) ! Default to zero
       pd3=aux_s%n(3)
    else
       pn3=1
       pd3=1.
    end if
    
    allocate(pdip1(pn1*pn2*pn3))

    ! Assumes 3D array as input
    nullify(paux1)
    paux1 => aux1_in
    paux1 = paux1/sum(abs(paux1))
  end subroutine

  ! --------------------------------------------------------------
  function xsergey_dip_lop(adj,add,mm,dd) result(stat)
    ! Find 3d dip(p) and coh field by sergey method
    logical, intent(in)         :: adj,add
    real                        :: mm(:),dd(:)
    real, allocatable, dimension(:,:,:) :: aux3,field3,xres3,dip3
    integer                     :: nwind(3), ierr,stat

    stat=1
    allocate(dip3(pn1,pn2,pn3))
    allocate(field3(pn1,pn2,pn3),xres3(pn1,pn2,pn3),aux3(pn1,pn2,pn3))


    ! Get SERGEYs slopes from auxiliary 
    aux3 = reshape(paux1,(/pn1,pn2,pn3/))
    ierr=find_sergey_fast((/pn1,pn2,pn3/),paux1,pniter,peps1,peps2,pdip1,.false.,.false.)

    ! Reshape 3D <=> 1D
    dip3   = reshape(pdip1,(/pn1,pn2,pn3/))
    field3 = reshape(mm,(/pn1,pn2,pn3/))
    xres3  = reshape(dd,(/pn1,pn2,pn3/))

    ! Apply dip to model
    call wavekill_3p(adj,add,1., dip3, field3, xres3)
    mm= reshape(field3,(/pn1*pn2*pn3/))
    dd= reshape(xres3,(/pn1*pn2*pn3/))

    deallocate(aux3,field3,xres3,dip3)
    stat=0
  end function

  ! --------------------------------------------------------------
  ! ========================  XPUCK_DIP   ========================
  ! --------------------------------------------------------------
  subroutine xpuck_dip_init(aux_s,aux1_in,smooth)
   ! Initialize Xpuck function
    ! aux_s types 3D argument passes the dimensions 
    ! aux_in passes the soft data
    type(sep3d)                :: aux_s
    real, dimension(:), target :: aux1_in
    integer                    :: smooth(3)
    
    ! Set smoothing factor on Z axis for Puck2d 
    pnwind=smooth
    !    call xsergey_dip_init(aux_s,aux1_in,pniter,peps1,peps2)
    ! Set size of aux
    pn1=aux_s%n(1)
    pn2=aux_s%n(2)
    pd1=aux_s%d(1)
    pd2=aux_s%d(2)

    if (sep3d_ndims(aux_s) >2) then
       pn3=aux_s%n(3) ! Default to zero
       pd3=aux_s%n(3)
    else
       pn3=1
       pd3=1.
    end if
    
    allocate(pdip1(pn1*pn2*pn3))
    ! Assumes 3D array as input
    nullify(paux1)
    paux1 => aux1_in
    paux1 = paux1/sum(abs(paux1))

  end subroutine


  ! --------------------------------------------------------------
  function xpuck_dip_lop(adj,add,mm,dd) result(stat)
    ! Find 3d dip(p) and coh field by puck method
    real                  :: mm(:),dd(:)
    logical, intent(in)   :: adj,add
    integer               :: nwind(3), ierr,stat
    real, allocatable, dimension(:,:,:) :: aux3,field3,xres3,dip3,j_coh3,j_res3


    stat=1
    allocate(aux3(pn1,pn2,pn3),field3(pn1,pn2,pn3),xres3(pn1,pn2,pn3),dip3(pn1,pn2,pn3))
    allocate(j_coh3(pn1,pn2,pn3),j_res3(pn1,pn2,pn3))

    ! Get Puck_slopes from auxiliary data
    aux3 = reshape(paux1,(/pn1,pn2,pn3/))
    ierr=find_puck_dip_fast(aux3,pnwind,dip3,j_coh3,j_res3)

    ! Reshape 3D<=>1D  & 
    ! Get Residual
    field3 = reshape(mm,(/pn1,pn2,pn3/))
    xres3  = reshape(dd,(/pn1,pn2,pn3/))
    call wavekill_3p(adj,add,1., dip3, field3, xres3)
    mm    = reshape(field3,(/pn1*pn2*pn3/))
    dd    = reshape(xres3 ,(/pn1*pn2*pn3/))
    pdip1 = reshape(dip3  ,(/pn1*pn2*pn3/))

    deallocate(aux3,field3,xres3,dip3)
    deallocate(j_coh3,j_res3)
    stat=0
  end function


  ! ==============================================================
  ! --------------------------------------------------------------
  ! ==============================================================


  subroutine wavekill( adj,add, aa, pp, uu, vv )
    ! Apply (aa.d/dx + pp.d/dt) uu = vv
    logical, intent(in)         :: adj,add
    real    :: aa, pp, s11, s12, s21, s22, uu(:,:,:), vv(:,:,:)
    integer :: i1,i2,n1,n2,n3,i3

    ! Adjoint Null
    if (.not. add) then
       if (adj) then
          uu=0.
       else
          vv=0.
       endif
    endif

    ! Compute and print Sij
    s11 = -aa-pp; s12 =  aa-pp; 
    s21 = -aa+pp; s22 =  aa+pp
    
    ! Get dimensions
    n1=size(uu,1); 
    n2=size(uu,2); 
    n3=size(uu,3)
    
    ! Loop over axis and apply filter
    if (adj) then
       do i3=1, n3 !xline
          do i1= 1, n1-1 !time 
             do i2= 1, n2-1  !inline 
                uu(i1  ,i2  ,i3) = uu(i1  ,i2  ,i3)+ s11*vv(i1,i2,i3)
                uu(i1  ,i2+1,i3) = uu(i1  ,i2+1,i3)+ s12*vv(i1,i2,i3)
                uu(i1+1,i2  ,i3) = uu(i1+1,i2  ,i3)+ s21*vv(i1,i2,i3)
                uu(i1+1,i2+1,i3) = uu(i1+1,i2+1,i3)+ s22*vv(i1,i2,i3)
             end do
          end do
       end do
    else
       do i3=1, n3 !xline
          do i1= 1, n1-1 !time 
             do i2= 1, n2-1  !inline 
                vv(i1,i2,i3) = vv(i1,i2,i3)+&
                     uu(i1  ,i2,i3)*s11 + uu(i1  ,i2+1,i3)*s12 +&
                     uu(i1+1,i2,i3)*s21 + uu(i1+1,i2+1,i3)*s22
             end do
          end do
       end do
       ! Set boundaries
       vv(n1,:,:)=0. !vv(n1-1,:,:)
       vv(:,n2,:)=0. !vv(:,n2-1,:)
    end if

  end subroutine 

  ! --------------------------------------------------------------
  subroutine wavekill_3p(adj,add,aa, pp, uu, res)
    ! Find res = (aa.d/dx + pp.d/dt) with 3D pp field
    logical, intent(in)         :: adj,add
    integer :: i1,i2,n1,n2,n3,i3
    real    :: aa,uu(:,:,:), res(:,:,:), pp(:,:,:)
    real, dimension(:,:,:), allocatable :: s11, s12, s21, s22


    ! Adjoint Null
    if (.not. add) then
       if (adj) then
          uu=0.
       else
          res=0.
       endif
    endif
   
    ! Get dimensions
    n1=size(uu,1); 
    n2=size(uu,2); 
    n3=size(uu,3)


    
    ! Compute and print Sij
    allocate(s11(n1,n2,n3),s12(n1,n2,n3),s21(n1,n2,n3),s22(n1,n2,n3))
    s11 = -aa-pp; s12 =  aa-pp; 
    s21 = -aa+pp; s22 =  aa+pp

    

    ! Loop over axis and apply filter
    if (adj) then
       do i3=1, n3 !xline
          do i1= 1, n1-1 !time 
             do i2= 1, n2-1  !inline 
                uu(i1  ,i2  ,i3) = uu(i1  ,i2  ,i3)+ s11(i1,i2,i3)*res(i1,i2,i3)
                uu(i1  ,i2+1,i3) = uu(i1  ,i2+1,i3)+ s12(i1,i2,i3)*res(i1,i2,i3)
                uu(i1+1,i2  ,i3) = uu(i1+1,i2  ,i3)+ s21(i1,i2,i3)*res(i1,i2,i3)
                uu(i1+1,i2+1,i3) = uu(i1+1,i2+1,i3)+ s22(i1,i2,i3)*res(i1,i2,i3)
             end do
          end do
       end do
    else
       do i3=1, n3 !xline
          do i1= 1, n1-1 !time 
             do i2= 1, n2-1  !inline 
                res(i1,i2,i3) = res(i1,i2,i3)+&
                     uu(i1  ,i2,i3)*s11(i1,i2,i3) + uu(i1  ,i2+1,i3)*s12(i1,i2,i3) +&
                     uu(i1+1,i2,i3)*s21(i1,i2,i3) + uu(i1+1,i2+1,i3)*s22(i1,i2,i3)
             end do
          end do
       end do
       ! Set boundaries
       res(n1,:,:)=0. !res(n1-1,:,:)
       res(:,n2,:)=0. !res(:,n2-1,:)
    end if

    deallocate(s11,s12,s21,s22)

  end subroutine 



  ! ==============================================================
  ! --------------------------------------------------------------
  ! ==============================================================

  subroutine find_puck_dips(field,dip1,coh1)
    ! Find SCALAR dips(p) and coh by puck method
    real :: field(:,:,:)
    real :: dip1,coh1
    real :: dx(size(field,1),size(field,2),size(field,3))
    real :: dt(size(field,1),size(field,2),size(field,3))
    real :: xx,xt,tt
    
    ! Get d/dt,d/dx
    call wavekill(.false.,.false.,1.,0.,field,dx)
    call wavekill(.false.,.false.,0.,1.,field,dt)
    
    xx=1e-30+sum(dx*dx)
    xt=1e-30+sum(dx*dt)
    tt=1e-30+sum(dt*dt)
    
    ! Find SCALAR dip (p) and coherency
    coh1=sqrt((xt/tt)*(xt/xx))
    dip1=-xt/tt
                
  end subroutine 

  ! --------------------------------------------------------------
  logical function find_puck_dip_fast(field,nwind,dip3,coh3,res3)
    ! Find 3d dip(p) and coh field by puck method
        integer,intent(in)  :: nwind(3)
        real   ,intent(in)  :: field(:,:,:)
        real   ,intent(out) :: dip3(:,:,:),coh3(:,:,:) ,res3(:,:,:)
        real   ,pointer, dimension(:,:,:)  :: dx,dt,xt,tt,xx
        real                :: ipct
        integer             :: b(3),e(3),i1,i2,i3,npct,i
        integer             :: ierr,n1,n2,n3

        find_puck_dip_fast=.false.
        ! Get dimensions and allocate
        n1=size(field,1); 
        n2=size(field,2); 
        n3=size(field,3)
        allocate(dx(n1,n2,n3),dt(n1,n2,n3),xt(n1,n2,n3),stat=ierr)
        if (ierr.ne.0) then; 
           write(0,*) "trouble allocating dx and dt"; 
           return; 
        end if
        ! Get d/dx & d/dt
        call wavekill(.false.,.false.,1.,0.,field,dx);
        call wavekill(.false.,.false.,0.,1.,field,dt);
        xt=dt*dx; 
        tt=>dt; tt=dt*dt; 
        xx=>dx; xx=dx*dx; 
        
        ! Smoothing over each axes
        call tri_smooth(xx,nwind); 
        call tri_smooth(tt,nwind); 
        call tri_smooth(xt,nwind)
        
        ! Compute dip and coherency
        where (tt/=0) dip3=-xt/tt
        where (tt/=0 .and. xx/=0) coh3=sqrt((xt/tt)*(xt/xx))
        
        call wavekill_3p(.false.,.false.,1., dip3,field, res3)
        deallocate(dx,dt,xt)
        find_puck_dip_fast=.true.

      end function 

  ! --------------------------------------------------------------
  logical function find_puck_dip_slow(field,nwind,dip3,coh3,res3,verb)
        real,intent(in)     :: field(:,:,:)
        integer,intent(in)  :: nwind(3)
        logical, intent(in) :: verb
        real,intent(out)    :: dip3(:,:,:),coh3(:,:,:) ,res3(:,:,:)
        integer             :: b(3),e(3),i1,i2,i3,npct,i
        real                :: ipct

        find_puck_dip_slow=.false.

        ! Init progress report variables
        npct=size(field,2)*size(field,3); 
        ipct=.0001; 
        i=1

        ! Loop over xline
        do i3=1,size(field,3)
           ! set min & max index3
           b(3)=max(1,i3-nwind(3))
           e(3)=min(size(field,3),i3+nwind(3))

           ! Loop over inline
           do i2=1,size(field,2)
              ! set min & max index2
              b(2)=max(1,i2-nwind(2))
              e(2)=min(size(field,2),i2+nwind(2))
              ! Print progress percentage
              i= i+1
              if (real(i)/real(npct) > ipct .and. verb) then
                 ipct=ipct+.05
                 write(0,*) "Finished ",nint(real(ipct)/real(npct)*100.), "%"
              end if

              ! Loop over time/depth
              do i1=1,size(field,1)
                 ! set min & max index1
                 b(1)=max(1,i1-nwind(1))
                 e(1)=min(size(field,1),i1+nwind(1))
                 ! actual dip finder kernel
                 call find_puck_dips(field(b(1):e(1),b(2):e(2),b(3):e(3)),&
                      dip3(i1,i2,i3),coh3(i1,i2,i3))
              end do

           end do
        end do
        call wavekill_3p(.false.,.false.,1., dip3,field, res3)
        find_puck_dip_slow=.true.

  end function 

  ! --------------------------------------------------------------
  logical function find_sergey_slow(n,field,niter,nl_iter,eps,dip,verb)
        real,intent(in) :: field(:)
        logical, intent(in) :: verb
        real :: eps,mean
        real,pointer ::  dip(:)
        integer, intent(in) :: niter,nl_iter
        integer :: n(3),ierr,n123,ndim,iter
        real, allocatable  :: u1(:),u2(:),dp(:)
        logical            :: dave
                find_sergey_slow=.false.
                n123=product(n)
                allocate(u1(n123),u2(n123),dp(n123),stat=ierr)
                if (ierr.ne.0) then
                write(0,*) "trouble allocating temp arrays"
                return
                end if
                call grad3_init(n)
                call  allpass_3d_init(n(1),n(2),n(3),dip,dave)
                call  allder_3d_init(n(1),n(2),n(3),dip,dave)
                ndim=0
                if (n(1) >1) then
                        ndim= ndim+1
                end if
                if (n(2) >1) then
                        ndim= ndim+1
                end if
                if (n(3) >1) then
                        ndim= ndim+1
                end if
                do iter=1,nl_iter 
                        if (verb) then
                                write(0,*) "working on non-linear iteration",iter
                        end if
                        ierr=allpass_3d_lop(.false.,.false.,field,u2)
                        ierr=allder_3d_lop(.false.,.false.,field,u1)
                        mean=sqrt(sum(u1*u1))
                        u1=u1/mean
                        u2=u2/mean
                        call divide_apply_cg(.false.,n123,u2,u1,dp,eps,niter,grad3_op,ndim*n123)
                        dip=dip+dp
                end do 
                deallocate(dp,u1,u2)
                call allpass_3d_close()
                call allder_3d_close()
                call grad3_close()
                find_sergey_slow=.true.
  end function 

  ! --------------------------------------------------------------
  logical function find_sergey_fast(n,field,niter,eps1,eps2,dip,verb,dave)
        logical, intent(in) :: verb
        real,pointer :: field(:)
        real :: eps1,eps2,mean
        real,pointer ::  dip(:) 
        integer, intent(in) :: niter
        integer :: n(3),ierr,n123,ndim,iter
        real, allocatable  :: u1(:),u2(:),dp(:)
        logical            :: dave
                find_sergey_fast=.false.
                n123=product(n)
                ! write(0,*) "dave",dave
                allocate(u1(n123),u2(n123),dp(n123),stat=ierr)
                if (ierr.ne.0)then; write(0,*) "trouble allocating temp arrays"; return; end if
                call  allpass_3d_init(n(1),n(2),n(3),dip,dave)
                ! write(0,*) "dave",dave
                call   allder_3d_init(n(1),n(2),n(3),dip,dave)
                ! write(0,*) "dave",dave
                dip=0.

                do iter=1,niter 
                        if (verb) write(0,*) "working on non-linear iteration111",iter
                        ! write(0,*) "dave",dave
                        ierr=allpass_3d_lop(.false.,.false.,field,u2)
                        ! write(0,*) "dave",dave
                        ierr=allder_3d_lop(.false.,.false.,field,u1)
                        mean=sqrt(sum(u1*u1))
                        u1=u1/mean
                        u2=u2/mean
                        call divide_apply_tri(n,u2,u1,dp,eps1,eps2)
                        dip=dip+dp
                end do 
                deallocate(dp,u1,u2)
                call allpass_3d_close()
                call allder_3d_close()
                find_sergey_fast=.true.
  end function 


end module 
