program fit_mcmillan
  use bmad
  use multipole_opt
  use calcf_interface
  implicit none

  type (lat_struct) lat
  type (coord_struct), allocatable :: orbit(:)
  type (coord_struct) co
  type (ele_struct) ele_map, ele_mpole
  type (em_field_struct) field, field_mpole

  integer lun
  integer nargs, iargc, ios
  integer ix
  integer track_state
  integer i
  integer nparams/40/,nconstraints/242/
  integer n
  integer max_loops
 
  character*120 line, lat_file, lat_file_name
  character*120 string
  character*3 cmax_loops

  logical err_flag, local_ref_frame/.false./
  real(rp) y, s_pos, x
  real(rp), allocatable::param(:)
  real(rp) grid_half_width/0.022/
  real(rp), allocatable::yfit(:)
  real(rp) chis

nargs = command_argument_count()
  if (nargs == 1) then
    call get_command_argument(1, lat_file)
  elseif(nargs == 2)then
    call get_command_argument(1,lat_file)
    call get_command_argument(2,cmax_loops)
    read(cmax_loops,*)max_loops
    print *,'max_loops = ', max_loops
  else if(lat_file_name /= '')then
     lat_file = lat_file_name
  else
    lat_file = 'bmad.'
    print '(a,$)','  Lattice file name?  (Default = bmad.)  '
    read(5,'(a)') line
    call string_trim(line, line, ix)
    lat_file = line
    if (ix==0) lat_file = 'bmad.'
  endif
    print *, ' lat_file = ', lat_file
   call bmad_parser(lat_file, lat)
   call reallocate_coord (orbit, lat%n_ele_max)
   ele_mpole = lat%ele(2)
   n=nparams/2
   allocate(param(1:2*n))
   allocate(yfit(1:nconstraints))

   print *,' number of a_pole_elec = ', size(ele_mpole%a_pole_elec)
   param(1:n) = ele_mpole%a_pole_elec(1:n)
   param(n+1:2*n) = ele_mpole%b_pole_elec(1:n)

   call calcf(lat,nconstraints,nparams, chis, yfit,grid_half_width) 
   print '(a,i10,a,es12.4,a,es12.4)',' nconstraints', nconstraints,' chis = ', chis, ' sqrt(chis/nconstraints)=', sqrt(chis/nconstraints)
   if(max_loops > 0)call opt_mpoles(lat,nconstraints,nparams,grid_half_width,param,max_loops)

 end


subroutine calcf(lat,nconstraints,nparams, chis, yfit,grid_half_width)
! compute the electric field over a grid (+-4.5 X +-4.5cm) using the grid map quadrupole element (which is ele(1)) in this lattice)
! then compute the electric field over the same grid using the multipole quadrupole element (ele(2) in this lattic)
! then take the sum of the squares of the difference of Ex and Ey over the entire grid.
! That sum is chis2

  use bmad

  implicit none

  type (lat_struct) lat
  type (coord_struct), allocatable :: orbit(:)
  type (coord_struct), save :: co
  type (ele_struct), save :: ele_map
  type(ele_struct), save :: ele_mpole
  type (em_field_struct), save :: field
  type (em_field_struct) field_mpole
  type(em_field_struct),allocatable,save:: grid(:,:)
  type(em_field_struct),allocatable, save::grid_m(:,:), delta_grid(:,:)

  integer nparams, nconstraints, i, ix, iy
  integer n
  integer, save :: steps, steps_half

  logical err_flag, local_ref_frame/.false./
  logical first/.true./

  real(rp)  s_pos, x,y
  real(rp) chis 
  real(rp), allocatable:: yfit(:)
  real(rp), optional::grid_half_width  ! something less than 0.042 m 
  real(rp), save:: gh_width
  real(rp) delta

  if(present (grid_half_width))gh_width=grid_half_width
   ele_map = lat%ele(1)
   ele_mpole = lat%ele(2)
  if(first)then
   call init_coord(co)
   print *,' element_map = ', ele_map%name
   print *,' element_mpole = ', ele_mpole%name
   steps = sqrt(float(nconstraints/2))
   if(steps**2 /= nconstraints/2)then
     print *,' nconstraints is not a perfect square'
     stop
   endif
   steps_half = (steps - 1)/2
   if(steps_half*2 /= steps-1)then
    print *,' nconstraints is not a perfect square of an odd number '
    stop
   endif
   allocate(grid(1:steps,1:steps))
   allocate(grid_m(1:steps,1:steps))
   allocate(delta_grid(1:steps,1:steps))

   s_pos = ele_map%s/2   

   delta = gh_width/steps_half
 !  print '(3(a,i12),2(a,es12.4))',' nconstaints = ',nconstraints,' steps = ', steps, 'steps_half = ', steps_half, ' delta = ', delta,' grid_half_width = ', grid_half_width


   y=-gh_width- delta
   iy=0
   do while(iy < steps)
    co%vec=0
     y=y+delta
     co%vec(3) = y
     iy=iy+1
     x=-gh_width-delta
     ix=0
     do while(ix < steps) 
       x = x+delta
       co%vec(1) =x
       ix=ix+1
       call em_field_calc(ele_map, lat%param, s_pos, co,local_ref_frame,field,err_flag=err_flag)
!      call em_field_calc(ele_mpole, lat%param, s_pos, co,local_ref_frame,field_mpole,err_flag=err_flag)
!      write(11, '(9es12.4)')co%vec(1),co%vec(3),s_pos, field%E, field_mpole%E
!   print *,' ix = ',ix,' iy = ', iy, ' x = ', x,' y = ',y
      grid(ix,iy) =field 
     end do
   end do
   first = .false.
  endif !first

! find mupltipoles in multipole element
open(unit=11)
  chis = 0
   i=0
   s_pos = ele_map%s/2
   delta = gh_width/steps_half
!print *,' delta = ', delta,' gh_width = ', gh_width,' steps_half = ', steps_half
   y=-gh_width- delta
   iy=0
   do while(iy < steps)
    co%vec=0
     y=y+delta
     co%vec(3) = y
     iy=iy+1
     x=-gh_width-delta
     ix=0
     do while(ix< steps) 
       x = x+delta
       co%vec(1) =x
       ix=ix+1

!   call em_field_calc(ele_map, lat%param, s_pos, co,local_ref_frame,field,err_flag=err_flag)
    call em_field_calc(ele_mpole, lat%param, s_pos, co,local_ref_frame,field_mpole,err_flag=err_flag)
    grid_m(ix,iy)=field_mpole

   write(11, '(9es12.4)')co%vec(1),co%vec(3),s_pos, grid(ix,iy)%E, grid_m(ix,iy)%E

    delta_grid(ix,iy)%E = grid(ix,iy)%E-grid_m(ix,iy)%E
    i=i+1
    yfit(i) = delta_grid(ix,iy)%E(1)
   i=i+1
    yfit(i) = delta_grid(ix,iy)%E(2)
    chis = chis + yfit(i)**2
 if(i > nconstraints)then
  print *,' i greater than nconstraints in calcf: i-',i
  stop
endif

   end do
  end do
   close(unit=11)
return
end subroutine

subroutine calcf_dyda(lat,nconstraints,nparams,chis, yfit, dyda)
use bmad
use calcf_interface,except_dummy =>calcf_dyda
implicit none
type(lat_struct) lat
  real(rp), allocatable :: yfit(:)
  real(rp), dimension(:,:) ::  dyda
  real(rp) fp, f0, xx
  real(rp) dx
  real(rp) chis
  real(rp), dimension(:), allocatable, save :: yfit0(:)
  integer nconstraints, nparams
  integer n, way, IGRAD
  integer ix_mpole/2/
  data way/1/, IGRAD/8/

  allocate(yfit0(1:nconstraints))
  call calcf(lat, nconstraints,nparams,chis,yfit)
  yfit0=yfit
  f0=chis

  way = -way
  dx = 0.5**IGRAD*way
  do n = 1, nparams
   if(n<= nparams/2)then
     xx = lat%ele(ix_mpole)%a_pole_elec(n)
     lat%ele(ix_mpole)%a_pole_elec(n) = xx+dx
   endif
   if(n > nparams/2)then
     xx = lat%ele(ix_mpole)%b_pole_elec(n-nparams/2)
     lat%ele(ix_mpole)%b_pole_elec(n-nparams/2) = xx+dx
   endif

     call calcf(lat, nconstraints,nparams,chis,yfit)
     fp = chis
     if(n<=nparams/2)lat%ele(ix_mpole)%a_pole_elec(n) = xx
     if(n > nparams/2)lat%ele(ix_mpole)%b_pole_elec(n-nparams/2) = xx

     dyda(1:nconstraints,n) = (yfit(1:nconstraints) - yfit0(1:nconstraints))/dx 
     
  end do
  deallocate(yfit0)

  return
end subroutine calcf_dyda