module dimwgh2_
  !! image perturbation weight

  use utilities
  implicit none

  type(paxis) :: laz,lahx,lahy,lamx,lamy

  real, private :: pi

  real,    dimension(:),         pointer :: kmx,kmy,khx,khy, kz
  real,    dimension(:,:),       pointer :: ksx,ksy,krx,kry
  real,    dimension(:,:,:,:),   pointer :: kss,krr
  real,    dimension(:,:,:,:),   pointer ::  uu,wgh
  real,    dimension(:,:,:,:),   pointer :: tms,tmr
contains
  !----------------------------------------------------------------
  subroutine dimwgh2_init(az,ahx,ahy,amx,amy)
    type(paxis) :: az,ahx,ahy,amx,amy
    integer :: iz,ihx,ihy,imx,imy
    integer :: er

    laz=az
    lahx=ahx
    lahy=ahy
    lamx=amx
    lamy=amy

    pi = acos(-1.)

    allocate(kmx(amx%n),stat=er)
    if(er/=0) call seperr('cant allocate kmx')
    call klinemap(kmx,amx%n,amx%d)
    allocate(kmy(amy%n),stat=er)
    if(er/=0) call seperr('cant allocate kmy')
    call klinemap(kmy,amy%n,amy%d)
    allocate(khx(ahx%n),stat=er)
    if(er/=0) call seperr('cant allocate khx')
    call klinemap(khx,ahx%n,ahx%d)

    allocate(khy(ahy%n),stat=er)
    if(er/=0) call seperr('cant allocate khy')
    call klinemap(khy,ahy%n,ahy%d)

    allocate(kz(az%n),stat=er)
    if(er/=0) call seperr('cant allocate kz')
    call klinemap( kz, az%n, az%d)

    !! ksx(hx,mx), krx(hx,mx)
    allocate(ksx(ahx%n, amx%n),stat=er)
    if(er/=0) call seperr('cant allocate ksx')
    allocate(krx(ahx%n, amx%n),stat=er)
    if(er/=0) call seperr('cant allocate krx')

    !! ksy(hy,my), kry(hy,my)
    allocate(ksy(ahy%n, amy%n),stat=er)
    if(er/=0) call seperr('cant allocate ksy')
    allocate(kry(ahy%n, amy%n),stat=er)
    if(er/=0) call seperr('cant allocate kry')

    !! ks(hx,hy,mx,my)
    allocate(kss(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate kss')

    !! kr(hx,hy,mx,my)
    allocate(krr(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate krr')

    !!  uu(hx,hy,mx,my)
    allocate( uu(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate  uu')

    !! wgh(hx,hy,mx,my)
    allocate(wgh(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate wgh')

    !! tms(hx,hy,mx,my)
    allocate(tms(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate tms')
    !! tmr(hx,hy,mx,my)
    allocate(tmr(ahx%n,ahy%n,amx%n,amy%n),stat=er)
    if(er/=0) call seperr('cant allocate tmr')

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

    !! ksx = kmx - khx
    !! krx = kmx + khx
    do ihx=1,ahx%n
       do imx=1,amx%n
          ksx(ihx,imx) = kmx(imx)-khx(ihx)
          krx(ihx,imx) = kmx(imx)+khx(ihx)
       end do
    end do

    !! ksy = kmy - khy
    !! kry = kmy + khy
    do ihy=1,ahy%n
       do imy=1,amy%n
          ksy(ihy,imy) = kmy(imy)-khy(ihy)
          kry(ihy,imy) = kmy(imy)+khy(ihy)
       end do
    end do

    !! kss = sqrt(ksx^2 + ksy^2) = |ks|
    !! krr = sqrt(krx^2 + kry^2) = |kr|
    do ihy=1,ahy%n
       do ihx=1,ahx%n
          do imy=1,amy%n
             do imx=1,amx%n
                kss(ihx,ihy,imx,imy) = &
                ksx(ihx,imx)**2      + &
                ksy(ihy,imy)**2

                krr(ihx,ihy,imx,imy) = &
                krx(ihx,imx)**2      + &
                kry(ihy,imy)**2
             end do
          end do
       end do
    end do
    kss = sqrt(kss)
    krr = sqrt(krr)

  end subroutine dimwgh2_init
  !----------------------------------------------------------------  
  subroutine dimwgh2(cdata)
    complex, dimension(:,:,:,:,:), pointer :: cdata
    integer :: iz
    real    :: kz2
    real    :: epst

    cdata(1,:,:,:,:)=0.0
    do iz=2,laz%n
       kz2=kz(iz)**2

       !! mu^2
       uu = &
       ( 4*kz2 + (krr-kss)**2) * &
       ( 4*kz2 + (krr+kss)**2) / 16

       where( uu > 1.*epsilon(1.))
          tms=kss**2/ uu * kz2
          tmr=krr**2/ uu * kz2
       end where

       epst=0.99
       where(tms<epst)
          tms=1./sqrt(1.-tms)
       elsewhere
          tms=1.
       end where
       where(tmr<epst)
          tmr=1./sqrt(1.-tmr)
       elsewhere 
          tmr=1.
       end where

       if(kz2>1.*epsilon(1.)) then
          wgh=0.5*sqrt(uu)/sqrt(kz2+10*epsilon(1.))*(tms + tmr)
       else
          wgh=1.0
       end if

!       wgh = wgh / maxval(wgh)

       cdata(iz,:,:,:,:) = &
       cdata(iz,:,:,:,:) * wgh 

    end do
  end subroutine dimwgh2
  !----------------------------------------------------------------
  subroutine klinemap(k,n,d)
    real, dimension(:), pointer :: k
    integer, intent(in) :: n
    real,    intent(in) :: d
    integer             :: i,m

    if(n/=1) then
       m = n/2

       do i=1,m
          k(i  ) =    i-1
          k(i+m) = -m+i-1
       end do

    else
       k=0.
    end if

    k = k / (n*d)  !! 1/length
    k = 2 * pi * k !! radians
  end subroutine klinemap
  !----------------------------------------------------------------
end module dimwgh2_


