!One-Way Modeling /Migration....modified from Yaxiun's
module oneway_mod
  use history_mod 
  use io_mod
  use array_mod
  use param_mod 
  use wavelet_mod
  use myfft_mod
  use map_mod
  use grid_mod
  use vel_mod
  use ctaper_mod
  use ssf_mod
  implicit none

contains
!=======================================================================================
  integer function oneway_init() result(stat)
    call from_param("adj", adj, .false. )
    call vel_ref_input_history()                                !read vel history
    if (.not. adj)then
        call wavelet_input()                                    !read wavelethistory
        call param_init()                                       !initialize modeling oneway parameters 
        if (.not. only_green) call data_output_history()        !output the data history
        call data_array_alloc()                                 !allocate memory 
        call propagation_grid()                                 !propagation grid
        call velocity_grid()                                    !velocity grid
        if (.not. only_green) call data_io_init()               !data to zero
        if (lgreen    ) call green_output_history()             !output the greens history
  !      if (lgreen    ) call greens_io_init()                   !greens to zero
        if (wfld_d_out) call wfld_output_history(wfld_d_)       !output the wavefield history
        if (wfld_d_out) call wfld_io_init(wfld_d_)              !down wavefield to zero
        call k2_map()
    else
        call wavelet_input()                                    !read wavelet history
        call param_init()                                       !initialize migration parameters 
        call mig_output_history()                               !output migration history
        call mig_array_alloc()                                  !allocate memory 
        call propagation_grid()                                 !propagation grid
        call velocity_grid()                                    !velocity grid
        call mig_io_init()                                      !mig to zero
        if (wfld_d_out) call wfld_output_history(wfld_d_)       !output down wavefield history
        if (wfld_u_out) call wfld_output_history(wfld_u_)       !output up wavefield history
        if (wfld_d_out) call wfld_io_init(wfld_d_)              !down wavefield to zero
        if (wfld_u_out) call wfld_io_init(wfld_u_)              !up   wavefield to zero
        call k2_map()
    endif
    stat = 0
  end function
!=======================================================================================
  integer function oneway_clos() result(stat)
        call array_dealloc()
        stat = 0
  end function
!!=======================================================================================
  integer function oneway_oper() result(stat)
    integer                                        :: index_aw
    integer                                        :: ixg, iyg, niw
    real,    dimension(:)          ,allocatable    :: vref
    real,    dimension(:,:)        ,allocatable    :: vel_slice,ref_slice
    integer, dimension(:,:)        ,allocatable    :: ind_vref
    real                                           :: w 
    complex , dimension(:,:)       ,allocatable    :: tmp_wfld_d
    complex , dimension(:,:)       ,allocatable    :: tmp_wfld__d
    complex , dimension(:,:)       ,allocatable    :: tmp_wfld_u
    complex , dimension(:,:)       ,allocatable    :: tmp_wfld__u
    complex , dimension(:,:)       ,allocatable    :: wfld_slice_d
    complex , dimension(:,:)       ,allocatable    :: wfld_slice_u
    complex , dimension(:,:)       ,allocatable    :: split 
    real    , dimension(:,:,:)     ,allocatable    :: tmp_mig
    real    , dimension(:,:,:,:,:) ,allocatable    :: tmp_migh
    complex , dimension(:,:,:)     ,allocatable    :: green, green_s, green_r
    complex , dimension(:,:,:,:,:) ,allocatable    :: wfld_d 
    complex , dimension(:,:,:    ) ,allocatable    :: wfld_u 
    real    , dimension(:)         ,allocatable    :: kz2  
    integer                                        :: isz, isx, isy, iisz, iisx, iisy
    integer                                        :: igz, igx, igy, iigz, iigx, iigy
    real                                           :: z, x, y  
    real                                           :: sx, sy, gx, gy
    real                                           :: kz,wkbj
    integer                                        :: icount 
    complex                                        :: phase_shaping 
    icount=0;    
    open(unit=10, file=report_, status='unknown')

    nnode =omp_get_num_procs(); thread=omp_get_thread_num();  nodes=min(nnode-1,nodes) 
    call omp_set_num_threads(nodes); write(*,*) 'using', nodes, 'of' , nnode
!need to remove this...prob. with reading velocitie slices
    allocate(vel_slice(pgrid%ax%n,pgrid%ay%n)); allocate(ref_slice(pgrid%ax%n,pgrid%ay%n))
    if (.not. adj) then
        write(0,*)'Greens function modeling......'
        do iz=1,vgrid%az%n; call data_io_vel(vel_slice, iz); call data_io_ref(ref_slice, iz); vel(iz,:,:)=vel_slice; refl(iz,:,:)=ref_slice; enddo
    else
        write(0,*)'Migration........'
        do iz=1,vgrid%az%n; call data_io_vel(vel_slice, iz); vel(iz,:,:)=vel_slice; enddo
        allocate(mig(igrid%az%n,igrid%ax%n,igrid%ay%n),tmp_mig(pgrid%az%n,pgrid%ax%n,pgrid%ay%n)); mig=0.; tmp_mig=0.
        !allocate(migh(igrid%az%n,igrid%ax%n,igrid%ay%n,igrid%ahx,igrid%ahy%n), & 
        !     tmp_migh(igrid%az%n,igrid%ax%n,igrid%ay%n,igrid%ahx,igrid%ahy%n)); migh =0.; tmp_migh=0.
    endif
!$OMP parallel do private(inode,iw,w,index_aw,ixg,iyg,wfld_d,wfld_u,iz,vel_slice,ind_vref,vref,nref,green,greens,tmp_wfld_d,tmp_wfld__d,tmp_wfld_u,tmp_wfld__u,kz2,iv,ikx,iky,ix,iy,data,sx,isx,iisx,sy,isy,iisy,gx,igx,iigx,gy,igy,iigy,tmp_mig, tmp_migh,split, kz, wfld_slice_u,wfld_slice_d) schedule(dynamic,1) 
  do iw=1, dgrid%aw%n                                                            ! freq. count.
       w = dgrid%aw%o + (iw-1)*dgrid%aw%d; 
       inode=omp_get_thread_num();
       if(abs(csou(iw))<=mcsou*.005 .and. iw <= 0.05*dgrid%aw%n)cycle
       if(abs(csou(iw))<=mcsou*.005 .and. iw >= 0.15*dgrid%aw%n)cycle
       stat = ssf_init(kx2, ky2, pgrid%az%d)

       if (allocated(wfld_slice_d)) deallocate(wfld_slice_d);  allocate(wfld_slice_d(pgrid%ax%n,pgrid%ay%n))
       if (allocated(tmp_wfld_d)) deallocate(tmp_wfld_d); allocate(tmp_wfld_d(pgrid%ax%n,pgrid%ay%n))
       if (allocated(tmp_wfld__d)) deallocate(tmp_wfld__d); allocate(tmp_wfld__d(pgrid%ax%n,pgrid%ay%n))
 if (.not. adj)then
       if(only_green )then;
                if ((mod(iw,jw)/=1 .and. jw/=1) .or. icount>ggrid%aw%n)then; 
                        cycle; 
                else; 
                        icount=icount+1; 
                        write(* ,*)'iw.:', iw ,' of ',dgrid%aw%n,' @ node ', inode        
                endif; 
       else
                write(* ,*)'iw.:', iw ,' of ',dgrid%aw%n,' @ node ', inode        
       endif
       index_aw = .5 + (w - ggrid%aw%o)/ggrid%aw%d; 
                  allocate(wfld_d(         imax_z,      pgrid%ax%n,      pgrid%ay%n,ggrid%ax_surf%n, ggrid%ay_surf%n)); wfld_d=0.;
       if (lgreen)allocate(greens(ggrid%az_modl%n, ggrid%ax_modl%n, ggrid%ay_modl%n,ggrid%ax_surf%n, ggrid%ay_surf%n));
       do ixg=1,dgrid%agx%n;    ! surf. shot/geoph. pos. loop
       do iyg=1,dgrid%agy%n;
                if (spike)then
                        wfld_d(imin_z,ixg+(pgrid%index_datpro_gxmin-1),iyg+(pgrid%index_datpro_gymin-1),ixg,iyg)=1;
                else
                        wfld_d(imin_z,ixg+(pgrid%index_datpro_gxmin-1),iyg+(pgrid%index_datpro_gymin-1),ixg,iyg)=csou(iw);
                endif;
                do iz=imin_z,pgrid%az%n;   ! depth loop
                        if (iz /=imin_z) then 
                                call vel_alloc(vel_slice,ind_vref,vref); vel_slice = vel(iz,:,:); call velocity(vref, vel_slice, ind_vref, nref, vnum);
                                tmp_wfld_d=wfld_slice_d; 
                                wfld_slice_d=0.
                                do iv=1, nref                                                                           !vel. loop  
                                        tmp_wfld__d=tmp_wfld_d
                                        stat = ssf_oper(adj, -1, w, vref(iv), vel_slice, tmp_wfld__d)
                                        where (ind_vref == iv)
                                             wfld_slice_d = wfld_slice_d + tmp_wfld__d
                                        end where
                                enddo
                       else
                                wfld_slice_d=wfld_d(1,:,:,ixg,iyg); 
                       endif   

                       if(lgreen .and. ((mod(iw,jw)==1 .or. jw==1) .and. int(iw/jw)<=ggrid%aw%n)) greens(iz,:,:,ixg,iyg)=&
                                                                wfld_slice_d(ggrid%index_modpro_minx:ggrid%index_modpro_maxx, & 
                                                                             ggrid%index_modpro_miny:ggrid%index_modpro_maxy)
                enddo !end depth
       enddo
       enddo !end geoph.

       if(lgreen .and. ((mod(iw,jw)==1 .or. jw==1) .and. icount<=ggrid%aw%n)) then 
       !$OMP CRITICAL
                call data_io_greens(greens, index_aw)
       !$OMP END CRITICAL 
       endif

       if (.not. only_green) then
           !call data_compute
       endif !end if not only green...compute one way data...demigration

       if(wfld_d_out) wfld_out_d(:,ggrid%index_modpro_minx:ggrid%index_modpro_maxx, & 
                                   ggrid%index_modpro_miny:ggrid%index_modpro_maxy,iw,1) = greens(:,:,:,1,1)
       deallocate(greens,tmp_wfld_d,tmp_wfld__d,wfld_slice_d,wfld_d)
  else
       write(* ,*)'iw.:', iw ,' of ',dgrid%aw%n,' @ node ', inode         
       write(10,*)'iw.:', iw ,' of ',dgrid%aw%n,' @ node ', inode       
       index_aw = 1.5 + (w - ggrid%aw%o)/dgrid%aw%d;
       allocate(wfld_d(imax_z, pgrid%ax%n, pgrid%ay%n, dgrid%asx%n, dgrid%asy%n)); wfld_d=0.;
       allocate(wfld_u(imax_z, pgrid%ax%n, pgrid%ay%n)); wfld_u=0.
       allocate(data(dgrid%agx%n, dgrid%agy%n, dgrid%asx%n, dgrid%asy%n)); data=0.
       allocate(tmp_mig(pgrid%az%n,pgrid%ax%n,pgrid%ay%n)); tmp_mig=0.
!       allocate(tmp_migh(pgrid%az%n,pgrid%ax%n,pgrid%ay%n)); tmp_migh=0.
       if (allocated(wfld_slice_u)) deallocate(wfld_slice_u);  allocate(wfld_slice_u(pgrid%ax%n,pgrid%ay%n))
       if (allocated(tmp_wfld_u)) deallocate(tmp_wfld_u); allocate(tmp_wfld_u(pgrid%ax%n,pgrid%ay%n))
       if (allocated(tmp_wfld__u)) deallocate(tmp_wfld__u); allocate(tmp_wfld__u(pgrid%ax%n,pgrid%ay%n))
       if (allocated(split    )) deallocate(split   ); allocate(split(pgrid%ax%n,pgrid%ay%n))

       if (pgrid%ay%n== 1 .or. pgrid%ax%n==1) then
       !!..2D: 45 deg. rotation i^(-1/2)
                phase_shaping = cmplx(cos(PI/4.), -sin(PI/4.)) / sqrt(w)           
       else
       !!..3D: 90 deg. rotation i^(-1)
                phase_shaping = cmplx(0., - 1. * w)       
       end if

       !$OMP CRITICAL
                call data_io_data("reed", data, iw)
       !$OMP END CRITICAL
       do isx=1,dgrid%asx%n;  sx = dgrid%asx%o + (isx-1)*dgrid%asx%d ; iisx = 1.5 + (sx - pgrid%ax%o)/pgrid%ax%d !hsot loop
       do isy=1,dgrid%asy%n;  sy = dgrid%asy%o + (isy-1)*dgrid%asy%d ; iisy = 1.5 + (sy - pgrid%ay%o)/pgrid%ay%d 
                wfld_d(imin_z,iisx,iisy,isx,isy)=(csou(iw));
                wfld_u(imin_z,pgrid%index_datpro_gxmin:pgrid%index_datpro_gxmax, &
                         pgrid%index_datpro_gymin:pgrid%index_datpro_gymax)=data(:,:,isx,isy)
                do iz=imin_z,pgrid%az%n;   ! depth loop
                        if (iz /=imin_z) then 
                                call vel_alloc(vel_slice,ind_vref,vref); vel_slice = vel(iz,:,:); call velocity(vref, vel_slice, ind_vref, nref, vnum);
                                tmp_wfld_d=wfld_slice_d; wfld_slice_d=0.
                                tmp_wfld_u=wfld_slice_u; wfld_slice_u=0.
                                do iv=1, nref                                                                           !vel. loop  
                                        tmp_wfld__d=tmp_wfld_d; 
                                        tmp_wfld__u=tmp_wfld_u;
                                        stat = ssf_oper(adj, -1, w, vref(iv), vel_slice, tmp_wfld__d)
                                        stat = ssf_oper(adj, +1, w, vref(iv), vel_slice, tmp_wfld__u)
                                        where (ind_vref == iv)
                                             wfld_slice_d = wfld_slice_d + tmp_wfld__d
                                             wfld_slice_u = wfld_slice_u + tmp_wfld__u
                                        end where
                                enddo
                       else
                           wfld_slice_d=wfld_d(1,:,:,isx,isy); 
                           wfld_slice_u=wfld_u(1,:,:        ); 
                       endif    
                       tmp_mig(iz,:,:)=tmp_mig(iz,:,:)+real(wfld_slice_u(:,:)*conjg(wfld_slice_d(:,:))*phase_shaping);
                enddo !end depth
       enddo
       enddo !end shot
       !$OMP CRITICAL
       if (nostack)then; 
           mig=tmp_mig; call mig_io_mig(mig, index_aw);  
       else; 
           write(*,*)'adding freq:', index_aw, 'sum_mig:',sum(mig), 'sum tmp_mig:',sum(tmp_mig); mig=mig+tmp_mig; call mig_io_mig(mig, 1); 
       endif
       !$OMP END CRITICAL
       deallocate(wfld_d, wfld_u, tmp_mig)
       endif
       enddo
!$OMP end parallel do

       if (.not. adj) then; if(wfld_d_out) call wfld_out(wfld_out_d,wfld_d_); 
       else
       mig(:,:,:)=tmp_mig(:,:,:)
       if(wfld_d_out) call wfld_out(wfld_out_d,wfld_d_); if(wfld_u_out) call wfld_out(wfld_out_u,wfld_u_);
       stat = 0
       deallocate(mig,tmp_mig)    
       endif

       stat=0;

  end function

!!!!!!!!!!!!!!!!  ###########
  end module
