program nlin_dir

use sep 
use OW_type
use OW_grabvel2_areal
use OW_exec_time
use OW_parms_areal
use OW_wei_areal
use OW_mpi
use steepest_mod
use base_spline_mod


implicit none

real                                          :: start, finish, start1, finish1, start2, finish2
integer                                       :: ni,i,ishot,j,k,titer,stat1
integer                                       :: ir,is,iww
complex,allocatable                           :: ws(:,:,:),wr(:,:,:)
real, allocatable                             :: velint(:),vel0(:),pick1(:),velout(:)
double precision, allocatable                 :: x(:), g(:),fun(:),fn(:)
real, allocatable                             :: l(:), u(:)
integer                                       :: mpierr,i_r,ierr,status,shtida,ils,imin

real                                          :: lambda,norm_dx,minxx
double precision                              :: ft,f2,gt,gn,g2,ngrad0,ngrad1,teta,fmin
double precision, allocatable                 :: xn(:)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

 call MPI_INIT(mpierr)
 if (mpierr .ne. MPI_SUCCESS) then
    write(0,*) 'Error starting MPI program. Terminating.'
    call MPI_ABORT(MPI_COMM_WORLD, wpar%i_mpi, mpierr)
 end if
 call sep_init()
 call owwei_mpi_init() 
 call set_no_putch()
 if (wpar%i_mpi==0) call owparam()
 call owparm_bcast1()
write(0,*) "Main: Parameter broadcast"

 allocate(x(img%z%n*img%xm%n*img%ym%n));x=dble(0.)
 allocate(g(img%z%n*img%ym%n*img%xm%n));g=dble(0.)
 allocate(pick1(img%xm%n*img%ym%n))
 allocate(xn(img%z%n*img%xm%n*img%ym%n));xn=dble(0.)
 allocate(vel0(nzold*img%ym%n*img%xm%n));vel0=0
 allocate(velout(nzold*img%ym%n*img%xm%n));velout=0
 allocate(fun(max_iters));fun=0
 allocate(fn(nls));fn=0 

write(0,*) "Main: vectors allocated"
! call init_files()
!write(0,*) "Main: files initialized"

 if (wpar%i_mpi==0) then
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    ! Read velocities
    allocate(velint(v%z%n*v%xm%n*v%ym%n));velint=0
    call sreed("vel",velint,v%z%n*v%xm%n*v%ym%n*4)
    call owgrab_velim(velint,vel0,ni)
    deallocate(velint)
    !----------------------------------------------
    ! Horizon with no velocity update
write(0,*) "Before reading horizon"

    if (exist_file("horv0")) then
       open(51,file=hor1,status='old')
       do i=1,img%xm%n*img%ym%n
          read(51,*,end=111) pick1(i)
       end do
111    close(51)
    else
       pick1=0
    end if
write(0,*) "Before sending picks"
    do i_r=1,wpar%n_mpi-1
       write(0,*) "SENDING picks and velocity to node ",i_r
       call MPI_SEND(pick1, img%xm%n*img%ym%n,MPI_REAL,i_r,105,MPI_COMM_WORLD,ierr)
       call MPI_SEND(vel0, nzold*img%ym%n*img%xm%n,MPI_REAL,i_r,110,MPI_COMM_WORLD,ierr)
    end do
 else
    call MPI_RECV(pick1, img%xm%n*img%ym%n,MPI_REAL,0,105,MPI_COMM_WORLD,status,ierr)
    call MPI_RECV(vel0, nzold*img%ym%n*img%xm%n,MPI_REAL,0,110,MPI_COMM_WORLD,status,ierr)
    write(0,*) "Receiving picks and velocity at node ",wpar%i_mpi
 end if

 allocate(ws(sou%hx%n,sou%hy%n,wblock%ns),wr(sou%hx%n,sou%hy%n,wblock%ns))
! SEND/RECV SHOTS

    do is=1,shtblk
       call owwei_init(is,1,1)
       if (is == shtblk) then
          if (shtres.eq.0) then
             shtida = shtblm
          else
             shtida = shtres
          end if
       else
          shtida = shtblm
       end if
       do iww=1, wblock%nb
          if (wpar%i_mpi==0) then
             do i_r=1,shtida
                ws=0;wr=0
                write(0,*) "Sent shot ",(is-1)*shtblm+i_r, " to node ", i_r," freq.block = ",iww
                call gathers_io(iww,(is-1)*shtblm+i_r,wr(:,:,1:wblock%blen(iww)),"rec","reed")
                call gathers_io(iww,(is-1)*shtblm+i_r,ws(:,:,1:wblock%blen(iww)),"sou","reed")
                call MPI_SEND(wr,sou%hx%n*sou%hy%n*wblock%blen(iww),MPI_COMPLEX,i_r,120,MPI_COMM_WORLD,ierr)
                call MPI_SEND(ws,sou%hx%n*sou%hy%n*wblock%blen(iww),MPI_COMPLEX,i_r,130,MPI_COMM_WORLD,ierr)
                call MPI_SEND(i_r,1,MPI_INTEGER,i_r,131,MPI_COMM_WORLD,ierr)
             end do
          else
             ws=0;wr=0
             call MPI_RECV(wr,sou%hx%n*sou%hy%n*wblock%blen(iww),MPI_COMPLEX,0,120,MPI_COMM_WORLD,status,ierr)
             call MPI_RECV(ws,sou%hx%n*sou%hy%n*wblock%blen(iww),MPI_COMPLEX,0,130,MPI_COMM_WORLD,status,ierr)
             call MPI_RECV(i_r,1,MPI_INTEGER,0,131,MPI_COMM_WORLD,status,ierr)
             call gathers_io(iww,is,wr(:,:,1:wblock%blen(iww)),"rec","rite")
             call gathers_io(iww,is,ws(:,:,1:wblock%blen(iww)),"sou","rite")
             write(0,*) "RECV shot ",(is-1)*shtblm+i_r," in node ", wpar%i_mpi," freq.block = ",iww
          end if
       end do
    end do
 deallocate(wr,ws)
write(0,*) "==================================="
 call MPI_BARRIER(MPI_COMM_WORLD,ierr)
if (wpar%i_mpi==0) then
write(0,*) "==================================="
write(0,*) "Beginning of the optimization loops"
write(0,*) "==================================="
end if
!--------------------------------------------
 ! Initial velocity model = background velocity
 do k=nzmin,nzmax;do j=1,img%ym%n;do i=1,img%xm%n
    x((k-nzmin)*img%ym%n*img%xm%n+(j-1)*img%xm%n+i)=dble(vel0((k-1)*img%ym%n*img%xm%n+(j-1)*img%xm%n+i))
 end do;end do;end do

 call to_history("n4",1,"Velout")
 call srite("Velout",vel0,4*nzold*img%xm%n*img%ym%n)

 xn=x
 minxx=sngl(minval(x))
 if (bounds) then
    allocate(l(img%z%n*img%ym%n*img%xm%n));l=0.
    allocate(u(img%z%n*img%ym%n*img%xm%n));u=0.
    if (cbounds) then
       l=vmin
       u=vmax
    else
       l=p_lo*1/x
       u=p_up*1/x
    end if
 end if 

 call steep_mpi_init(pick1,minxx)
 if (scaling) call scale_init()
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! Main code loop starts here  
!! 

 fun=0;gn=0;g=0.
 call fgeval_init(1,sngl(xn))
! call init_files()
 call scale_init1(.true.)
 call MPI_BARRIER(MPI_COMM_WORLD,ierr)
 call fgeval(vel0=sngl(xn),ff=fun(1),gg=gn,grad=g,ngrad=ngrad0)
 if (wpar%i_mpi==0) then
    write(0,*) 'Iteration = ',1, " fn = ",fun(1),0
 end if
 lambda=1.
 i=1
800 x=xn+lambda*g
    if (any(mask=(abs(1/x - 1/xn) > deltav/xn)).and.i < 40) then
        lambda=factor*lambda
        x=xn+lambda*g
        i=i+1
        if (i > 40) goto 1000
        goto 800
    end if
1000 g=g*lambda
! x=xn+lambda*g
! do while(any(mask=(abs(1/x - 1/xn) > deltav/xn)).or.i < 40)
!    lambda=factor*lambda
!    x=xn+lambda*g
!    i=i+1
! end do
! g=g*lambda
 write(0,*) "lambda = ",lambda,i
 call to_history("n4",1,"grad_")
 call grad_io("grad_","rite",1,sngl(g))

 do iters =2, max_iters
    fn=0.;gn=0
    do ils=1,nls
       x = xn + alpha*ils*g
       write(0,*) ils,"Max/min vel.",int(maxval(1/x)),int(minval(1/x))
!       call init_files()
       call scale_init1(.false.)
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       call fgeval(vel0=sngl(x),ff=fn(ils),gg=gn)
    end do
    if (wpar%i_mpi==0) write(0,*) iters," fn_vector = ", fn
    fmin = fn(1)
    imin = 1
    do ils=1,nls
       if (fn(ils).le.fmin) then
          fmin = fn(ils)
          imin = ils
       end if
    end do

    if (fmin <= fun(iters-1)) then
       x = xn + alpha*imin*g
!       alpha = alpha*imin/(0.5*(nls+1))
       alpha = alpha * 0.5
       xn = x
    else
       alpha = alpha * 0.5
    end if
    do k=nzmin,nzmax;do j=1,img%ym%n;do i=1,img%xm%n
       vel0((k-1)*img%ym%n*img%xm%n+(j-1)*img%xm%n+i)=sngl(xn((k-nzmin)*img%ym%n*img%xm%n+(j-1)*img%xm%n+i))
    end do;end do;end do
    call to_history("n4",iters,"Velout")
    call srite("Velout",vel0,4*nzold*img%xm%n*img%ym%n)
    g=0.
    call fgeval_init(iters)
!    call init_files()
    call scale_init1(.false.)
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    call fgeval(vel0=sngl(xn),ff=fun(iters),gg=gn,grad=g,ngrad=ngrad1)
    write(0,*) iters,ngrad0,ngrad1

    if (ngrad0 > tiny(1.)) then
       g=g*ngrad1/ngrad0*lambda
    else
       g=g*lambda
    end if
    call to_history("n4",iters,"grad_")
    call grad_io("grad_","rite",iters,sngl(g))
    ! ****************************************************************** 
    ! Line search complete 
    ! ****************************************************************** 
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)

    if (wpar%i_mpi==0) then
       write(0,*) 'Iteration = ',iters, " fn = ",fun(iters),alpha
    end if

 end do

01  format(a9,i7,a6,i9)
02  format(a9,i7,2(a6,i9))
03  format(a9,i7)
04  format(a12,i4,a7,i11)
1212 close (3)
 if (scaling) call scale_close()
 if (allocated(x)) deallocate(x)
 if (allocated(g)) deallocate(g)
 if (allocated(xn)) deallocate(xn)
 if (allocated(vel0)) deallocate(vel0)
 call steep_mpi_close()
 call sep_close()
 call MPI_FINALIZE(mpierr)
 stop
end program
