program steepest_ls3D

! 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
real, allocatable                             :: velint(:),vel0(:),pick1(:),velout(:)
double precision, allocatable                 :: x(:), g(:)
real, allocatable                             :: l(:), u(:)
integer                                       :: mpierr,i_r,ierr,status

real                                          :: lambda,lam1,lam2,lam3,lamc,lam4
double precision                              :: ft,fn,f2,gt,gn,g2,h20,h320,h32
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()
 call owparam()
 call owparm_bcast1()
 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(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

 call init_files()

 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

    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

    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

 !----------------------------------------------
 ! Initial velocity model = background velocity
 do k=nzmin,nzmax;do j=1,img%ym%n;do i=1,img%xm%n
!    x((i-1)*img%ym%n*nz+(j-1)*nz+k-nzmin+1)=dble(vel0((i-1)*img%ym%n*nzold+(j-1)*nzold+k))
    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
 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)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!! Main code loop starts here  
!! changed or to and
 do while (err > tola.or.iters <= max_iters) 
    fn=0;gn=0;g=0.
    call fgeval_init(iters)
    call fgeval(vel0=sngl(x),ff=fn,gg=gn,grad=g,riteD=.true.,wsave=.true.)
    if (wpar%i_mpi==0) then
       call to_history("n6",iters,"yline")
    end if
    dx=-(g/maxval(abs(g)))*minval(x) 	! Update the Newton search direction 
    err=norm(dx)				! Error 
    if (iters==1) tola=tol(1)*err+tol(2) 	! Set tolerance for stopping criterium 
    if (wpar%i_mpi==0) write(0,*) minval(dx),maxval(dx)
    ! ******************** 
    ! Line Searching below 
    ! ******************** 
    lam3=1.
    gn=fn*fn
900 x=xn+lam3*dx
    if (any(mask=(1/x < l.or.1/x > u))) 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
    ft=0;gt=0
    call fgeval(vel0=sngl(x),ff=ft,gg=gt,riteD=.false.,wsave=.false.)
    if (wpar%i_mpi==0) write(0,*) "==> ",gt,gn,lam3,lambda_min

    do while (gt > gn.and.lam3 > lambda_min) 
901    lam3=factor*lam3
       x=xn+lam3*dx
       if (any(mask=(1/x < l.or.1/x > u))) 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 fgeval(vel0=sngl(x),ff=ft,gg=gt,riteD=.false.,wsave=.false.)
       if (wpar%i_mpi==0) write(0,*) "01) fn = ",fn," ft = ",ft
    end do     

    lam2=factor*lam3; 
    x=xn+lam2*dx 
    f2=0;g2=0
    call fgeval(vel0=sngl(x),ff=f2,gg=g2,riteD=.false.,wsave=.false.)
    h20=(g2-gn)/lam2
    h32=(gt-g2)/(lam3-lam2)
    h320=(h32-h20)/lam3
    lamc=0.5*(lam2-h20/h320); 
    x=xn+lamc*dx
    f2=0;g2=0
    call fgeval(vel0=sngl(x),ff=f2,gg=g2,riteD=.false.,wsave=.false.)
    if (wpar%i_mpi==0) write(0,*) "02) fn = ",fn," ft = ",ft," f2 = ",f2
    if (g2 < gt) then
       lambda=lamc
       fn=f2
       gn=g2
    else 
       lambda=lam3
       fn=ft
       gn=gt
    end if   
    ! ****************************************************************** 
    ! Line search complete 
    ! ****************************************************************** 
    xn=xn+lambda*dx ! Update Newton iterate
    x=xn 
    if (wpar%i_mpi==0) then
       write(0,*) 'Iteration = ',iters,' error=',err
       write(0,*) "03) fn = ",fn
       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))
!          vel0((i-1)*img%ym%n*nzold+(j-1)*nzold+k)=sngl(x((i-1)*img%ym%n*nz+(j-1)*nz+k-nzmin+1))
       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 if
    iters=iters+1;            !Update iterations counter 
 end do

 if (allocated(x)) deallocate(x)
 if (allocated(g)) deallocate(g)
 if (allocated(xn)) deallocate(xn)
 if (allocated(dx)) deallocate(dx)
 if (allocated(vel0)) deallocate(vel0)
 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)
 call steep_mpi_close()
 call sep_close()
 call MPI_FINALIZE_(mpierr)
end program
