module steepest_mod

use sep
use OW_grabvel2_areal
use OW_exec_time
use OW_parms_areal
use OW_taper_areal
use OW_wei_areal
use lloyd_vel
use base_spline_mod
use OW_mpi

implicit none
integer, public, save                  :: iter,realz
real, private, allocatable             :: taper(:,:),pick1(:)
real, private, allocatable             :: image(:,:,:,:,:),image0(:,:,:,:,:),g_ir(:,:),grad0(:)
real, private, allocatable             :: img1(:,:,:,:),img0(:,:,:,:),img_ir(:,:,:,:,:)
double precision, private, allocatable :: gspl(:)

contains 
!---------------------------------------------------------------------
subroutine steep_mpi_init(pick_in)
integer                  :: i, i_r
real                     :: pick_in(:)

    allocate(pick1(img%xm%n*img%ym%n))
    pick1=pick_in
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Set up the taper
    allocate(taper(kxm%n,kym%n));taper=0.
    call owtaper(taper)
    allocate(image0(kxm%n,nsecdimx,kym%n,nsecdimy,zblock%ns));image0=0
    allocate(image(kxm%n,nsecdimx,kym%n,nsecdimy,zblock%ns));image=0
    allocate(g_ir(kxm%n*nz*kym%n,wpar%n_mpi-1));g_ir=0.
    allocate(grad0(kxm%n*nz*kym%n));grad0=0
    allocate(img0(kxm%n,nsecdimx,kym%n,nsecdimy));img0=0
    allocate(img1(kxm%n,nsecdimx,kym%n,nsecdimy));img1=0
    allocate(img_ir(kxm%n,nsecdimx,kym%n,nsecdimy,wpar%n_mpi-1));img_ir=0
    if (axk%n.ne.99999 .and. ayk%n.ne.99999 .and. azk%n.ne.99999) allocate(gspl(axk%n*ayk%n*azk%n))

end subroutine
!---------------------------------------------------------------------
subroutine steep_mpi_close()
  if (allocated(taper)) deallocate(taper)
  if (allocated(pick1)) deallocate(pick1)
  if (allocated(image)) deallocate(image)
  if (allocated(image0)) deallocate(image0)
  if (allocated(img1)) deallocate(img1)
  if (allocated(img0)) deallocate(img0)
  if (allocated(img_ir)) deallocate(img_ir)
  if (allocated(grad0)) deallocate(grad0)
  if (allocated(g_ir)) deallocate(g_ir)
  if (allocated(gspl)) deallocate(gspl)
end subroutine
!---------------------------------------------------------------------
subroutine fgeval_init(iter_in)
 integer                     :: iter_in
 iter=iter_in

 if (nreal.ne.1) then
    realz=mod(iter,nreal)
 else
    realz=1
 end if

end subroutine

!---------------------------------------------------------------------
subroutine init_files()

integer            :: i,j,k
real,allocatable   :: a(:)

allocate(a(kxm%n*kym%n*nsecdimx*nsecdimy)); a=0
call sseek("bkgimg",0,0)
call sseek("Dimage",0,0)
call sseek("grad",0,0)
do i=1,img%z%n
   call srite("bkgimg",a,4*kxm%n*kym%n*nsecdimx*nsecdimy)
   call srite("Dimage",a,4*kxm%n*kym%n*nsecdimx*nsecdimy)
   call srite("grad",a(1:kxm%n*kym%n),4*kxm%n*kym%n)
end do
deallocate(a)
end subroutine
!---------------------------------------------------------------------
subroutine fgeval(vel0,ff,gg,grad)

double precision,optional          :: grad(:)
double precision                   :: ff,gg,aa,bb,fun
real                               :: vel0(:)
!---
integer                            :: ierr,idip,stat1,i_r,status,izb,iz, ihy,iy,ihx,ix
real                               :: start2,finish2

    do idip=1,shtblk
       call feval_exe(idip,vel0)
    end do

    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    call owexec_time(start2)

!    do izb=1,zblock%nb
!       image0=0;image=0
!       if(wpar%i_mpi > 0) call image_io("bkgimg","reed",izb,image(:,:,:,:,1:zblock%blen(izb)),.false.)
!       do iz=1,zblock%blen(izb)
!          call MPI_BARRIER(MPI_COMM_WORLD,ierr)
!          call MPI_ALLREDUCE(image(:,:,:,:,iz),image0(:,:,:,:,iz),&
!                kxm%n*nsecdimx*kym%n*nsecdimy,MPI_REAL,MPI_SUM,MPI_COMM_WORLD,ierr)
!       end do
!       call image_io("bkgimg","rite",izb,image0(:,:,:,:,1:zblock%blen(izb)))
!    end do
    do izb=1,img%z%n
       img0=0;img1=0;img_ir=0
       if(wpar%i_mpi == 0) write(0,*) "passei 000", izb
       if(wpar%i_mpi > 0) call image_io2("bkgimg","reed",izb,img1)
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if(wpar%i_mpi == 0) write(0,*) "passei 001", izb
       do i_r=1,wpar%n_mpi-1
          if (i_r==wpar%i_mpi) then
             call MPI_SEND(img1,size(img1),MPI_REAL,0,i_r,MPI_COMM_WORLD,ierr)
          else if(wpar%i_mpi==0) then
             call MPI_RECV(img_ir(:,:,:,:,i_r),size(img1),MPI_REAL,i_r,i_r,MPI_COMM_WORLD,status,ierr)
          end if
       end do
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if(wpar%i_mpi == 0) write(0,*) "passei 002", izb
       do i_r=1,wpar%n_mpi-1
          img0=img0+img_ir(:,:,:,:,i_r)
       end do
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if(wpar%i_mpi == 0) write(0,*) "passei 003", izb
       if (wpar%i_mpi==0) then
          do i_r=1,wpar%n_mpi-1
             call MPI_SEND(img0,size(img0),MPI_REAL,i_r,7,MPI_COMM_WORLD,ierr)
          end do
       else
          call MPI_RECV(img0,size(img0),MPI_DOUBLE_PRECISION,0,7,MPI_COMM_WORLD,status,ierr)
       end if
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if(wpar%i_mpi == 0) write(0,*) "passei 004", izb
       call image_io2("bkgimg","rite",izb,img0)
       if(wpar%i_mpi == 0) write(0,*) "passei 005", izb
    end do

    call owexec_time(finish2)
    if (wpar%i_mpi==1) write(0,*) "FINISHED COPYING bimg ",int((finish2-start2))," s"

    if (scaling) then
       if (wpar%i_mpi==1) write(0,*) "@@1 "
          call ow_scale1(iter)
       if (wpar%i_mpi==1) write(0,*) "@@2 "
    end if
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    call owexec_time(start2)
    ff=dble(0.);gg=ff
    call owdso_nozb0(pick1)
!    call owdso_nozb1(pick1)
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)

    if (wpar%i_mpi==1) write(0,*) "@@3 "
    do izb=1,zblock%nb
       image=0
       write(0,*) "@@3 ",izb,wpar%i_mpi
       if (wpar%i_mpi==1) then
          call image_io("Dimage","reed",izb,image(:,:,:,:,1:zblock%blen(izb)),.false.)
          call MPI_SEND(image(:,:,:,:,1:zblock%blen(izb)),kxm%n*nsecdimx*kym%n*nsecdimy*zblock%blen(izb),&
               MPI_REAL,0,9,MPI_COMM_WORLD,ierr)
       else if(wpar%i_mpi==0) then
          call MPI_RECV(image(:,:,:,:,1:zblock%blen(izb)),kxm%n*nsecdimx*kym%n*nsecdimy*zblock%blen(izb),&
               MPI_REAL,1,9,MPI_COMM_WORLD,status,ierr)
          call image_io("Dimage","rite",izb,image(:,:,:,:,1:zblock%blen(izb)))
       end if
    end do
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)

    if (wpar%i_mpi==1) write(0,*) "@@4 "
    ff=0;gg=0
    do izb=1,zblock%nb
       image=0
       write(0,*) "@@4 ",izb,wpar%i_mpi
       call image_io("Dimage","reed",izb,image(:,:,:,:,1:zblock%blen(izb)),.false.)
       do iz=1,zblock%blen(izb)
          do ihy=1,nsecdimy
             do iy=1,img%ym%n
                do ihx=1,nsecdimx
                   do ix=1,img%xm%n
                      ff=ff+image(ix,ihx,iy,ihy,iz)*image(ix,ihx,iy,ihy,iz)
                   end do
                end do
             end do
          end do
       end do
    end do
    gg = ff*ff
    if (wpar%i_mpi==1) write(0,*) "@@5 "

    call owexec_time(finish2)
    if (wpar%i_mpi==1) write(0,*) "FINISHED DSO IN ",int((finish2-start2))," s"
    write(0,*) "ff= ",ff," gg= ",gg,wpar%i_mpi
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    if (present(grad)) then
       do idip=1,shtblk
          call geval_exe(idip,vel0);
       end do
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       grad=0;

       if(wpar%i_mpi > 0) call grad_io("grad","reed",grad0,.false.)
       do i_r=1,wpar%n_mpi-1
          if (i_r==wpar%i_mpi) then
             call MPI_SEND(grad0,size(grad),MPI_REAL,0,i_r,MPI_COMM_WORLD,ierr)
          else if(wpar%i_mpi==0) then
             call MPI_RECV(g_ir(:,i_r),size(grad),MPI_REAL,i_r,i_r,MPI_COMM_WORLD,status,ierr)
          end if
       end do
!       write(0,*) "passei01",maxval(grad0),minval(grad0),wpar%i_mpi
       grad=0.
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if (wpar%i_mpi==0) then
          do i_r=1,wpar%n_mpi-1
             grad=grad+g_ir(:,i_r)
          end do
       end if
       if (wpar%i_mpi==0) then
          do i_r=1,wpar%n_mpi-1
             call MPI_SEND(grad,size(grad),MPI_DOUBLE_PRECISION,i_r,11,MPI_COMM_WORLD,ierr)
          end do
       else
          call MPI_RECV(grad,size(grad),MPI_DOUBLE_PRECISION,0,11,MPI_COMM_WORLD,status,ierr)
       end if
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)

       grad=-grad
       if (chop) then
          aa=(1-chophi)*maxval(grad)
          bb=(1-choplo)*minval(grad)
          where(grad>aa) grad=aa
          where(grad<bb) grad=bb
       end if
!       if (wpar%i_mpi==0) write(0,*) "02 ",minval(grad),maxval(grad)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       if (axk%n.ne.99999 .and. ayk%n.ne.99999 .and. azk%n.ne.99999) then
           !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
           ! B-spline representation of grad
              stat1 = base_spln3d_init(img%xm,img%ym,img%z,axk,ayk,azk)
              call owexec_time(start2)
              gspl=dble(0.)
              stat1 = base_spln3d_oper(.true. , .false., gspl, grad)
              grad=dble(0.)
              stat1 = base_spln3d_oper(.false. ,.false., gspl, grad)
              call owexec_time(finish2)
              write(3,*) "SMOOTHED GRADIENT IN ",int((finish2-start2))," s",wpar%i_mpi
       end if
!       call MPI_BARRIER(MPI_COMM_WORLD,ierr)

       if (wpar%i_mpi==0) write(0,*) "03 ",minval(grad),maxval(grad)
       call oweditvel(pick1,grad)
       if (wpar%i_mpi==0) write(0,*) "05 ",minval(grad),maxval(grad)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Somente teste
       grad=sign(sqrt(abs(grad)),grad)
! apagar depois
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
       call grad_io("grad","rite",sngl(grad))

    end if
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)

end subroutine
!---------------------------------------------------------------------
subroutine feval_exe(idip,vel0)

real                                          :: start, finish, start1, finish1, vel0(:)
integer                                       :: idip,ireal
real, allocatable                             :: vref(:,:,:),dcig(:,:,:,:,:)
double precision                              :: fun
integer                                       :: izb,i_r,status,ierr

       if (extrap(1:4).eq."pspi") then
          if (lloyd) then 
            allocate(vref(nref,2,img%z%n))
            call lloyd_vel1(vel0,vref)
!            if (wpar%i_mpi==0) call srite("vreff",vref,4*nref*2*img%z%n)
          end if
       end if
       !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
       ! Computing bkg.image
       call owexec_time(start)
       call owwei_init(idip,1,iter)
       if (lloyd) then
!          call owwei(taper=taper,vel=vel0,sref=vref)
          call owwei_nozb(taper=taper,vel=vel0,sref=vref)
       else
!          call owwei(taper=taper,vel=vel0)
          call owwei_nozb(taper=taper,vel=vel0)
       end if
       call owexec_time(finish)
       write(3,112) "COMPUTED BACKGROUND IMAGE IN ",(finish-start)/60," min"
       if (allocated(vref)) deallocate(vref)

 111 format(a7,i4,a8,i4,a4,f6.2,a4)
 112 format(a29,f10.2,a4)
end subroutine
!---------------------------------------------------------------------
 subroutine geval_exe(idip,vel0)

 real                                          :: start, finish, start1, finish1, start2, finish2,vel0(:)
 integer                                       :: ni,i,ishot,j,k,stat1,idip
 real, allocatable                             :: vref(:,:,:)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Computing the gradient
!
       if (extrap(1:4).eq."pspi") then
          if(lloyd) then
            allocate(vref(nref,2,img%z%n))
            call lloyd_vel1(vel0,vref)
          end if
       end if
       start1=0.;finish1=0.;

       call owexec_time(start1)
       call owwei_init(idip,1,iter)
       if (lloyd) then
          call owwemva(taper,vel0,iter,vref)
       else
          call owwemva(taper,vel0,iter)
       end if

       call owexec_time(finish1)
       if (wpar%i_mpi==1) write(3,112) "COMPUTED THE GRADIENT IN ",(finish1-start1)/60," min OF ITERATION # ", iter
       if (allocated(vref)) deallocate(vref)
 111 format(a26,i4,a8,i4,a4,f6.2,a4)
 112 format(a29,f10.2,a20,i4)
 end subroutine
!---------------------------------------------------------------------

double precision function norm(x)
  double precision                  :: x(img%xm%n,img%ym%n,img%z%n)
  integer                           :: i,j,k
  norm=0.

  do i=1,img%z%n
     do j=1,img%ym%n
        do k=1,img%xm%n
           norm=norm+x(k,j,i)**2
        end do
     end do
  end do
  norm=sqrt(norm)
end function
!---------------------------------------------------------------------
end module
