! Find all files 'all_phase_space_at_#####_ns.dat'
! Fit f(x) = a+b*x+c**2 to momentum vs yamp = sqrt(y**2/betay+yp**2*betay)
     
program fit_p_vs_y
  use nr
  use sim_utils

  implicit none
  
  interface
     subroutine funcs(x, arr)
       use nrtype
       use precision_def
       implicit none
       real(rp), intent(in) :: x
       real(rp), dimension(:), intent(out) :: arr
     end subroutine funcs
  end interface
  type phase_space_data_struct 
     integer  muon
     real(rp)  jx ,  jy ,  x ,  xp ,  y ,  yp ,  vec5 ,  pz ,  t 
     real(rp) s , sx , sy , sz , phi 
  end type phase_space_data_struct
  type time_struct
     type(phase_space_data_struct) p(1600000)
     integer i, n
     real(rp) time, offset
     character*300  phase_space_file   
  end type time_struct

  type(time_struct), allocatable :: ts(:)
  
  character*600 string, new_string
  character*20 dir_min, dir_max

  
  real(rp), allocatable :: amp_y(:),y(:),yp(:),pz(:), sig(:)
  real(rp), allocatable :: a(:)
  logical, allocatable :: maska(:)
  real(rp), allocatable :: covar(:,:)
  real(rp) chisq
!  real(rp) jx,jy,x,xp,vec5,t,s,sx,sy,sz,phi
  real(rp) betay/20./
  real(rp) chi,z(1:10)
  real(rp) average_y, average_yp

  
  integer n,i,j, idum
  integer num_files
  integer lun, reason,lun6
  integer nargs
  integer muon, npoints,deg/3/
  integer n1, n2

 nargs = command_argument_count()  

   string = 'ls all_phase_space_at*.dat | wc -l > num_files.dat'  
   call execute_command_line(string)

   lun=lunget()
   open(unit=lun,file='num_files.dat', status='old')

   string='ls all_phase_space_at_*.dat_trimmed > out.dat'
   call execute_command_line(string)

   read(lun,*)num_files
   num_files = num_files + 1
   allocate(ts(num_files))

   open(unit=lun,file='out.dat', status='old')

   n=1

   ts(1)%phase_space_file = 'all_Energy_vs_time_spin.dat'
   ts(1)%phase_space_file = 'all_MARK_INFLECTOR_DS_phase_space.dat_trimmed'
   ts(1)%n = 1
   ts(1)%time = 0
   ts(1)%offset = 0
   
  do while(.true.)
    read(lun, '(a)', IOSTAT = reason)new_string
    if(reason < 0)exit
    n=n+1
    ts(n)%phase_space_file=trim(new_string)
    read(ts(n)%phase_space_file(20:24),*) ts(n)%time
    ts(n)%n = n
    ts(n)%offset = 0
    print '(i10,1x,a,es12.4)',n, ts(n)%phase_space_file, ts(n)%time
  end do
  close(lun)

!_______________________________________________________________________

  n=0
  open(lun, file = ts(1)%phase_space_file)     ! this loop is to count the number of lines in the phase space files. We will assume they are all the same
  do while(.true.)
     read(lun,'(a)', IOSTAT = reason)new_string
     if(reason < 0)exit
     n=n+1
  end do
  close(unit=lun)
 
 npoints = n  !number of lines in phase space files
 deg=3
 allocate( amp_y(1:npoints+10),y(1:npoints+10),yp(1:npoints+10),pz(1:npoints+10), sig(1:npoints+10))
 allocate( a(1:3))
 allocate( maska(1:3))
 allocate (covar(1:3,1:3))



 maska=.true.  
 n1=num_files - 1
 do n=1,num_files
    i=0
    open(lun, file=ts(n)%phase_space_file)
    do while(.true.)
       read(lun,'(a)', IOSTAT = reason)string
       if(reason < 0) exit
       if(index(string,'!')/=0 .or. index(string,'muon')/=0 .or. index(string,'AtStart')/=0)cycle
       i=i+1
       if(i > 1600000)then
          print *, ' More than 1600000 particles, array out of bounds'
          print '(a,i10,a)',' There are ', npoints,' muons in the file'
          stop
       endif
       if(ts(n)%offset == 0)then
          read(string,*)ts(n)%p(i)%muon, ts(n)%p(i)%jx, ts(n)%p(i)%jy, ts(n)%p(i)%x, ts(n)%p(i)%xp, ts(n)%p(i)%y, ts(n)%p(i)%yp, ts(n)%p(i)%vec5,&
               ts(n)%p(i)%pz, ts(n)%p(i)%t, ts(n)%p(i)%s, ts(n)%p(i)%sx,ts(n)%p(i)%sy,ts(n)%p(i)%sz,ts(n)%p(i)%phi
       else
          read(string,*)ts(n)%p(i)%muon, z(1:10),  ts(n)%p(i)%x, ts(n)%p(i)%xp, ts(n)%p(i)%y, ts(n)%p(i)%yp,  ts(n)%p(i)%pz, ts(n)%p(i)%t, &
               ts(n)%p(i)%sx,ts(n)%p(i)%sy,ts(n)%p(i)%sz,ts(n)%p(i)%phi
       endif
       
       if(ts(n)%p(i)%muon == 0)then
          print '(a)',' muon num = 0'
          print '(a)', string
       endif
       
       y(i) = ts(n)%p(i)%y
       yp(i)=ts(n)%p(i)%yp
       pz(i) = ts(n)%p(i)%pz
       sig(i) = 1.
    end do
    average_y = sum(y(:))/i
    average_yp = sum(yp(:))/i
    amp_y(:) = sqrt((y(:)-average_y)**2/betay + (yp(:)-average_yp)**2*betay)
    ts(n)%i = i
    call lfit(amp_y(1:i), pz(1:i),sig(1:i),a,maska, covar, chisq,funcs)
    chi=0
    do j=1,i
       chi = chi + (pz(j)-(a(1)+a(2)*amp_y(j)+a(3)*amp_y(j)**2))**2
    end do
    
    print '(a,a5,3es12.4,a,es12.4)',trim(ts(n)%phase_space_file), '  a = ',a,' chisq = ', chisq
    if(n == 1)then
       lun6 = lunget()
       open(unit=lun6,file='fitted_p_vs_y_at_times.dat')
       write(lun6,'(13a12)')'time','a','b','c','chisq',' covar(1,1)', 'sig_a','covar(2,2)','sig_b','covar(3,3)','sig_c','sqrt(chi)','sqrt(chi/ndf)'
    endif
    
    write(lun6,'(5es12.4,6es12.4,2es12.4)')ts(n)%time,a,chisq, (covar(j,j),(chisq*covar(j,j)/i)**.5, j=1,3), sqrt(chi), sqrt(chi/i)

 end do

 do n=1,num_files
    if(n /= n1)call sort_by_muon_num(ts(n1), ts(n))
 end do
  
end program

subroutine funcs(x, arr)
  use nrtype
  use precision_def
  implicit none
  real(rp), intent(in) :: x
  real(rp), dimension(:), intent(out) :: arr

  arr(1) = 1
  arr(2) = x
  arr(3) = x**2

  return
end subroutine funcs

subroutine sort_by_muon_num(ts1,ts2)
  use sim_utils

  implicit none
  
  type phase_space_data_struct 
     integer  muon
     real(rp)  jx ,  jy ,  x ,  xp ,  y ,  yp ,  vec5 ,  pz ,  t 
     real(rp) s , sx , sy , sz , phi 
  end type phase_space_data_struct
  type time_struct
     type(phase_space_data_struct) p(1600000)
     integer i, n
     real(rp) time, offset
     character*300  phase_space_file   
  end type time_struct

  type(time_struct) ts1,ts2

  integer i,j
  integer, save :: lun1, lun2
  integer matches, mismatch
  logical, allocatable, save :: found(:)
  logical first/.true./
  integer, save :: lun1_last,lun2_last

  if(.not. allocated(found))allocate(found(1:ts1%i))
  found = .false.
  
  lun1 = 600 + ts1%n
  lun1_last = lun1
  if(first)open(unit=lun1, file =trim(ts1%phase_space_file)//'_sorted')
  lun2 = 600 + ts2%n
  if(lun2 == lun2_last)return
  lun2_last = lun2

  open(unit=lun2, file = trim(ts2%phase_space_file)//'_sorted')
  matches = 0
  mismatch=0
  do i=1,ts1%i
     do j=1,ts2%i
        if(ts1%p(i)%muon == ts2%p(j)%muon .and. abs(ts1%p(i)%pz-ts2%p(j)%pz)<5.e-5)then
if(ts2%p(j)%muon ==0)then
  print *,'something is wrong',' ts1%p(i)%muon = ', ts1%p(i)%muon,'  ts2%p(j)%muon =', ts2%p(j)%muon
endif
        if(first)then
           write(lun1,'(i10,14es12.4)')ts1%p(i)%muon, ts1%p(i)%jx, ts1%p(i)%jy, ts1%p(i)%x, ts1%p(i)%xp, ts1%p(i)%y, ts1%p(i)%yp, ts1%p(i)%vec5, ts1%p(i)%pz, ts1%p(i)%t, &
                ts1%p(i)%s,ts1%p(i)%sx,ts1%p(i)%sy,ts1%p(i)%sz,ts1%p(i)%phi
           endif
           write(lun2,'(i10,14es12.4)')ts2%p(j)%muon, ts2%p(j)%jx, ts2%p(j)%jy, ts2%p(j)%x, ts2%p(j)%xp, ts2%p(j)%y, ts2%p(j)%yp, ts2%p(j)%vec5, ts2%p(j)%pz, ts2%p(j)%t, &
                ts2%p(j)%s,ts2%p(j)%sx,ts2%p(j)%sy,ts2%p(j)%sz,ts2%p(j)%phi
           matches = matches + 1
           found(i)=.true.
           exit
        endif
     end do
     if(.not. found(i))then
        mismatch = mismatch + 1
        print '(2a,a,i10)','Problem with missing muon in ',trim(ts1%phase_space_file),'No match for ',ts1%p(i)%muon
     endif
     
  end do
  print '(2a,i10,a,i10)',' number of matches ',trim(ts2%phase_space_file), matches,'   mismatches =', mismatch
  close(unit=lun1)
  first = .false.
  close(unit=lun2)
  
end subroutine sort_by_muon_num