program RWE2D_wavepath
  use sep
  use util_
  use RWE2D_wavepath_wemig
  use C2R2D
  use sep_timers_mod
  implicit none
  !! ax is the Cartesian length dimension associated with Velocity 
  !! az is the Cartesian depth  dimension associated with Velocity 
  !! aw is the frequency axis
  !! as is the shot axis
  !! ag is the offset axis in Ray space
  !! at is the extrapolation direction in Ray space
  !! ar is the number of references
  type(myaxis) :: ax,az,aw,ar,as,at,ag

  integer :: ix,iz,ii,iw2,ishot(1),iis,mxl,mxr,oldn,ig,it,numit,in,nw,izcut(1),ixcut(1),ib
  integer :: nref,nsx,nsz,norm,cnx,cnz,migapp,ptap,nzcut,nxcut,nblock
  real    :: forward,xmin,zmin,xmax,zmax,dxx,dzz,omin,omax,pmax,zcut,jscale,njs
  real    :: rms,rmsold,www,r0
  real,    allocatable,dimension(:,:)  :: Vel,gnorm,Velin,dVel,Vtmp,CDnew,CDold,dVelold
  real,    allocatable,dimension(:,:)  :: gradient,Steplength,gtmp,Dscale,Sfilter
  real,    pointer    ,dimension(:,:)  :: HessScale,Htmp
  real,    allocatable,dimension(:,:,:):: raysin
  real,    allocatable,dimension(:,:):: Dwfld
  complex, pointer    ,dimension(:,:)  :: data,swf,model,resid,inject
  complex, allocatable,dimension(:,:,:):: swftmp,datatmp,wfld,SGF,Hwfld

  logical                              :: verbose,kin,adj,wf,is1D,usecg,forwardmodel
  character(len=128) :: name

  !! Geometry
  integer, allocatable :: mask(:,:)
  real,    allocatable :: refs(:,:,:), fields(:,:,:)

  !! Timers
  integer :: t1,t2,t3,t4,t5,t6,t7,t8,t9,t10,t11,t12,t13,t14,t15
  logical :: logic


  !call sep_init(SOURCE)
  call sep_init()
  call getch_add_string("head=/dev/null")
  call getch_add_string("noheader=y")
  call set_no_putch()
  call sep_begin_prog()

  write(0,*) 'STARTING PROGRAM'

  call from_param("forward",forward,0.) !! Forward scattering option
  call from_param("nref",nref,30)        !! # reference velocities
  call from_param("verbose",verbose,.false.)
  call from_param("kinematic",  kin,.true.)
  call from_param("norm",norm,1)
  call from_param("nsx",nsx,3)
  call from_param("nsz",nsz,3)
  call from_param("omin",omin,2000.)
  call from_param("omax",omax,2400.)
  call from_param("wf",wf,.false.)
  call from_param("pmax",pmax,0.01)
  call from_param("numit",numit,1)
  call from_param("ptap",ptap,10)
  call from_param("zcut",zcut,0.)
  call from_param("jscale",jscale,1.)
  call from_param("is1D",is1D,.true.)
  call from_param("usecg",usecg,.true.)
  call from_param("nblock",nblock,1)
  call from_param("forwardmodel",forwardmodel,.true.)
  !! 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)
  !! Do only 1 frequency
  nw = aw%n
  aw%n=1


  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


  !! . . Create offset axis
  call aouaxis(az,1,"image")
  call aouaxis(ax,2,"image")
  call auxputch("esize","i",4,"image")
  call auxputch("n3","i",nw,"image")
  call auxputch("d3","r",1.,"image")
  call auxputch("o3","r",0.,"image")

  !  call aouaxis(ag,1,"DH.H")
  !  call aouaxis(at,2,"DH.H")!

  !  call aouaxis(at,1,"gnorm.H")
  !  call aouaxis(ag,2,"gnorm.H")

  !  call auxputch("esize","i",8,"DH.H")
  !  call auxputch("n3","i",2,"DH.H")
  !  call auxputch("d3","r",1.,"DH.H")
  !  call auxputch("o3","r",0.,"DH.H")

  call aouaxis(az,1,"HessScale.H")
  call aouaxis(ax,2,"HessScale.H")
  call aouaxis(az,1,"Steplength.H")
  call aouaxis(ax,2,"Steplength.H")
  call auxputch("n3","i",numit,"Steplength.H")
  call auxputch("n4","i",nw,"Steplength.H")

  call aouaxis(az,1,"CDnew.H")
  call aouaxis(ax,2,"CDnew.H")
  call auxputch("n3","i",numit,"CDnew.H")
  call auxputch("n4","i",nw,"CDnew.H")

  call aouaxis(az,1,"Dvel.H")
  call aouaxis(ax,2,"Dvel.H")
  call auxputch("n3","i",numit,"Dvel.H")
  call auxputch("n4","i",nw,"Dvel.H")

  !! . . Output Imaged WF in RWE coordinates
  call aouaxis(at,1,"Rimage")
  call aouaxis(ag,2,"Rimage")
  call auxputch("n3","i",numit,"Rimage")
  call auxputch("n4","i",nw,"Rimage")
  call auxputch("esize","i",4,"Rimage")

  call aouaxis(at,1,"Dwfld.H")
  call aouaxis(ag,2,"Dwfld.H")
  call auxputch("n3","i",numit,"Dwfld.H")
  call auxputch("n4","i",nw   ,"Dwfld.H")

  call aouaxis(at,1,"Sfilter.H")
  call aouaxis(ag,2,"Sfilter.H")
  call auxputch("n3","i",numit,"Sfilter.H")
  call auxputch("n4","i",nw   ,"Sfilter.H")


  call aouaxis(at,1,"Perturb.H")
  call aouaxis(aw,2,"Perturb.H")
  call auxputch("n3","i",3 ,"Perturb.H")
  call auxputch("o3","r",0.,"Perturb.H")
  call auxputch("d3","r",1.,"Perturb.H")
  call auxputch("n4","i",as%n,"Perturb.H")
  call auxputch("n5","i",numit,"Perturb.H")
  call auxputch("n6","i",1+nw/nblock, "Perturb.H")

  call auxputch("esize","i",8,"Perturb.H")

  !! Input data files
  allocate( swftmp(ag%n,nw,as%n),datatmp(at%n,nw,as%n),Dscale(nw,as%n) )

  !! Along data line
  allocate(  model(at%n,nblock),data(at%n,nblock),inject(at%n,nblock),resid(at%n,nblock) )

  !! Extrapolation
  allocate( swf(ag%n,nblock) )

  !! Velocity and Step
  allocate( dVel(az%n,ax%n),Vtmp(az%n,ax%n),CDold(az%n,ax%n),CDnew(az%n,ax%n) )
  allocate( dVelold(az%n,ax%n),gtmp(az%n,ax%n),Sfilter(at%n,ag%n) )
  allocate( Dwfld(at%n,ag%n),HessScale(az%n,ax%n),Htmp(az%n,ax%n) )
  allocate( Vel(at%n,ag%n),Velin(az%n,ax%n),gradient(at%n,ag%n),Steplength(az%n,ax%n) )

  !! Wavefield
  allocate( Hwfld(ag%n,at%n,nblock),wfld(ag%n,at%n,nblock),SGF(ag%n,at%n,nblock))

  !! Geometry
  allocate( mask(ag%n ,at%n), refs(nref,6,at%n),fields(ag%n,6,at%n) )
  allocate( raysin(at%n,ag%n,2),gnorm(at%n,ag%n) )

  !! . . Read in Rays and Velocity
  call sep_read(raysin,"rays")
  call sep_read(Velin ,"vel ")
  write(0,*) 'VELOCITY MIN/MAX', minval(Velin),maxval(Velin)

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

  !! . . Initialize Routines
  call geometry_init(at%n,ag%n,at%d,ag%d,kin,nref,verbose)
  call wemig_init(ag,at,aw,ar,forward,kin,nref,verbose,ptap,nzcut,nblock)
  call C2R_init(0,0,norm,ax%d,az%d,ax%n,az%n,ag%n,at%n,ax%o,az%o)
  call wftaper_init(omin,omax)

  !! read in Source and data wavefields
  call sep_read(swftmp ,"swf")
  call sep_read(datatmp,"data")


  do iw2=1,nw
     do ii=1,as%n
        write(0,*) maxval(abs(swftmp(:,iw2,ii))),maxval(abs(datatmp(:,iw2,ii)))
        Dscale(iw2,ii) = jscale

        datatmp(:,iw2,ii)=datatmp(:,iw2,ii)/&
                    maxval(abs(datatmp(:,iw2,ii)))/sqrt( dble(iw2-1)*aw%d+aw%d)
        swftmp(:,iw2,ii)=  swftmp(:,iw2,ii)/&
                    maxval(abs( swftmp(:,iw2,ii)))/sqrt( dble(iw2-1)*aw%d+aw%d)
        !  swftmp(:,iw2,ii)=jscale* swftmp(:,iw2,ii)/(abs(datatmp(1,iw2,ii)))
        write(0,*) maxval(abs(swftmp(:,iw2,ii))),maxval(abs(datatmp(:,iw2,ii)))

     end do
  end do
  datatmp = datatmp*jscale
  swftmp =  swftmp*jscale
  CDold=0.;CDnew=0.

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !!
  !! . . MAIN PROGRAM
  !!
  njs = jscale
  logic=init_sep_timers()
  logic=setup_next_timer("Frequencies",t1)
  logic=setup_next_timer("NL Interations",t2)
  logic=setup_next_timer("Shots",t3)
  logic=setup_next_timer("Vel update",t4)
  logic=setup_next_timer("StepCalc",t5)
  logic=setup_next_timer("geometry",t6)
  logic=setup_next_timer("Interp1",t7)
  logic=setup_next_timer("DataAssign",t8)
  logic=setup_next_timer("FP1",t9)
  logic=setup_next_timer("Reshape",t10)
  logic=setup_next_timer("Residual",t11)
  logic=setup_next_timer("BPResid",t12)
  logic=setup_next_timer("DVelcalc",t13)
  logic=setup_next_timer("Interp2",t14)
  logic=setup_next_timer("Output",t15)



  call ray_location(raysin)
  do iw2=1,nw,nblock !! . . Loop over input frequencies
     call start_timer_num(t1)
     !     raysin(:,:,1)=raysin(:,:,1)+as%o
     www = 2.*3.1415926*( (iw2-1.)*aw%d+aw%o )
     rms = 0.;rmsold=0;r0=0.
     do in=1,numit !! . . Number of non-linear iterations
        call start_timer_num(t2)
 
        !        if (in .gt. 1 .and. rms .gt. 1.05*rmsold) cycle
        !! . . Forward Modelling Section (x2)
        dVel=0.;
        do ii=1,as%n   !! . . Loop over shot profiles

           !! . . Initalization and info
           Vel=0.; swf = 0.;wfld=0.;data=0.;inject=0.;SGF=0.;gtmp=0.;resid=0.
           gradient=0.;model=0.
           call start_timer_num(t3)
           call info_one(ii,iw2,www,as%n,nw,in,numit)

           !! . . Calculate velocity on rays
           call start_timer_num(t7)
           call C2R_run(Vel,raysin,Velin,gnorm,.false.)
           call stop_timer_num(t7)
           call start_timer_num(t6)
           call geometry_calc(raysin,Vel,fields,refs,mask,gnorm)
           call stop_timer_num(t6)

           !! . . Assign starting wavefield
           call start_timer_num(t8)
           swf  =  swftmp(:,iw2:iw2+nblock,ii)   
           data = datatmp(:,iw2:iw2+nblock,ii)
           call stop_timer_num(t8)

           !! . . Forward propagation of swf to generate modelled wavefield
           call start_timer_num(t9)    
           call wemig(swf,inject,wfld,.true.,ii,fields,refs,mask,gnorm,iw2)
           call stop_timer_num(t9)
           
           call start_timer_num(t10)    
           do ib=1,nblock
              wfld(:,:,ib) = 4.*wfld(:,:,ib)/((www+dble(ib-1)*aw%d)**0.55)
           end do

           model=wfld(nzcut,:,:)
           call wftaper(model,data);

           SGF   =  wfld
           call stop_timer_num(t10)

           !! . . Calculate residuals 
           call start_timer_num(t11)    
           call Diffdata(model,data,resid)
           rms = sqrt(sum(abs(conjg(resid(:,1))*resid(:,1))))
           if (in .eq. 1) r0=sqrt(sum(abs(conjg(data(:,1))*data(:,1))))
           write(0,*) '-  RESIDUALS L2: ',rms
           write(0,*) '-  RESIDUALS Change (in %): ',100*(rms-r0)/r0
           call stop_timer_num(t11)

           !! . . Back Propagate Residuals
           call start_timer_num(t12)    
           wfld=0.;swf=0.;!! . . Back propagate residuals
           call wemig(swf,resid,wfld,.false.,ii,fields,refs,mask,gnorm,iw2)
           call stop_timer_num(t12)

           !! . . Calculate Gradient
           call start_timer_num(t13)    
           call Perturbation_calc(SGF,wfld,gradient,iw2) ; !! . . Image Vel Perturb
           call stop_timer_num(t13)

           !! . . Interpolate to Cartesian
           call start_timer_num(t14)    
           call C2R_init(nsx,nsz,norm,ax%d,az%d,ax%n,az%n,ag%n,at%n,ax%o,az%o)
           call C2R_run(gradient,raysin,gtmp,gnorm,.true.)
           call stop_timer_num(t14)

           call PixelBalance(gtmp)
           call CartSurface_filter(gtmp,www)

           !! . . If 1D set values equal to ray center
           call start_timer_num(t15)    
           if (is1D .eq. .true.) then
              do iz=1,az%n
                 gtmp(iz,:)=sum(gtmp(iz,:))!/dble(ax%n)
              end do
              write(0,*) '-  1D assumption used:'
           end if


           !! . . Normalize by magnitude of data . . balance out
           dVel=dVel+gtmp
           call srite("Dvel.H",dVel,size(dVel)*4)
           !raysin(:,:,1)=raysin(:,:,1)+as%d
           call stop_timer_num(t15)

           call stop_timer_num(t3)

        end do !! . . End Shot-point loop

        !      call sep_write(gradient,"Rimage");    !! Write out RWE-based image field
        call start_timer_num(t4)

        !!------------------------------
        !!
        !! . . Update the Velocity model
        !!
        !!------------------------------

        !        if (in .gt. 1 .and. rms .gt. 1.05*rmsold) cycle

        if (usecg .eq. .true.) then
           write(0,*) '-  Using Conjugate Gradient '
           if (in .eq. 1) then
              CDnew = dVel
           else
              where (abs(dVel/dVelold) .gt. 0.1)
                 CDnew=dVel+dVel*(dVel-dVelold)/(dVelold**2+0.1)*CDold
              elsewhere
                 CDnew=dVel
              end where
           end if
        else
           write(0,*) '-  Using Gradient '
           CDnew = dVel
        end if
        write(0,*) '-  PERTURBATION', minval(dVel),maxval(dVel)
        CDold = CDnew;
        dVelold = dVel

        call srite("CDnew.H",CDnew,size(CDnew)*4)

        Vtmp=0.; 

!        Vtmp = Velin + pmax*dVel* maxval(Velin)/maxval(abs(dVel))
        Vtmp = Velin + pmax*dVel

        call stop_timer_num(t4)

        !!------------------------------
        !!
        !! . . Calculate Conjugate Direction
        !!
        !!------------------------------

        !raysin(:,:,1)=raysin(:,:,1)-(as%n)*as%d
        Dwfld=0.
        do ii=1,as%n   !! . . Loop over shot points
           Vel=0.
           call start_timer_num(t5)
           write(0,*) '-  Hessian Modelling '

           call C2R_run(Vel,raysin,Vtmp,gnorm,.false.)

           call geometry_calc(raysin,Vel,fields,refs,mask,gnorm)
           write(0,*) '-  MIN/MAX Vel: ',minval(Vel),maxval(Vel)

           !! . . Input starting wavefield
           swf = 0.; Hwfld=0.; data=0.;inject=0.
           swf = swftmp(:,iw2:iw2+nblock,ii)

           !! . . Forward propagation of swf in pert. vel. to generate wavefield
           !!call wemig(swf,inject,Hwfld,.true.,ii,fields,refs,mask,gnorm,iw2 )
           call Diffwavefield(SGF,Hwfld,Dwfld)
           !call sep_write(Dwfld,'Dwfld.H')

           !! . . Interpolate Perturbation to Cartesian
           call C2R_init(nsx,nsz,norm,ax%d,az%d,ax%n,az%n,ag%n,at%n,ax%o,az%o)         
           call C2R_run(Dwfld,raysin,HessScale,gnorm,.true.)
           !raysin(:,:,1)=raysin(:,:,1)+as%d
           call stop_timer_num(t5)
        end do !! . . End Shot-point loop

        !  raysin(:,:,1)=raysin(:,:,1)-(as%n)*as%d

        where ( abs(HessScale/(pmax**2)+dVel**2 ) .eq. 0.0 )
           Steplength = 0.
        elsewhere
           Steplength = dVel**2 / (HessScale/(pmax**2)+dVel**2 )
        end where

        !        if (in .eq. 1) then  

        Velin = Velin + Steplength*CDnew
        !        else
!                   Velin = Velin + (r0-rms)/abs(r0-rms)*Steplength*CDnew
        !          write(0,*) '-  (RMSOLD-RMS)/RMS: ',  (rmsold-rms)/abs(rmsold-rms)
        !        end if

        call PixelBalance(Velin)

        call info_two(Velin)
        call stop_timer_num(t2)
        rmsold=rms
        HessScale=0.; Steplength=0.;dVel=0.; Dwfld=0.;
     end do
     call sep_write(Velin,"image")
     call stop_timer_num(t1)

  end do
  write(0,*) 'Interpolated Perturb.: ',minval(Velin),maxval(Velin)
  call print_timers()

  call sep_end_prog()
  call exit()
end program RWE2D_wavepath
