program RWE2D_WI
  use sep
  use util_
  use RWE2D_WI_wemig
  use C2R2D

  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,iw,ishot(1),iis,mxl,mxr,oldn,ig,it,numit,in
  integer :: nref,nsx,nsz,norm,cnx,cnz,migapp,ptap
  real    :: forward,xmin,zmin,xmax,zmax,dxx,dzz,omin,omax,pmax
  real,    allocatable,dimension(:,:)  :: Vel,gnorm,Velin,dVel,Vtmp,CDnew,CDold 
  real,    allocatable,dimension(:,:)  :: gradient,Steplength,gtmp,Dwfld
  real,    pointer    ,dimension(:,:)  :: HessScale,Htmp,dVelold
  real,    allocatable,dimension(:,:,:):: raysin      

  complex, pointer    ,dimension(:,:)  :: swf,model
  complex, pointer    ,dimension(:,:)  :: data,inject,resid
  complex, allocatable,dimension(:,:,:):: SGF,wfld,swftmp,datatmp

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

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

  !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,16)        !! # 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)
  write(0,*) 'READ IN PARAMETERS'

  !! 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)
  write(0,*) 'INPUT AXIS'
  write(0,*) 'DEPTH',at%n,at%o,at%d
  write(0,*) 'X1   ',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",1,"image")
  call auxputch("d3","r",1.,"image")
  call auxputch("o3","r",0.,"image")

  call aouaxis(az,1,"HessScale.H")
  call aouaxis(ax,2,"HessScale.H")
  call aouaxis(az,1,"Steplength.H")
  call aouaxis(ax,2,"Steplength.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("esize","i",4,"Rimage")

  call aouaxis(at,1,"Dwfld.H")
  call aouaxis(ag,2,"Dwfld.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("esize","i",8,"Perturb.H")

  allocate( swf(ag%n,aw%n),model(at%n,aw%n),dVel(az%n,ax%n),Vtmp(az%n,ax%n),CDold(az%n,ax%n) )
  allocate( data(at%n,aw%n),wfld(ag%n,at%n,aw%n),SGF(ag%n,at%n,aw%n),inject(at%n,aw%n))
  allocate( Dwfld(at%n,ag%n),swftmp(ag%n,aw%n,as%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) )
  allocate( raysin(at%n,ag%n,2),gnorm(at%n,ag%n),resid(at%n,aw%n),datatmp(at%n,aw%n,as%n) )
  allocate( mask(ag%n ,at%n), refs(nref,6,at%n),fields(ag%n,6,at%n),CDnew(az%n,ax%n) )
  allocate( dVelold(az%n,ax%n),gtmp(az%n,ax%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)

  !! . . 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)
  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")

  CDold=0.;CDnew=0.

  raysin(:,:,1)=raysin(:,:,1)+as%o-as%d
  do in=1,numit
     !!-------------------------------------------------
     !!
     !! . . Forward Modelling Section (x2)
     !!
     dVel=0.; dVelold=0.;gradient=0.;
     do ii=1,as%n   !! . . Loop over shot points

        write(0,*) '-------------------------------------------'
        write(0,*) '-'
        write(0,*) '-  Forward Modelling '
        write(0,*) '-  WORKING ON SHOT: ', ii, ' of ', as%n
        write(0,*) '-  Iteration      : ', in, ' of ',numit 
        write(0,*) '-'
        write(0,*) '------------------------------------------'

        Vel=0.; swf = 0.;wfld=0.;data=0.;inject=0.;SGF=0.;gtmp=0.;resid=0.
        gradient=0.
        raysin(:,:,1)=raysin(:,:,1)+as%d

        adj=.false.; !! Calculate velocity on rays
        call C2R_run(Vel,raysin,Velin,gnorm,adj)
        write(0,*) 'VEL  IN: ',minval(Velin),maxval(Velin)
        write(0,*) 'VEL OUT: ',minval(Vel  ),maxval(Vel  )

        !! . . Set up geometry
        call geometry_calc(raysin,Vel,fields,refs,mask,gnorm)
        !     call geometry_close()

        !! . . Assign starting wavefield
        swf  =  swftmp(:,:,ii);   data = datatmp(:,:,ii)
        write(0,*) 'SWF  MIN/MAX', minval(abs(swf )),maxval(abs(swf ))
        write(0,*) 'DATA MIN/MAX', minval(abs(data)),maxval(abs(data))

        !! . . Forward propagation of swf to generate modelled wavefield
        adj=.true.
        call wemig(swf,inject,wfld,adj,ii,fields,refs,mask,gnorm )
        write(0,*) 'FORWARD PROPAGATED SWF TO GENERATE MODEL WAVEFIELD'

        !! . . Get modelled wavefield, apply taper and find residuals
        SGF = wfld  !! . . Source Greens function
        model=wfld(ag%n/2,1:at%n,:)
        
        call wftaper(model,data);

        do iw=1,1!,aw%n
           model(:,iw)=model(:,iw)*maxval(abs(data(:,iw)))/maxval(abs(model(:,iw)))
        end do

        call Diffdata(model,data,resid)

        adj=.false.;wfld=0.;swf=0.;!! . . Back propagate residuals
        call wemig(swf,resid,wfld,adj,ii,fields,refs,mask,gnorm)

        call Perturbation_calc(SGF,wfld,gradient) ; !! . . Image Vel Perturb

        adj=.true.; !! . . 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(gradient,raysin,gtmp,gnorm,adj)
        dVel=dVel+1000.*gtmp
     end do !! . . End Shot-point loop
     write(0,*) 'Min/Max Gradient:',minval(gradient),maxval(gradient)

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

     !!------------------------------
     !!
     !! . . Update the Velocity model
     !!
     !!------------------------------
       
     where (dVel .ne. 0.)
        CDnew = dVel - dVel*(dVel-dVelold)/(dVel**2)*CDold
     end where

     where (dVel .eq. 0.)
        CDnew=0.
     end where

     CDold=0.! CDnew;
     dVelold=dVel
     ! dVel is the calculated perturbation
     ! Vtmp is the velocity used for 3rd updating
     Vtmp=0.;  !! Perturbed velocity

!     Vtmp = Velin + dVel*pmax*maxval(Velin)
      Vtmp = Velin + pmax*dVel * maxval(Velin)/maxval(abs(dVel))
    write(0,*) 'Cartesian domain Perturb.: ',minval(Vtmp),maxval(Vtmp) 
     !  call sep_write(dVel,"image")

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

     raysin(:,:,1)=raysin(:,:,1)-as%n*as%d
     Dwfld=0.
     do ii=1,as%n   !! . . Loop over shot points
        write(0,*) '-------------------------------------------'
        write(0,*) '-'
        write(0,*) '-  Hessian Modelling '
        write(0,*) '-  WORKING ON SHOT:', ii, ' of ', as%n
        write(0,*) '-  Iteration      :', in, ' of ',numit 
        write(0,*) '-'
        write(0,*) '------------------------------------------'

        adj=.false. ; !! Calculate velocity on rays
        raysin(:,:,1)=raysin(:,:,1)+as%d

        Vel=0.
        call C2R_run(Vel,raysin,Vtmp,gnorm,adj)
        write(0,*) 'VEL  IN: ',minval(Vtmp),maxval(Vtmp)
        write(0,*) 'VEL OUT: ',minval(Vel  ),maxval(Vel  )

        !! . . Set up geometry
        call geometry_calc(raysin,Vel,fields,refs,mask,gnorm)
        call geometry_close()

        !! . . Input starting wavefield
        swf = 0.; wfld=0.; data=0.;inject=0.
        swf = swftmp(:,:,ii)

        !! . . Forward propagation of swf in pert. vel. to generate wavefield
        adj=.true.;
        call wemig(swf,inject,wfld,adj,ii,fields,refs,mask,gnorm )

        call Diffwavefield(SGF,wfld,Dwfld)

        adj=.true.; !! . . 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)
        Htmp=0.
        call C2R_run(Dwfld,raysin,Htmp,gnorm,adj)
        HessScale=HessScale+Htmp
     end do !! . . End Shot-point loop
     raysin(:,:,1)=raysin(:,:,1)-as%n*as%d

     if (minval(HessScale) .lt. 0.) HessScale=HessScale-2.*minval((HessScale))

     call sep_write(HessScale,'HessScale.H')
     call sep_write(Dwfld,'Dwfld.H')

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

     where ((HessScale/(pmax**2)+dVel**2 ) .ne. 0.) 
        Steplength = dVel**2 / (HessScale/(pmax**2)+dVel**2 )
     end where

     call sep_write(Steplength,'Steplength.H')

     write(0,*) 'dVel: ', minval(dVel),maxval(dVel)
     write(0,*) 'Hess: ', minval(10000*HessScale),maxval(10000*HessScale)

     !! . . Step length 
     write(0,*) 'STEP LENGTH: ',minval(Steplength),maxval(Steplength)

     Velin = Velin + Steplength * CDnew !dVel
     HessScale=0.; Steplength=0.;dVel=0.; Dwfld=0.;
  end do

  write(0,*) 'Interpolated Perturb.: ',minval(Velin),maxval(Velin)
  call sep_write(Velin,"image")

  call sep_end_prog()
  call exit()
end program RWE2D_WI
