program nlin_dir
! line_search = Option switch for line seaching (1=2-pt,0=simple) 
!      = 0 = none 
!      = 1 = simple 
!      = 2 = 2-point parabolic 
!      = 3 = 2-point parabolic 

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(:), g0(:)
real, allocatable                             :: l(:), u(:)
integer                                       :: mpierr,i_r,ierr,status,shtida

real                                          :: lambda,lam1,lam2,lam3,lamc,lam4,norm_g0,norm_g1,norm_dx
double precision                              :: ft,fn,f2,gt,gn,g2,h20,h320,h32,ggn,teta,minxx
double precision, allocatable                 :: xn(:), dx(:)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

 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(g0(img%z%n*img%ym%n*img%xm%n));g0=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(dx(img%z%n*img%xm%n*img%ym%n));dx=dble(0.)
 allocate(vel0(nzold*img%ym%n*img%xm%n));vel0=0
 allocate(velout(nzold*img%ym%n*img%xm%n));velout=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)
write(0,*) "==================================="
write(0,*) "Beginning of the optimization loops"
write(0,*) "==================================="
!--------------------------------------------
 ! 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

 if (wpar%i_mpi==0) then
    call to_history("n4",1,"Velout")
    call srite("Velout",vel0,4*nzold*img%xm%n*img%ym%n)
 end if

 xn=x
 minxx=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)
 if (scaling) call scale_init()
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! Main code loop starts here  
!! 
 do while (iters <= max_iters) 
    fn=0;gn=0;g=0.
    call fgeval_init(iters)
    call init_files()
    call scale_init1(.true.)
    call fgeval(vel0=sngl(x),ff=fn,gg=gn,grad=g)

    if (multsc) then
       if (iters == 1.or.iters == 4.or.iters == 6) then
          dx=(g/maxval(abs(g)))*minxx 	! Update the Newton search direction 
          g0=dx
       else if (iters==2.or.iters==3.or.iters==5.or.iters>6) then 
          norm_g1=norm(g)
          if (norm_g0.gt.tiny(1.)) then
             dx=(g/maxval(abs(g)))*minxx*norm_g1/norm_g0
          else
             dx=(g/maxval(abs(g)))*minxx
          end if
          if (dot_product(g0,g0).gt.dble(tiny(1.))) then
             teta=max(dble(0.),dot_product(dx,dx-g0)/dot_product(g0,g0))
          else
             write(0,*) "Gradient too small. Optimization stopped" 
             goto 1212
          end if
          g0=dx
          dx=dx+teta*g0
          do while (norm(dx).gt.norm_dx)
             dx=dx*.9
          end do
       end if
    else
       if (iters == 1) then
          dx=(g/maxval(abs(g)))*minxx 	! Update the Newton search direction 
          g0=dx
       else 
          norm_g1=norm(g)
          if (norm_g0.gt.tiny(1.)) then
             dx=(g/maxval(abs(g)))*minxx*norm_g1/norm_g0
          else
             dx=(g/maxval(abs(g)))*minxx
          end if
          if (dot_product(g0,g0).gt.dble(tiny(1.))) then
             teta=max(dble(0.),dot_product(dx,dx-g0)/dot_product(g0,g0))
          else
             write(0,*) "Gradient too small. Optimization stopped" 
             goto 1212
          end if
          g0=dx
          dx=dx+teta*g0
          do while (norm(dx).gt.norm_dx)
             dx=dx*.9
          end do
       end if
    end if
    norm_g0=norm(g)
    norm_dx=norm(dx)
    call to_history("n4",iters,"grad_")
    call grad_io("grad_","rite",iters,sngl(dx))
    write(0,*) "DX ... ",norm_g0,norm_g1," proc = ",wpar%i_mpi
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    ! ******************** 
    ! Line Searching below 
    ! ******************** 
    lam3=1.

900 x=xn+lam3*dx
    if (any(mask=(1/x < l.or.1/x > u)).and.(norm(1/x).gt.1.2*norm(1/xn).or.norm(1/x).lt.0.8*norm(1/xn))) then
       lam3=factor*lam3
       if (wpar%i_mpi==0) write(0,*) "1) Bound reached - recomputing lam3", lam3
       if (lam3.lt.tiny(1.)) then
          if (wpar%i_mpi==0) write(0,*) "lam3 < tiny(1.) ==> progr.STOPPED"
          if (wpar%i_mpi==0) write(0,*) "Problem with gradient values and/or crongrats, optimization converged!!!"  
          goto 1212
       end if
       goto 900
    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(x((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,"vtemp")
    call srite("vtemp",vel0,4*nzold*img%xm%n*img%ym%n)
    ft=0;gt=0
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    call init_files()
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    call scale_init1(.false.)
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    call fgeval(vel0=sngl(x),ff=ft,gg=gt)

    iww=0
    lam2=lam3
    write(0,*) " @@@ gt= ",gt, " gn= ",gn," proc = ",wpar%i_mpi
    do while (gt > gn.and.iww<=5) 
901    lam3=factor*lam3
       x=xn+lam3*dx
       if (any(mask=(1/x < l.or.1/x > u)).and.(norm(1/x).gt.1.2*norm(1/xn).or.norm(1/x).lt.0.8*norm(1/xn))) then
          if (wpar%i_mpi==0) write(0,*) "2) Bound reached - recomputing lam3", lam3
          if (lam3.lt.tiny(1.)) then
             if (wpar%i_mpi==0) write(0,*) "lam3 < tiny(1.) ==> progr.STOPPED"
             if (wpar%i_mpi==0) write(0,*) "Problem with gradient values and/or crongrats, optimization converged!!!"  
             goto 1212
          end if
          goto 901
       end if
       ft=0;gt=0
       call MPI_BARRIER(MPI_COMM_WORLD,ierr)
       call init_files()
       call scale_init1(.false.)
       call fgeval(vel0=sngl(x),ff=ft,gg=gt)
       write(0,*) "01) fn = ",fn," ft = ",ft," proc = ",wpar%i_mpi,iww
       iww=iww+1
    end do     
    if (iww==6 .and. gt > gn) then
       lambda=0
    else
       lambda=lam3
    end if

    ! ****************************************************************** 
    ! Line search complete 
    ! ****************************************************************** 
    call MPI_BARRIER(MPI_COMM_WORLD,ierr)
    if (lambda==0) then
       write(0,*) "YOU'RE IN PROBLEM: Line search did not converge!!!"
       goto 1212
    end if
    xn=xn+lambda*dx ! Update Newton iterate
    x=xn 
    iters=iters+1;            !Update iterations counter 
    if (wpar%i_mpi==0) then
       write(0,*) 'Iteration = ',iters,wpar%i_mpi
       write(0,*) "03) fn = ",fn,wpar%i_mpi
    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(x((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)

 end do

 if (wpar%i_mpi==0) then
    ! Success statement based on error and iters ******** 
    if (err > err1.or.iters == max_iters) then
       write(0,*) 'Convergence may not have occurred; error is greater than 20% grad_0 or max_iters was reached' 
    else if (err < err1) then
       write(0,*)  'The error is less than 20% grad_0; convergence has likely been achieved' 
    end if 
 end if
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(g0)) deallocate(g0)
 if (allocated(xn)) deallocate(xn)
 if (allocated(dx)) deallocate(dx)
 if (allocated(vel0)) deallocate(vel0)
 call steep_mpi_close()
 call sep_close()
 call MPI_FINALIZE(mpierr)
 stop
end program
