module geodesic_lm_mod

use bmad
use bmadz_struct
use bmadz_interface
use bunchcross_mod
use zquad_lens_mod
use pretz_mod
use twiss_max_mod
use constraints_mod
use nonlin_mod
use opt_mod
use geodesic_lm !In sim_utils/geodesic_lm

type (indep_var_struct), save  :: indep_var_o, sext_var_o
type (lat_struct), save ::        lat_o
type (constraint_struct), save :: sp_o
type (db_struct), save ::       cesr_o
type (pc_struct), save ::         pc_o
type (twiss_max_struct), save ::  twiss_max_o
type (pretz_struct), save ::      pretz_o
type (global_struct), save ::     global_o
type (zquad_struct), save ::      quad_o
type (nonlin_ele_struct), save :: nonlin_o
type (moment_struct), save ::     moment_o
type (coord_struct), allocatable, save :: co_o(:)

contains


!-----------------------------------------------------------------
!-----------------------------------------------------------------
!-----------------------------------------------------------------

subroutine geo_func(mm, nn, a, y_fit)
!subroutine mrqmin_funcs(x,a,yfit,dyda)

  implicit none
  
  type (indep_var_struct) comp_indep_var
  type (an_indep_var)  thevar
  type (ele_struct)  ele, ele2

  integer :: mm, nn
  real(rp), dimension(nn) :: a
  !real(rp), dimension(:,:) :: dyda
  real(rp), dimension(mm) :: y_fit
  real(rp), pointer :: a_ptr

  integer n
  integer ma
  integer ix_ele, attrib, num


  ma = indep_var_o%n_var
  do n = 1, indep_var_o%n_var !number of parameters (n)
    a_ptr => var_attribute(lat_o, indep_var_o%v(n))
    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)
  enddo

  call calcf(lat_o, sp_o, pc_o, quad_o, co_o, global_o,  &
       twiss_max_o, pretz_o, nonlin_o, indep_var_o)
  call fom(lat_o, cesr_o, sp_o, pc_o, quad_o, co_o, global_o,  &
       twiss_max_o, pretz_o, nonlin_o, moment_o, indep_var_o, sext_var_o)

  y_fit(1:sp_o%n_constraint) = sp_o%c(1:sp_o%n_constraint)%contribution
  !this is an array of constraints' contributions to the FOM. Constraint means
  !ideal point. Best fit is where this y_fit = 0. (residuals).

  !call bmad_calc_dyda(lat_o, sp_o, pc_o, quad_o, co_o, global_o, &
  !     twiss_max_o, pretz_o, cesr_o, nonlin_o, indep_var_o, &
  !     dyda, moment_o, sext_var_o)

  return
!end subroutine mrqmin_funcs
end subroutine geo_func

!-----------------------------------------------------------------
!-----------------------------------------------------------------
!-----------------------------------------------------------------

subroutine jacobian(mm, nn, x, fjac)
    implicit none
    integer :: mm, nn
    real(8) :: x(nn), fjac(mm, nn)

    call bmad_calc_dyda(lat_o, sp_o, pc_o, quad_o, co_o, global_o, &
       twiss_max_o, pretz_o, cesr_o, nonlin_o, indep_var_o, &
       fjac, moment_o, sext_var_o)

end subroutine jacobian


!-------------------------------------------------------------------------
!-------------------------------------------------------------------------
!-------------------------------------------------------------------------

subroutine avv(mm, nn, x, v, acc)
    implicit none
    integer :: mm, nn
    real(8) :: x(nn), v(nn), acc(nn)
    !dummy subroutine for calculating second directional derivative
    !as long as geodesic_lm_param%analytic_acc = .false. in
    !sim_utils/geodesic_lm/geodesic_lm_mod.f90 this will not be called
end subroutine

!-----------------------------------------------------------------
!-----------------------------------------------------------------
!-----------------------------------------------------------------


!Deals with geolevmar overstepping bounds
subroutine callback(mm, nn, x, fvec, fjac, accepted, info)
    implicit none
    integer mm, nn, info, accepted
    real(8) x(nn), fvec (mm)
    real(8) fjac(mm,nn)
    !Trivial subroutine as bounds are reflected in cost
end subroutine

!-------------------------------------------------------------------------
!-------------------------------------------------------------------------
!-------------------------------------------------------------------------


  subroutine geodesic_lm_opt (lat, indep_var, cesr, con, pc, quad,  &
       co, global, twiss_max, pretz, nonlin, moment, sext_var)

    implicit none
    
    type (indep_var_struct)  indep_var, sext_var
    type (indep_var_struct)  comp_indep_var
    type (an_indep_var)  thevar
    type (lat_struct)  lat
    type (constraint_struct)  con
    type (db_struct) cesr
    type (pc_struct) pc
    type (twiss_max_struct) twiss_max
    type (pretz_struct)  pretz
    type (global_struct)  global
    type (zquad_struct)  quad
    type (coord_struct), allocatable :: co(:)
    type (nonlin_ele_struct)  nonlin
    type (moment_struct)  moment

    real(rp) demerit
    real(rp) v0(indep_var_maxx), vdel(indep_var_maxx)
    real(rp) fig_start, fig_end !start and end costs
    real(rp) p(indep_var_maxx) !parameter vector
    real(rp) fig
    real(rp) comp_vec(indep_var_maxx+3)
    
    real(rp), allocatable :: y_fit(:), a(:)
    real(rp), pointer :: a_ptr

    !real(rp), allocatable, save :: x(:), y(:), sig(:), a(:)
    !real(rp), dimension(:,:), allocatable, save :: covar, alpha
    !real(rp), save :: chisq, alambda
    
    integer i, n
    integer :: num_param, num_residuals
    !integer,save ::loops
    integer ix_ele, attrib, num
    !logical, allocatable, save :: maska(:)
    !logical,save :: finished !,nit_needed = .true.
    character*60 new_bmad_file

    !!geolevmar diagnostic quantities
    real(rp), allocatable :: fjac(:,:), dtd(:,:)
    integer :: info, niters, nfev, njev, naev, converged


    !!

!  gen = con%n_cycles

    call reallocate_coord( co_o, lat%n_ele_max )
    call allocate_nonlin_ele( nonlin_o, lat%n_ele_max )
    call allocate_moment( moment_o, lat%n_ele_max )

 
    sp_o = con
    lat_o = lat
    cesr_o = cesr
    pc_o = pc
    twiss_max_o = twiss_max
    pretz_o = pretz
    global_o = global
    quad_o = quad
    nonlin_o = nonlin
    co_o = co
    indep_var_o = indep_var
    sext_var_o = sext_var
    
    do n = 1, indep_var%n_var
      p(n) = var_attribute(lat, indep_var%v(n))
    !  v0(n) = p(n)
    !  vdel(n) = indep_var%v(n)%delta
    enddo


    !call calcf(lat, con, pc, quad, co, global, twiss_max, pretz, nonlin,  &
    !                                                               indep_var)
    !call fom(lat, cesr, con, pc, quad, co, global, twiss_max, pretz,  &
    !                                 nonlin, moment, indep_var, sext_var)

    lat_o = lat

    !fig_start = con%figure_of_merit
    !call showme_topten(con, lat, 0)

    !if (con%plot) call plotdo_bmadz('/xs', lat, co, pc)

    !finished = .false.

    !initilization of variables, arrays
    !if (init_needed) then
      !loops = 0
      !n = con%n_constraint
          !if (any(con%c(1:n)%weight == 0.)) then
      !  do i=1,n
      !    if(con%c(i)%weight == 0) &
      !         print '(a12,1x,i3,1x,a16)',' Constraint ', i, ' has zero weight'
      !  end do
      !  print *,' Marquardt requires all nonzero weights'
      !  stop
      !endif
      !if (allocated(x)) deallocate(x,y,sig,a,covar,alpha,maska)
      !allocate(x(1:n), y(1:n), sig(1:n))
      !allocate(a(1:ma), maska(1:ma))
      !allocate(covar(1:ma,1:ma), alpha(1:ma,1:ma))

      !forall (i = 1:con%n_constraint) x(i) = i 
      !y(1:con%n_constraint) = 0.
      !sig(1:con%n_constraint) = 1./con%c(1:con%n_constraint)%weight
      !a(1:indep_var%n_var) = p(1:indep_var%n_var)
      !alambda = -1
      !forall (i = 1:ma) maska(i) = .true.
      !init_needed = .false.

    !endif

    !if (alambda > 1e10) then
    !  print *, 'MARQUARDT: AT MINIMUM'
    !  finished = .true.
    !  init_needed = .true.
    !  alambda = 0
    !  call mrqmin(x, y, sig, a, maska, covar, alpha, chisq, mrqmin_funcs, alambda)
    !else

    !  loops = loops +1
    !  if (loops == con%n_loops) alambda = 0
    !  call mrqmin(x, y, sig, a, maska, covar, alpha, chisq, mrqmin_funcs, alambda)
    !endif
  
    !print *, 'chisq = ', chisq

    num_param = indep_var%n_var !(n) number of parameters for geolevmar
    num_residuals = con%n_constraint !(m) number of residuals for geolevmar

    !allocate(a(1:num_param))
    allocate(y_fit(1:num_residuals))
    !y_fit line probably not necessary
    y_fit(1:con%n_constraint) = sp_o%c(1:con%n_constraint)%contribution
    allocate(dtd(1:num_param,1:num_param))
    allocate(fjac(1:num_residuals,1:num_param))

    !*********************Notes on geolevmar***********************************
    !fjac = array containing jacobian
    !info = set to nonzero in callback if the programs should need to terminate
    !dtd = levenburg marquardt damping matrix of size (num_param, num_param)
    !niters =  the number of steps taken by geolevmar
    !nfev = number of times geo_func was evaluated
    !njev = number of times jacobian was evaluated
    !naev = number of times AVV was calculated
    !converged = integer classifying reason geolevmar stopped optimization. See
    !  sim_utils/geodesic_lm/leastsq.f90 to decode converged
    !**************************************************************************
    !p is initial parameter starting point

    nfev=1000*con%n_loops
    njev=100*con%n_loops
    naev=100*con%n_loops

    !open(unit=10,file='~/geoinfo', access='APPEND')

    call run_geodesic_lm(geo_func, jacobian, Avv, p, y_fit, fjac, callback, info,&
         dtd, con%n_loops, nfev, njev, naev, converged)

    !close(10)

    !updates model with optimized parameters
    do n = 1, indep_var%n_var
      thevar = indep_var_o%v(n)
      print *, n, ' ', thevar%ele_name, thevar%val_name, ' ', p(n)
      a_ptr => var_attribute(lat, indep_var_o%v(n))
      a_ptr = p(n)
      !fill in best fit parameters 'p'
      if(con%sol_com >=0)then
        ix_ele = comp_indep_var%v(n)%rindex
        attrib = comp_indep_var%v(n)%val_id
        call update_ele(comp_vec(n), lat, ix_ele, attrib, num) 
      endif
    enddo

    !Not related to geolevmar, will leave in for future use:
    if (con%sol_com >= 0 .and. con%sol_com <= 2) then
      comp_vec(1:indep_var%n_var)= a(1:indep_var%n_var)
      call compensation_indep_var(con%sol_com, con%compensating_ele, &
           lat, indep_var, comp_indep_var, comp_vec)
      do n = indep_var%n_var+1,2+con%sol_com
        thevar = comp_indep_var%v(n)
        print *, n, ' ', thevar%ele_name, thevar%val_name, ' ', comp_vec(n)
        a_ptr => var_attribute(lat, comp_indep_var%v(n))
        a_ptr = comp_vec(n)
        ix_ele = comp_indep_var%v(n)%rindex
        attrib = comp_indep_var%v(n)%val_id
        call update_ele (comp_vec(n), lat, ix_ele, attrib, num) 
      enddo
    endif

    !Calculate new function value and costs
    call calcf(lat, con, pc, quad, co, global, twiss_max, pretz, &
         nonlin, indep_var)

    call fom(lat, cesr, con, pc, quad, co, global, twiss_max, pretz,  &
         nonlin, moment, indep_var, sext_var)

    !Prepare output file
    if (con%new_bmad_file<=' ') then
      new_bmad_file = 'new_bmad_input.'
    else
      new_bmad_file = con%new_bmad_file
    endif

    if (con%sol_com >= 0 .and. con%sol_com <= 2) then
      comp_vec(1:indep_var%n_var)= a(1:indep_var%n_var)
      call compensation_indep_var(con%sol_com, con%compensating_ele, &
           lat, indep_var, comp_indep_var, comp_vec)
      call opt_out(lat%input_file_name,new_bmad_file,comp_indep_var, &
           comp_vec(1:comp_indep_var%n_var))
    else
      call opt_out(lat%input_file_name,new_bmad_file,indep_var, a)  
    endif

    fig_end = con%figure_of_merit !final cost

    print '(1x,a12,e12.4)' , 'fig_start = ', fig_start
    print '(1x,a10,e12.4,6x,a11,e12.4,6x,a9,i3)', &
         'fig_end = ', fig_end!, ' alambda = ',alambda,'loops = ',loops

    return
  end subroutine geodesic_lm_opt

end module geodesic_lm_mod