program RWE2D
  !! !!!!!!!!!!!!!!!!!!!
  !!
  !! 		Written by Jeff Shragge
  !!				Feb 18/2008
  !!				contact: jshragge@gmail.com
  !!
  !! !!!!!!!!!!!!!!!!!!!!!!!!!!

  use sep           !! Basic SEPlib I/O 
  use util_         !! handles axis I/O and aux file I/O
  use RWE2D_wemig   !! Migration driver module
  use C2R2D         !! Ray-to-Cartesian interpolation 
  use off2ang2D

  implicit none
  type(myaxis) :: aE1,aE3,aW,aS1,aX1,aX3,aR1,aH1,aA1,aZZ

  integer :: iX1,iX3,ii,iw,iis,iE1,iE3,iS1,method,ih,an
  integer :: nsx,nsy,nsz,norm,itest,nbytes,i1,i3,ncx,ncz,nh
  real    :: r0,eps,rextra,minang,maxang,rscale,atloc,ad,ao
  real    :: ddx1,ddx3,dd1,dd3,rro,rrd,xloc
  character (len=128) :: name
  integer, allocatable :: iix1(:)
  real,    allocatable :: ox1(:)
  real,    allocatable,dimension(:,:)  :: vtmp
  real,    allocatable,dimension(:,:)  :: RVel,CVel,gnorm,fields
  real,    allocatable,dimension(:,:,:):: CImg,Rimg,Image
  real,    allocatable,dimension(:,:,:):: raysin
  complex, allocatable,dimension(:,:)  :: swfsave,rwfsave,rwf,swf  
  logical :: verbose,adj,newimage,planewave

  !! !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !!
  !! . . HEADERS IN/OUT 
  !!
  call sep_init()
  call getch_add_string("head=/dev/null")
  call getch_add_string("noheader=y")
  call set_no_putch()
  call auxinout("image")
  call sep_begin_prog()
  call from_param("newimage",newimage,.true.)
  call from_param("planewave",planewave,.false.)
  allocate( iix1(2),ox1(2) )

  !! 	Read in input dimensions from the "rwf" tag
  !!	Determines the radial steps grid parameters
  call RWF_Info(aE1,aW,aS1,aR1)

  !! . . Read in input dimensions of the "vel" tag
  !! . . Determines the global image size
  call VEL_Info(aX1,aX3)

  !! . . Obtain information on offsets
  call from_param("nh",aH1%n,1)
  aA1%n=aH1%n

  !! . . See if Image file exists. If not, then write one out..
  name = "image"
  if ( newimage ) then
     write(0,*) 'WRITING OUT NEW IMAGE VOLUME'
     call IMAGE_Out_Info(aX1,aX3,aH1)
     allocate ( Image( aX3%n,aX1%n,1) )
     do ih=1,aH1%n
        call srite('image', Image, 4*aX1%n*aX3%n)
     end do
     call sseek('image', 0 , 0 ) 
     deallocate( Image )
  end if

  call sseek('image', 0 , 0 )

  !! . . Read in the parameters
  call Read_In_Parameters(aE3,verbose,norm,nsx,nsz,eps,minang,maxang,method,newimage,nh)
  write(0,*) 'nh,aH1%n',nh,aH1%n
  !! . . Initialize propagation module
  call wemig_axes_init(aE1,aE3,aW,aX1,aX3,aS1,aR1,aH1,aA1)

  !! . . Set scaling of the coordinate system.  
  !! . . Cart is 1. EC is (nr-1)*dr/2. (half-width)
  if ( method .eq. 1 ) then    !! CARTESIAN
     call wemig_param_init(verbose,1.,minang,maxang,eps,method)
     r0 = 1.
  else !! ELLIPTIC
     r0 = dble(aE3%n-1.)*aR1%d/2.
     call wemig_param_init(verbose,r0,minang,maxang,eps,method)
  end if

  !! . . END OF HEADERS IN/OUT
  !!
  !! ------------------------------------------------------------
  !!   
  !! . . GEOMETRY SECTION
  !! . . Generate Coordinate System
  allocate( raysin(aE1%n,aE3%n,2),fields(aE1%n,aE3%n) )
  allocate( gnorm(aE1%n,aE3%n) )
  if (method .eq. 1) then      !! CARTESIAN
     call  CART_Setup(raysin,fields,gnorm,aE3%n,aE3%o,aE3%d)
     aE1%o=aR1%o; aE1%d=aR1%d; aH1%d=aE1%d ; aH1%o=-(aH1%n-1)/2*aH1%d
  else if (method .eq. 2) then !! ELLIPTIC
     call ELLIP_Setup(raysin,fields,gnorm,aE1%n,aE1%o,aE1%d)
     aH1%d=aE1%d; aH1%o=-(aH1%n-1)/2*aH1%d
  end if

  !! . . Initialize Angle Gathers
  aZZ%n=aE3%n; aZZ%d=aE3%d*r0; aZZ%o=aE3%o
  if (aH1%n .gt. 1) call off2ang2D_init(aE3,aE1,aH1,aA1,.true.,.false.)

  !! . . Reinitialize the axes because E1 or E3 could have changed
  call wemig_axes_init(aE1,aE3,aW,aX1,aX3,aS1,aR1,aH1,aA1)

  !! . . Write out parameters to header files  
  call write_out(aX1,aX3,aE1,aE3,aS1,aR1,aH1,aA1)
  call sep_write(raysin,"rays")
  call Geometry_QC(raysin,aE1,aE3,aS1,aW,aR1,aH1,aA1)
  an=aA1%n; ad=aA1%d; ao=aA1%o
  call IMAGE_Out_Info(aX1,aX3,aA1)
  aA1%n=an; aA1%d=ad; aA1%o=ao

  !! . . END OF GEOMETRY SECTION
  !!
  !! ---------------------------------------------------------------
  !!    
  !! . . START SOURCE WAVEFIELD SECTION
  !!
  !! . . Use only one swf for all imaging ... and shift it around

  if (planewave) then
     allocate( swfsave(aE1%n,aW%n),swf(aE1%n,aW%n) )
  else
     allocate( swfsave(aE1%n,aW%n),swf(aE1%n,aW%n) )
     call sreed("swf", swfsave , 8*aE1%n*aW%n) 
     !! . . Position wavefield on surface
     if ( method .eq. 1) then
        swf = swfsave
     else if (method .eq. 2) then
        call ELLIP_Position_XWF(swfsave,swf,raysin,aS1%o)
     end if
     deallocate( swfsave )
  end if
  !! . . END SOURCE WAVEFIELD SECTION
  !!
  !!----------------------------------------------------------------
  !!
  !! . . START MIGRATION SECTION
  !!
  !! . . Allocate space for receiver wavefield and local velocities
  allocate ( rwf (aE1%n, aW%n) , rwfsave (aE1%n, aW%n) )
  allocate (RVel (aE1%n,aE3%n) ,    RImg (aE1%n,aE3%n,nh) )
  !! . . Loop over shot points: inline and streamers
  INLINE_SHOT: do iS1 = 1,aS1%n   

     if (verbose) write(0,*) '        * RAYS at X:',minval(raysin(:,1,1)),maxval(raysin(:,1,1))
     if (verbose) write(0,*) '        * RAYS at Z:',minval(raysin(:,1,2)),maxval(raysin(:,1,2))

     if (planewave ) then
        swf=0. ; swfsave=0.
        call sreed("swf", swfsave , 8*aE1%n*aW%n) 
        !! . . Position wavefield on surface
        if ( method .eq. 1) then
           swf = swfsave
        else if (method .eq. 2) then
           call ELLIP_Position_XWF(swfsave,swf,raysin,(iS1-1.)*aS1%d+aS1%o)
        end if
     end if

     !! RECEIVER WAVEFIELD section
     !! Read in next Receiver wavefield
     rwf=0.;rwfsave=0.
     call sreed ("rwf", rwfsave , 8*aE1%n*aW%n ) !! Read in only current receiver wavefield
     if (method .eq. 1) then
        rwf = rwfsave
     else if (method .eq. 2) then
        call ELLIP_Position_XWF(rwfsave,rwf,raysin,(iS1-1.)*aS1%d+aS1%o)
     end if

     !! . . VELOCITY SECTION
     call Get_Vel_Indicies(raysin,iix1,ox1)
     ncx = iix1(2)-iix1(1)+1
     ncz = aX3%n

     allocate (  vtmp(ncz,ncx) ,CVel(ncz,ncx) )
     call Zero_arrays(RVel,RImg,ncz,ncx,nh,aE1,aE3,aH1)

     !! Read in velocity subset
     name="vel"
     call sseek("vel",0,0)
     nbytes = 4 * ( ncz*(iix1(1)-1) )
     call sseek("vel",nbytes,0)
     call sreed( name, CVel , ncz*ncx*4)

     !! . . Do Cart -> RWE Velocity Interpolation
     xloc = aX1%o + ( iix1(1) - 1) * aX1%d
     call C2R2D_init(0,0,norm,aX1%d,aX3%d,ncx,ncz,aE1%n,aE3%n,xloc,0.,nh)
     call C2R2D_VELOCITY (RVel,raysin,CVel,gnorm)  
     call Velocity_Check(CVel,RVel)

     RImg=0.
     call wemig( rwf,swf,RImg,iS1,RVel,fields )   

     !! . . Generate ADCIGS
     if (aH1%n .gt. 1) 	call off2ang2D_run(Rimg,aH1,aA1)

     !! . . Do RWE -> CART Image Interpolation
     xloc = aX1%o + ( iix1(1) - 1 ) * aX1%d
     allocate ( CImg(ncz,ncx,nh) , Image(ncz,ncx,nh) )
     call C2R2D_init(nsx,nsz,norm,aX1%d,aX3%d,ncx,ncz,aE1%n,aE3%n,xloc,0.,nh)
     call C2R2D_IMAGE (RImg,raysin,CImg,gnorm)

     !! ------------------------------------------------------------------
     !! . . 
     !! . . IMAGE WRITE OUT
     !! . . 
     !! ------------------------------------------------------------------
     !! . . Add single image to global image
     !! . . Read in image subset
     if (verbose) write(0,*) '        * Reading in Image volume'
     if (verbose) write(0,*) '        * RIMG/CIMG max: ',maxval(RImg),maxval(CImg)

     call sseek('image', 0 , 0 ) !! Rewind
     do ih=1,aH1%n
        nbytes = 4*( ncz*(iix1(1)-1) + (ih-1)*aX1%n*aX3%n  )
        call sseek("image",nbytes,0)
        call sreed('image', Image(:,:,ih), 4*ncz*ncx)
     end do
     
     !! . . Update local Cart image
     if (verbose) write(0,*) '        * Updating Image'
     !$OMP PARALLEL DO PRIVATE(ih,i1,i3)
     do ih=1,nh
        do i1=1,ncx
           do i3=1,ncz
              Image(i3,i1,ih) = Image(i3,i1,ih) + CImg(i3,i1,ih) 
           end do
        end do
     end do
     !$OMP END PARALLEL DO

     !! Write out image
     call sseek('image', 0 , 0 )
     do ih=1,aH1%n
        nbytes = 4*( ncz*(iix1(1)-1) + (ih-1)*aX1%n*aX3%n )
        call sseek("image",nbytes,0)
        call srite("image", Image(:,:,ih), 4*ncz*ncx)
     end do

     !! Advance rays to next step
     !$OMP PARALLEL DO PRIVATE(iE3,iE1)
     do iE3=1,aE3%n
        do iE1=1,aE1%n
           raysin(iE1,iE3,1)=raysin(iE1,iE3,1)+aS1%d
        end do
     end do
     !$OMP END PARALLEL DO

     !! . . deallocate variables associated with temp Cart image
     deallocate ( CVel , CImg , Image , vtmp )

  end do INLINE_SHOT

  if (nh .gt. 1) call off2ang2D_close()
  call sep_end_prog()
  call exit()

end program RWE2D
