!!.. This module contains various kinds of Hessian approximations
!!-------------------------------------------------------------------
!!.. Author: Yaxun Tang, tang@sep.stanford.edu
!!-------------------------------------------------------------------
module wem_hesapp_mod
  use  base_struct_mod
  implicit none


  integer                       , private :: indx, indy, indz
  type(hes_grid)                , private :: hgrid
  type(out_grid)                , private :: ogrid
  type(pro_grid)                , private :: pgrid 
  character(len=256)            , private :: filetag
  integer, external             , private :: sreed_window, srite_window

contains


  integer function wem_hesapp_init(hgrid_in, ogrid_in, pgrid_in, filetag_in) result(stat)
    type(hes_grid)        :: hgrid_in
    type(out_grid)        :: ogrid_in
    type(pro_grid)        :: pgrid_in
    character(len=*)      :: filetag_in

    stat = 1

    indx = ogrid_in%index_outpro_xmin
    indy = ogrid_in%index_outpro_ymin
    indz = ogrid_in%index_outpro_zmin

    hgrid = hgrid_in
    ogrid = ogrid_in
    pgrid = pgrid_in

    filetag = filetag_in

    stat = 0

  end function

  integer function wem_hesapp_clos() result(stat)
    stat = 1

    indx = 0; indy = 0; indz = 0

    stat = 0
  end function

  integer function wem_hesapp_oper(hess, wfds, wfdr, ws) result(stat)
    real   , dimension(:,:,:) :: hess
    complex, dimension(:,:,:) :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:,:) :: wfdr
    real                      :: ws
    stat = 1

    call hesapp_kernel(hess, wfds, wfdr, ws)

    stat = 0
  end function


  !!.. Hessian via random phase encoding
  subroutine hesapp_kernel(hess, wfds, wfdr, ws)
    real   , dimension(:,:,:) :: hess
    complex, dimension(:,:,:) :: wfds !!.. wfds(nx, ny, nz)
    complex, dimension(:,:,:) :: wfdr
    complex                   :: soufnc
    real                      :: ws
 
    integer                   :: ix_modl, iy_modl, iz_modl
    integer                   :: ix_offd, iy_offd, iz_offd
    integer                   :: jx_offd, jy_offd, jz_offd
    integer                   :: index_x_modl, index_y_modl, index_z_modl
    integer                   :: index_x_offd, index_y_offd, index_z_offd
    real                      :: x_modl, y_modl, z_modl
    real                      :: x_offd, y_offd, z_offd
    real                      :: hx_offd, hy_offd, hz_offd
    real                      :: wt_modl, wt_offd
    complex                   :: greens_sou, greens_rec


    do iz_offd=1, hgrid%az_offd%n; hz_offd = hgrid%az_offd%o+(iz_offd-1)*hgrid%az_offd%d; jz_offd = nint(hz_offd/pgrid%az%d)
    do iy_offd=1, hgrid%ay_offd%n; hy_offd = hgrid%ay_offd%o+(iy_offd-1)*hgrid%ay_offd%d; jy_offd = nint(hy_offd/pgrid%ay%d)
    do ix_offd=1, hgrid%ax_offd%n; hx_offd = hgrid%ax_offd%o+(ix_offd-1)*hgrid%ax_offd%d; jx_offd = nint(hx_offd/pgrid%ax%d)

       call hesapp_fileio(hess, hgrid, ogrid, ix_offd, iy_offd, iz_offd, filetag, "reed")

       do iz_modl=1, ogrid%az%n; index_z_modl = max(1, min(indz+iz_modl-1,pgrid%az%n))      ; z_modl = pgrid%az%o + (index_z_modl-1)*pgrid%az%d
                                 index_z_offd = max(1, min(index_z_modl+jz_offd,pgrid%az%n)); z_offd = pgrid%az%o + (index_z_offd-1)*pgrid%az%d
       do iy_modl=1, ogrid%ay%n; index_y_modl = max(1, min(indy+iy_modl-1,pgrid%ay%n))      ; y_modl = pgrid%ay%o + (index_y_modl-1)*pgrid%ay%d
                                 index_y_offd = max(1, min(index_y_modl+jy_offd,pgrid%ay%n)); y_offd = pgrid%ay%o + (index_y_offd-1)*pgrid%ay%d
       do ix_modl=1, ogrid%ax%n; index_x_modl = max(1, min(indx+ix_modl-1,pgrid%ax%n))      ; x_modl = pgrid%ax%o + (index_x_modl-1)*pgrid%ax%d
                                 index_x_offd = max(1, min(index_x_modl+jx_offd,pgrid%ax%n)); x_offd = pgrid%ax%o + (index_x_offd-1)*pgrid%ax%d

          greens_sou =       wfds(index_x_modl,index_y_modl,index_z_modl)  &
                     * conjg(wfds(index_x_offd,index_y_offd,index_z_offd))

          select case (hgrid%hestyp)
            case (1) !!.. source intensity, assuming receiver side Green's function to be 1
                 greens_rec = cmplx(1., 0.)
            case (2) !!.. approximate the receiver side Green's function as a aperture dependent constant (not ready yet)
                 if (pgrid%ay%n>1 .and. pgrid%ax%n>1) then
                    wt_modl = rec_green3d(x_modl, y_modl, pgrid%gxmin, pgrid%gxmax, pgrid%gymin, pgrid%gymax, z_modl)
                    wt_offd = rec_green3d(x_offd, y_offd, pgrid%gxmin, pgrid%gxmax, pgrid%gymin, pgrid%gymax, z_offd)
                 else
                    wt_modl = rec_green2d(x_modl, pgrid%gxmin, pgrid%gxmax, z_modl)
                    wt_offd = rec_green2d(x_offd, pgrid%gxmin, pgrid%gxmax, z_offd)
                 end if
                 greens_rec = wt_modl
            case (3) !!.. Hessian via phase-encoding 
                 greens_rec = conjg(wfdr(index_x_modl,index_y_modl,index_z_modl)) &
                            *       wfdr(index_x_offd,index_y_offd,index_z_offd)
          end select

          hess(ix_modl,iy_modl,iz_modl) = hess(ix_modl,iy_modl,iz_modl) + real(greens_sou*greens_rec)

       end do
       end do   
       end do

       call hesapp_fileio(hess, hgrid, ogrid, ix_offd, iy_offd, iz_offd, filetag, "rite")

    end do
    end do
    end do

    return
  end subroutine

  subroutine hesapp_fileio(hess, hgrid, ogrid, index_x_offd, index_y_offd, index_z_offd, filetag, flag)
    real, dimension(:,:,:) :: hess
    type(hes_grid)         :: hgrid
    type(out_grid)         :: ogrid
    integer                :: index_x_offd, index_y_offd, index_z_offd
    character(len=*)       :: filetag
    character(len=4)       :: flag

    integer, dimension(6)  :: n_cube, j_wind, n_wind, f_wind
    integer                :: n_dim_cube, err

    n_dim_cube = 6
    n_cube = (/hgrid%ax_modl%n,hgrid%ay_modl%n,hgrid%az_modl%n,hgrid%ax_offd%n,hgrid%ay_offd%n,hgrid%az_offd%n/)
    j_wind = (/1,1,1,1,1,1/)
    n_wind = (/ogrid%ax%n,ogrid%ay%n,ogrid%az%n,1,1,1/)
    f_wind = (/ogrid%index_outimg_xmin-1,ogrid%index_outimg_ymin-1,ogrid%index_outimg_zmin-1,index_x_offd-1,index_y_offd-1,index_z_offd-1/)

    if (flag == "rite") then
       err = srite_window(filetag,n_dim_cube,n_cube,n_wind,f_wind,j_wind,4,hess)
    else if (flag == "reed") then
       hess = 0.
       err = sreed_window(filetag,n_dim_cube,n_cube,n_wind,f_wind,j_wind,4,hess)
    end if

    return
  end subroutine

!--------------------------------------------------------------------------
! 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
    if (abs(zs-0.)<0.000001) then
       rec_green2d = 1.
    else 
       a = (gx_max - mx) / zs
       b = (gx_min - mx) / zs
       rec_green2d = asinh(a) - asinh(b)
    end if
  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
