module findmatch_mod
  
  use helix
  use helicon
  use polydiv
  use nchconest
  use nchelicon
  
  use nhelix
  use ncnhconest
  use ncnhelicon
  
  use cgstep_mod
  use solver2_mod
  use verbosity_mod
  
  use mkwallwt
  use patch
  use tent
  use random
    use dotest

! Weight = L1
!  use irls
!  use irls_priv

  implicit none
  
  integer, private                     :: nf,ndata
  type (filter), private               :: acorr
  real, private, dimension(:), pointer :: xcorr,space,matrix

contains

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! Find a single matched-filter     !!
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  ! Minimizes residual:
  ! wght ( data * match - targt ) = 0
  subroutine findmatch(data,targt,match,niter,dtest,init,noise,wght, &
        & nmem_in,nfreq_in)
    optional                              ::wght,nmem_in,nfreq_in
    real, dimension(:), pointer           :: data,targt,wght
    integer                               :: niter,nmem_in,nfreq_in
    logical,                     optional :: init
    logical                                    :: dtest
    real,                        optional :: noise
    type (filter) :: match
    integer       :: stat,i,nmem,nfreq
    logical       :: cginv,initchol
    real          :: sd,rnd

    cginv=(niter.ge.0)
 
    if (present(nfreq_in)) then 
       nmem=nmem_in
       nfreq=nfreq_in
    end if

    if (present(noise)) then
       sd=noise*sqrt(dot_product(data,data))
       do i=1,size(data)
          rnd=sd*random_normal()
          data(i)=data(i)+rnd
       end do
    end if

    if (present(nmem_in)) then
       call cgfindmatch(data,targt,match,niter,dtest,nmem,nfreq)
    else
        if (present(wght)) then 
               call cgfindmatch(data,targt,match,niter,dtest,weight=wght)
        else
write(0,*)'before cgfindmatch'
               call cgfindmatch(data,targt,match,niter,dtest)
write(0,*)'after cgfindmatch'
        end if
    end if
  end subroutine findmatch

  !! Find match with conjugate gradients (slow)
  subroutine cgfindmatch(data,targt,match,niter,dtest1,nmem,nfreq,weight)
    optional                                 :: nmem,nfreq,weight
    real, dimension(:), pointer           :: data,targt,weight
    logical                                    :: dtest1
    integer                               :: nmem,nfreq
    type (filter) :: match
    integer       :: stat,niter

    if (verb>1) write(0,*) ' Using CG to solve system- there may be better way'

    call nchconest_init(data,match)

if (dtest1) then
    write(0,*)'*****************************************'
    write(0,*)'Dot product test for nchconest_lop'
    call dot_est(nchconest_lop, size(match%flt), size(targt))
    write(0,*)'*****************************************'
end if

    if (present(nfreq)) then
       write(0,*) 'Using L1 norm with nonlinear CG'
       write(0,*) 'nmem',nmem
       write(0,*) 'nfreq',nfreq
!       call solver(nchconest_lop,cgstep,match%flt,targt,niter,wght=l1_priv, &
       call solver(nchconest_lop,cgstep,match%flt,targt,niter, &
            &     nmem=nmem,nfreq=nfreq)
    else

        if (present(weight)) then
write(0,*)''
write(0,*)''
write(0,*)'ANTES weight',weight
write(0,*)''
write(0,*)''
               call solver(nchconest_lop,cgstep,match%flt,targt,niter,wt=weight,verb=.true.)
write(0,*)''
write(0,*)''
write(0,*)'DESPUES weight',weight
write(0,*)''
write(0,*)''
        else
               call solver(nchconest_lop,cgstep,match%flt,targt,niter,verb=.true.)
        end if
    end if

    call cgstep_close()

  end subroutine cgfindmatch

  !! Calculates:
  !! output = input * match
  subroutine applymatch(input,output,match,dtest)
    
    real, dimension(:), pointer :: input, output
    type (filter)               :: match
    integer stat
    logical            :: dtest    !! dot product test
    
    call nchelicon_init(match)

if (dtest) then
    write(0,*)'*****************************************'
    write(0,*)'Dot product test for nchelicon_lop'
    call dot_est(nchelicon_lop, size(input), size(output))
    write(0,*)'*****************************************'
end if

    stat=nchelicon_lop(.false.,.false.,input,output)
 
  end subroutine applymatch

  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
  !! Non-stationary matched filtering !!
  !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

  !! Finds non-stationary matched-filter
  subroutine findnsmatch(data,targt,match,reg,mode,niter,eps,dtest,&
        & wght,nfreq_in,nmem_in)
    real                        :: eps
    real, dimension(:), pointer           :: data,targt
    type (nfilter)                        :: match
    type ( filter)                        :: reg
    integer                               :: mode
    integer                               :: niter
    integer, optional                     :: nfreq_in,nmem_in
    logical                                :: dtest
    real, dimension(:), pointer, optional :: wght

    real, dimension(:), pointer :: coeffs
    integer :: stat,nhmax,i,nfreq,nmem
    logical :: sverb

    sverb=(verb>2)
    nhmax=size(match%hlx(1)%lag)  ! assumes all filters are same length

    write(0,*) 'epsilon = ',eps

    allocate (coeffs(nhmax*size(match%hlx)))
    coeffs=0.

    call ncnhconest_init(data,match,nhmax)       ! initialize hconest

if (dtest) then
    write(0,*)'*****************************************'
    write(0,*)'Dot product test for ncnhconest_lop'
    call dot_est(ncnhconest_lop, size(coeffs), size(targt))
    write(0,*)'*****************************************'
end if

    select case(mode)

    case (1)  ! Regularized solver

       if (verb>2) write(0,*) '  using regularized conjugate gradients'
       call helicon_init(reg)                       ! initialize regularizer
       !!call ncnhconest_init(data,match,nhmax)       ! initialize hconest

       nfreq=-1
       nmem=5
       if (present(nfreq_in)) nfreq=nfreq_in
       if (present(nmem_in)) nmem=nmem_in

       ! solve problem
       if (nfreq>0) then
          write(0,*) 'nfreq',nfreq,'nmem',nmem
          call solver_reg(ncnhconest_lop,cgstep,helicon_lop,size(coeffs), &
          &   coeffs,targt,niter,eps,x0=coeffs,nmem=nmem, &
          &   nfreq=nfreq,verb=sverb)
!          &   nfreq=nfreq,wght=l1_priv,verb=sverb)
       else
            if (present(wght)) then 
                  call solver_reg(ncnhconest_lop,cgstep,helicon_lop,size(coeffs), &
                  &   coeffs,targt,niter,eps,x0=coeffs,verb=sverb,wt=wght)
            else
                  call solver_reg(ncnhconest_lop,cgstep,helicon_lop,size(coeffs), &
                  &   coeffs,targt,niter,eps,x0=coeffs,verb=sverb)
            end if
       end if

    case (2)  ! Preconditioned solver

       if (verb>2) write(0,*) '  using preconditioned conjugate gradients'
       call polydiv_init(size(coeffs),reg)          ! initialize preconditioner
       !!call ncnhconest_init(data,match,nhmax)       ! initialize hconest

       nfreq=-1
       nmem=5
       if (present(nfreq_in)) nfreq=nfreq_in
       if (present(nmem_in)) nmem=nmem_in

       ! solve problem
       if (nfreq>0) then
          write(0,*) 'nfreq',nfreq,'nmem',nmem
          call solver_prec(ncnhconest_lop,cgstep,polydiv_lop,size(coeffs), &
          &   coeffs,targt,niter,eps,x0=coeffs,nmem=nmem, &
          &   nfreq=nfreq,verb=sverb)
!          &   nfreq=nfreq,wght=l1_priv,verb=sverb)
       else
            if (present(wght)) then
                  call solver_prec(ncnhconest_lop,cgstep,polydiv_lop,size(coeffs),&
                    & coeffs,targt,niter,eps,x0=coeffs,verb=sverb,wt=wght)
            else
                  call solver_prec(ncnhconest_lop,cgstep,polydiv_lop,size(coeffs),&
                    & coeffs,targt,niter,eps,x0=coeffs,verb=sverb)
            end if
       end if

    case (3)  ! Simple solver

       if (verb>2) write(0,*) '  CG with no regularization - mileage may vary'
       !!call ncnhconest_init(data,match,nhmax)       ! initialize hconest

        if (present(wght)) then
               call solver(ncnhconest_lop,cgstep,coeffs,targt,niter,&
                & x0=coeffs,verb=sverb,wt=wght)
        else    
               call solver(ncnhconest_lop,cgstep,coeffs,targt,niter, &
                & x0=coeffs,verb=sverb)
        end if

    case default
       call seperr('Mode not implemented in findnsmatch')
    end select

    if (verb>2) write(0,*) '  Done with solver'

    !! Copy coefficients into NS filter type
    do i=1,size(match%hlx)
       match%hlx(i)%flt(1:nhmax)=coeffs(1+(i-1)*nhmax:i*nhmax)
    end do

    deallocate(coeffs)
    call cgstep_close()

  end subroutine findnsmatch
  
  !! Applies non-stationary matched-filter
  subroutine applynsmatch(input,output,match,dtest)
    real, dimension(:), pointer :: input, output
    type (nfilter)              :: match
    integer                     :: stat
    logical                     :: dtest

    call ncnhelicon_init(match)

if (dtest) then
    write(0,*)'*****************************************'
    write(0,*)'Dot product test for ncnhelicon_lop'
    call dot_est(ncnhelicon_lop, size(input), size(output))
    write(0,*)'*****************************************'
end if

    stat=ncnhelicon_lop(.false.,.false.,input,output)
    call ncnhelicon_close()

  end subroutine applynsmatch

  !! Flags edges of cube as missing data
  !! Eliminates helical wrap around
  subroutine flagedges(missing,nwind)
    logical, dimension(:), pointer :: missing
    integer, dimension(:), pointer :: nwind
    integer :: i,j,k
    k=1
    missing=.false.
    do i=1,size(nwind)-1
       k=k*nwind(i)
       do j=1,nwind(i+1)
          missing(j*k)=.true.
       end do
    end do
  end subroutine flagedges
  
end module findmatch_mod
