module wei_global_calc_mod
  use  wei_global_defi_mod
  use sep
  implicit none

contains

!=================================================================================
! calculating the global parameters, including the propagation grid, 
! image grid and the output image grid
!---------------------------------------------------------------------------------
  subroutine wei_global_calc(ish)
    integer                                 :: ish
    integer                                 :: ishx 
    integer                                 :: ishy 
    real, allocatable, dimension(:,:), save :: dgrid_header_rec_save
    real                             , save :: dgrid_gxmin_save
    real                             , save :: dgrid_gymin_save
    write(0,*)"in wei_global_calc"

    !!.. locate the shot location
    !!.. if the source header is presented, read the source
    !!.. location from the header
    if(dgrid%hdr_sou) then

       dgrid%sx = dgrid%header_sou(1)
       dgrid%sy = dgrid%header_sou(2)

    !!.. if the source header is not provided, calculate from the history file
    else

       ishx = mod(ish-1, dgrid%asx%n)+1
       ishy = int(1. + real(ish-1)/real(dgrid%asx%n))
       dgrid%sx = dgrid%asx%o + (ishx-1)*dgrid%asx%d
       dgrid%sy = dgrid%asy%o + (ishy-1)*dgrid%asy%d

    end if

    !!.. if the input receiver coordinate is in offset space
    if (offset) then
       if (ish == 1) then
           if (dgrid%hdr_rec) then
              allocate(dgrid_header_rec_save(2, dgrid%ntr))
              dgrid_header_rec_save(1,:) = dgrid%header_rec(1,:)
              dgrid_header_rec_save(2,:) = dgrid%header_rec(2,:) 
           end if
           dgrid_gxmin_save = dgrid%agx%o
           dgrid_gymin_save = dgrid%agy%o
       end if

       if (dgrid%hdr_rec) then
          dgrid%header_rec(1,:) = dgrid_header_rec_save(1,:) + dgrid%sx
          dgrid%header_rec(2,:) = dgrid_header_rec_save(2,:) + dgrid%sy
       else
          dgrid%agx%o = dgrid_gxmin_save + dgrid%sx
          dgrid%agy%o = dgrid_gymin_save + dgrid%sy 
       end if
    end if

    !!.. define the propagation grid
    call wei_propagation_grid()

    !!.. define the input velocity grid, vgrid defines the entire velocity
    !!.. but the input size of velocity is equal to pgrid
    call wei_velocity_grid()

    !!.. compute the x and y index of the output image on the current propagation grid
    call wei_imgout_grid()

    !!.. compute the wavenumber parameters
    if(associated(kx2)) deallocate(kx2)
    if(associated(ky2)) deallocate(ky2)
    allocate(kx2(pgrid%ax%n))
    allocate(ky2(pgrid%ay%n))
    call k2_map(kx2, pgrid%ax%n, pgrid%ax%d)
    call k2_map(ky2, pgrid%ay%n, pgrid%ay%d)

  end subroutine
!=================================================================================
!.. calculate parameters for the wavefield propagation grid
!---------------------------------------------------------------------------------
  subroutine wei_propagation_grid()
    real                :: pgrid_xmin, pgrid_ymin
    real                :: pgrid_xmax, pgrid_ymax
    integer             :: itr
    !!.. the spacing of the propagation grid is assumed to
    !!.. be the same as the receiver spacing
    pgrid%ax%d = dgrid%agx%d * pgrid%jx
    pgrid%ay%d = dgrid%agy%d * pgrid%jy

    !!.. the source location for a shot gather
    pgrid%sx = dgrid%sx
    pgrid%sy = dgrid%sy

    write(0,*)"pgrid:", pgrid%sx, pgrid%sy


    !!.. the minimum and maximum coordinates for the recorded receiver wavefield
    !!.. for the case of unsorted data
    if (dgrid%hdr_rec) then
       pgrid%gxmin = minval(dgrid%header_rec(1, :))
       pgrid%gymin = minval(dgrid%header_rec(2, :))
       pgrid%gxmax = maxval(dgrid%header_rec(1, :))
       pgrid%gymax = maxval(dgrid%header_rec(2, :))
   !!.. for the case of regular data
    else
       pgrid%gxmin = dgrid%agx%o
       pgrid%gymin = dgrid%agy%o
       pgrid%gxmax = dgrid%agx%o + (dgrid%agx%n-1)*dgrid%agx%d
       pgrid%gymax = dgrid%agy%o + (dgrid%agy%n-1)*dgrid%agy%d
    end if
    !!.. the coordinate for the depth direction
    pgrid%az = vgrid%az

    !!.. compute the migration aperture
    !!.. for simplicity, currently I am just using the maximum distance
    !!.. between the source and the receiver (better make it adaptive)
    pgrid_xmin = pgrid%gxmin; pgrid_xmax = pgrid%gxmax
    pgrid_ymin = pgrid%gymin; pgrid_ymax = pgrid%gymax

    if (pgrid%sx .ge. pgrid%gxmax) pgrid_xmax = pgrid%sx
    if (pgrid%sy .ge. pgrid%gymax) pgrid_ymax = pgrid%sy

    if (pgrid%sx .le. pgrid%gxmin) pgrid_xmin = pgrid%sx
    if (pgrid%sy .le. pgrid%gymin) pgrid_ymin = pgrid%sy

    !!.. add the padding on both x and y direction
    pgrid_xmin = pgrid_xmin - pgrid%nxpad_beg*pgrid%ax%d
    pgrid_xmax = pgrid_xmax + pgrid%nxpad_end*pgrid%ax%d

    pgrid_ymin = pgrid_ymin - pgrid%nypad_beg*pgrid%ay%d
    pgrid_ymax = pgrid_ymax + pgrid%nypad_end*pgrid%ay%d

    !!.. make sure the propagation grid is within the velocity model 
    pgrid_xmin = max(pgrid_xmin, vgrid%ax%o)
    pgrid_xmax = min(pgrid_xmax, vgrid%ax%o + (vgrid%ax%n-1)*vgrid%ax%d)

    pgrid_ymin = max(pgrid_ymin, vgrid%ay%o)
    pgrid_ymax = min(pgrid_ymax, vgrid%ay%o + (vgrid%ay%n-1)*vgrid%ay%d)

    !!.. propagation grid 
    pgrid%ax%o = pgrid_xmin
    pgrid%ay%o = pgrid_ymin
    pgrid%ax%n = 1.5 + (pgrid_xmax - pgrid_xmin)/pgrid%ax%d
    pgrid%ay%n = 1.5 + (pgrid_ymax - pgrid_ymin)/pgrid%ay%d

    !!.. determine the source location on the propagation grid
    pgrid%index_datpro_sx = 1.5 + (pgrid%sx - pgrid%ax%o) / pgrid%ax%d
    pgrid%index_datpro_sy = 1.5 + (pgrid%sy - pgrid%ay%o) / pgrid%ay%d

    !!.. determine the index of receivers on the propagation grid
    !!.. for the case of unsorted data, all indexes should be calculated to 
    !!.. map the receivers onto the propagation grid
    if (dgrid%hdr_rec) then
       do itr=1, dgrid%ntr
          pgrid%index_datpro_hrec(1, itr) = 1.5 + (dgrid%header_rec(1, itr) - pgrid%ax%o) / pgrid%ax%d
          pgrid%index_datpro_hrec(2, itr) = 1.5 + (dgrid%header_rec(2, itr) - pgrid%ay%o) / pgrid%ay%d
       end do
    !!.. for the case of regular data, only the first receiver location on the 
    !!.. propagation grid need to be calculated
    else
       pgrid%index_datpro_gxmin = 1.5 + (dgrid%agx%o - pgrid%ax%o) / pgrid%ax%d
       pgrid%index_datpro_gymin = 1.5 + (dgrid%agy%o - pgrid%ay%o) / pgrid%ay%d
    end if

write(0,*)"in propagation grid", pgrid%index_datpro_sx, pgrid%index_datpro_sy, pgrid%ax%d, pgrid%ay%d

  end subroutine
!====================================================================================
!.. calculate the output image grid
!------------------------------------------------------------------------------------
  subroutine wei_imgout_grid()

    real  :: aix_max, aiy_max
    real  :: aix_min, aiy_min
    real  :: apx_max, apy_max
    real  :: apx_min, apy_min
    real  :: aix_length, aiy_length
    real  :: apx_length, apy_length

    ogrid%ax  = igrid%ax
    ogrid%ay  = igrid%ay
    ogrid%ahx = igrid%ahx
    ogrid%ahy = igrid%ahy
    ogrid%az  = igrid%az

    aix_min = igrid%ax%o
    apx_min = pgrid%ax%o
    aix_max = igrid%ax%o + (igrid%ax%n-1)*igrid%ax%d
    apx_max = pgrid%ax%o + (pgrid%ax%n-1)*pgrid%ax%d

    aiy_min = igrid%ay%o
    apy_min = pgrid%ay%o
    aiy_max = igrid%ay%o + (igrid%ay%n-1)*igrid%ay%d
    apy_max = pgrid%ay%o + (pgrid%ay%n-1)*pgrid%ay%d


    if (pgrid%ax%n > 1) then
       if(apx_min < aix_min) then
          if(apx_max < aix_min) then
             ogrid%ax%n = 0
          else if(apx_max >= aix_min .and. apx_max < aix_max) then
             ogrid%ax%o = igrid%ax%o
             ogrid%ax%n = 1.5 + (apx_max-ogrid%ax%o)/ogrid%ax%d
          else
             ogrid%ax%o = igrid%ax%o
             ogrid%ax%n = 1.5 + (aix_max-ogrid%ax%o)/ogrid%ax%d
          end if
       else if(apx_min >= aix_min .and. apx_min <= aix_max) then
          if(apx_max < aix_max) then
             ogrid%ax%o = pgrid%ax%o
             ogrid%ax%n = 1.5 + (apx_max-ogrid%ax%o)/ogrid%ax%d
          else 
             ogrid%ax%o = pgrid%ax%o
             ogrid%ax%n = 1.5 + (aix_max-ogrid%ax%o)/ogrid%ax%d
          end if
       else
          ogrid%ax%n = 0
       end if
    end if

    if (pgrid%ay%n > 1) then
       if(apy_min < aiy_min) then
          if(apy_max < aiy_min) then
             ogrid%ay%n = 0
          else if(apy_max >= aiy_min .and. apy_max < aiy_max) then
             ogrid%ay%o = igrid%ay%o
             ogrid%ay%n = 1.5 + (apy_max-ogrid%ay%o)/ogrid%ay%d
          else
             ogrid%ay%o = igrid%ay%o
             ogrid%ay%n = 1.5 + (aiy_max-ogrid%ay%o)/ogrid%ay%d
          end if
       else if(apy_min >= aiy_min .and. apy_min <= aiy_max) then
          if(apy_max < aiy_max) then
             ogrid%ay%o = pgrid%ay%o
             ogrid%ay%n = 1.5 + (apy_max-ogrid%ay%o)/ogrid%ay%d
          else
             ogrid%ay%o = pgrid%ay%o
             ogrid%ay%n = 1.5 + (aiy_max-ogrid%ay%o)/ogrid%ay%d
          end if
       else
          ogrid%ay%n = 0
       end if
    end if


    !!.. computing the first and the last x, y index of the output image on the final image grid
    ogrid%index_outimg_xmin = 1.5 + (ogrid%ax%o - igrid%ax%o)/igrid%ax%d
    ogrid%index_outimg_ymin = 1.5 + (ogrid%ay%o - igrid%ay%o)/igrid%ay%d
    ogrid%index_outimg_xmax = ogrid%index_outimg_xmin + ogrid%ax%n - 1
    ogrid%index_outimg_ymax = ogrid%index_outimg_ymin + ogrid%ay%n - 1

    !!.. computing the first and the last x, y index of the output image on the propagation grid
    ogrid%index_outpro_xmin = 1.5 + (ogrid%ax%o - pgrid%ax%o)/pgrid%ax%d
    ogrid%index_outpro_ymin = 1.5 + (ogrid%ay%o - pgrid%ay%o)/pgrid%ay%d
    ogrid%index_outpro_xmax = ogrid%index_outpro_xmin + ogrid%ax%n - 1
    ogrid%index_outpro_ymax = ogrid%index_outpro_ymin + ogrid%ay%n - 1

  end subroutine
!===================================================================================
  subroutine wei_velocity_grid()
    real  :: vgrid_xmin, vgrid_xmax
    real  :: vgrid_ymin, vgrid_ymax
    real  :: pgrid_xmin, pgrid_xmax
    real  :: pgrid_ymin, pgrid_ymax

    vgrid_xmin = vgrid%ax%o
    vgrid_ymin = vgrid%ay%o
    vgrid_xmax = vgrid%ax%o + (vgrid%ax%n-1)*vgrid%ax%d
    vgrid_ymax = vgrid%ay%o + (vgrid%ay%n-1)*vgrid%ay%d

    pgrid_xmin = pgrid%ax%o
    pgrid_ymin = pgrid%ay%o
    pgrid_xmax = pgrid%ax%o + (pgrid%ax%n-1)*pgrid%ax%d
    pgrid_ymax = pgrid%ay%o + (pgrid%ay%n-1)*pgrid%ay%d   

    if(pgrid_xmin < vgrid_xmin) then
       write(0,*)"pgrid_xmin, vgrid_xmin", pgrid_xmin, vgrid_xmin
       call seperr("propagation grid out of velocity grid: pgrid_xmin < vgrid_xmin")
    end if
    if(pgrid_xmax > vgrid_xmax) then
       write(0,*)"pgrid_xmax, vgrid_xmax", pgrid_xmax, vgrid_xmax
       call seperr("propagation grid out of velocity grid: pgrid_xmax > vgrid_xmax")
    end if
    if(pgrid_ymin < vgrid_ymin) then
       write(0,*)"pgrid_ymin, vgrid_ymin", pgrid_ymin, vgrid_ymin
       call seperr("propagation grid out of velocity grid: pgrid_ymin < vgrid_ymin")
    end if
    if(pgrid_ymax > vgrid_ymax) then
       write(0,*)"pgrid_ymax, vgrid_ymax", pgrid_ymax, vgrid_ymax
       call seperr("propagation grid out of velocity grid: pgrid_ymax > vgrid_ymax")
    end if

    !!.. computing the first and the last x,y index of the input velocity on the whole velocity grid
    vgrid%index_provel_xmin = 1.5 + (pgrid%ax%o - vgrid%ax%o)/vgrid%ax%d
    vgrid%index_provel_ymin = 1.5 + (pgrid%ay%o - vgrid%ay%o)/vgrid%ay%d
    vgrid%index_provel_xmax = vgrid%index_provel_xmin + vgrid%ax%n - 1
    vgrid%index_provel_ymax = vgrid%index_provel_ymin + vgrid%ay%n - 1

  end subroutine 
!===================================================================================
  subroutine wei_block_calc()
    integer  :: iwb, izb

    wblock%bres = .false.
    zblock%bres = .false.

    wblock%ds = dgrid%aw%d
    zblock%ds = vgrid%az%d

    wblock%nb = dgrid%aw%n / wblock%ns
    zblock%nb = vgrid%az%n / zblock%ns

    wblock%nb_res = mod(dgrid%aw%n, wblock%ns)
    zblock%nb_res = mod(vgrid%az%n, zblock%ns)

    if(wblock%nb_res /= 0) then
       wblock%nb = wblock%nb + 1
       wblock%bres = .true.
    end if

    if(zblock%nb_res /= 0) then
       zblock%nb = zblock%nb + 1
       zblock%bres = .true.
    end if

    allocate(wblock%blen(wblock%nb), wblock%os(wblock%nb), wblock%indx(wblock%nb))
    allocate(zblock%blen(zblock%nb), zblock%os(zblock%nb), zblock%indx(zblock%nb))

    wblock%blen = wblock%ns
    zblock%blen = zblock%ns

    if(wblock%bres) wblock%blen(wblock%nb) = wblock%nb_res
    if(zblock%bres) zblock%blen(zblock%nb) = zblock%nb_res

    do iwb=1, wblock%nb
       wblock%os(iwb)   = dgrid%aw%o + (iwb-1)*wblock%ns*dgrid%aw%d
       wblock%indx(iwb) = 1 + (iwb-1)*wblock%ns
    end do

    do izb=1, zblock%nb
       zblock%os(izb)   = vgrid%az%o + (izb-1)*zblock%ns*vgrid%az%d
       zblock%indx(izb) = 1 + (izb-1)*zblock%ns
    end do
write(0,*)"zblock%indx", zblock%indx
write(0,*)"zblock%nb", zblock%nb
  end subroutine
!===================================================================================
  subroutine wei_radians_calc()
    !!.. convert frequencies to radians
    if (.not.radian) then
       dgrid%aw%d  = dgrid%aw%d * 2.0 * PI
       dgrid%aw%o  = dgrid%aw%o * 2.0 * PI
    end if

  end subroutine
!===================================================================================
  subroutine wei_counter_init()
    nstep = vgrid%az%n * dgrid%aw%n
    rstep = nstep / num_report
    istep = 0
    astep = 1
  end subroutine
!===================================================================================
  subroutine wei_wfld_alloc(nws_in)
    integer  :: nws_in
    allocate(csou(nws_in))
    allocate(crec(dgrid%ntr, nws_in))
    allocate(wfld_d(pgrid%ax%n, pgrid%ay%n, nws_in))
    allocate(wfld_u(pgrid%ax%n, pgrid%ay%n, nws_in))

    csou = 0.
    crec = 0.
    wfld_d = 0.
    wfld_u = 0.

    memory = 0
    memory = memory + 8*(size(csou)+size(crec)+size(wfld_d)+size(wfld_u))
  end subroutine
!===================================================================================
  subroutine wei_modl_alloc(nzs_in)
    integer  :: nzs_in
    allocate(model(ogrid%ax%n, ogrid%ay%n, ogrid%ahx%n, ogrid%ahy%n, nzs_in))
    allocate(vmod(pgrid%ax%n, pgrid%ay%n, nzs_in))
    allocate(vref(nref, nzs_in))
    model = 0.
    vmod  = 0.
    vref  = 0.
    memory = memory + 4*(size(model)+size(vmod)+size(vref))

    if (hessian_type > 0) then
       allocate(normd(ogrid%ax%n, ogrid%ay%n, ogrid%ahx%n, ogrid%ahy%n, nzs_in)) 
       normd = 0.
       memory = memory + 4*size(normd)
    end if
  end subroutine
!===================================================================================
  subroutine wei_wfld_dealloc()
    deallocate(csou)
    deallocate(crec)
    deallocate(wfld_d)
    deallocate(wfld_u)
  end subroutine
!===================================================================================
  subroutine wei_modl_dealloc()
    deallocate(model)
    deallocate(vmod)
    deallocate(vref)
    if (hessian_type > 0) deallocate(normd)
  end subroutine
!===================================================================================
!.. initialize downward and upward wavefields according to the propagation grid
!-----------------------------------------------------------------------------------
  subroutine wei_wfld_d_init(nws_in)
    integer  :: nws_in
    integer  :: iws, itr, i1, i2, ind1, ind2
    real, dimension(:,:), pointer :: random

!    allocate(random(nws_in, 2))
!    call sep_read(random, "phase_random")

!write(0,*)"nws_in", nws_in
!write(0,*)"index", pgrid%index_datpro_sx 
    wfld_d = 0.
    !!.. locate the source wavefield onto the propagation grid 
    do iws=1, nws_in
       wfld_d(pgrid%index_datpro_sx, pgrid%index_datpro_sy, iws) = csou(iws)!*cmplx(cos(random(iws,1)), sin(random(iws,1)))
    end do

!    deallocate(random)
  end subroutine

!====================================================================================
!.. forward: wfld_u => crec      in case of demigration
!.. adjoint: crec   => wfld_u    in case of migration
!------------------------------------------------------------------------------------
  subroutine wei_wfld_u_init(adj, add, nws_in)
    logical  :: adj, add
    integer  :: nws_in
    integer  :: iws, itr, i1, i2, ind1, ind2
    real     :: gx, gy

    if (.not. add) then
       if (adj) then
          wfld_u = 0.
       else
          crec = 0.
       end if
    end if 

    if (dgrid%hdr_rec) then
       do itr=1, dgrid%ntr
          i1 = pgrid%index_datpro_hrec(1,itr)
          i2 = pgrid%index_datpro_hrec(2,itr) 
          do iws=1, nws_in
             if (adj) then
                wfld_u(i1, i2, iws) = wfld_u(i1, i2, iws) + crec(itr, iws)
             else
                crec(itr, iws) = crec(itr, iws) + wfld_u(i1, i2, iws)
             end if
          end do
       end do
    else

       do iws=1, nws_in
          do i2=1, dgrid%agy%n
             gy = dgrid%agy%o + (i2-1)*dgrid%agy%d
             ind2 = 1.5 + (gy - pgrid%ay%o)/pgrid%ay%d
             do i1=1, dgrid%agx%n
                gx = dgrid%agx%o + (i1-1)*dgrid%agx%d
                ind1 = 1.5 + (gx - pgrid%ax%o)/pgrid%ax%d
                itr  = i1 + (i2-1)*dgrid%agx%n
                if (adj) then
                   wfld_u(ind1, ind2, iws) = wfld_u(ind1, ind2, iws) + crec(itr, iws)
                else
                   crec(itr, iws) = crec(itr, iws) + wfld_u(ind1, ind2, iws)
                end if  
             end do
          end do
       end do


    end if

    
  end subroutine

!============================================================================
!
  subroutine k2_map(k,n,d)
    real,  dimension(:) :: k
    integer, intent(in) :: n
    real,    intent(in) :: d
    integer             :: i,m
    k=0.
    if(n/=1) then
       m = n/2
       do i=1,m+1
          k(i)=i-1
       end do
       do i=m+2,n
          k(i)=i-n-1
       end do
    else
       k=0. !! make sure k=0 if n=1
    end if

    k = k / (n*d)      !! 1/length
    k = 2 * PI * k     !! radians
    k = k * k          !! k squared
  end subroutine 


end module






