! Routine to find the twiss at start of injection line that correspond to twiss specified at inflector
subroutine optimize_twiss(lat, twiss, twiss_opt)
 use bmad
 use muon_mod
 use nr
 implicit none
 type (lat_struct) lat
! type (coord_struct), allocatable :: co(:)
 type (g2twiss_struct) twiss, twiss_opt
 real(rp) x0(8), y1(8), y0(8), target(8)
 real(rp) ytemp(8)
 integer track_state
 integer n, iteration
 integer i
 logical err_flag, err
 real(rp) dydx(8,8), dx(8), b(8,8), M(8,8), Identity(8,8)
 real(rp) delta/0.001_rp/
 real(rp) tol/1.e-4/

! call reallocate_coord (orb, lat%branch(0)%n_ele_max)

!initial guess for twiss parameters at start of injection line

 x0 = 0.
 x0(1) = 50.
 x0(2) = 10.
 x0(3) = -10.
 x0(4) = 0.

 target(1) = twiss%betax
 target(2) = twiss%betay
 target(3) = twiss%alphax
 target(4) = twiss%alphay
 target(5) = twiss%etax
 target(6) = twiss%etapx
 target(7) = twiss%etay
 target(8) = twiss%etapy

print '(a,1x,8es12.4)','target', target

! n = lat%branch(0)%n_ele_track-3 !Mark inflector ds
 n=-1

 do i=1,lat%branch(0)%n_ele_track
  if (index(lat%branch(0)%ele(i)%name, 'INFA')/=0)n=i
 end do
 if(n>0)then
  print *,' element at end of injection line = ', lat%branch(0)%ele(n)%name
 else
  print *,' Optimize incident: end element not found '
 endif

! call lat_make_mat6(lat, -1, co, ix_branch=0, err_flag = err_flag)

 lat%branch(0)%ele(0)%a%beta = x0(1)
 lat%branch(0)%ele(0)%b%beta = x0(2)
 lat%branch(0)%ele(0)%a%alpha = x0(3)
 lat%branch(0)%ele(0)%b%alpha = x0(4)
 lat%branch(0)%ele(0)%x%eta   = x0(5)
 lat%branch(0)%ele(0)%x%etap  = x0(6) 
 lat%branch(0)%ele(0)%y%eta   = x0(7) 
 lat%branch(0)%ele(0)%y%etap  = x0(8) 

 call twiss_propagate_all(lat, ix_branch=0)
                do i=1,lat%branch(0)%n_ele_track
                    write(74,'(i10,1x,3es12.4)')i,lat%branch(0)%ele(i)%s, lat%branch(0)%ele(i)%a%beta, lat%branch(0)%ele(i)%b%beta
                end do
               

 y0(1) =  lat%branch(0)%ele(n)%a%beta 
 y0(2) =  lat%branch(0)%ele(n)%b%beta
 y0(3) = lat%branch(0)%ele(n)%a%alpha
 y0(4) = lat%branch(0)%ele(n)%b%alpha
 y0(5) = lat%branch(0)%ele(n)%x%eta
 y0(6) =  lat%branch(0)%ele(n)%x%etap
 y0(7) = lat%branch(0)%ele(n)%y%eta
 y0(8) = lat%branch(0)%ele(n)%y%etap

! compute jacobian

iteration = 1
!print *
! print '(i10,1x,8es12.4)',iteration, dx
! print '(a10,1x,8es12.4)','start', x0
! print '(a10,1x,8es12.4)','end', y0
! print '(a10,1x,8es12.4)','diff', (target - y0)

do while(any(abs (y0(1:6) - target(1:6) ) > tol))
 do i=1,8 ! compute derivates

  if(i == 1)then
    lat%branch(0)%ele(0)%a%beta =  x0(i) + delta
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%a%beta =  x0(i)
  endif
  if(i == 2)then
    lat%branch(0)%ele(0)%b%beta =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%b%beta =  x0(i)
  endif
  if(i == 3)then
    lat%branch(0)%ele(0)%a%alpha =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%a%alpha =  x0(i)
  endif
  if(i == 4)then
    lat%branch(0)%ele(0)%b%alpha =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%b%alpha =  x0(i)
  endif
  if(i == 5)then
    lat%branch(0)%ele(0)%x%eta =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%x%eta =  x0(i)
  endif
  if(i == 6)then
    lat%branch(0)%ele(0)%x%etap =  x0(i)+ delta 
    call twiss_propagate_all(lat,ix_branch = 0) 
    lat%branch(0)%ele(0)%x%etap =  x0(i)
  endif  
  if(i == 7)then
    lat%branch(0)%ele(0)%y%eta =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0)
    lat%branch(0)%ele(0)%y%eta =  x0(i)
  endif
  if(i == 8)then
    lat%branch(0)%ele(0)%y%etap =  x0(i) + delta 
    call twiss_propagate_all(lat,ix_branch = 0)
    lat%branch(0)%ele(0)%y%etap =  x0(i)
  endif

  y1(1) =  lat%branch(0)%ele(n)%a%beta 
  y1(2) =  lat%branch(0)%ele(n)%b%beta
  y1(3) = lat%branch(0)%ele(n)%a%alpha
  y1(4) = lat%branch(0)%ele(n)%b%alpha
  y1(5) = lat%branch(0)%ele(n)%x%eta
  y1(6) =  lat%branch(0)%ele(n)%x%etap
  y1(7) = lat%branch(0)%ele(n)%y%eta
  y1(8) = lat%branch(0)%ele(n)%y%etap

  dydx(1:8,i) = (y1(1:8) - y0(1:8))/delta

 end do  !derivatives loop

 b(1:8,1) = y0
 b(1:8,2:8) = 0.
! target = y0 + dydx * dx => dx = M * (target - y0)
 M = dydx
 call gaussj(M(1:6,1:6),b(1:6,1:6))
 dx(1:6) = matmul(M(1:6,1:6),(target(1:6)-y0(1:6)))/50
 dx(7:8) = 0.
 x0 = x0 + dx
 
  iteration = iteration + 1

 lat%branch(0)%ele(0)%a%beta = x0(1)
 lat%branch(0)%ele(0)%b%beta = x0(2)
 lat%branch(0)%ele(0)%a%alpha = x0(3)
 lat%branch(0)%ele(0)%b%alpha = x0(4)
 lat%branch(0)%ele(0)%x%eta   = x0(5)
 lat%branch(0)%ele(0)%x%etap  = x0(6) 
 lat%branch(0)%ele(0)%y%eta   = x0(7) 
 lat%branch(0)%ele(0)%y%etap  = x0(8) 
 call twiss_propagate_all(lat, ix_branch = 0)
 y0(1) =  lat%branch(0)%ele(n)%a%beta 
 y0(2) =  lat%branch(0)%ele(n)%b%beta
 y0(3) = lat%branch(0)%ele(n)%a%alpha
 y0(4) = lat%branch(0)%ele(n)%b%alpha
 y0(5) = lat%branch(0)%ele(n)%x%eta
 y0(6) =  lat%branch(0)%ele(n)%x%etap
 y0(7) = lat%branch(0)%ele(n)%y%eta
 y0(8) = lat%branch(0)%ele(n)%y%etap

!print *
! print '(i10,1x,8es12.4)',iteration, dx
! print '(a10,1x,8es12.4)','start', x0
! print '(a10,1x,8es12.4)','end', y0
! print '(a10,1x,8es12.4)','diff', (target - y0)

  if(iteration > 1000)then
   print *, iteration ,  ' iterations in optimize twiss return'
   return
  endif
 end do

print *,' Twiss at start of line that yield selected twiss at mid point of inflector' 
print '(10x,1x,8a12)','beta x','beta y','alpha x','alpha y','eta x','etap x','eta y','etap_y'
 print '(i10,1x,8es12.4)',iteration, dx
 print '(a10,1x,8es12.4)','start', x0
 print '(a10,1x,8es12.4)','end', y0
 print '(a10,1x,8es12.4)','diff', (target - y0)


return

end subroutine