module srm_storm2_
  use utilities
  use pctrist_

  implicit none
  integer, parameter, private:: dbl=selected_int_kind(15)

  real,    dimension(:),   allocatable, private :: kzi,khxi,khyi,kmxi,kmyi
  complex, dimension(:,:), allocatable, private :: model, data
  real,    dimension(:,:), allocatable, private :: str

  real,    private :: eps
  logical, private :: push
contains
  !----------------------------------------------------------------
  subroutine storm_srm_init(nth,az,ahx,ahy,amx,amy,eps_in,push_in,scale)
    integer,     intent(in) :: nth
    type(paxis), intent(in) :: az,ahx,ahy,amx,amy
    real,        intent(in) :: eps_in
    logical,     intent(in) :: push_in
    logical,     intent(in) :: scale
    type(paxis)             :: akz

    integer :: iz,ihx,ihy,imx,imy
    real    :: pi
    integer :: stat

    push = push_in
    eps  = eps_in

    pi = acos(-1.)

    allocate( data(nth,az%n), stat=stat)
    if(stat /=0) call seperr("cant allocate data")
    allocate(model(nth,az%n), stat=stat)
    if(stat /=0) call seperr("cant allocate model")
    allocate(  str(nth,az%n), stat=stat)
    if(stat /=0) call seperr("cant allocate str")

    allocate( kzi( az%n), stat=stat)
    if(stat /=0) call seperr("cant allocate kzi")
    allocate(khxi(ahx%n), stat=stat)
    if(stat /=0) call seperr("cant allocate khxi")
    allocate(khyi(ahy%n), stat=stat)
    if(stat /=0) call seperr("cant allocate khyi")
    allocate(kmxi(amx%n), stat=stat)
    if(stat /=0) call seperr("cant allocate kmxi")
    allocate(kmyi(amy%n), stat=stat)
    if(stat /=0) call seperr("cant allocate kmyi")

    !! kz
!!$    do iz=1,az%n
!!$       kzi(iz) = az%o + (iz-1)*az%d
!!$    end do
!!$    kzi=2*pi*kzi
!!$    akz%o = 2*pi*az%o
!!$    akz%d = 2*pi*az%d
!!$    akz%n = az%n

    !! kz
    do iz = 1,az%n
       kzi(iz) = (-az%n/2+iz-1)
    end do
    kzi   = 2*pi* kzi / (az%n*az%d)

    akz%o = kzi(1)
    akz%d = 2*pi / (az%n*az%d)
    akz%n = az%n

    !! khx
    khxi=0.
    if(ahx%n>1) then
       do ihx = 1,ahx%n/2
          khxi(ihx        ) = (         ihx-1)
          khxi(ihx+ahx%n/2) = (-ahx%n/2+ihx-1)
       end do
       khxi = 2*pi* khxi / (ahx%n*ahx%d)
    end if

    !! khy
    khyi=0.
    if(ahy%n>1) then
       do ihy = 1,ahy%n/2
          khyi(ihy        ) = (         ihy-1)
          khyi(ihy+ahy%n/2) = (-ahy%n/2+ihy-1)
       end do
       kmyi = 2*pi* khyi / (ahy%n*ahy%d)
    end if

    !! kmx
    kmxi=0.
    if(amx%n>1) then
       do imx = 1,amx%n/2
          kmxi(imx        ) = (         imx-1)
          kmxi(imx+amx%n/2) = (-amx%n/2+imx-1)
       end do
       kmxi = 2*pi* kmxi / (amx%n*amx%d)
    end if

    !! kmy
    kmyi=0.
    if(amy%n>1) then
       do imy = 1,amy%n/2
          kmyi(imy        ) = (         imy-1)
          kmyi(imy+amy%n/2) = (-amy%n/2+imy-1)
       end do
       kmyi = 2*pi* kmyi / (amy%n*amy%d)
    end if

!!$    !! khx
!!$    khxi=0.
!!$    if(ahx%n>1) then
!!$       do ihx=1,ahx%n
!!$          khxi(ihx) = ahx%o + (ihx-1)*ahx%d
!!$       end do
!!$       !! khx is not centered
!!$       do ihx=ahx%n/2+1,ahx%n
!!$          khxi(ihx) = khxi(ihx) - (ahx%n)*ahx%d
!!$       end do
!!$       khxi=2*pi*khxi
!!$    end if
!!$
!!$    !! khy
!!$    khyi=0.
!!$    if(ahy%n>1) then
!!$       do ihy=1,ahy%n
!!$          khyi(ihy) = ahy%o + (ihy-1)*ahy%d
!!$       end do
!!$       !! khy is not centered
!!$       do ihy=ahy%n/2+1,ahy%n
!!$          khyi(ihy) = khyi(ihy) - (ahy%n)*ahy%d
!!$       end do
!!$       khyi=2*pi*khyi
!!$    end if
!!$
!!$    !! kmx
!!$    kmxi=0.
!!$    if(amx%n>1) then
!!$       do imx=1,amx%n
!!$          kmxi(imx) = amx%o + (imx-1)*amx%d
!!$       end do
!!$       !! kmx is not centered
!!$       do imx=amx%n/2+1,amx%n
!!$          kmxi(imx) = kmxi(imx) - (amx%n)*amx%d
!!$       end do
!!$       kmxi=2*pi*kmxi
!!$    end if
!!$
!!$    !! kmy
!!$    kmyi=0.
!!$    if(amy%n>1) then
!!$       do imy=1,amy%n
!!$          kmyi(imy) = amy%o + (imy-1)*amy%d
!!$       end do
!!$       !! kmy is not centered
!!$       do imy=amy%n/2+1,amy%n
!!$          kmyi(imy) = kmyi(imy) - (amy%n)*amy%d
!!$       end do
!!$       kmyi=2*pi*kmyi
!!$    end if

    call pctrist_init(nth,akz,akz,push,eps)

  end subroutine storm_srm_init
  !----------------------------------------------------------------
  subroutine storm_srm_compute(ith,isl,f,ratio,az,ahx,ahy,imx,imy,depthfreeze)
    integer,           intent(in) :: ith
    integer(kind=dbl), intent(in) :: isl
    real,        intent(in) :: ratio
    integer,     intent(in) ::    imx,imy
    type(paxis), intent(in) :: az,ahx,ahy
    integer                 :: iz,ihx,ihy
    complex, dimension(:,:,:,:), pointer :: f
    logical,     intent(in) :: depthfreeze

    real    :: ro
    real    :: kz,khx,khy,kmx,kmy
    real    :: mks,mkg
    real    :: kzh,kzm
    real    :: zzs,zzg

    kmx = kmxi(imx)
    kmy = kmyi(imy)

    if(push) then
       ro =   ratio
    else
       ro = 1/ratio
    end if

    do ihy = 1, ahy%n;      khy = khyi(ihy)
       do ihx = 1, ahx%n;   khx = khxi(ihx)
          mks=sqrt( (kmx-khx)**2 + (kmy-khy)**2 ) !! | ks |
          mkg=sqrt( (kmx+khx)**2 + (kmy+khy)**2 ) !! | kg |

          do iz = 1, az%n;  kz  = kzi (iz )
             !! dispersion relation - doc in SEP100

             kzh = (4*kz**2 + (mkg-mks)**2)/4
             kzm = (4*kz**2 + (mkg+mks)**2)/4 

             if(depthfreeze) then
                zzs =         kzh*kzm - kz**2 * mks**2 / ro**2
                zzg =         kzh*kzm - kz**2 * mkg**2 / ro**2
             else
                zzs = ro**2 * kzh*kzm - kz**2 * mks**2
                zzg = ro**2 * kzh*kzm - kz**2 * mkg**2
             end if

             str(ith,iz) = az%o
             if (zzs>0..and. zzg>0.) &
             str(ith,iz)=0.5/kz*(sqrt(zzs)+sqrt(zzg))
          end do
          call pctrist_define(ith,str)

          data(ith,:) = f(:,ihx,ihy,isl)
          call pctrist_apply(ith,model,data)
          f(:,ihx,ihy,isl) = model(ith,:)

       end do
    end do

  end subroutine storm_srm_compute
  !----------------------------------------------------------------
  subroutine storm_srm_nomove(ith,isl,f,ratio,az,ahx,ahy,imx,imy)
    integer,           intent(in) :: ith
    integer(kind=dbl), intent(in) :: isl
    real,        intent(in) :: ratio
    integer,     intent(in) ::    imx,imy
    type(paxis), intent(in) :: az,ahx,ahy
    integer                 :: iz,ihx,ihy
    complex, dimension(:,:,:,:), pointer :: f

    real    :: ro
    real    :: kz,kmx,kmy
    real    :: zzm
    real    :: mkm2

    if(push) then
       ro =   ratio
    else
       ro = 1/ratio
    end if
    ro=1/ro !! undo zero offset migration

    kmx = kmxi(imx)
    kmy = kmyi(imy)
    mkm2= kmx**2 + kmy**2

    do ihy = 1, ahy%n
       do ihx = 1, ahx%n
          do iz = 1, az%n;  kz  = kzi (iz )
             zzm = ro**2 * ( kz**2 + mkm2) - mkm2

             str(ith,iz) = az%o
             if (zzm>0.) str(ith,iz)= abs(kz)/kz * sqrt(zzm)
          end do
          call pctrist_define(ith,str)

          data(ith,:) = f(:,ihx,ihy,isl)
          call pctrist_apply(ith,model,data)
          f(:,ihx,ihy,isl) = model(ith,:)

       end do
    end do

  end subroutine storm_srm_nomove

  !----------------------------------------------------------------
  subroutine storm_srm_close()
    integer :: stat
    call pctrist_close()

    deallocate(kzi,  stat=stat)
    if(stat /=0) call seperr("cant deallocate kzi")
    deallocate(khxi, stat=stat)
    if(stat /=0) call seperr("cant deallocate khxi")
    deallocate(khyi, stat=stat)
    if(stat /=0) call seperr("cant deallocate khyi")
    deallocate(kmxi, stat=stat)
    if(stat /=0) call seperr("cant deallocate kmxi")
    deallocate(kmyi, stat=stat)
    if(stat /=0) call seperr("cant deallocate kmyi")

    deallocate(data,  stat=stat)
    if(stat /=0) call seperr("cant deallocate data")
    deallocate(model, stat=stat)
    if(stat /=0) call seperr("cant deallocate model")
    deallocate(str,   stat=stat)
    if(stat /=0) call seperr("cant deallocate str")

  end subroutine storm_srm_close
  !----------------------------------------------------------------
end module srm_storm2_












