module wei_cig_mod

  use wei_type_mod

  implicit none

  integer, private :: nmx, npx
  integer, private :: nmy, npy
  integer, private :: nhx
  integer, private :: nhy

  integer, private :: indx
  integer, private :: indy
  
  real   , private :: dpx, dmx 
  real   , private :: dpy, dmy
  real   , private :: dhx
  real   , private :: dhy

  real   , private :: ohx, opx, opy
  real   , private :: ohy, omx, omy

  real   , private :: gx_min, gy_min
  real   , private :: gx_max, gy_max

  logical, private :: decon
  real   , private :: eps
  integer, private :: hessian_type

  real, parameter, private :: PI = 3.14159265359
contains
!===========================================================================
  integer function wei_cig_init(decon_in, eps_in, hessian_type_in, ogrid_in, pgrid_in) result(stat)
    logical               :: decon_in
    real                  :: eps_in
    integer               :: hessian_type_in
    type(out_grid)        :: ogrid_in
    type(pro_grid)        :: pgrid_in 
    

    nmx = ogrid_in%ax%n ; dmx = ogrid_in%ax%d ; omx = ogrid_in%ax%o
    nmy = ogrid_in%ay%n ; dmy = ogrid_in%ay%d ; omy = ogrid_in%ay%o
    nhx = ogrid_in%ahx%n; dhx = ogrid_in%ahx%d; ohx = ogrid_in%ahx%o
    nhy = ogrid_in%ahy%n; dhy = ogrid_in%ahy%d; ohy = ogrid_in%ahy%o
    npx = pgrid_in%ax%n ; dpx = pgrid_in%ax%d ; opx = pgrid_in%ax%o
    npy = pgrid_in%ay%n ; dpy = pgrid_in%ay%d ; opy = pgrid_in%ay%o

    indx = ogrid_in%index_outpro_xmin
    indy = ogrid_in%index_outpro_ymin

    decon   = decon_in
    eps     = eps_in
    hessian_type = hessian_type_in

    gx_min = pgrid_in%gxmin
    gx_max = pgrid_in%gxmax

    gy_min = pgrid_in%gymin
    gy_max = pgrid_in%gymax

write(0,*)"init", gx_min, gx_max, gy_min, gy_max
    stat = 0
  end function
!=============================================================================
  integer function wei_cig_clos() result(stat)
    nmx = 0 ; dmx = 0.; omx = 0.
    nmy = 0 ; dmy = 0.; omy = 0.
    nhx = 0 ; dhx = 0.; ohx = 0.
    nhy = 0 ; dhy = 0.; ohy = 0.
    dpx = 0.; dpy = 0.
    indx = 0
    indy = 0
    stat = 0
    decon = .false.
    hessian_type = 0
  end function
!=============================================================================
  integer function wei_cig_oper(adj, add, ws, zs, wfld_slice_d, wfld_slice_u, model_slice, normd_slice) result(stat)
    optional                                            :: normd_slice
    logical,                              intent(in)    :: adj, add
    real                                                :: ws, zs
    complex, dimension(:,:)    , pointer, intent(inout) :: wfld_slice_d
    complex, dimension(:,:)    , pointer, intent(inout) :: wfld_slice_u
    real   , dimension(:,:,:,:), pointer, intent(inout) :: model_slice
    real   , dimension(:,:,:,:), pointer, intent(inout) :: normd_slice
    integer :: ihx, ihy, iws, imx, imy
    integer :: jhx, jhy, jxd, jxu, jyd, jyu
    real    :: hx, hy
    real    :: denom
    complex :: phase_shaping


    integer :: ipx, ipy
    real    :: wt, summ, dis
    real    :: mx, my, px, py, xd, yd
    stat = 1
    wt = 1.
    if (.not.add) then
       if (adj) then
          model_slice = 0.
          if (present(normd_slice)) normd_slice = 0.
       else
          wfld_slice_u = 0.
       end if
    end if
    
    if (nmy == 1 .or. nmx ==1) then
       !!.. in case of 2D migration, 45 degree rotation i^(-1/2)
       phase_shaping = cmplx(cos(PI/4.), -sin(PI/4.)) / sqrt(ws)           
    else
       !!.. in case of 3D migration, 90 degree rotation i^(-1)
       phase_shaping = cmplx(0., - 1. * ws)       
    end if


    do ihy=1, nhy
       hy = ohy + (ihy-1)*dhy
       jhy = nint(hy / dpy)

       do ihx=1, nhx
          hx = ohx + (ihx-1)*dhx
          jhx = nint(hx / dpx)

          do imy=1, nmy
             jyd = (imy+indy-1) - jhy
             jyu = (imy+indy-1) + jhy
             if (jyd .lt.1 .or. jyd.gt.size(wfld_slice_d,2) .or. &
                 jyu .lt.1 .or. jyu.gt.size(wfld_slice_u,2)) cycle
             
             !jyd = max(1, min((imy+indy-1) - jhy, size(wfld_slice_d, 2)))
             !jyu = max(1, min((imy+indy-1) + jhy, size(wfld_slice_u, 2)))

             do imx=1, nmx
                jxd = (imx+indx-1) - jhx
                jxu = (imx+indx-1) + jhx
                if (jxd .lt.1 .or. jxd.gt.size(wfld_slice_d,1) .or. &
                    jxu .lt.1 .or. jxu.gt.size(wfld_slice_u,1)) cycle

                !jxd = max(1, min((imx+indx-1) - jhx, size(wfld_slice_d, 1)))
                !jxu = max(1, min((imx+indx-1) + jhx, size(wfld_slice_u, 1)))

                !!.. if migration, the output is model_slice
                if (adj) then
                   if (decon) then
                      denom = cabs(wfld_slice_d(jxd, jyd))
                      model_slice(imx, imy, ihx, ihy)          &
                      = model_slice(imx, imy, ihx, ihy)        &
                      + conjg(wfld_slice_d(jxd, jyd))          &
                      * wfld_slice_u(jxu, jyu)/(denom+eps*eps) &
                      * phase_shaping
                   else
                      model_slice(imx, imy, ihx, ihy)       &
                      = model_slice(imx, imy, ihx, ihy)     &
                      + conjg(wfld_slice_d(jxd, jyd))       &
                      *       wfld_slice_u(jxu, jyu)        &
                      * phase_shaping 
                   end if
                !!.. if demigration, the output is wfld_slice_u
                else
                   wfld_slice_u(jxu, jyu)                      &
                   = wfld_slice_u(jxu, jyu)                    &
                   + wfld_slice_d(jxd, jyd)                    &
                   * conjg(phase_shaping)                      &
                   * cmplx(model_slice(imx, imy, ihx, ihy), 0.)      
                end if

                if (hessian_type == 1) then   !!.. only the shot illumination
                   normd_slice(imx, imy, ihx, ihy)          &
                   = normd_slice(imx, imy, ihx, ihy)        &
                   + cabs(wfld_slice_d(jxd, jyd)) * cabs(phase_shaping)

                else if (hessian_type == 2) then !!.. source intensity multiplies receiver intensity
                                                 !!.. let data(omega) = 1 to get the diagonal of Hessian
                   normd_slice(imx, imy, ihx, ihy)          &
                   = normd_slice(imx, imy, ihx, ihy)        &
                   + cabs(wfld_slice_d(jxd, jyd)) * cabs(wfld_slice_u(jxu, jyu)) * cabs(phase_shaping)

                else if (hessian_type == 3) then !!.. approximate receiver side Green's function with a constant
                   xd = omx + (jxd-1)*dmx
                   yd = omy + (jyd-1)*dmy
                   if (nmy > 1 .and. nmy > 1) then
                      wt = rec_green3d(xd, yd, gx_min, gx_max, gy_min, gy_max, zs)
                   else
                      wt = rec_green2d(xd, gx_min, gx_max, zs)
                   end if
                   normd_slice(imx, imy, ihx, ihy)          &
                   = normd_slice(imx, imy, ihx, ihy)        &
                   + cabs(wfld_slice_d(jxd, jyd)) * cabs(phase_shaping) * wt 

                else
         
                end if

             end do
          end do
       end do
    end do
    stat = 0
  end function

!--------------------------------------------------------------------------
! Approximate receiver side Green's function for constant velocity and
! high frequency and long spread
!--------------------------------------------------------------------------
  real function rec_green3d(mx, my, gx_min, gx_max, gy_min, gy_max, zs)
    real  :: mx, my, gx_min, gy_min, gx_max, gy_max, zs
    real  :: a, b, c, d, tmp1, tmp2
!!.. upper bound
    a = 1.
    b = mx
    c = (gy_max-my)**2 + zs*zs
    d = gy_max - my
    tmp1 = integral_oper(gx_min, gx_max, a, b, c, d)

!!.. lower bound
    a = 1.
    b = mx
    c = (gy_min-my)**2 + zs*zs
    d = gy_min - my
    tmp2 = integral_oper(gx_min, gx_max, a, b, c, d)

    rec_green3d = tmp1 - tmp2 

  end function
!--------------------------------------------------------------------------
  real function rec_green2d(mx, gx_min, gx_max, zs)
    real  :: mx, gx_min, gx_max, zs
    real  :: a, b

    a = (gx_max - mx) / zs
    b = (gx_min - mx) / zs
    rec_green2d = asinh(a) - asinh(b)
  end function
!--------------------------------------------------------------------------
  real function asinh(x)
    real  :: x
    asinh = log(x + sqrt(x*x+1.))
  end function 
!--------------------------------------------------------------------------
! solve for ingeration
!--------------------------------------------------------------------------
  real function integral_oper(x_min, x_max, a, b, c, d)
    real :: x_min, x_max, a, b, c, d
    integral_oper = integral_kern(x_max, a, b, c, d) -  &
                    integral_kern(x_min, a, b, c, d)
  end function
!--------------------------------------------------------------------------
! 3-D integral kernel for int_a^b log(sqrt(a*(x-b)^2+c)+d)  
!--------------------------------------------------------------------------
  real function integral_kern(x, a, b, c, d)
    real :: x, a, b, c, d
!In[8]:= Integrate[Log[Sqrt[a*(x-b)^2+c]+d],x]
!
!                2         Sqrt[a] (b - x)
!Out[8]= ((-c + d ) ArcTan[---------------] +
!                                     2
!                           Sqrt[c - d ]
!
!             2                 Sqrt[a] d (b - x)
!       (c - d ) ArcTan[---------------------------------] -
!                                 2                    2
!                       Sqrt[c - d ] Sqrt[c + a (b - x) ]
!
!                 2
!       Sqrt[c - d ] (Sqrt[a] x +
!
!                                                    2
!          Sqrt[a] (b - x) Log[d + Sqrt[c + a (b - x) ]] +
!
!                                     2
!          d Log[2 (Sqrt[c + a (b - x) ] + Sqrt[a] (b - x))])) /
!
!                        2
!     (Sqrt[a] Sqrt[c - d ])


     integral_kern = ((-c + d*d)*atan(sqrt(a)*(b-x)/sqrt(c-d*d)) +                       &
                     (c - d*d)*atan(sqrt(a)*d*(b-x)/sqrt(c-d*d)/sqrt(c+a*(b-x)*(b-x))) - &
                     sqrt(c - d*d)*(sqrt(a)*x +                                          &
                     sqrt(a)*(b - x)*Log(d + sqrt(c + a*(b-x)*(b-x) )) +                 &
                     d*log(2*(sqrt(c+a*(b-x)*(b-x))+sqrt(a)*(b-x))))) /                  &
                     (sqrt(a)*sqrt(c-d*d))

  end function
end module












