program OWlbfgsc

use sep 
use OW_type
use OW_grabvel_areal
use OW_exec_time
use OW_parms_areal
use OW_taper_areal
use OW_wei_areal
!use OW_slant
!use OW_dotst
use solver_lbfgsb_mod
use base_spline_mod
use fx2d

implicit none

character*80                                  :: hor1,log,listagem
real                                          :: start, finish, start1, finish1, start2, finish2
real                                          :: o3,a,b,vmin,vmax,stat
integer                                       :: ixs,ino,ni,i,st,idip,n3,t3,iz,j,k,iter,ix,ish,titer
real, allocatable                             :: velint(:),taper(:),pick1(:)
real, allocatable                             :: deltaI(:,:,:),scatS(:),deltaS(:)
real, allocatable                             :: R_image(:,:,:),B_image(:,:,:)
real, allocatable                             :: vel(:),vref(:,:,:),vel0(:)
complex, allocatable                          :: recv(:,:),source(:,:),source1(:)
logical                                       :: cond,bounds
type(axis)                                    :: axk, ayk, azk
real                                          :: axkd, aykd, azkd, axko, ayko, azko
integer                                       :: axkn, aykn, azkn
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Variables of L-BFGSC
!
 integer                       :: nmax, mmax ! nmax - dimension of the problem and mmax - max.# of ltd.memory corrections.
 
!     Declare the variables needed by the code.
!     A description of all these variables is given at the end of 
!     the driver.
 
 character*60                  :: task, csave
 logical                       :: lsave(4)
 integer                       :: n, m, iprint,subiter
 integer,allocatable           :: nbd(:), iwa(:)
 integer                       :: isave(44)
 double precision              :: f, factr, pgtol
 double precision              :: dsave(29)
 double precision, allocatable :: x(:), l(:), u(:), g(:),wa(:)
 real, allocatable             :: spl(:),spl2(:)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 call sep_init()
 call sep_begin_prog()
 start=0.;finish=0.
 call owexec_time(start)
 call owparam()

 call from_param("oxk", axko, 0.); call from_param("dxk", axkd, 1.); call from_param("nxk", axkn, 99999 )
 call from_param("oyk", ayko, 0.); call from_param("dyk", aykd, 1.); call from_param("nyk", aykn, 99999 )
 call from_param("ozk", azko, 0.); call from_param("dzk", azkd, 1.); call from_param("nzk", azkn, 99999 )

 nmax=img%z%n*img%xm%n
 mmax=5

 call from_param("log",log,"log")
 open(unit=2,file=log)
 call from_param("list",listagem,"listagem")
 open(unit=3,file=listagem)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Memory allocation
!
 allocate(x(nmax),l(nmax),u(nmax),g(nmax),nbd(nmax))
 allocate(spl(nmax)); allocate(spl2(nmax))
 allocate(wa(2*mmax*nmax+4*nmax+12*mmax*mmax+12*mmax))
 allocate(iwa(3*nmax))
 allocate(deltaS(img%z%n*kxm%n));deltaS=0.
 allocate(scatS(img%z%n*img%xm%n));scatS=0.
 allocate(deltaI(img%z%n,nsecdim,kxm%n));deltaI=0.
 allocate(R_image(img%z%n,nsecdim,img%xm%n));R_image=0.
 allocate(B_image(img%z%n,nsecdim,img%xm%n));B_image=0.
 allocate(taper(kxm%n))
 allocate(velint(v%z%n*v%xm%n),vel(img%z%n*kxm%n),vel0(img%z%n*img%xm%n))
 allocate(vref(nref,2,v%z%n))

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Set up the taper
!
 call owtaper(tapertr,taper)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Read velocities
!

 call sreed("vel",velint,v%z%n*v%xm%n*4)
 if(exist_file("vref")) then
    call sreed("vref",vref,v%z%n*2*nref*4)
 end if
 call owgrab_vel(velint,vel0,.true.,ni)

 deallocate(velint)

!     We wish to have output at every iteration.
 iprint = 99

! Tolerances and constraints
 factr  = 1.0d0+3
 pgtol  = 1.0d-16
 call from_param("bounds",bounds,.false.)
 if (bounds) then
    call from_param("vmin",vmin)
    call from_param("vmax",vmax)
    l=dble(1/vmax)
    u=dble(1/vmin)
    nbd=2
 else
    nbd=0
 end if 

!----------------------------------------------
! Horizon with no velocity update
if (exist_file("horv0")) then
   call from_param('horv0',hor1)
   open(51,file=hor1,status='old')
   allocate(pick1(img%xm%n));pick1=9999.
   do i=1,img%xm%n
      read(51,*,end=1) pick1(i)
   end do
1 close(51)
end if
!----------------------------------------------
! Initial velocity model = background velocity
 x=vel0

 iter=0
 subiter=1
 titer = 1
 task = 'START'

 cond=.true.

! do while(cond .or. iter <= niter)
 do while(iter <= niter)
    cond=.false.

    call setulb(nmax,mmax,x,l,u,nbd,f,g,factr,pgtol,wa,iwa,task,iprint,csave,lsave,isave,dsave)
 
    if (task(1:2) .eq. 'FG') then
!        the minimization routine has returned to request the
!        function f and gradient g values at the current x.
!        Compute function value f for the sample problem.
!
!     x is the updated slowness
!

       cond=.true.

!       if (exist_file("horv0")) then
!          call oweditvel(pick1,x,vel0)
!       end if

       subiter=subiter+1

       call to_history("n4",titer,"Velout")
       call srite("Velout",sngl(1/x),4*img%z%n*img%xm%n)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Computing bkg.image
!
       start1=0.;finish1=0.;
       B_image=0.;R_image=0.;
       call sseek("rec",0,0)
       call sseek("sou",0,0)
       call owexec_time(start1)
       if (not(areal)) then
          if (allocated(source1)) deallocate(source1); allocate(source1(sou%w%n))
          source1=0;
          if (allocated(source)) deallocate(source); allocate(source(rec%h%n,sou%w%n))
          source=0;
          call sreed("sou",source1,8*sou%w%n)
          source(abs(floor(rec%h%o/rec%h%d)) + 1,:)=source1
       end if

       do idip=1,rec%xd%n
          start=0.;finish=0. 
          call owexec_time(start)
          if (areal) then
             if (allocated(source)) deallocate(source); allocate(source(sou%h%n,sou%w%n))
             source=0;
             call sreed("sou",source,8*sou%h%n*sou%w%n)
             vel=sngl(x)
             ni=1
          else
             call owgrab_vel(sngl(x),vel,.false.,ni,idip)
          end if
             
          if (allocated(recv))   deallocate(recv);   allocate(recv(rec%h%n,rec%w%n))
          recv=0;deltaI=0

          call sreed("rec",recv,8*rec%h%n*rec%w%n)
          if (exist_file("vref")) then
             call owwei_init()
             call owwei(taper,recv,source,vel,deltaI,vref)
          else
             call owwei_init()
             call owwei(taper,recv,source,vel,deltaI)
          end if
          if (areal) then
             B_image=B_image+deltaI
          else
             B_image(:,:,ni:ni+kxm%n-1)=B_image(:,:,ni:ni+kxm%n-1)+deltaI
          end if
          call owexec_time(finish)

          write(3,111) "EXTRAPOLATED WAVEFIELDS = ",idip," OUT OF ",rec%xd%n," IN ",(finish-start)/60," min"
       end do

       call owexec_time(finish1)
       write(3,112) "COMPUTED BACKGROUND IMAGE IN ",(finish1-start1)/60," min OF ITERATION # ", iter
       call to_history("n4",titer,"bkgimg")
       if (timeshift) then
          call srite("bkgimg",B_image,4*img%z%n*img%th%n*img%xm%n)
       else
          call srite("bkgimg",B_image,4*img%z%n*img%sh%n*img%xm%n)
       end if
!
!
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Computing DSO
!
!       if (timeshift) then
!          if (allocated(R_image)) deallocate(R_image) ;allocate(R_image(img%z%n,img%sh%n,img%xm%n));R_image=0.
!          do ixm=1,img%xm%n
!          stat = owslant_init(sngl(1/x))
!          call owslant_exe(.true.,.false.,B_image,R_image)
!          stat = owslant_close()
!          deallocate(B_image);allocate(B_image(img%z%n,img%sh%n,img%xm%n))
!          B_image=R_image
!       end if
       if (zweight) then
          call owzwht(.false., B_image, R_image)
          B_image=R_image
       end if
       call owdso(.false. , B_image, R_image)
       if (nxfilt.ne.99999) then
          stat = fx2d_lop (img%z%n,nsecdim,img%xm%n,5,50,5,int(1.5*img%z%n/(5+1)),int(1.5*img%xm%n/(50-5+1)),20.,R_image,B_image)
!          call owxwht(.false., B_image, R_image)
!          call owxwht(.true. , B_image, R_image)
          R_image=B_image
       end if       
       call owdso(.true.  , B_image, R_image)
       if (zweight) then
          R_image=B_image
          call owzwht(.true. , B_image, R_image)
       end if

       call to_history("n4",titer,"Dimage")
       call srite("Dimage",B_image,4*img%z%n*nsecdim*img%xm%n)
!
!
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Computing the gradient
!

       start1=0.;finish1=0.;

       call sseek("rec",0,0)
       call sseek("sou",0,0)
       call owexec_time(start1)
       if (not(areal)) then
          if (allocated(source1)) deallocate(source1); allocate(source1(sou%w%n))
          source1=0;
          if (allocated(source)) deallocate(source); allocate(source(rec%h%n,sou%w%n))
          source=0;
          call sreed("sou",source1,8*sou%w%n)
          source(abs(floor(rec%h%o/rec%h%d)) + 1,:)=source1
       end if

       start1=0.;finish1=0.;scatS=0.
       call owexec_time(start1)
       do idip=1,rec%xd%n
          start=0.;finish=0. 
          call owexec_time(start)
          if (areal) then
             if (allocated(source)) deallocate(source); allocate(source(sou%h%n,sou%w%n))
             source=0;
             call sreed("sou",source,8*sou%h%n*sou%w%n)
             vel=x
             ni=1
          else
             call owgrab_vel(sngl(x),vel,.false.,ni,idip)
          end if
             
          if (allocated(recv))   deallocate(recv);   allocate(recv(rec%h%n,rec%w%n))
          recv=0;deltaS=0
          call sreed("rec",recv,8*rec%h%n*rec%w%n)

          call owwemva_init(idip)
          if (exist_file("vref")) then
             call owwemva(.true.,taper,recv,source,vel,deltaS,B_image(:,:,ni:ni+kxm%n-1),vref)
          else
             call owwemva(.true.,taper,recv,source,vel,deltaS,B_image(:,:,ni:ni+kxm%n-1))
          end if
          if (areal) then
             scatS=scatS+deltaS
          else
             do j=ni,ni+kxm%n-1;do k=1,img%z%n
                scatS((j-1)*img%z%n+k)=scatS((j-1)*img%z%n+k)+deltaS((j-ni)*img%z%n+k)
             end do;end do
          end if
          
          call owexec_time(finish)
          write(3,111) "EXTRAPOLATED WAVEFIELDS = ",idip," OUT OF ",rec%xd%n," IN ",(finish-start)/60," min"
       end do

       g=dble(scatS)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! Smoothing/scaling the gradient

       if (axk%n .ne. 99999.and.azk%n .ne. 99999) then
!         stat = base_spln3d_init(img%xm,img%ym,img%z,axk,ayk,azk)
          if (titer <= 5) then
             axk%d=axkd*(6-titer)
             ayk%d=aykd*(6-titer)
             azk%d=azkd
             axk%n=axkn/(6-titer)
             ayk%n=aykn/(6-titer)
             azk%n=azkn
             axk%o=axko
             ayk%o=ayko
             azk%o=azko
          else
             axk%d=axkd
             ayk%d=aykd
             azk%d=azkd
             axk%n=axkn
             ayk%n=aykn
             azk%n=azkn
             axk%o=axko
             ayk%o=ayko
             azk%o=azko
          end if
          stat = base_spln2d_init(img%z,img%xm,azk,axk)

          call owexec_time(start2)
!          spl2=sngl(g)
!          stat = fx2d_lop (img%z%n,1,img%xm%n,5,50,5,int(1.5*img%z%n/(5+1)),&
!                           int(1.5*img%xm%n/(50-5+1)),40.,spl2,spl)
!          g=dble(spl)
          stat = base_spln2d_oper(.true. , .false., spl, g)
          stat = base_spln2d_oper(.false., .false., spl, g)
          call owexec_time(finish2)
          write(3,112) "SMOOTHED GRADIENT IN ",(finish2-start2)/60
       end if

       g = g/maxval(abs(g))
!       if (titer < 30 ) then
          g=-g*(1-exp(-float(titer)/damp))
!       else
!          g=-g*(1-exp(-float(titer)))
!       end if
!
!
       if (exist_file("horv0")) then
          call oweditvel(pick1,g)
       end if
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
       call owexec_time(finish1)

       write(3,112) "COMPUTED THE GRADIENT IN ",(finish1-start1)/60," min OF ITERATION # ", iter

       call to_history("n4",titer,"grad")
       call srite("grad",sngl(g),4*img%z%n*img%xm%n)

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!  Computing the objective function
!
       a = 0.
       do i=1,img%xm%n
          do j=1,nsecdim
             if (maxval(abs(B_image(:,j,i))) > a) a=maxval(abs(B_image(:,j,i)))
          end do
       end do
       B_image=B_image/a
       f=dble(0.)
       do i=1,img%xm%n
          do k=1,img%z%n
             f = f + dot_product(B_image(k,:,i),B_image(k,:,i))
          end do
       end do

       if (timeshift) then
          if (allocated(B_image)) deallocate(B_image) ;allocate(B_image(img%z%n,img%th%n,img%xm%n));B_image=0.
       end if

       write(2,11) " iter = ",iter," subiter = ",subiter,"titer = ",titer," dotprod grad = ", &
                     dot_product(g,g) , " Obj.function = ", f
       titer=titer+1

    end if

    if (task(1:5) .eq. 'NEW_X')  then
       iter=iter+1
       subiter=1
       cond=.true.
    end if

 end do
 111 format(a26,i4,a8,i4,a4,f6.2,a4)
 112 format(a29,f10.2,a20,i4)

 deallocate(x,l,u,g,nbd)
 deallocate(spl)
 deallocate(wa)
 deallocate(iwa)
 deallocate(source)
 deallocate(recv)
 deallocate(deltaS)
 deallocate(scatS)
 deallocate(deltaI)
 deallocate(R_image)
 deallocate(B_image)
 deallocate(taper)
 deallocate(vel)
 deallocate(vref)
 close(2)
 close(3)
11 format(a8,i4,a11,i5,a8,i4,2(a16,f12.5))
 call sep_end_prog()

end program OWlbfgsc
