module RWE2D_WI_wemig
  use sep
  use util_ 
  use fft 
  use geometry2D 
  use sep_timers_mod

  implicit none
  integer, private :: nt,ng,nw,nr,nloop,nref,ntap,ptap
  real,    private :: dt,dg,dw,ot,og,ow,okg,dkg,forward,eps,pi
  logical, private :: kin,verbose

  integer, allocatable,private :: allmask(:)
  real,    allocatable,private :: msk(:),mtt(:),tap(:),ikz2d(:),Afilt(:),iikz2d(:)
  real,    allocatable,private :: allrefs(:,:),allfields(:,:),witap(:)
  complex, allocatable,private :: rwfl(:),rwtt(:),rdax(:)

  real,    private :: sixth
  real,    private :: c1,c2

  complex, allocatable, private :: lk(:),rk(:),ck(:)
  complex, allocatable, private :: a(:),b(:),c(:),u(:)
  real,                 private :: kmu,   knu,    kro 
  real,    allocatable, private :: mu(  :),nu(  :),ro(  :)
  real,    allocatable, private :: m0(:),n0(:),r0(:)

contains

  !----------------------------------------------------------------  
  subroutine wemig_init(ag,at,aw,ar,forward_in,kin_in,nref_in,verbose_in,ptap_in)
    type(myaxis)     :: ag,at,aw,ar
    integer          ::                               nref_in,           ptap_in
    real             ::             forward_in
    logical          ::                        kin_in,        verbose_in

    integer :: ig
    pi=acos(-1.)
    forward=forward_in
    nt=at%n;    ng=ag%n;    nw=aw%n
    dt=at%d;    dg=ag%d;    dw=aw%d * 2.*pi
    ot=at%o;    og=ag%o;    ow=aw%o * 2.*pi
    nref=nref_in
    ptap=ptap_in

    call from_param("ntap",ntap,5)
    call from_param("nloop",nloop,3)
    okg=  -pi/    dg
    dkg=2.*pi/(ng*dg)


    call from_param("c1",c1,0.5)
    call from_param("c2",c2,0.)
    call from_param("sixth",sixth,0.122996)

    kin=kin_in
    verbose=verbose_in;    
    eps=0.0000001

    allocate(rdax(ng))
    allocate(ikz2d(ng))
    allocate(iikz2d(ng))
    allocate(rwfl(ng))
    allocate(rwtt(ng),msk(ng));
    allocate(mtt(ng),tap(ng))
    allocate(allmask  (ng)   )
    allocate(allfields(ng,6) ) 
    allocate(allrefs(nref,6))
    allocate(Afilt(ng))
    allocate(witap(nt))
    msk=1.;mtt=1.;tap=1.;ikz2d=1.;iikz2d=0.;witap=1.
    call rwetaper_init()

    allocate(lk(ng),rk(ng),ck(ng))
    allocate(a(ng),b(ng-1),c(ng-1),u(ng))
    kmu = 1./ dt         !! 1/ dt
    knu = 1./(2.*dg**2)  !! 1/( 2 dg^2)
    kro = 1./(dt*dg**2)  !! 1/(dt dg^2)
    allocate(m0(ng),n0(ng),r0(ng))
    allocate(mu(   ng),nu(   ng),ro(   ng))

    do ig=0,ng-1
       Afilt(ig+1)=cos( dble(ig)*pi/dble(ng-1)-pi/4. ) !0.4636
!        Afilt(ig+1)=(cos( dble(ig)*pi/dble(ng-1)-3.81/180.*pi))
   end do

    Afilt=Afilt/(Afilt(ng/2))
    where (Afilt .ge. 1)
        Afilt=1.
    end where

    where (Afilt .le. 0.)
       Afilt=0.
    end where
    
    Afilt=4.*Afilt
   ! call auxputch("n1","i",ng,"Afilt.H")
   ! call auxputch("o1","r",og,"Afilt.H")
   ! call auxputch("d1","r",dg,"Afilt.H")
   ! call srite("Afilt.H",Afilt,ng*4)
  end subroutine wemig_init

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

  !! WE migration driver
  subroutine wemig(rwf,injector,wfld,adj,iss,fields,refs,mask,gnorm)
    integer          :: iw,it,ig,ir,iss,mask(:,:)
    real             :: rarg,w
    real             :: refs(:,:,:),fields(:,:,:),gnorm(:,:)
    complex,pointer  :: rwf(:,:),rax(:),injector(:,:)
    complex          :: wfld(:,:,:)
    logical          :: logic,adj

    call wavefield_write1()
    call wavefield_write2()
    do iw=1,1!nw  !! Frequency Loop           

       w=ow+(iw-1)*dw
       rax => rwf(:,iw)
       rax = rax*conjg(sqrt(cmplx(0.,-w))) !*Afilt
       
       !! . . Initial state
       call rwetaper(rax) 
       wfld(:,1,iw)=rax

       if (adj) then
          rarg= -1.
          do it=2,nt
             rax(ng/2) = rax(ng/2)+injector(it-1,iw)
             call wemig_ssf(rax,w,mask(:,it),fields(:,:,it),refs(:,:,it),rarg)
             rwtt=rax;rax=0. 
             call fth(.false.,.false.,rwtt)
             do ir=1,nref  
                if (refs(ir,3,it) .lt. 0.000000001) cycle 
                rwfl=rwtt
                call wemig_phs(rwfl,w,it,ir,refs(:,:,it),rarg)
                call fth( .true.,.false.,rwfl)
!                call wemig_psc_coef(fields(:,3,it),fields(:,4,it),&
 !                    refs(ir,3,it),refs(ir,4,it))
 !               call wemig_fds(rwfl,w,it,rarg)
                call wemig_ref(rwfl,it,ir,mask(:,it),rax)
             end do
             call rwetaper(rax) 
              wfld(1:ng,it,iw)=rax

         end do !! . . z-loop

       else !! Propagating backwards

          rarg=1.
          rax=rax!*sqrt(cmplx(0.,-w)) !*Afilt
          do it=nt-1,1,-1
             rax(ng/2)=rax(ng/2)+injector(it+1,iw)
             call wemig_ssf(rax,w,mask(:,it),fields(:,:,it),refs(:,:,it),rarg)
             rwtt=rax;rax=0. 
             call fth(.false.,.false.,rwtt)
             do ir=1,nref  
                if (refs(ir,3,it) .lt. 0.000000001) cycle 
                rwfl=rwtt
                call wemig_phs(rwfl,w,it,ir,refs(:,:,it),rarg)
                call fth( .true.,.false.,rwfl)
                call wemig_psc_coef(fields(:,3,it),fields(:,4,it),&
                               refs(ir,3,it), refs(ir,4,it))
                call wemig_fds(rwfl,w,it,rarg)
                call wemig_ref(rwfl,it,ir,mask(:,it),rax)
             end do
             call rwetaper(rax)  
             wfld(1:ng,it,iw)=rax   
          end do !! . . z-loop
       end if

       !write(0,*) "PSC",iw,nw,w/(2.*pi),maxval(abs(rwtt))

    end do  !! . . w-loop
    if (adj .eq. .true. ) call srite("wfld1.H",wfld,size(wfld)*8)
    if (adj .eq. .false.) call srite("wfld2.H",wfld,size(wfld)*8)

  end subroutine wemig

  !----------------------------------------------------------------
  !! split-step Fourier
  subroutine wemig_ssf(rdax,w,allmask,allfields,allrefs,rdir) 
    integer :: ig,allmask(:)
    real    :: w,rdir,allfields(:,:),allrefs(:,:)
    complex :: rdax(:)

    do ig=1,ng 
       ikz2d(ig)=-dt*( &
          w**2*allrefs(allmask(ig),3 )* &
         (allfields(ig,3 )-allrefs(allmask(ig),3))- &
          allrefs(allmask(ig),6)*(allfields(ig,6)-allrefs(allmask(ig),6)))/&
          sqrt(abs(allrefs(allmask(ig),3)**2*w**2-allrefs(allmask(ig),6)**2)+eps) 

!       iikz2d(ig)= -dt*( allrefs(allmask(ig),2)*&
!         (allfields(ig,2)-allrefs(allmask(ig),2)))

    end do

!    rdax = rdax*cmplx( cos(ikz2d) ,rdir*sin(ikz2d) )*exp(iikz2d)
    rdax = rdax*cmplx( cos(ikz2d) ,rdir*sin(ikz2d) )*exp(rdir*dt*allfields(:,2))

  end subroutine wemig_ssf
  !----------------------------------------------------------------  

  !! Fourier-domain phase shift
  subroutine wemig_phs(rdax,w,it,ir,allrefs,rdir)
    integer :: ig,it,ir
    real    :: w,kg,rdir,carg,carg1,allrefs(:,:),pkz
    complex :: kz,rdax(:)

    carg1= (allrefs(ir,3)* w)**2 - allrefs(ir,6)**2
 
    do ig=1,ng
       kg=okg+(ig-1)*dkg

       carg = carg1 - (allrefs(ir,4)*kg)**2 
       
       if ( carg .lt. 0) then

          pkz = -dt*abs(sqrt(-carg))
          rdax(ig)=rdax(ig)*exp(pkz)
        
       else

          pkz =-dt*( sqrt(carg) )
          rdax(ig)=rdax(ig)*cmplx(cos(pkz),rdir*sin(pkz))

       end if

       pkz = -dt*kg*allrefs(ir,1) !Was +1
       rdax(ig)=rdax(ig)*cmplx(cos(pkz),rdir*sin(pkz))

    end do

  end subroutine wemig_phs


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

  subroutine wemig_ref(rwfl,it,ir,allmask,rdax)
    complex :: rdax(:),rwfl(:)
    integer :: it,ir,ig,ii,allmask(:)

    mtt=0.
    where( allmask .eq. ir) 
       mtt=1.
    end where

    do ii=1,nloop !! Smoothing Iterations
       msk(1 ) = mtt(1   )
       msk(ng) = mtt(ng-1)
       do ig=2,ng-1
          msk(ig)=( mtt(ig-1)+mtt(ig+1) )/2.
       end do
       mtt=msk
    end do

    rdax = rdax + msk*rwfl

  end subroutine wemig_ref

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

  subroutine rwetaper_init()
    integer :: itap,jtap,ig

    if (ntap .gt. 1) then

       do itap=1,ntap
          jtap = abs(ntap-itap)
          tap(   itap) = cos(pi/2.* dble(jtap)/dble(ntap-1))
       end do

       do itap=1,ntap
          jtap = abs(ntap-itap)
          tap(ng-itap+1) = cos(pi/2.* dble(jtap)/dble(ntap-1))
       end do

    end if

  end subroutine rwetaper_init
  !----------------------------------------------------------------

  subroutine wftaper_init(omin,omax)
    integer :: itap,jtap,ig,imin,imax
    real :: omin,omax
    real,allocatable:: wfline(:)

    allocate(wfline(ptap))

    imin=floor( (omin-ot)/dt )
    imax=ceiling ( (omax-ot)/dt )

    imin=max(imin,1)
    imax=min(imax,nt)
    write(0,*) 'MIN',omin,imin,ptap
    write(0,*) 'MAX',omax,imax,ptap
    witap=0.
    witap(imin:imax)=1.

    do ig=1,ptap
       wfline(ig)=1-cos(pi/2.*dble(ig-1)/dble(ptap-1))
    end do

    do ig=1,ptap
       witap(imin-1+ig)=wfline(ig)
       witap(imax+1-ig)=wfline(ig)
    end do

     call auxputch("n1","i",nt,"wftaper.H")
     call auxputch("o1","r",ot,"wftaper.H")
     call auxputch("d1","r",dt,"wftaper.H")
     call srite("wftaper.H",witap,size(witap)*4)
    deallocate(wfline)
  end subroutine wftaper_init
  !----------------------------------------------------------------

  subroutine wftaper(model,data)
    complex :: model(:,:),data(:,:)
    integer :: iw

    do iw=1,nw
       model(:,iw)=model(:,iw)*witap
       data (:,iw)= data(:,iw)*witap
    end do

  end subroutine wftaper
  !----------------------------------------------------------------

  subroutine residtaper(resid)
    complex :: resid(:,:)
    integer :: iw

    do iw=1,nw
       resid(:,iw)= resid(:,iw)*witap
    end do

  end subroutine residtaper
  !----------------------------------------------------------------

  subroutine rwetaper(dax)
    complex, pointer  :: dax(:)
    dax = dax * tap
  end subroutine rwetaper

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

  subroutine Diffdata(model,data,resid)
    integer  :: iw
    complex :: resid(:,:)
    complex, pointer  :: model(:,:),data(:,:)
    do iw=1,1!nw
!       where ( abs(model(:,iw)) .lt. 0.01*maxval(abs(model(:,iw))) )
!          resid(:,iw)=0.
!       end where
       
!       where ( abs(model(:,iw)) .ge. 0.01*maxval(abs(model(:,iw))) )
       resid(:,iw) =  (model(:,iw) - data(:,iw))*2.71/maxval(abs(data(:,iw)))!/&
!                 sqrt(conjg(model(:,iw))*model(:,iw))
!       end where

       write(0,*) 'D,M,R MAX: ',maxval(abs(data(:,iw))),&
       maxval(abs(model(:,iw))),maxval(abs(resid(:,iw))),&
       maxval(abs(data(:,iw)))/maxval(abs(model(:,iw)))
    end do

    call residtaper(resid)

    write(0,*) 'CALCULATED PERTURBATION'
    write(0,*) 'Model: ',minval(abs(model)),maxval(abs(model))
    write(0,*) 'Data : ',minval(abs(data )),maxval(abs(data ))
    write(0,*) 'Resid: ',minval(abs(resid)),maxval(abs(resid))

    call srite("Perturb.H",model,size(model)*8)
    call srite("Perturb.H",data ,size(data )*8)
    call srite("Perturb.H",resid,size(resid)*8)

  end subroutine Diffdata
 !----------------------------------------------------------------

  subroutine Diffwavefield(SGF,wfld,Dwfld)
    integer  :: iw,it,ig
    real     :: Dwfld(:,:)
    complex  :: SGF(:,:,:), wfld(:,:,:)

    do iw=1,1!nw
       do ig=1,ng
          do it=1,nt
             Dwfld(it,ig) = Dwfld(it,ig) +&
                 cabs( SGF(ig,it,iw) - wfld(ig,it,iw) )**2
!                  conjg(SGF(ig,it,iw) - wfld(ig,it,iw))*(SGF(ig,it,iw) - wfld(ig,it,iw))
         enddo
       end do
    end do

    write(0,*) 'HESSIAN WAVEFIELD PERTURB: ',minval(Dwfld),maxval(Dwfld)
  end subroutine Diffwavefield
  !----------------------------------------------------------------

  subroutine data_write1()
      call auxputch("n1","i",nt  ,"data1.H")
      call auxputch("n2","i",nw  ,"data1.H")
      call auxputch("d1","r",dt  ,"data1.H")
      call auxputch("d2","r",dw  ,"data1.H")
      call auxputch("o1","r",ot  ,"data1.H")
      call auxputch("o2","r",ow  ,"data1.H")
      call auxputch("esize","i",8,"data1.H")
      call auxputch("n3","i",1   ,"data1.H")
      call auxputch("d3","r",1.  ,"data1.H")
      call auxputch("o3","r",0.  ,"data1.H")
  end subroutine data_write1
   !----------------------------------------------------------------

  subroutine wavefield_write1()
      call auxputch("n1","i",ng,   "wfld1.H")
      call auxputch("n2","i",nt  , "wfld1.H")
      call auxputch("n3","i",nw,   "wfld1.H")
      call auxputch("d1","r",dg,   "wfld1.H")
      call auxputch("d2","r",dt,   "wfld1.H")
      call auxputch("d3","r",dw,   "wfld1.H")
      call auxputch("o1","r",og,   "wfld1.H")
      call auxputch("o2","r",ot,   "wfld1.H")
      call auxputch("o3","r",ow,   "wfld1.H")
      call auxputch("esize","i",8, "wfld1.H")
  end subroutine wavefield_write1
   !----------------------------------------------------------------

  subroutine wavefield_write2()
      call auxputch("n1","i",ng,   "wfld2.H")
      call auxputch("n2","i",nt,   "wfld2.H")
      call auxputch("n3","i",nw,   "wfld2.H")
      call auxputch("d1","r",dg,   "wfld2.H")
      call auxputch("d2","r",dt,   "wfld2.H")
      call auxputch("d3","r",dw,   "wfld2.H")
      call auxputch("o1","r",og,   "wfld2.H")
      call auxputch("o2","r",ot,   "wfld2.H")
      call auxputch("o3","r",ow,   "wfld2.H")
      call auxputch("esize","i",8, "wfld2.H")
  end subroutine wavefield_write2

  !----------------------------------------------------------------
  subroutine wemig_psc_coef(aa,bb,a0,b0)
    real    :: aa(:),bb(:)
    real    :: a0,  b0
    integer :: it

    !write(0,*) "compute PSC coefficients"
!    n0 = a0*(c1*(aa/a0-1.)-(bb/b0-1.))*(b0/a0)**2
    n0 = -c1 * (b0/a0)**2 *(aa-a0)+ b0/a0   *(bb-b0)
    r0 = 3*c2* (b0/a0)**2 
    m0 = 1.

    m0 = m0 * kmu
    n0 = n0 * knu
    r0 = r0 * kro

  end subroutine wemig_psc_coef
 !----------------------------------------------------------------

  subroutine wemig_ffd_coef(aa,bb,a0,b0)
    real    :: aa(:),bb(:)
    real    :: a0,b0
    integer :: it
    real,    allocatable :: d1(:),d2(:)
    real    :: tt

    allocate(d1(ng),d2(ng))

    tt = b0
 
!    d1=aa*(bb/aa)**2-a0*(b0/a0)**2
!    d2=aa*(bb/aa)**2-a0*(b0/a0)**4

    d1=a0*((b0/tt)/a0)**2 * ( ( bb/b0)**2* a0/aa    - 1. )
    d2=a0*((b0/tt)/a0)**4 * ( ( bb/b0)**4*(a0/aa)**3- 1. )

    where(abs(d1)<eps ) 
       d1=eps
    end where

    m0 =      1    
    n0 = c1 * d1 ! * tt**2
    r0 = c2 !* d2/d1 * tt**2

    m0 = m0 * kmu
    n0 = n0 * knu
    r0 = r0 * kro

    deallocate(d1,d2)
  end subroutine wemig_ffd_coef
  !----------------------------------------------------------------
  subroutine wemig_coef_report()

    write(0,*) minval(m0),"< mu <",maxval(m0)
    write(0,*) minval(n0),"< nu <",maxval(n0)
    write(0,*) minval(r0),"< ro <",maxval(r0)

  end subroutine wemig_coef_report
  !----------------------------------------------------------------  
  !! finite-differences solver
  subroutine wemig_fds(dax,w,it,dir)
    real                  :: w,dir
    integer               :: it
    complex, dimension(:) :: dax
    integer               :: ig

    mu = m0
    nu = n0 / w
    ro = r0 / w**2

    lk = cmplx(ro+sixth*mu, nu)
    rk = cmplx(ro+sixth*mu,-nu)
    ck = cmplx(         mu, 0.)

    !!             |                      |
    !!     b       |          a           |     c
    !!  ro + i nu  |  mu - 2(ro + i nu )  |  ro + i nu
    !!             |                      |
    a = ck - 2.*lk  
    b = lk(1  :ng-1)
    c = lk(1+1:ng  )

    do ig=2,ng-1
       u(ig)=           rk(ig) *(dax(ig-1)+dax(ig+1)) & 
       +    (ck(ig)-2.0*rk(ig))* dax(ig)
       !!             |                      | 
       !!  ro - i nu  |  mu - 2 (ro - i nu)  | ro - i nu
       !!             |                      | 
    end do
    ig= 1; u(ig)=(ck(ig)-2.0*rk(ig))*dax(ig) + rk(ig)*dax(ig+1)
    ig=ng; u(ig)=(ck(ig)-2.0*rk(ig))*dax(ig) + rk(ig)*dax(ig-1)

    call thr(a,b,c,u,ng)
    dax = u
  end subroutine wemig_fds
  !----------------------------------------------------------------
  !! tridiagonal solver
  subroutine thr(a,b,c,v,n)
    integer :: n
    complex :: a(n),b(n-1),c(n-1),v(n)
    integer :: i

    b(1)=b(1)/a(1)
    v(1)=v(1)/a(1)
    do i=2,n
       a(i)=a(i)-c(i-1)*b(i-1)
       b(i)=b(i)/a(i)
       v(i)=(v(i)-c(i-1)*v(i-1))/a(i)
    end do
    do i=n-1,1,-1
       v(i)=v(i)-b(i)*v(i+1)
    end do
  end subroutine thr
  !----------------------------------------------------------------
  subroutine attribute(x)
    complex :: x(:)
    write(0,*) minval( real(x)),maxval( real(x))
    write(0,*) minval(aimag(x)),maxval(aimag(x))
  end subroutine attribute

  !----------------------------------------------------------------
  subroutine Perturbation_calc(SGF,wfld,gradient)
    integer             :: iw,it,ig
    real                :: w,mv
    complex,intent(in ) :: SGF(:,:,:),wfld(:,:,:)
    real   ,intent(out) :: gradient(:,:)

    do iw=1,1!nw
       w=ow+(iw-1)*dw
       mv=maxval(abs(SGF(:,:,iw)))
       do ig=1,ng
          do it=1,nt
             gradient(it,ig) = gradient(it,ig) &
                  - w**2 * real( (SGF(ig,it,iw)) * conjg(wfld(ig,it,iw)) )!!/&
!!                        real( conjg(SGF(ig,it,iw))*SGF(ig,it,iw)+0.01*mv) 
          end do
       end do
    end do

  end subroutine Perturbation_calc


end module  RWE2D_WI_wemig
