module RWE2D_WI_routines
  use sep
  use util_ 
  use fft 
  use geometry2D_TOPO 
  use RWE2D_WI_wemig
  use C2R2D

  implicit none
  integer, private :: nt,ng,nw,nr,ntap,ptap,nzcut
  integer, private :: nblock,ns,nsx,nsz,norm,nx,nz
  real,    private :: dt,dg,dw,ot,og,ow,okg,dkg,eps,pi
  real,    private :: os,ds,dx,dz,ox,oz,omin,omax,zcut
  logical, private :: verbose

  real,    allocatable,private :: flds(:,:,:),witap(:)
  real,    allocatable,private :: gnorm(:,:),RGrad(:,:),TGrad(:,:),RVel(:,:)
  complex, allocatable,private :: inject(:,:),SGF(:,:),model(:,:),resid(:,:)
  complex, allocatable,private :: rwf(:,:),swf(:,:),wfld(:,:)

contains

  subroutine routines_init(ag,at,ar,as,ax,az,aw,raysin)
    type(myaxis)        :: ag,at,ar,as,ax,az,aw
    integer :: ig,nzcut_in,izcut(1)
    real :: raysin(:,:,:)

    write(0,*) 'IN INIT'
    pi=acos(-1.)
    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
    os=as%o;    ds=as%d;    ns=as%n
    ox=ax%o;    dx=ax%d;    nx=ax%n
    oz=az%o;    dz=az%d;    nz=az%n

    verbose=.false.
    write(0,*) nt,ng,nw,ns,nx,nz
    write(0,*) 'allocate flds'
    allocate (flds(ng,4,nt))
     write(0,*) 'shit flds'
   allocate (witap(nt) )
     write(0,*) 'allocated witap'
   allocate(resid(nt,1))

    write(0,*) 'shit'
    allocate (gnorm(nt,ng),RGrad(nt,ng),inject(nt,1),SGF(ng,nt),wfld(ng,nt))
write(0,*) 'balls'
    allocate (RVel(nt,ng),TGrad(nz,nx),model(nt,1),swf(ng,1),rwf(nt,1) )
    
    write(0,*) 'ALLOCATED ARRAYS'
    call from_param("omin",omin,2000.)
    call from_param("omax",omax,2400.)
    call from_param("verbose",verbose,.false.)
    call from_param("norm",norm,1)
    call from_param("nsx",nsx,3)
    call from_param("nsz",nsz,3)
    call from_param("zcut",zcut,0.)
    call from_param("ptap",ptap,20)
    call from_param("ntap",ntap,10)
    write(0,*) 'INPUT VARIABLES'

    call wftaper_init(omin,omax)
    write(0,*) 'WFTAPER'
    call C2R_init(nsx,nsz,norm,dx,dz,nx,nz,ng,nt,ox,oz)
    call wemig_init(ag,at,aw,ar,verbose,nzcut)
    !! Where to cut

    izcut = minloc(abs(raysin(1,:,2) - zcut))
    nzcut= izcut(1)
    write(0,*) 'CUTTING AT ',nzcut,' of ', ag%n 

    call wemig_init(ag,at,aw,ar,verbose,nzcut)

  end subroutine routines_init

  !********************************************************

  subroutine Calc_Gradient(CVel,raysin,data,source,w,CGrad,OF)
    integer ::it,ig,ix,iz,is
    real :: w,OF
    real  :: CVel(:,:),raysin(:,:,:),CGrad(:,:)
    complex :: data(:,:),source(:,:)
    logical :: adj

    !! Init Geometry each time (to prevent memory leak)
    call geometry_init(nt,ng,dt,dg,verbose)

    shots: do is=1,ns   !! . . Loop over shot profiles

       !! . . Initalization and info
       !     call WI_ZERO(RVel,gnorm,fields,swf,rwf,inject,wfld,model,resid,SGF,RGrad,TGrad)
       RVel=0.
       gnorm=0.
       flds=0.
       swf=0.
       rwf=0.
       inject=0.
       wfld=0.
       model=0.
       resid=0.
       SGF=0.
       RGrad=0.
       TGrad=0.

       !! !!!!!!!!!!!!!!!!!!!!!!!!!!!!
       !!
       !! GEOMETRY
       !! . . Calculate velocity on rays
       !! Input  :: CVel, raysing
       !! Output :: RVel,gnorm
       call C2R_run(RVel,raysin,CVel,gnorm,.false.)

       !! . . Calculate geometry of coordinates
       !! Input  :: raysing,Rvel,gnorm
       !! Output :: fields
       call geometry_calc(raysin,RVel,flds,gnorm)
       !!
       !! END GEOMETRY
       !!
       !! !!!!!!!!!!!!!!!!!!!!!!!!!

       !! Get copies of source and data for this shot
       do ig=1,ng
          swf(ig,1)=source(ig,is)
       end do
       do ig=1,nt
          rwf(it,1)=  data(it,is)
       end do

       !! . . Forward propagation of swf to generate modeled wavefield
       !! Input  :: swf,inject,adj,ii,fields,gnorm,w
       !! Output :: wfld
       adj=.true.
       call wemig(swf,inject,wfld,adj,is,flds,gnorm,w,RVel)

       !! Cut out modeled data from proper location
       do it=1,nt
          model(it,1)=wfld(nzcut,it)
       end do

       !! Taper data
       !! IN/OUT :: model,data
       call wftaper(model,rwf)

       !! . . Calculate residuals 
       do it=1,nt
          resid(it,1) = model(it,1)-rwf(it,1)
       end do

       do it=1,nt
          OF = OF + 1./2.*real ( conjg(model(it,1)-rwf(it,1))*(model(it,1)-rwf(it,1)))
       end do

       !! Store temp G(x,s) Wavefield
       do ig=1,ng
          do it=1,nt
             SGF(ig,it) = wfld(ig,it)
          end do
       end do

       !! . . Back Propagate Residuals
       !! Input  :: swf,resid,is,fields,gnorm,w
       !! Output :: wfld
       adj=.false.
       wfld=0.;
       swf=0.;
       call wemig(swf,resid,wfld,adj,is,flds,gnorm,w,RVel)

       !! . . Calculate Gradient
       do ig=1,ng
          do it=1,nt
             RGrad(it,ig)=-w**2*real(SGF(ig,it)*wfld(ig,it))
          end do
       end do

       !! . . Interpolate to Cartesian
       !! Input  :: RGrad,raysin,gnorm
       !! Output :: CGrad
       call C2R_init(nsx,nsz,norm,dx,dz,nx,nz,ng,nt,ox,oz)
       call C2R_run(RGrad,raysin,TGrad,gnorm,.true.)

       !! . . Update CGrad
       do iz=1,nz
          do ix=1,nx
             CGrad(iz,ix) = CGrad(iz,ix)+TGrad(iz,ix)
          end do
       end do

    end do shots  !! . . End Shot-point loop

    call geometry_deallocate()

  end subroutine Calc_Gradient

  !*****************************************************************

  subroutine Calc_ObjFun(CVel,raysin,data,source,w,OF)
    integer :: ig,it,is
    real :: w,OF
    real  :: CVel(:,:),raysin(:,:,:)
    complex :: data(:,:),source(:,:)
    logical :: adj

    call geometry_init(nt,ng,dt,dg,verbose)
    call C2R_init(nsx,nsz,norm,dx,dz,nx,nz,ng,nt,ox,oz)
    call  wftaper_init(omin,omax)
    !! Value of the Objective Function
    OF=0.

    do is=1,ns   !! . . Loop over shot profiles
       !! . . Initalization and info
       RVel=0.; 
       wfld=0.;
       data=0.;
       inject=0.;
       resid=0.
       model=0.; 
       swf=0.; 
       rwf=0.;

       !! Get copies of source and data for this shot
       do ig=1,ng
          swf(ig,1)=source(ig,is)
       end do
       do ig=1,nt
          rwf(it,1)=data(it,is)
       end do

       !! . . Calculate velocity on rays
       !! Input  :: CVel, raysing
       !! Output :: RVel,gnorm
       call C2R_run(RVel,raysin,CVel,gnorm,.false.)

       !! . . Calculate geometry of coordinates
       !! Input  :: raysing,Rvel,gnorm
       !! Output :: fields
       call geometry_calc(raysin,RVel,flds,gnorm)

       !! . . Forward propagation of swf to generate modeled wavefield
       !! Input  :: swf,inject,adj,is,fields,gnorm,w
       !! Output :: wfld
       adj=.true.
       call wemig(swf,inject,wfld,adj,is,flds,gnorm,w,RVel)

       !! Cut out modeled data from proper location
       do it=1,nt
          model(it,1)=wfld(nzcut,it)
       end do

       !! Taper data
       !! IN/OUT :: model,data
       call wftaper(model,rwf)

       !! . . Calculate  
       do it=1,nt
          OF = OF + 1./2.*real ( conjg(model(it,1)-rwf(it,1))*(model(it,1)-rwf(it,1)))
       end do

    end do !! . . End Shot-point loop

    call geometry_deallocate()

  end subroutine Calc_ObjFun

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

  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)

    witap=0.
    witap(imin:nt)=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)
    end do
    witap(nt-4)=0.8;
    witap(nt-3)=0.6;
    witap(nt-2)=0.4;
    witap(nt-1)=0.2;
    witap(nt  )=0.0

    deallocate(wfline)
  end subroutine wftaper_init
  !----------------------------------------------------------------

  subroutine wftaper(model,data)
    complex :: model(:,:),data(:,:)
    integer :: it
    do it=1,nt
       model(it,1)=model(it,1)*witap(it)
    end do
    do it=1,nt
       data (it,1)=data (it,1)*witap(it)
    end do
  end subroutine wftaper
  !----------------------------------------------------------------

  subroutine ray_location(raysin)
    real :: raysin(:,:,:)
    write(0,*) '***************************************'
    write(0,*) '*                                     *'
    write(0,*) '* RAYS STARTING AT: ',minval(raysin(:,:,1)),   '*'
    write(0,*) '*                                     *'
    write(0,*) '***************************************'

  end subroutine ray_location

  !----------------------------------------------------------------
!! Taper Initialization Program

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

  subroutine routine_info(at,ag,as,aw,ax,az)
    type(myaxis)       :: at,ag,as,aw,ax,az
    write(0,*) 'DEPTH',at%n,at%o,at%d
    write(0,*) 'H-POS',ag%n,ag%o,ag%d
    write(0,*) 'SHOTS',as%n,as%o,as%d
    write(0,*) 'FREQ ',aw%n,aw%o,aw%d
    write(0,*) 'VEL X',ax%n,ax%o,ax%d
    write(0,*) 'VEL Z',az%n,az%o,az%d

  end subroutine routine_info

  subroutine Get_WI_axes(ag,aw,as,at,az,ax)
    type(myaxis)      :: ag,aw,as,at,az,ax
    character(len=128) :: name

    !! Read in input dimensions
    name="swf"
    call ainaxis(ag,1,name) 
    call ainaxis(aw,2,name) 
    call ainaxis(as,3,name) 

    name="rays"
    call ainaxis(at,1,name)

    !!    Read in velocity file
    name="vel"
    call ainaxis(az,1,name)
    call ainaxis(ax,2,name)

  end subroutine Get_WI_axes

end module  RWE2D_WI_routines
