!+
! Subroutine : CUSTOM_SET_TUNE (phi_a_set, phi_b_set, dk1, lat, orb, ok, match_in)
!
! Description: Subroutine to Q_tune a lat. Program will set the tunes to
!              within 0.001 radian (0.06 deg).
!
! Arguments  :  
!         Input:
!            phi_a_set       -- Real: Horizontal set tune (radians)
!            phi_b_set       -- Real: Vertical set tune (radians)
!            dk1(n_ele_max) -- Real: Relative amount to vary a quad in tuning.
!                       dk1(i) relates to lat%ele(i). Those quads with a
!                       positive dk1(i) will be varied as one group and the
!                       quads with negative dk1(i) will be varied as another
!                       group.
!            match_in        -- Logical, optional, if true insert match element to qtune
!
!        Output:
!            lat      -- lat_struct: Q_tuned lat
!            orb(0:) -- Coord_struct: New closed orbit.
!            ok        -- Logical: Set True if everything is ok. 
!                         False otherwise.        
!-

subroutine custom_set_tune (phi_a_set, phi_b_set, dk1, lat, orb, ok, match_in)

  use bmad

  implicit none

  type (lat_struct) lat
  type (lat_struct), save :: lat_save
  type (ele_struct) ave, ele_adjacent, match_ele
  type (coord_struct), allocatable :: orb(:)

  real(rp) phi_a_set, phi_b_set, dphi_a, dphi_b
  real(rp) phi_a, phi_b, d_xx, d_xy, d_yx, d_yy, det
  real(rp) l_beta_a, l_beta_b, dk_x, dk_y
  real(rp) eps_rel(4), eps_abs(4)
  real(rp), allocatable :: dk1(:)

  integer i, j, dk11
  integer ix, status
  integer i_dim
  integer ix_ele

  logical ok, err
  logical, optional :: match_in
  logical first_time/.true./
  logical match

  character(20) :: r_name='custom_set_tune'

  lat_save = lat

  eps_rel(:) = 0.000001
  eps_abs(:) = 0.000001

  dk11 = nint(dk1(1))

! q_tune

  match = logic_option(.false., match_in)
  if(dk1(1)>0 .or. match)print *,' Use match element to q_tune '  !dk1(1)>0 means that there is already a qtuneing match element

    if(match .and. first_time) then
       dk11 = lat%n_ele_track + 1
       call init_ele(match_ele)
       match_ele%name = 'QTUNE_E'
       match_ele%key = match$
!       match_ele%value(match_end_orbit$) = .true.
       match_ele%value(x1_limit$) = 0
       match_ele%value(y1_limit$) = 0
       match_ele%aperture_type = rectangular$
       match_ele%s = lat%ele(dk11-1)%s
       ix_ele = dk11 -1
       call insert_element(lat, match_ele, ix_ele)
       call element_locator('QTUNE_E',lat,ix_ele)
      dk11=ix_ele
      lat%ele(ix_ele)%value(beta_a0$) = lat%ele(ix_ele-1)%a%beta 
      lat%ele(ix_ele)%value(beta_b0$) = lat%ele(ix_ele-1)%b%beta 
      lat%ele(ix_ele)%value(beta_a1$) = lat%ele(ix_ele-1)%a%beta 
      lat%ele(ix_ele)%value(beta_b1$) = lat%ele(ix_ele-1)%b%beta 
      lat%ele(ix_ele)%value(alpha_a0$) = lat%ele(ix_ele-1)%a%alpha 
      lat%ele(ix_ele)%value(alpha_b0$) = lat%ele(ix_ele-1)%b%alpha 
      lat%ele(ix_ele)%value(alpha_a1$) = lat%ele(ix_ele-1)%a%alpha 
      lat%ele(ix_ele)%value(alpha_b1$) = lat%ele(ix_ele-1)%b%alpha 
      lat%ele(ix_ele)%value(eta_x0$) = lat%ele(ix_ele-1)%x%eta 
      lat%ele(ix_ele)%value(eta_y0$) = lat%ele(ix_ele-1)%y%eta 
      lat%ele(ix_ele)%value(etap_x0$) = lat%ele(ix_ele-1)%x%etap 
      lat%ele(ix_ele)%value(etap_y0$) = lat%ele(ix_ele-1)%y%etap 
      lat%ele(ix_ele)%value(eta_x1$) = lat%ele(ix_ele-1)%x%eta 
      lat%ele(ix_ele)%value(eta_y1$) = lat%ele(ix_ele-1)%y%eta 
      lat%ele(ix_ele)%value(etap_x1$) = lat%ele(ix_ele-1)%x%etap 
      lat%ele(ix_ele)%value(etap_y1$) = lat%ele(ix_ele-1)%y%etap
      call lattice_bookkeeper(lat) 
       call lat_make_mat6(lat,-1)
       dk1(1) =2
       call element_locator(match_ele%name,lat,ix_ele)       
!       call type_ele(lat%ele(ix_ele))  
       first_time = .false.
!       return
    endif

  do i = 1, 20

!    call closed_orbit_from_tracking (lat, orb, 4, eps_rel, eps_abs, orb(0))
    call element_locator('PATCH_RF_W1',lat,ix)
    i_dim = 4
    if(ix > 0)then
      i_dim = 6
      print *,' RF frequency shift. i_dim = ', i_dim
      print '(1x,a13,1x,a12,e12.4)', ' PATCH_RF_W1 ' ,' z_offset = ', lat%ele(ix)%value(z_offset$)
    endif
    call closed_orbit_calc( lat, orb, i_dim, err_flag = err)
    print *,' custom_set_tune: return from first closed_orbit_calc '
    ok = .not. err
    if (.not. ok) then
       lat = lat_save
       return
    endif
    call track_all (lat, orb)
    call lat_make_mat6 (lat, -1, orb)


    call twiss_at_start(lat, status = status)
    ok = (status == ok$)
    if(.not. lat%param%stable .or. .not. ok)then
      lat = lat_save
      return
    endif   

    call twiss_propagate_all(lat)

    phi_a = lat%ele(lat%n_ele_track)%a%phi
    phi_b = lat%ele(lat%n_ele_track)%b%phi
    dphi_a = phi_a_set - phi_a 
    dphi_b = phi_b_set - phi_b 
!    if (abs(dphi_a) < 0.0014 .and. abs(dphi_b) < 0.0014) return
    if (abs(dphi_a) < 0.0002 .and. abs(dphi_b) < 0.0002) return
    if(abs(dphi_a) > twopi .or. abs(dphi_b) > twopi)then
       ok=.false.
       lat = lat_save
       return
    endif

    call out_io(s_info$,r_name,'CURRENT TUNE: \2f12.4\ ', r_array = (/ phi_a/twopi, phi_b/twopi /) )
    print *,' phi_a_set, dphi_a', phi_a_set, dphi_a
    if(i == 20) exit
    d_xx = 0
    d_xy = 0
    d_yx = 0
    d_yy = 0

    if(abs(dk1(1)) > 1)then  !use match element to qtune
      ele_adjacent = lat%ele(dk11-1)
      lat%ele(dk11)%value(beta_a0$) = ele_adjacent%a%beta 
      lat%ele(dk11)%value(beta_b0$) = ele_adjacent%b%beta 
      lat%ele(dk11)%value(beta_a1$) = ele_adjacent%a%beta 
      lat%ele(dk11)%value(beta_b1$) = ele_adjacent%b%beta 
      lat%ele(dk11)%value(alpha_a0$) = ele_adjacent%a%alpha 
      lat%ele(dk11)%value(alpha_b0$) = ele_adjacent%b%alpha 
      lat%ele(dk11)%value(alpha_a1$) = ele_adjacent%a%alpha 
      lat%ele(dk11)%value(alpha_b1$) = ele_adjacent%b%alpha 
      lat%ele(dk11)%value(dphi_a$) =  lat%ele(dk11)%value(dphi_a$) + dphi_a
      lat%ele(dk11)%value(dphi_b$) =  lat%ele(dk11)%value(dphi_b$) + dphi_b

      lat%ele(dk11)%value(eta_x0$) = ele_adjacent%x%eta 
      lat%ele(dk11)%value(eta_y0$) = ele_adjacent%y%eta 
      lat%ele(dk11)%value(etap_x0$) = ele_adjacent%x%etap 
      lat%ele(dk11)%value(etap_y0$) = ele_adjacent%y%etap 
      lat%ele(dk11)%value(eta_x1$) = ele_adjacent%x%eta 
      lat%ele(dk11)%value(eta_y1$) = ele_adjacent%y%eta 
      lat%ele(dk11)%value(etap_x1$) = ele_adjacent%x%etap 
      lat%ele(dk11)%value(etap_y1$) = ele_adjacent%y%etap
      call set_flags_for_changed_attribute (lat%ele(dk11)) 
     else 

    do j = 1, lat%n_ele_max
      if (dk1(j) == 0) cycle
      call twiss_at_element (lat%ele(j), average = ave)
      l_beta_a =  abs(dk1(j)) * ave%a%beta * ave%value(l$) / 2
      l_beta_b = -abs(dk1(j)) * ave%b%beta * ave%value(l$) / 2

      if (dk1(j) > 0) then
        d_xx = d_xx + l_beta_a
        d_yx = d_yx + l_beta_b
      else
        d_xy = d_xy + l_beta_a
        d_yy = d_yy + l_beta_b
      endif
    enddo

    det = d_xx * d_yy - d_xy * d_yx
    dk_x = (d_yy * dphi_a - d_xy * dphi_b) / det
    dk_y = (d_xx * dphi_b - d_yx * dphi_a) / det

! put in the changes

    do j = 1, lat%n_ele_max
      if (dk1(j) == 0) cycle
      if (dk1(j) > 0) then
        lat%ele(j)%value(k1$) = lat%ele(j)%value(k1$) + abs(dk1(j)) * dk_x
      else
        lat%ele(j)%value(k1$) = lat%ele(j)%value(k1$) + abs(dk1(j)) * dk_y
      endif
      call set_flags_for_changed_attribute(lat%ele(j), lat%ele(j)%value(k1$))
    enddo
   endif

   call lattice_bookkeeper(lat)

  enddo


  call out_io(s_error$,r_name,'CANNOT GET TUNE RIGHT.', 'CURRENT TUNE: \2f12.4\ ', &
       'SET TUNE:     \2f\ ', r_array = (/ phi_a/twopi,  phi_b/twopi, phi_a_set/twopi, phi_b_set/twopi /) )


   ok = .false.
   lat = lat_save

!  call err_exit

end subroutine