module fgevaluation
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 base_spline_mod
!
use helderiv
use identity_mod
use find_xz_dip_mod
use solver_reg_l2_mod
use cgstep_mod
!
implicit none
real, private, save, allocatable                        :: taper(:),pick(:)
contains
!===========================================
!subroutine owfgeval_init()
!end subroutine
!===========================================
subroutine taper_set()
 allocate(taper(kxm%n))
 call owtaper(tapertr,taper)
end subroutine
!===========================================
subroutine taper_close()
 deallocate(taper)
end subroutine
!===========================================
subroutine wolfe_init(pick_in)
real            :: pick_in(:)
allocate(pick(size(pick_in)))
pick=pick_in
end subroutine
!===========================================
subroutine wolfe_close()
deallocate(pick)
end subroutine
!===========================================
subroutine wolfe_search2(x,p,gamma,c1,c2,f0,df0,iter,xplus,alpha,fval,dfval,nfe,nge)

double precision                       :: gamma
double precision                       :: x(:),p(:),c1,c2,f0,df0(:)
double precision                       :: alpha,fval,dfval(:)
integer                                :: nfe, nge,iter

!! internal variables
double precision           :: slope0,slope,xplus(size(x)),old_alpha,old_fval,old_dfval(size(df0)),old_slope
double precision           :: temp,alpha_lo,alpha_hi,f_lo,f_hi,df_lo(size(df0)),slope_lo,a,b,dalpha,c
integer                    :: lsiter,iter_cnt,firsttrip,ls_iter
real                       :: B_image(img%xm%n,nsecdim,img%z%n),D_image(img%xm%n,nsecdim,img%z%n),x1(size(x)),x0(size(x))
!real, pointer              :: z0(:)
real                       :: z0(img%xm%n,img%z%n)
integer                    :: i

ls_iter=5 
!! [ALPHA,FVAL,DFVAL,NFE,NGE] = wolfe_search2(F,DF,PARAMS,X,P,ALPHA0,C1,C2,F0,DF0,TRACE)
!! Performs search for point satisfying the two strong Wolfe conditions:
!!
!!	f(x+alpha.p) <= f(x) + c1.alpha p'.grad f(0)		(WC1)
!!	|p'.grad f(x+alpha.p)| <= c2|p'.grad f'(x)|		(WC2b)
!!
!! ALPHA0 is the initial value of alpha in the search.  If ALPHA0
!!	satisfies the Wolfe conditions, then the search stops there. 
!! F0 is f(x) at the initial x.
!! DF0 is grad f(x) at the initial x.
!! DF(PARAMS,X) is the gradient of F(PARAMS,X) with
!!	respect to X
!! Note that PARAMS is the parameter for F: F(PARAMS,X)
!! If TRACE is non-zero then print out information about the process
!!
!! The algorithm follows that of Wright & Nocedal "Numerical Optimization",
!! pp. 60-62, section 3.4
!!
!! Returns
!!	ALPHA	-- step length parameter
!!	FVAL	-- function value at end
!!	DFVAL	-- gradient value at end
!!	NFE	-- # function evaluations
!!	NGE	-- # gradient evaluations

!!
!! Check parameters
!!
alpha = 0
fval = f0
dfval = df0
nfe = 0
nge = 0

if (not(0<c1.and.c1<c2.and.c2<1)) then
  write(4,*) 'wolfe_search2: Error: Need 0 < c1 < c2 < 1'
  write(4,*) 'wolfe_search2: c1 = ',c1,' c2 = ', c2
  return
end if

! Initialization

slope0 = dot_product(p,df0)
if (slope0 >= 0) then
  write(4,*) 'wolfe_search2: Error: Need a descent direction'
  write(4,*) 'wolfe_search2: p*grad f(x) = ',slope0
  return
end if
!
! Braketing phase
!
write(4,*) 'wolfe_search2: Bracketing phase: alpha = ', gamma
alpha = gamma
old_alpha = 0
old_fval = f0
old_dfval = df0
old_slope = slope0
write(4,*) 'wolfe_search2: f(x) = ',f0,' p*grad f(x) = ',slope0

! Main loop
firsttrip = 1
lsiter=1
do while (lsiter < ls_iter )
  B_image=0
  D_image=0
  fval=0d+0
  dfval=0d+0
  xplus = x+alpha*p
  if (allocated(pick)) then
     call oweditvel(pick,xplus,sngl(x))
  end if

  where(xplus.gt.(1+maxvchg)*maxval(x)) xplus=(1+maxvchg)*maxval(x)
  where(xplus.lt.(1+maxvchg)*minval(x)) xplus=(1+maxvchg)*minval(x)
  call owfeval(xplus,fval,B_image,D_image)

!  if (smooth) then
!     allocate(z0(img%xm%n*img%z%n))
!!     do i=1,img%xm%n
!!        z0((i-1)*img%z%n+1:i*img%z%n)=D_image(i,nsecdim/2,:)
!!     end do
!!     x0=reshape(transpose(reshape(sngl(xplus),(/img%xm%n,img%z%n/))),(/img%xm%n*img%z%n/))
!!     call helderiv_init ( img%z%n,serg_na,serg_epsh)
!!     call xsergey_dip_init(z0,serg_niter,serg_eps1,serg_eps2)
!!     call solver_reg_l2(m=x1,d=x0, Fop=I_lop, Aop=xsergey_dip_lop, stepper=cgstep, &
!!          nAop=img%z%n*img%xm%n,niter=dip_niter,eps=dip_eps,m0=x0,verb=.true.)
!!     xplus=reshape(transpose(reshape(dble(x1),(/img%z%n,img%xm%n/))),(/img%xm%n*img%z%n/))
!!     call helderiv_close()
!     do i=1,img%z%n
!        z0((i-1)*img%xm%n+1:i*img%xm%n)=D_image(:,nsecdim/2,i)
!     end do
!     call helderiv_init ( img%xm%n,serg_na,serg_epsh)
!     call xsergey_dip_init(z0,serg_niter,serg_eps1,serg_eps2)
!     call solver_reg_l2(m=x1,d=sngl(xplus), Fop=I_lop, Aop=xsergey_dip_lop, stepper=cgstep, &
!          nAop=img%z%n*img%xm%n,niter=dip_niter,eps=dip_eps,m0=sngl(xplus),verb=.true.)
!     xplus=dble(x1)
!     call helderiv_close()
!     deallocate(z0)
!  end if

  nfe = nfe + 1
  write(4,*) lsiter,'wolfe_search2: alpha = ',alpha,' f(x+alpha*p) = ',fval
  if (fval > f0 + c1*alpha*slope0 .or. fval >= old_fval ) then
     write(4,*) lsiter,'wolfe_search2: (WC1) failed or f increased'
     exit
  end if
  write(4,*) lsiter,'wolfe_search2: (WC1) holds & f decreased'

  call owgeval(xplus,dfval,B_image)!,lsiter,iter)
  if (envelope) then
     do i=1,img%z%n
        z0(:,i)=D_image(:,nsecdim/2,i)
     end do       
     call owgeval(xplus,dfval,B_image,z0)!,lsiter,iter)
  else
     call owgeval(xplus,dfval,B_image)!,lsiter,iter)
  end if
  nge = nge + 1
  slope = dot_product(p,dfval)

  write(4,*) lsiter,'wolfe_search2: dot(p,grad f(x+alpha*p)) = ', slope
  if ( abs(slope) <= c2*abs(slope0) ) then
     write(4,*) lsiter,'wolfe_search2: (WC2) holds'
    return
  end if
  if ( slope >= 0 ) then
    write(4,*) lsiter,'wolfe_search2: d2f(alpha) >= 0'
    exit
  end if

  ! Update variables -- note no upper limit on alpha
  temp = alpha
  ! alpha = alpha + old_alpha
  alpha = 2*alpha
  old_alpha = temp
  old_fval = fval
  old_slope = slope
  lsiter=lsiter+1
end do

write(4,*) 'wolfe_search2: Entering zoom phase.'

!
! "Zoom" phase
!
alpha_lo = old_alpha
alpha_hi =     alpha
f_lo     = old_fval
f_hi     =     fval
df_lo    = old_dfval
slope_lo = old_slope
!!df_hi   =     dfval
!!slope_hi=     slope

write(4,*) 'wolfe_search2: zoom phase: alpha_lo = ',alpha_lo,' alpha_hi = ',alpha_hi
iter_cnt = 0
lsiter=1
do while (lsiter < ls_iter )

!  ! form quadratic interpolant of function values on alpha_lo, alpha_hi
!  ! and the derivative at alpha_lo...
!  ! and find min of interpolant within interval [alpha_lo,alpha_hi]
!  a = f_lo
!  b = slope_lo
  dalpha = alpha_hi-alpha_lo
!  c = (f_hi - f_lo - dalpha*slope_lo)/dalpha**2
!  ! iter_cnt
!  if ( c <= 0 .or. mod(iter_cnt,3) == 2 ) then
    ! Use bisection
    alpha = alpha_lo + 0.5*dalpha
    write(4,*) lsiter,'wolfe_search2: using bisection: c = ', c
!  else
!    ! Use min of quadratic
!    alpha = alpha_lo - 0.5*b/c
!    write(4,*) lsiter,'wolfe_search2: using quadratic: a = ', a,' b = ',b, ' c = ', c
!  end if
  write(4,*) lsiter,'wolfe_search2: alpha = ', alpha

  ! main part of loop
  xplus = x + alpha*p
  if (allocated(pick)) then
     call oweditvel(pick,xplus,sngl(x))
  end if

  D_image=0
  B_image=0
  fval=0d+0
  dfval=0d+0
  call owfeval(xplus,fval,B_image,D_image)!,lsiter,iter)

!  if (smooth) then
!     allocate(z0(img%xm%n*img%z%n))
!!     do i=1,img%xm%n
!!        z0((i-1)*img%z%n+1:i*img%z%n)=D_image(i,nsecdim/2,:)
!!     end do
!!     x0=reshape(transpose(reshape(sngl(xplus),(/img%xm%n,img%z%n/))),(/img%xm%n*img%z%n/))
!!     call helderiv_init ( img%z%n,serg_na,serg_epsh)
!!     call xsergey_dip_init(z0,serg_niter,serg_eps1,serg_eps2)
!!     call solver_reg_l2(m=x1,d=x0, Fop=I_lop, Aop=xsergey_dip_lop, stepper=cgstep, &
!!          nAop=img%z%n*img%xm%n,niter=dip_niter,eps=dip_eps,m0=x0,verb=.true.)
!!     xplus=reshape(transpose(reshape(dble(x1),(/img%z%n,img%xm%n/))),(/img%xm%n*img%z%n/))
!!     call helderiv_close()
!     do i=1,img%z%n
!        z0((i-1)*img%xm%n+1:i*img%xm%n)=D_image(:,nsecdim/2,i)
!     end do
!     call helderiv_init ( img%xm%n,serg_na,serg_epsh)
!     call xsergey_dip_init(z0,serg_niter,serg_eps1,serg_eps2)
!     call solver_reg_l2(m=x1,d=sngl(xplus), Fop=I_lop, Aop=xsergey_dip_lop, stepper=cgstep, &
!          nAop=img%z%n*img%xm%n,niter=dip_niter,eps=dip_eps,m0=sngl(xplus),verb=.true.)
!     xplus=dble(x1)
!     call helderiv_close()
!     deallocate(z0)
!  end if

  nfe = nfe + 1
  write(4,*) lsiter,'wolfe_search2: f(x+alpha*p) = ', fval
  if (fval > f0 + c1*alpha*slope0 .or. fval >= f_lo ) then
     write(4,*) lsiter,'wolfe_search2: zoom: (WC1) fails or f increased'
     write(4,*) lsiter,'wolfe_search2: zoom: update alpha_hi'
     alpha_hi = alpha
     f_hi   = fval
  else
    write(4,*) lsiter,'wolfe_search2: zoom: (WC1) holds and f decreased'

    if (envelope) then
       do i=1,img%z%n
          z0(:,i)=D_image(:,nsecdim/2,i)
       end do       
       call owgeval(xplus,dfval,B_image,z0)!,lsiter,iter)
    else
       call owgeval(xplus,dfval,B_image)!,lsiter,iter)
    end if
    nge = nge + 1
    slope = dot_product(p,dfval)
    write(4,*) lsiter,'wolfe_search2: dot(p,grad f(x+alpha*p)) = ', slope
    if (abs(slope) <= c2*abs(slope0)) then
       write(4,*) lsiter,'wolfe_search2: zoom: (WC2b) holds'
       return
    end if
    if ( slope*dalpha >= 0 ) then
       write(4,*) lsiter,'wolfe_search2: zoom: alpha_hi <- alpha_lo & update alpha_lo'
       alpha_hi = alpha_lo
       f_hi     = f_lo
       !!! df_hi  = df_lo
       !!! slope_hi= slope_lo
       alpha_lo = alpha
       f_lo     = fval
       df_lo    = dfval
       slope_lo = slope
    else
       write(4,*) lsiter,'wolfe_search2: zoom: update alpha_lo'
       alpha_lo = alpha
       f_lo     = fval
       df_lo    = dfval
       slope_lo = slope
    end if
  end if
  ! Update iteration count
  iter_cnt = iter_cnt + 1
  lsiter=lsiter+1
end do
return
end subroutine
!===========================================
subroutine owfeval(x,fun,R_image,B_image)!,subiter,iter)

double precision                              :: x(:), fun
real                                          :: start, finish, start1, finish1
real,optional                                 :: B_image(:,:,:)
real                                          :: R_image(:,:,:)
real,allocatable                              :: D_image(:,:,:)
integer                                       :: ni,i,ishot,j,k,iter,subiter,titer
real, allocatable                             :: deltaI(:,:,:),illum(:,:),illumination(:,:)
real, allocatable                             :: vel(:)
complex, allocatable                          :: recv(:,:),source(:,:),source1(:)

allocate(deltaI(kxm%n,nsecdim,img%z%n));deltaI=0.
allocate(illum(kxm%n,img%z%n),vel(img%z%n*kxm%n));illum=0.;vel=0
allocate(illumination(img%xm%n,img%z%n));illumination=0.
start1=0.;finish1=0.

if(not(present(B_image))) then
  allocate(D_image(img%xm%n,nsecdim,img%z%n));D_image=0
end if
!-------------------------------------------
! Computing background image
!
       call taper_set()
       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 ishot=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,ishot)
          end if
          if (allocated(recv))   deallocate(recv);   allocate(recv(rec%h%n,rec%w%n))
          recv=0;deltaI=0;illum=0.

          call sreed("rec",recv,8*rec%h%n*rec%w%n)
          call owwei_init()
          call owwei(taper,recv,source,vel,deltaI,illum)

          if (present(B_image)) then
             if (areal) then
                B_image=B_image+deltaI
                illumination=illumination+illum
             else
                B_image(ni:ni+kxm%n-1,:,:)=B_image(ni:ni+kxm%n-1,:,:)+deltaI
                illumination(ni:ni+kxm%n-1,:)=illumination(ni:ni+kxm%n-1,:)+illum
             end if
          else
             if (areal) then
                D_image=D_image+deltaI
                illumination=illumination+illum
             else
                D_image(ni:ni+kxm%n-1,:,:)=D_image(ni:ni+kxm%n-1,:,:)+deltaI
                illumination(ni:ni+kxm%n-1,:)=illumination(ni:ni+kxm%n-1,:)+illum
             end if
          end if
          call owexec_time(finish)

          write(3,111) "EXTRAPOLATED WAVEFIELDS = ",ishot," 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

!-------------------------------------------
!
! Computing DSO
!

       if (present(B_image)) then
          R_image=B_image
!          if (envelope) then
!             call owenvelope(R_image)
!          end if
          if (zwgt.ne.99999.) then
             call owzlinwht(R_image)
          end if
          call owdso(.false. , R_image)
          call owdso(.true.  , R_image)
       else
          R_image=D_image
!          if (envelope) then
!             call owenvelope(R_image)
!          end if
          if (zwgt.ne.99999.) then
             call owzlinwht(R_image)
          end if
          call owdso(.false. , R_image)
          call owdso(.true.  , R_image)
       end if
       if (allocated(pick)) then
          do j=1,img%sh%n
             call oweditvel_sngl(pick,R_image(:,j,:))
          end do
      end if

!       call to_history("n4",titer,"Dimage")
       call srite("Dimage",R_image,4*img%z%n*nsecdim*img%xm%n)

!-------------------------------------------
!  Computing the objective function
!
       fun=0.
       do i=1,img%z%n
          do j=1,img%sh%n
             do k=1,img%xm%n
                fun = fun + R_image(k,j,i)*R_image(k,j,i)
             end do
          end do
       end do
       
       write(2,113) " iter = ",iter," titer= ",titer," J(x) = ", fun

if (allocated(deltaI)) deallocate(deltaI)
if (allocated(illum)) deallocate(illum)
if (allocated(illumination)) deallocate(illumination)
if (allocated(vel)) deallocate(vel)
if (allocated(recv)) deallocate(recv)
if (allocated(source)) deallocate(source)
if (allocated(source1)) deallocate(source1)
if (allocated(D_image)) deallocate(D_image)
 111 format(a26,i4,a8,i4,a4,f6.2,a4)
 112 format(a29,f10.2,a20,i4)
 113 format(2(a8,i4),a8,e25.16)
 call taper_close()
end subroutine
!===========================================
subroutine owgeval(x,g,B_image,z0)

logical                                       :: cond,bounds,cbounds
double precision                              :: g(:),x(:)
real,optional                                 :: z0(:,:)
real                                          :: B_image(:,:,:)
real                                          :: start, finish, start1, finish1, start2, finish2
integer                                       :: ni,i,ishot,iz,j,k,iter,subiter,titer,i1,i2,i3,stat1
real, allocatable                             :: scatS(:),deltaS(:),vel(:)
complex, allocatable                          :: recv(:,:),source(:,:),source1(:)
double precision,allocatable                  :: gspl(:)


       call taper_set()
       if (axkn.ne.99999 .or. azkn.ne.99999) then
          stat1 = base_spln2d_init(img%xm,img%z,axk,azk)
          allocate(gspl(axk%n*azk%n));gspl=dble(0.)
!          stat1 = base_spln3d_init(img%xm,img%ym,img%z,axk,ayk,azk)
       end if

       start1=0.;finish1=0.;g=0.
       allocate(vel(img%z%n*kxm%n));vel=0
       allocate(deltaS(img%z%n*kxm%n));deltaS=0
       allocate(scatS(img%xm%n*img%z%n));scatS=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 ishot=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,ishot)
          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(ishot)
          call owwemva(.true.,taper,recv,source,vel,deltaS,B_image(ni:ni+kxm%n-1,:,:))
          if (areal) then
             scatS=scatS+deltaS
          else
             do j=ni,ni+kxm%n-1;do k=1,img%z%n
                scatS((k-1)*img%xm%n+j)=scatS((k-1)*img%xm%n+j)+deltaS((k-1)*kxm%n+j-ni+1)
             end do;end do
          end if
          call owexec_time(finish)
          write(3,111) "EXTRAPOLATED WAVEFIELDS = ",ishot," OUT OF ",rec%xd%n," IN ",(finish-start)/60," min"
       end do

       g=dble(scatS)

       if (envelope) then
          call owenvelope1(z0)
          do i=1,img%z%n
             g((i-1)*img%xm%n+1:i*img%xm%n) = g((i-1)*img%xm%n+1:i*img%xm%n) / z0(:,i)
          end do
       end if


       write(0,*) "main => ",maxval(g),minval(g)

       if (axkn.ne.99999 .or. azkn.ne.99999) then
          call owexec_time(start2)
          !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
          ! B-spline representation of grad
          gspl=dble(0.)
          stat1 = base_spln2d_oper(.true. , .false., gspl, g)
          stat1 = base_spln2d_oper(.false. ,.false., gspl, g)
          call owexec_time(finish2)
          !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!  
       end if

       write(0,*) "main => ",maxval(g),minval(g)
       write(3,112) "SMOOTHED GRADIENT IN ",(finish2-start2)/60

       g=g/maxval(abs(g))

       call owexec_time(finish1)
       write(3,112) "COMPUTED THE GRADIENT IN ",(finish1-start1)/60," min OF ITERATION # ", iter
       call taper_close()
       if (allocated(deltaS)) deallocate(deltaS)
       if (allocated(scatS)) deallocate(scatS)
       if (allocated(vel)) deallocate(vel)
       if (allocated(recv)) deallocate(recv)
       if (allocated(source)) deallocate(source)
       if (allocated(source1)) deallocate(source1)
       if (allocated(gspl)) deallocate(gspl)
 111 format(a26,i4,a8,i4,a4,f6.2,a4)
 112 format(a29,f10.2,a20,i4)
 113 format(2(a8,i4),a8,a5,a8,e15.6,2(a8,I5))
end subroutine
!===========================================
!======================================================================
subroutine owenvelope1( B_image)

  integer                   :: i1,i2,i3
  real                      :: B_image(:,:)
  real                      :: bot(img%z%n), data(img%xm%n,img%z%n),a,b
  complex                   :: cdata(img%z%n,1)

  cdata=0.
  call owfftwr(cdata,1)
  do i2= 1, img%xm%n  
      cdata(:,1) = cmplx(B_image(i2,:),0.)
      call owfftexer(cdata(:,1),1)
      cdata(:,1)=cdata(:,1)/sqrt(1.*img%z%n)
      cdata(1,1)      = cdata(1,1)      / 2.
      cdata(1+img%z%n/2,1) = cdata(1+img%z%n/2,1) / 2.
      cdata(img%z%n/2+2:img%z%n,1) = 0.
      call owifftexer(cdata(:,1),1)
      cdata(:,1)=cdata(:,1)/sqrt(1.*img%z%n)
      bot = conjg( cdata(:,1)) * cdata(:,1)
      call triangle3( tc2, tc1, 1, img%z%n, bot, bot)
      data(i2,:) = sqrt( bot)
  end do 
  a=0
  do i2= 1, img%xm%n
     if (a < maxval(data(i2,:))) then
        a = maxval(data(i2,:))
     end if  
  end do
  B_image=(data+a)**2
end subroutine 
!======================================================================

end module
