module wei_S
  use wei_util
  use wei_type
  use wei_pro
  use wei_par
  use wei_xmap
  use wei_attr
  use wei_dump
  !---------------
  implicit none

  real,    dimension(:), pointer, private ::  sloline
  !! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  type(weslowmap), public :: slo
  !! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!

contains
  !----------------------------------------------------------------
  subroutine weiS_init(X)
    type(weslow), intent(inout) :: X
    character(len=128), parameter :: unit='weiS_init' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpar%impi==0) call weiS_his_init(X)
    if(wpar%impi/=0) call weiS_par_init(X)

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_init
  !----------------------------------------------------------------
  subroutine weiS_his_init(X)
    type(weslow), intent(inout) :: X
    character(len=128), parameter :: unit='weiS_his_init' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    call auxinout(X%tag)
    X%from = "history"

    call hinaxis(X%tag,X%amx,1); X%amx%l='mx'
    if(wpro%amx%n ==0) wpro%amx = X%amx

    call hinaxis(X%tag,X%amy,2); X%amy%l='my'
    if(wpro%amy%n ==0) wpro%amy = X%amy

    call hinaxis(X%tag,X%az ,3); X%az%l ='z_'
    if(wpro%az%n ==0) then
       wpro%az = X%az
    else
       if(.not. coax( X%az, wpro%az)) then
          call weiS_report(X)
          call seperr('S.az /= wpro.az')
       end if
    end if

    call from_param("velocity",X%velocity,.false.)
    call from_aux(trim(X%tag),"esize",X%esize,4)

    call weiS_blockcheck(X)
    call weiS_allocate(X)

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_his_init
  !----------------------------------------------------------------
  subroutine weiS_par_init(X)
    type(weslow), intent(inout) :: X
    character(len=128), parameter :: unit='weiS_par_init' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    call auxscr(X%tag)
    X%from = "parameter"

    X%amx = wpro%amx; 
    X%amy = wpro%amy;
    X%az  = wpro%az;
    X%esize=8
    X%velocity=.false.

    call houaxis(X%tag,X%amx,1)
    call houaxis(X%tag,X%amy,2)
    call houaxis(X%tag,X%az, 3)
    call to_history("esize",8,X%tag)

    call weiS_blockcheck(X)
    call weiS_allocate(X)
    call weiS_assign(X,(0.,0.))

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_par_init
  !----------------------------------------------------------------
  subroutine weiS_scr_init(X)
    type(weslow), intent(inout) :: X
    character(len=128), parameter :: unit='weiS_scr_init' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    call auxscr(X%tag)
    X%from = "parameter"

    X%amx = wpro%amx; call aouaxis(X%tag,X%amx,1)
    X%amy = wpro%amy; call aouaxis(X%tag,X%amy,2)
    X%az  = wpro%az;  call aouaxis(X%tag,X%az, 3)

    X%esize=8
    call auxputch("esize","i",8,X%tag)

    call weiS_blockcheck(X)
    call weiS_allocate(X)
    call weiS_assign(X,cmplx(0.))

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_scr_init
  !----------------------------------------------------------------
  subroutine weiS_blockcheck(X)
    type(weslow), intent(inout) :: X
    integer :: er

    if( (1.0*X%az%n)/wpro%nzs /= X%az%n/wpro%nzs ) then
       call weiS_report(X)
       call seperr('S.az.n is not a multiple of nzs')
    end if

  end subroutine weiS_blockcheck
  !----------------------------------------------------------------
  subroutine weiS_allocate(X)
    type(weslow), intent(inout) :: X
    integer :: er
    character(len=128), parameter :: unit='weiS_allocate' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    call weimem_add(X%esize,wpro%amx%n*wpro%amy%n*wpro%nzs,unit)
    !!--------------------
    allocate(X%bin(wpro%amx%n,wpro%amy%n,wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate S.bin')
    X%bin=(0.0,0.0)
    !!--------------------
    X%bites = X%esize * product( (/wpro%amx%n,wpro%amy%n,wpro%nzs/) )

    !! prestack binary array (PBA)
    ! OLD
    !    allocate(X%pba(wpro%amx%n,wpro%amy%n, wpro%ahx%n,wpro%ahy%n, wpro%nzs),stat=er)
    !    if(er/=0) call seperr('cant allocate X%pba')
    ! OLD

    call weimem_add(8,wpro%nzs,unit)
    allocate( sloline(wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate sloline')

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_allocate
  !----------------------------------------------------------------
  subroutine weiS_assign(X,v)
    type(weslow), intent(inout) :: X
    complex, intent(in) :: v
    integer :: izb
    character(len=128), parameter :: unit='weiS_assign' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    X%bin=v
    do izb=1,wpro%nzb
       call weiS_write(X,izb)
    end do

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_assign
  !----------------------------------------------------------------
  subroutine weiS_report(X)
    type(weslow), intent(in) :: X

    write(0,*)'        ************************************************'
    write(0,*)'        *       ',trim(X%tag),' from ',trim(X%from)
    write(0,'(a,i4,f10.2,f10.3)       ') "         *    Mx: ",X%amx%n, X%amx%o, X%amx%d
    write(0,'(a,i4,f10.2,f10.3)       ') "         *    My: ",X%amy%n, X%amy%o, X%amy%d
    write(0,'(a,i4,f10.2,f10.3,a,i3,a)') "         *     Z: ",X%az%n,  X%az%o,  X%az%d," ->",wpro%nzb," blocks"
    if(X%velocity) then
       write(0,'(a,i1,a)')               "         *    esize=",X%esize," velocity"
    else 
       write(0,'(a,i1,a)')               "         *    esize=",X%esize," slowness"
    end if
    if(associated(X%bin)) then
       write(0,'(a,f6.2,a,f6.2,a)        ') "         * block: ",X%bites/1024.," / ",8*size(X%bin)/1024.," Kbytes"
       write(0,*)'        ************************************************'
    end if

  end subroutine weiS_report
  !--------------------------------------------------------------
  subroutine weiS_read(X,k)
    type(weslow), intent(inout) :: X
    integer,      intent(in)    :: k
    integer :: stat
    real,allocatable  :: temp_buf(:,:,:)
    character(len=128), parameter :: unit='weiS_read' 
#ifdef DEBUG
    if(wpro%dbg>1) call in(unit)
#endif

    stat=sseek_block(X%tag,k-1,X%bites,0)
    if(X%esize==4) then
       call weimem_add(4,size(X%bin,1)*size(X%bin,2)*size(X%bin,3),unit)
       allocate(temp_buf(size(X%bin,1),size(X%bin,2),size(X%bin,3)),stat=stat)
       if(stat/=0) call erexit("cant allocate Re tmp slow")

       call sreed(X%tag, temp_buf, X%bites)
       if(X%velocity) temp_buf=1./temp_buf
       X%bin=cmplx(temp_buf,0.)

       call weimem_cut(4,size(X%bin,1)*size(X%bin,2)*size(X%bin,3),unit)
       deallocate(temp_buf)
    else
       call sreed(X%tag, X%bin, X%bites)
    end if

#ifdef DEBUG
    if(wpro%dbg>1) call out(unit)
#endif
  end subroutine weiS_read
  !--------------------------------------------------------------
  subroutine weiS_write(X,k)
    type(weslow), intent(inout) :: X
    integer,      intent(in)    :: k
    integer :: stat
    real,allocatable  :: temp_buf(:,:,:)
    character(len=128), parameter :: unit='weiS_write' 
#ifdef DEBUG
    if(wpro%dbg>1) call in(unit)
#endif

    stat=sseek_block(X%tag,k-1,X%bites,0)
    if(X%esize==4) then
       call weimem_add(4,size(X%bin,1)*size(X%bin,2)*size(X%bin,3),unit)
       allocate(temp_buf(size(X%bin,1),size(X%bin,2),size(X%bin,3)),stat=stat)
       if(stat/=0) call erexit("cant allocate Re tmp slow")

       temp_buf = real(X%bin)
       call srite(X%tag, temp_buf, X%bites)

       call weimem_cut(4,size(X%bin,1)*size(X%bin,2)*size(X%bin,3),unit)
       deallocate(temp_buf)
    else
       call srite(X%tag, X%bin, X%bites)
    end if

#ifdef DEBUG
    if(wpro%dbg>1) call out(unit)
#endif
  end subroutine weiS_write
  !----------------------------------------------------------------
  subroutine weiS_fkfx_allocate()
    integer :: er
    character(len=128), parameter :: unit='weiS_allocate' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    !! ref slowness (FKop)
    call weimem_add(4,wpro%nzs*(wpro%nfk+1),unit)
    allocate(slo%rs_s(wpro%nzs,wpro%nfk+1),stat=er) !! +1 for salt velocity
    if(er/=0) call seperr('cant allocate slo%rs_s')
    call weimem_add(4,wpro%nzs*(wpro%nfk+1),unit)
    allocate(slo%rs_r(wpro%nzs,wpro%nfk+1),stat=er) !! +1 for salt velocity
    if(er/=0) call seperr('cant allocate slo%rs_r')

    !! no ref slowness (FKop)
    call weimem_add(1,wpro%nzs,unit)
    allocate(slo%nr_s(wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%nr_s')
    call weimem_add(1,wpro%nzs,unit)
    allocate(slo%nr_r(wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%nr_r')

    !! interpolation maps
    call weimem_add(1,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%si_s(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)   
    if(er/=0) call seperr('cant allocate slo%si_s')
    call weimem_add(1,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%si_r(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)   
    if(er/=0) call seperr('cant allocate slo%si_r')

    !! prestack slowness
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%ss_s(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%ss_s')
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%ss_r(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%ss_r')

    !! prestack slowness
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%sw_s(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%sw_s')
    call weimem_add(4,wpro%amx%n*wpro%amy%n*wpro%ahx%n*wpro%ahy%n*wpro%nzs,unit)
    allocate(slo%sw_r(wpro%amx%n,wpro%amy%n,wpro%ahx%n,wpro%ahy%n,wpro%nzs),stat=er)
    if(er/=0) call seperr('cant allocate slo%sw_r')

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_fkfx_allocate
  !----------------------------------------------------------------
  subroutine weiS_interp(adj,X,line,xc,yc)
    type(weslow), intent(inout) :: X
    logical,                   intent(in) :: adj
    complex, dimension(:,:,:), pointer    :: cube
    real,    dimension(:),     pointer    :: line
    real,                      intent(in) :: xc,yc
    real    :: dx,dy
    integer :: ix,iy, jx,jy
    real    :: fx,fy
    real    :: w00,w01,w10,w11, sumw

    cube=>X%bin

    dx = xc
    dy = yc

    !! don't allow coordinates outside the slowness grid
    dx = min(dx,X%amx%o + (X%amx%n-1)* X%amx%d)
    dx = max(dx,X%amx%o)

    dy = min(dy,X%amy%o + (X%amy%n-1)* X%amy%d)
    dy = max(dy,X%amy%o)

    !! lower index
    ix = floor( 1+ (dx-X%amx%o) / X%amx%d )
    iy = floor( 1+ (dy-X%amy%o) / X%amy%d )

    !! keep indices in bounds
    ix = min(ix,X%amx%n)
    iy = min(iy,X%amy%n)
    ix = max(ix,        1  ); jx=min(ix+1,X%amx%n)
    iy = max(iy,        1  ); jy=min(iy+1,X%amy%n)

    !! linear weights
    fx = (dx - (X%amx%o + (ix-1)*X%amx%d) ) / X%amx%d
    fy = (dy - (X%amy%o + (iy-1)*X%amy%d) ) / X%amy%d
    fx = max(0., fx)
    fy = max(0., fy)

    !! 00   10       --> x
    !!    .         |
    !! 01   11      v
    !!               y
    w00=(1.-fx)*(1.-fy)
    w10=(   fx)*(1.-fy)
    w01=(1.-fx)*(   fy)
    w11=(   fy)*(   fy)

    !! normalize the sum of the weights
    sumw = w00+w10+w01+w11
    w00 = w00 / sumw
    w10 = w10 / sumw
    w01 = w01 / sumw
    w11 = w11 / sumw

    if(adj) then !! push
       cube(ix,iy,:) = cube(ix,iy,:) + w00 * line
       cube(jx,iy,:) = cube(jx,iy,:) + w10 * line
       cube(ix,jy,:) = cube(ix,jy,:) + w01 * line
       cube(jx,jy,:) = cube(jx,jy,:) + w11 * line
    else         !! pull
       line    = &
       w00*real(cube(ix,iy,:)) + &
       w10*real(cube(jx,iy,:)) + &
       w01*real(cube(ix,jy,:)) + &
       w11*real(cube(jx,jy,:))
    end if

  end subroutine weiS_interp
  !----------------------------------------------------------------
  !! pull slowness from ss and sr into bin
  subroutine weiS_pull(X)
    type(weslow), intent(inout) :: X
    integer :: imx,imy,ihx,ihy
    real    :: xc,yc
    character(len=128), parameter :: unit='weiS_pull' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    X%bin=(0.,0.)

    do ihy=1,wpro%ahy%n
       do ihx=1,wpro%ahx%n
          do imy=1,wpro%amy%n
             do imx=1,wpro%amx%n

                !! X%bin = f(X%ss,X%sr)
                xc=xmap%sx(imx,ihx)
                yc=xmap%sy(imy,ihy)
                sloline = slo%ss_s(imx,imy,ihx,ihy,:)
                call weiS_interp(.true.,X,sloline,xc,yc)

                xc=xmap%rx(imx,ihx)
                yc=xmap%ry(imy,ihy)
                sloline = slo%ss_r(imx,imy,ihx,ihy,:)
                call weiS_interp(.true.,X,sloline,xc,yc)

             end do
          end do
       end do
    end do

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_pull
  !----------------------------------------------------------------
  !! push slowness from bin to ss and sr
  subroutine weiS_push(X)
    type(weslow), intent(inout) :: X
    integer :: imx,imy,ihx,ihy
    real    :: xc,yc
    character(len=128), parameter :: unit='weiS_push' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    slo%ss_s = 0.
    slo%ss_r = 0.

    do ihy=1,wpro%ahy%n
       do ihx=1,wpro%ahx%n
          do imy=1,wpro%amy%n
             do imx=1,wpro%amx%n

                !! X%ss = f(X%bin,x,y)
                xc=xmap%sx(imx,ihx)
                yc=xmap%sy(imy,ihy)
                call weiS_interp(.false.,X,sloline,xc,yc)

                slo%ss_s(imx,imy,ihx,ihy,:) = &
                slo%ss_s(imx,imy,ihx,ihy,:) + sloline

                !----------------------------------------------------------------

                !! X%sr = f(X%bin,x,y)
                xc=xmap%rx(imx,ihx)
                yc=xmap%ry(imy,ihy)
                call weiS_interp(.false.,X,sloline,xc,yc)

                slo%ss_r(imx,imy,ihx,ihy,:) = &
                slo%ss_r(imx,imy,ihx,ihy,:) + sloline

             end do
          end do
       end do
    end do

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_push
  !----------------------------------------------------------------
  real function weiS_mean(X,k) result(smean)
    type(weslow), intent(inout) :: X
    integer :: k

    smean = sum(real(X%bin(:,:,k))) / (X%amx%n * X%amy%n)
  end function weiS_mean
  !----------------------------------------------------------------
  real function weiS_max(X,k) result(smax)
    type(weslow), intent(inout) :: X
    integer :: k

    smax = maxval(real(X%bin(:,:,k)))
  end function weiS_max
  !----------------------------------------------------------------
  real function weiS_min(X,k) result(smin)
    type(weslow), intent(inout) :: X
    integer :: k

    smin = minval(real(X%bin(:,:,k)))
  end function weiS_min
  !----------------------------------------------------------------
  subroutine weiS_mpicopy(X)
    type(weslow), intent(inout) :: X
    integer :: mpierr, msglen, msgsou
    integer :: izb
    character(len=128), parameter :: unit='weiS_mpicopy' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    msglen = size(X%bin,1)*size(X%bin,2)*size(X%bin,3)
    msgsou=0
    
    do izb=1,wpro%nzb
       if(wpar%impi==0) write(0,'(a,a,i3)') weitab,"S send izb=",izb
       if(wpar%impi==0) call weiS_read(X,izb)
       call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
       call MPI_BCAST(X%bin,msglen,MPI_COMPLEX,msgsou,MPI_COMM_WORLD,mpierr)
       if(wpar%impi/=0) call weiS_write(X,izb)
    end do
    
#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_mpicopy
  !----------------------------------------------------------------  
  subroutine weiS_mpisumm(X)
    type(weslow), intent(inout) :: X
    integer :: mpierr, msglen, msgrec
    integer :: izb
    complex, pointer :: tmp(:,:,:)
    character(len=128), parameter :: unit='weiS_mpisumm' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

        msglen = size(X%bin,1)*size(X%bin,2)*size(X%bin,3)
    msgrec=0

    if(wpar%impi==0) then
       call weimem_add(8,wpro%amx%n*wpro%amy%n*wpro%az%n,unit)
       allocate(tmp(wpro%amx%n,wpro%amy%n,wpro%az%n))
    else
       call weimem_add(8,1,unit)
       allocate(tmp(1,1,1))
    end if

    do izb=1,wpro%nzb
       if(wpar%impi==0) write(0,'(a,a,i3)') weitab,"S recv izb=",izb
       call weiS_read(X,izb)
       call MPI_BARRIER(MPI_COMM_WORLD,mpierr)
       call MPI_REDUCE(X%bin,tmp,msglen,MPI_COMPLEX,MPI_SUM,msgrec,MPI_COMM_WORLD,mpierr)
       if(wpar%impi==0) X%bin = tmp
       if(wpar%impi==0) call weiS_write(X,izb)
    end do

    if(wpar%impi==0) then
       call weimem_cut(8,wpro%amx%n*wpro%amy%n*wpro%az%n,unit)
       deallocate(tmp)
    else
       call weimem_cut(8,1,unit)
       deallocate(tmp)
    end if

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_mpisumm
  !----------------------------------------------------------------
  subroutine weiS_mpizero(X)
    type(weslow), intent(inout) :: X
    character(len=128), parameter :: unit='weiS_mpizero' 
#ifdef DEBUG
    if(wpro%dbg>0) call in(unit)
#endif

    if(wpar%impi/=0) call weiS_assign(X,cmplx(0.))

#ifdef DEBUG
    if(wpro%dbg>0) call out(unit)
#endif
  end subroutine weiS_mpizero
  !----------------------------------------------------------------
end module wei_S
