program RWE3D
  !! !!!!!!!!!!!!!!!!!!!
  !!
  !! 		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 RWE3D_wemig !! Migration driver module
  use C2R3D               !! Ray-to-Cartesian interpolation 
  use elliptic                !! Elliptic integrals for Ellipsoidal 
  implicit none
  type(myaxis) :: aE1,aE2,aE3,aW,aS1,aS2,aX1,aX2,aX3,aR1,aR2

  integer :: iX1,iX2,iX3,ii,iw,iis,iE1,iE2,iE3,iS1,iS2,ishot,stat,method
  integer :: nsx,nsy,nsz,norm,itest,nbytes,i1,i2,i3,ncx,ncy,ncz,nxint,nyint
  integer :: e1int1,e1int2,e2int1,e2int2
  real    :: r0,eps,rextra,minang,maxang,rscale,atloc
  real    :: ddx1,ddx2,ddx3,dd1,dd2,dd3,rro,rrd,amax,b,c
  character (len=128) :: name
  integer, allocatable :: iix1(:),iix2(:)
  real,    allocatable :: ox1(:),ox2(:)
  real,    allocatable,dimension(:,:)    :: vtmp
  real,    allocatable,dimension(:,:,:)  :: RVel,CVel,CImg,RImg,Image,gnorm
  real,    allocatable,dimension(:,:,:,:):: raysin,fields
  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,.false.)
  call from_param("planewave",planewave,.false.)
  allocate( iix1(2),iix2(2),ox1(2),ox2(2) )

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

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

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

  call IMAGE_Info(aX1,aX2,aX3)
  call sseek('image', 0 , 0 )

  !! . . Read in the parameters
  call Read_In_Parameters(aE3,verbose,norm,nsx,nsy,nsz,eps,&
  minang,maxang,method,newimage,nxint,nyint)

  !! . . Initialize propagation module
  call wemig_axes_init(aE1,aE2,aE3,aW,aX1,aX2,aX3,aS1,aS2,aR1,aR2)

  !! . . Set scaling of the coordinate system.  
  !! . . Cart is 1. EC is (nr-1)*dr/2. (half-width)
  if (method .le. 0 .or. method .ge. 4 ) then
     call seperr("IDIOT...USE method = 1, 2 or 3")
  end if

  if ( method .eq. 1 ) then
     call wemig_param_init(verbose,1.,minang,maxang,eps,method)
  else if (method .eq. 2) then
     r0 = dble(aE3%n-1.)*aR1%d/2.
     call wemig_param_init(verbose,r0,minang,maxang,eps,method)
  else if (method .eq. 3) then
     r0 =1.
     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,aE2%n,aE3%n,3),fields(aE1%n,aE2%n,aE3%n,3) )
  allocate( gnorm(aE1%n,aE2%n,aE3%n) )
  if (method .eq. 1) then
     call  CART_Setup(raysin,fields,gnorm,aE3%n,aE3%o,aE3%d)
     aE1%o=aR1%o; aE1%d=aR1%d
     aE2%o=aR2%o; aE2%d=aR2%d
  else if (method .eq. 2) then
     call CYNEL_Setup(raysin,fields,gnorm,aE2%n,aE2%o,aE2%d)
     aE1%o=aR1%o; aE1%d=aR1%d
  else if (method .eq. 3) then
     call ELLIP_Setup(raysin,fields,gnorm,aE1,aE2,aE3,aR1,aR2)	
  end if

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

  !! . . Write out parameters to header files  
  call write_out(aX1,aX2,aX3,aE1,aE2,aE3,aS1,aS2,aR1,aR2)
  call sep_write(raysin,"rays")
  call Geometry_QC(raysin,aE1,aE2,aE3,aS1,aS2,aW,aR1,aR2)

  !! . . 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,aE2%n,aW%n),swf(aE1%n,aE2%n,aW%n) )
  else
     allocate( swfsave(aE1%n,aE2%n,aW%n),swf(aE1%n,aE2%n,aW%n) )
     call sreed("swf", swfsave , 8*aE1%n*aE2%n*aW%n) 
     !! . . Position wavefield on surface
     if ( method .eq. 1) then
        swf = swfsave
     else if (method .eq. 2) then
        call CYNEL_Position_XWF(swfsave,swf,raysin)
     else if (method .eq. 3) then
        call ELLIP_Position_XWF(swfsave,swf,raysin,aR1,aR2)
     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,aE2%n, aW%n) , rwfsave (aE1%n,aE2%n, aW%n) )
  allocate (RVel (aE1%n,aE2%n,aE3%n) ,    RImg (aE1%n,aE2%n,aE3%n) )
  !! . . Loop over shot points: inline and streamers
  XLINE_SHOT: do iS2 = 1,aS2%n
     INLINE_SHOT: do iS1 = 1,aS1%n   
        if (planewave ) then
           swf=0. ; swfsave=0.
           call sreed("swf", swfsave , 8*aE1%n*aE2%n*aW%n) 
           !! . . Position wavefield on surface
           if ( method .eq. 1) then
              swf = swfsave
           else if (method .eq. 2) then
              call CYNEL_Position_XWF(swfsave,swf,raysin)
           else if (method .eq. 3) then
              call ELLIP_Position_XWF(swfsave,swf,raysin,aR1,aR2)
           end if
        end if


        !! RECEIVER WAVEFIELD section
        !! Read in next Receiver wavefield
        rwf=0.;rwfsave=0.
     	call sreed ("rwf", rwfsave , 8*aE1%n*aE2%n*aW%n ) !! Read in only current receiver wavefield
        if (method .eq. 1) then
           rwf = rwfsave
        else if (method .eq. 2) then
           call CYNEL_Position_XWF(rwfsave,rwf,raysin)
        else if (method .eq. 3) then
           call ELLIP_Position_XWF(rwfsave,rwf,raysin,aR1,aR2)
        end if
 
        !! ------------------------------------------------------------------
        !! . . 
        !! . . Read in velocity 
        !! . . 
        !! ------------------------------------------------------------------
        call Get_Vel_Indicies(raysin,iix1,iix2,ox1,ox2)
        ncx = iix1(2)-iix1(1)+1
        ncy = iix2(2)-iix2(1)+1
        ncz = aX3%n

        !! . . Allocate necessary arrays
        allocate (  CImg(ncz,ncx,ncy) , Image(ncz,ncx,ncy) )
        allocate (  vtmp(ncz,ncx)     ,  CVel(ncz,ncx,ncy) )
        call Zero_arrays(RVel,CVel,RImg,CImg,ncz,ncx,ncy,aE1,aE2,aE3)

        !! Read in minimal velocity subset
        name="vel"
        call sseek("vel",0,0)
        do iE2 = 1,ncy
           nbytes = 4*( ncz*aX1%n*(iix2(1)-1+(iE2-1)) + ncz*(iix1(1)-1) )
           call sseek("vel",nbytes,0)
           call sreed( name, CVel(:,:,iE2) , ncz*ncx*4)
        end do

        !! ------------------------------------------------------------------
        !! . . 
        !! . . Interpolate Vel from Cart to Riemannian
        !! . . 
        !! ------------------------------------------------------------------
     	call C2R3D_init(0,0,0,norm,aX1%d,aX2%d,aX3%d,&
        ncx,ncy,ncz,aE1%n,aE2%n,aE3%n,&
        aX1%o+(iix1(1)-1)*aX1%d,aX2%o+(iiX2(1)-1)*aX1%d,0.)
        call C2R3D_CART (RVel,raysin,CVel,gnorm,.false.)  

        !! . . Check to see if velocities are good
        call Velocity_Check(CVel,RVel)
        call Get_Vel_Indicies(raysin,iix1,iix2,ox1,ox2)

        !! ------------------------------------------------------------------
        !! . . 
        !! . . Propagation and imaging in Riemannian Space
        !! . . 
        !! ------------------------------------------------------------------
        ishot=iS1 + aS1%n*(iS2-1.)
        write(0,*) '        * Working on shot ',ishot ,' of ', aS1%n*aS2%n
     	!! . . Call propagation and imaging module
     	call wemig( rwf,swf,RImg,ishot,RVel,fields )   

        !! ------------------------------------------------------------------
        !! . . 
        !! . . Image Interpolation to Cartesian
        !! . . 
        !! ------------------------------------------------------------------
        do i1=1,nxint
           do i2=1,nyint

              !! Riemannian Cut indicies
              e1int1 = 1 + floor( aE1%n *dble(i1-1)/dble(nxint) )
              e1int2 =     floor( aE1%n *dble(i1  )/dble(nxint) )
              e2int1 = 1 + floor( aE2%n *dble(i2-1)/dble(nyint) )
              e2int2 =     floor( aE2%n *dble(i2  )/dble(nyint) )

              !! Cart Cut indicies
              call Get_Vel_Indicies(raysin(e1int1:e1int2,e2int1:e2int2,:,:),&
                                    iix1,iix2,ox1,ox2)
              
              !! Initialize interp routine
              call C2R3D_init(nsx,nsy,nsz,norm,&       !  Interp
                             aX1%d,aX2%d,aX3%d,&       !  Sampling
                             iix1(2)-iix1(1)+1,&       !  cnx: X-cart
                             iix2(2)-iix2(1)+1,&       !  cny: Y-cart
                             ncz,&                     !  cnz: Z-cart
                             e1int2-e1int1+1,&         !  rnx: X-RWE
                             e2int2-e2int1+1,&         !  rny: Y-RWE
                             aE3%n,&                   !  rnz: Z-RWE
                             aX1%o+(iix1(1)-1)*aX1%d,& ! xmin: X-min
                             aX2%o+(iix2(1)-1)*aX2%d,& ! ymin: Y-min
                             0.)                       ! zmin: Z-min

              !! Interp Image: RWE -> Cartesian
              call C2R3D_CART (  RImg(e1int1:e1int2,  e2int1:e2int2,:),&
                               raysin(e1int1:e1int2,  e2int1:e2int2,:,:),&
                               CImg (:,iix1(1):iix1(2),iix2(1):iix2(2)),&
                               gnorm,.true.)
           end do
        end do


        !! ------------------------------------------------------------------
        !! . . 
        !! . . 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 iE2 = 1,ncy
           nbytes = 4*( ncz*aX1%n*(iix2(1)-1+(iE2-1)) + ncz*(iix1(1)-1) )
           call sseek("image",nbytes,0)
           call sreed('image', Image(:,:,iE2) , 4*ncz*ncx)
        end do

        !! . . Update local Cart image
        if (verbose) write(0,*) '        * Updating Image'
        !$OMP PARALLEL DO PRIVATE(i2,i1,i3)
        do i2=1,ncy
           do i1=1,ncx
              do i3=1,ncz
                 Image(i3,i1,i2) = Image(i3,i1,i2) + CImg(i3,i1,i2) 
              end do
           end do
     	end do
        !$OMP END PARALLEL DO

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

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

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

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

     end do INLINE_SHOT

     !! . . Reset the inline axis
     !$OMP PARALLEL DO PRIVATE(iE3,iE2,iE1)
     do iE3=1,aE3%n
       	do iE2=1,aE2%n
           do iE1=1,aE1%n
              raysin(iE1,iE2,iE3,1)=raysin(iE1,iE2,iE3,1) - dble(aS1%n)*aS1%d
           end do
        end do
     end do
     !$OMP END PARALLEL DO

     !! . . Update the crossline axis
     !$OMP PARALLEL DO PRIVATE(iE3,iE2,iE1)
     do iE3=1,aE3%n
       	do iE2=1,aE2%n
           do iE1=1,aE1%n
              raysin(iE1,iE2,iE3,2)=raysin(iE1,iE2,iE3,2) + dble(iS2)*aS2%d
           end do
        end do
     end do
     !$OMP END PARALLEL DO

  end do XLINE_SHOT
  call sep_end_prog()
  call exit()

end program RWE3D
