module multipole_opt
use bmad
USE nrtype;USE nrutil, ONLY : nrerror, assert_eq, diagmult
use nr
use calcf_interface

type (lat_struct), save::lat_o
integer ix_mpole/2/

contains 

subroutine opt_mpoles(lat,nconstraints, nparams, grid_half_width,param,max_loops)

implicit none 

type (lat_struct) lat

 real(rp), allocatable, save :: x(:), y(:), sig(:), a(:)
 real(rp), dimension(:,:), allocatable, save :: covar, alpha
 real(rp), save :: chisq, alambda
 real(rp)  grid_half_width, chis
 real(rp), allocatable:: param(:), yf(:)
 logical, allocatable, save :: maska(:)
 integer nconstraints, nparams
 integer i
 integer n
 integer loops, max_loops


allocate(yf(1:nconstraints))

! n is the number of constraints which is the number of points in the grid X 2 or 2X (Number of grid points)
! ma is the number of parameters which is twice the number of multipoles (a and b for each)
! choose n to be a perfect square of an odd number (square grid) even about zero
! choose ma to be even 


 allocate(x(1:nconstraints), y(1:nconstraints), sig(1:nconstraints))
 allocate(a(1:nparams), maska(1:nparams))
 allocate(covar(1:nparams,1:nparams), alpha(1:nparams,1:nparams))


call write_mpoles(lat)

 forall(i=1:nconstraints) x(i)=i
 n=nparams
 y(:)=0.
 sig(1:)=1.
 a(:) = param(:)
 alambda=-1
 maska(:)=.true.
loops = 0
do while (loops < max_loops)
lat_o = lat
call calcf(lat, nconstraints, nparams,  chis, yf, grid_half_width=grid_half_width)
print *,' chis = ', chis
   loops = loops + 1
    if (alambda > 1e10) then
      print *, 'MARQUARDT: AT MINIMUM'
      alambda = 0
      call mrqmin(x, y, sig, a, maska, covar, alpha, chisq, mrqmin_funcs, alambda)
      exit
    else

      if (loops == max_loops) alambda = 0
      call mrqmin(x, y, sig, a, maska, covar, alpha, chisq, mrqmin_funcs, alambda)
    endif
  
  lat%ele(ix_mpole)%a_pole_elec(1:n/2) = a(1:n/2)
  lat%ele(ix_mpole)%b_pole_elec(1:n/2) = a(n/2+1:n)
  param = a
    print '(a,es12.4,a,i10,a,es20.12,a,10es16.8)', 'alambda =',alambda,'  loops = ', loops, '  chisq = ', chisq,'  b_pole_ele =', lat%ele(ix_mpole)%b_pole_elec(1:10)
end do
call calcf(lat, nconstraints, nparams,  chis, yf, grid_half_width=grid_half_width)
print '(a,i10,a,i10,a,es12.4,a,es12.4)',' nconstraints=', nconstraints, 'nparams = ', nparams, ' chis = ', chis, ' sqrt(chis/(nconstraints-nparams)=', sqrt(chis/(nconstraints-nparams))

call write_mpoles(lat)

deallocate(x,y,sig,a,covar,alpha,maska)

return
end

subroutine mrqmin_funcs(x,a,yfit,dyda)

  implicit none

  real(rp), dimension(:), intent(in) :: x,a
  real(rp), dimension(:,:), intent(out) :: dyda
  real(rp), dimension(:), intent(out) :: yfit

!  real(rp), dimension(:), intent(in) :: x,a
! real(rp), allocatable, intent(out) :: dyda(:,:)
  real(rp), allocatable :: yy(:)
  real(rp) chis
!  real(rp), pointer :: a_ptr

  integer n
  integer ma
  integer ix_ele, attrib, num
  integer nconstraints


!

!    a_ptr = a(n)
!    ix_ele = indep_var_o%v(n)%rindex
!    attrib = indep_var_o%v(n)%val_id
!    call update_ele(a(n), lat_o, ix_ele, attrib, num)

 n=size(a)
 nconstraints=size(x)
 allocate(yy(1:nconstraints))
 
  lat_o%ele(ix_mpole)%a_pole_elec(1:n/2) = a(1:n/2)
  lat_o%ele(ix_mpole)%b_pole_elec(1:n/2) = a(n/2+1:n)

!initialize fom

    call calcf(lat_o, nconstraints,n, chis, yy) 
!print '(a,es12.4,a,20es12.4)','chis = ',chis,' a(1:n) = ',a(1:n)
  call calcf_dyda(lat_o, nconstraints,n, chis, yy, dyda) 
  yfit = yy
  deallocate(yy)
  return
end subroutine mrqmin_funcs

subroutine write_mpoles(lat)
implicit none
 type (lat_struct) lat
 integer  i
write(12,'(a)')' a_pole_elec'
 do i=0,size(lat%ele(ix_mpole)%a_pole_elec)-1
 write(12,'(i10,es16.8)') i,lat%ele(ix_mpole)%a_pole_elec(i)
end do
write(12,'(a)')' b_pole_elec'
 do i=0,size(lat%ele(ix_mpole)%b_pole_elec)-1
 write(12,'(i10,es16.8)') i,lat%ele(ix_mpole)%b_pole_elec(i)
end do
return
end subroutine
end module multipole_opt