! get a list of directories and combine content of Time_Dep_Moments_at_END.dat into single file
! writes more information than compile_Time_dep.f90
! Command line reads 1- <location> like QUADM1, 2- <desired bin width> like 149.2e-9, 3- subdirectories to include, default is all
! If <location> = 'all' do all QUADM
program compile_all_time_dep

use precision_def
use nr
use bmad
use sim_utils
use compile_interface
use compile_mod
!use muon_mod
implicit none

TYPE (g2moment_struct), allocatable :: moments(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width
TYPE (g2moment_struct), allocatable :: moments_sum(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width 


type (running_moment_struct), allocatable :: running_moment(:)
type (running_moment_struct), allocatable :: running_moments_sum(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width 
 
integer, parameter:: n_time_bins=700000

 integer tbin, tbin_max/0/, lun
 integer i, n, lines/0/
 integer isign/1/
 integer m
 integer k, ix
 integer j
 integer nbins_max/590000/
 integer reason
 integer lun1
 integer nargs
 integer n_files/0/
 integer lun_0
 integer m_rev
 integer peak_bin
 integer nbins_to_combine
 integer jloc
 
 real(rp) two/2./, zero/0./
 complex(rp), allocatable ::data_fft_n(:), data_fft_x(:), data_fft_y(:), data_fft_xx(:), data_fft_yy(:), data_fft_xxx(:), data_fft_yyy(:)
 real(rp) time_bin/10.e-9/, desired_width/149.2e-9/
 real(rp), allocatable :: time(:)
 real(rp) freq
 real(rp) total_hits
 real(rp) time_off
 real(rp) qy,qx, betax, eta, beta_0
 real(rp) y_last, yp_last, betay,Cp, x_last, xp_last
 real(rp) x_average,pz,Ce
 real(rp) s
 real(rp) f_rev, Q_vert, qx_fft, qy_fft
 
 character*290 string, new_string
 character*100 dir, dir_file, dir_list/'all'/
 character*24000 file_string/' '/
 character*20 time_off_char
 character*20 location
 character*20 locations(12)/'QUADM1','QUADM2','QUADM3','QUADM4','QUADM5','QUADM6','QUADM7','QUADM8','QUADM9','QUADM10','QUADM11','QUADM12'/
 character*20 location_0
 character*200 cwd
 character*20 desired_width_word
 character*1 AS
 logical itexists
 logical end_flag/.false./, efield_end_flag/.false./
 
 allocate(moments(1:nbins_max), running_moment(1:nbins_max))
 allocate(running_moments_sum(1:nbins_max), moments_sum(1:nbins_max))
 allocate(time(1:nbins_max))

 location = 'END'
 nargs = cesr_iargc()
 if (nargs >= 1)then
    call cesr_getarg(1, location)
    if(nargs >= 2)then
      call cesr_getarg(2, desired_width_word)
      read(desired_width_word,*)desired_width
   endif
   if(nargs >= 3)call cesr_getarg(3, dir_list)
  endif
  print *, 'Using ', ' location = ', trim(location)
  print *, 'Using desired_width = ', desired_width
  print *,' directory list ',dir_list
  location_0 = location
  jloc=0
  do while (jloc <= size(locations) -1 .and. trim(location_0) == 'all') 
     jloc = jloc+1
     location = locations(jloc)
 
running_moments_sum(1:nbins_max)%n = 0
forall(i=1:6)moments_sum(1:nbins_max)%ave(i)=0
forall(i=1:6)running_moments_sum(1:nbins_max)%sum(i)=0
do j=1,6
   forall(i=1:6)running_moments_sum(1:nbins_max)%product(i,j)=0
end do
call getcwd(cwd)
string='ls > out.dat'
call execute_command_line(string)
lun=lunget()
open(unit=lun,file='out.dat', status='old')

do while(.true.)
 read(lun, '(a)', IOSTAT = reason)new_string
 if(reason < 0)exit
 if(  (index(new_string,'2020')==0) .and. (index(new_string,'2019')==0) .and. len(trim(new_string)) /= 7)cycle
 dir=trim(new_string)
 dir_file = trim(dir)//'/co_tunes_scraping_off.dat'
  print '(a)',dir_file
   inquire (file=trim(dir_file), exist=itexists)
  if (.not.itexists) cycle
  call get_tunes(dir,Qx,Qy,betax,betay,eta) 
  exit
end do
close(lun)

open(unit=lun,file='out.dat', status='old')
do while(.true.)
 read(lun, '(a)', IOSTAT = reason)new_string
 if(reason < 0)exit
 if(  (index(new_string,'2020')==0) .and. (index(new_string,'2019')==0) .and. len(trim(new_string)) /= 7)cycle
 if(trim(dir_list) /= 'all' .and. index(dir_list,trim(new_string)) == 0)cycle
  dir=trim(new_string)
  dir_file = trim(dir)//'/Time_Dep_Moments_at_'//trim(location)//'.dat'
!    dir_file = trim(dir)//'/Time_Dep_Moments_at_QUADM1.dat' 
  print '(a)',dir_file
 
  inquire (file=trim(dir_file), exist=itexists)
  if (.not.itexists) cycle
n_files = n_files+1



lines = 0
lun1=lunget()
   open(unit=lun1, file = trim(dir_file) )
   do while(.true.)
     read(lun1,'(a)', end=99)string
     if(index(string,'<')/= 0)cycle
      lines = lines + 1
      if(lines > nbins_max)then
         print *,' Moment array has size smaller than number of bins in data file'
         stop
      endif
!     read(string,*,end=99)tbin, time, number, orb%vec(1:6)
!      print '(a)',string
      read(string,'(i10)')i
      if(i > nbins_max)cycle
      read(string,'(i10,es18.8,i10,21(es12.4))',end=99)i, time(i),running_moment(i)%n, moments(i)%ave(:),&
                             (moments(i)%sigma(1+j,1+j),moments(i)%sigma(1+j,2+j),moments(i)%sigma(2+j,2+j), j=0,4,2), &
                             moments(i)%third(1:6)
     
      if(any(moments(i)%ave(:) ==0) .or. moments(i)%sigma(1,1) == 0 .or. moments(i)%sigma(3,3) ==0)then
       print *,' zero'
       cycle
      endif

      if(AS(1:1) /= 'A' .and. AS(1:1) /= 'S')then
        print '(a,$)', ' Data saved as Averages(A) or Sums(S)?'
        accept *,AS
      endif
   
      if(AS(1:1) == 'A')then
    
      moments_sum(i)%ave(:) = (moments_sum(i)%ave(:) * running_moments_sum(i)%n + moments(i)%ave(:) * running_moment(i)%n) &
                                                                                 /(running_moments_sum(i)%n + running_moment(i)%n)
      do j=0,4,2
         moments_sum(i)%sigma(1+j,1+j) =(moments_sum(i)%sigma(1+j,1+j) * running_moments_sum(i)%n + moments(i)%sigma(1+j,1+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
         moments_sum(i)%sigma(1+j,2+j) =(moments_sum(i)%sigma(1+j,2+j) * running_moments_sum(i)%n + moments(i)%sigma(1+j,2+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
         moments_sum(i)%sigma(2+j,2+j) =(moments_sum(i)%sigma(2+j,2+j) * running_moments_sum(i)%n + moments(i)%sigma(2+j,2+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
      end do
      running_moments_sum(i)%n = running_moments_sum(i)%n+running_moment(i)%n
      moments_sum(i)%n = running_moments_sum(i)%n

   else
      
      moments_sum(i)%ave(:) = (moments_sum(i)%ave(:) * running_moments_sum(i)%n + moments(i)%ave(:) * running_moment(i)%n) &
                                                                                 /(running_moments_sum(i)%n + running_moment(i)%n)
      do j=0,4,2
         moments_sum(i)%sigma(1+j,1+j) =(moments_sum(i)%sigma(1+j,1+j) * running_moments_sum(i)%n + moments(i)%sigma(1+j,1+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
         moments_sum(i)%sigma(1+j,2+j) =(moments_sum(i)%sigma(1+j,2+j) * running_moments_sum(i)%n + moments(i)%sigma(1+j,2+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
         moments_sum(i)%sigma(2+j,2+j) =(moments_sum(i)%sigma(2+j,2+j) * running_moments_sum(i)%n + moments(i)%sigma(2+j,2+j)*running_moment(i)%n)/ &
                                                                                  (running_moments_sum(i)%n + running_moment(i)%n)   
      end do
      running_moments_sum(i)%n = running_moments_sum(i)%n+running_moment(i)%n
      moments_sum(i)%n = running_moments_sum(i)%n

   else

      running_moments_sum(i)%sum(:)=running_moments_sum(i)%sum + moments(i)%ave(:)
      do j=0,4,2
         running_moments_sum(i)%product(1+j,1+j)= sunning_moments_sum(i)%product(1+j,1+j) + moments(i)%sigma(1+j,1+j)
         running_moments_sum(i)%product(1+j,2+j)= sunning_moments_sum(i)%product(1+j,2+j) + moments(i)%sigma(1+j,2+j)
         running_moments_sum(i)%product(2+j,2+j)= sunning_moments_sum(i)%product(2+j,2+j) + moments(i)%sigma(2+j,2+j)
      end do
      running_moments_sum(i)%n = running_moments_sum(i)%n + running_moment(i)%n
         

!      if((lines/10000)*10000 == lines)print *, lines
      tbin_max = max(tbin_max, i)

   end do

99 continue
   if(AS(1:1) == 'S')then
      moments_sum(i)%ave(:) = running_moments_sum(i)%sum/running_moments_sum(i)%n
      do j=0,4,2
         moments_sum(i)%sigma(1+j,1+j)= running_moments_sum(i)%product(1+j,1+j) /running_moments_sum(i)%n -  moments_sum(i)%ave(1+j)*moments_sum(i)%ave(1+j)
         moments_sum(i)%sigma(1+j,2+j)= running_moments_sum(i)%product(1+j,2+j) /running_moments_sum(i)%n -  moments_sum(i)%ave(1+j)*moments_sum(i)%ave(2+j)
         moments_sum(i)%sigma(2+j,2+j)= running_moments_sum(i)%product(2+j,2+j) /running_moments_sum(i)%n -  moments_sum(i)%ave(2+j)*moments_sum(i)%ave(2+j)
      end do
   endif
   

   close(lun1)

   n = log(float(tbin_max))/log(two)
   m = two**n
   
   call fft_x_y(time_bin, tbin_max, moments_sum, qx_fft, qy_fft,data_fft_n, data_fft_x, data_fft_y)

end do !cycle through directories
close(lun)
if(n_files == 0)then
   lun_0=lunget()
   open(unit=lun_0, access='sequential', file = 'NoData.dat',position='append') 
   write(lun_0,'(2a)')' No data in directory ',trim(cwd)
   close(unit=lun_0)
   goto 199
endif

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
   lun = lunget()
   open(unit = lun, file = 'EndTime_'//trim(location)//'.dat')
 print *,' EndTime_'//trim(location)//'.dat'
 write(lun,'(a)')'#! Data from files '
 write(lun,'(a)')'#!',dir_file

  write(lun, '(a9,1x,a12,3(3a14),9a14)')'time_bin','freq','Total Counts','Real fft','Imag fft',&
                                                 '<x>','Real fft','Imag fft',&
                                                 '<y>','Real fft','Imag fft', &
                                                 '<xx>','<xpx>','<pxpx>','<yy>','<ypy>','<pypy>','<zz>','<zpz>','<pp>'
  do i = 1, m
   freq=float(i)/time_bin/m/1.e6
   write(lun,'(i10,es12.4,(i14,2es14.6),6(3es14.6), 9es14.6)' ) i,freq,running_moments_sum(i)%n, data_fft_n(i),&
                          moments_sum(i)%ave(1), data_fft_x(i), &
                          moments_sum(i)%ave(3), data_fft_y(i), &
                           (moments_sum(i)%sigma(1+j,1+j),moments_sum(i)%sigma(1+j,2+j),moments_sum(i)%sigma(2+j,2+j), j=0,4,2)
   end do
close(lun)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
   lun = lunget()
   open(unit = lun, file = 'EndTime_rebin_'//trim(location)//'.dat')
 print *,' EndTime_rebin_'//trim(location)//'.dat'
 write(lun,'(a)')'#! Data from files '
 write(lun,'(a)')'#!',dir_file

  write(lun, '(a9,1x,a12,11a14)')'time_bin','Total Counts',&
                                                 '<x>',  '<y>',  '<xx>','<xpx>','<pxpx>','<yy>','<ypy>','<pypy>','<zz>','<zpz>','<pp>'

  nbins_to_combine = desired_width/time_bin
!   freq=float(i)/time_bin/m/1.e6
  do i = 1, m, nbins_to_combine
     write(lun,'(i10,i11,12es14.6)' ) (i+nbins_to_combine/2),sum(running_moments_sum(i:i+nbins_to_combine-1)%n),&
                         sum( moments_sum(i:i+nbins_to_combine-1)%ave(1))/nbins_to_combine, &
                          sum(moments_sum(i:i+nbins_to_combine-1)%ave(3)) /nbins_to_combine, &
                          (sum(moments_sum(i:i+nbins_to_combine-1)%sigma(1+j,1+j))/nbins_to_combine, &
                          sum(moments_sum(i:i+nbins_to_combine-1)%sigma(1+j,2+j))/nbins_to_combine, &
                          sum(moments_sum(i:i+nbins_to_combine-1)%sigma(2+j,2+j))/nbins_to_combine, j=0,4,2)
  end do
  close(lun)
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

199 continue
end do ! loop over locations
open(unit=lun, file='plot_width_cboamp_10-30us.dat')
write(lun, '(a13,a,a13,a,a16,a, a3,a9,a15,a1)')'location =  "',trim(location_0),'" ;  width= "',trim(desired_width_word),'" ;  dir_list= "',trim(dir_list),'" ;',' date = "',dir_list(1:15),'"'
print *,' write file plot_widt_cboamp_10-30us.dat '
close(lun)
end program compile_all_time_dep



subroutine get_tunes(dir, qx,qy,betax,betay,eta)
  use sim_utils
  implicit none
  logical itexists, twiss_found/.false./
 character*290 string, new_string, out_string
 character*100 dir
 character*50 dir_file
 character*10 word(1:5)/'Qx =','Qy =','Betax =','Betay =','eta ='/
character*10 test_word
 real(rp)  twiss(1:5), qx,qy,betax,betay,eta
 integer ix, reason, lun, lun1
 integer iword, w_len

 if(twiss_found) return

 ix=-1
 lun1=lunget()
 dir_file = trim(dir)//'/co_tunes_scraping_off.dat'
 inquire (file=dir_file, exist=itexists)
 if(.not.itexists) return
 open(lun1, file=trim(dir_file))
   do while(ix <= 0)
      read(lun1,'(a)', IOSTAT=reason)new_string
      if(reason < 0) exit
      do iword = 1,5
         w_len = len(trim(word(iword)))
         test_word = trim(word(iword))
        ix = index(new_string,trim(word(iword)))
        if(ix /=0)then
           call string_trim(new_string(ix+len(trim(word(iword))):), out_string,ix)
           read(out_string(1:ix),*)twiss(iword)
           print *,word(iword),twiss(iword)
        endif
        qx=twiss(1)
        qy=twiss(2)
        betax=twiss(3)
        betay=twiss(4)
        eta=twiss(5)
        
     end do
   
   end do
   twiss_found = .true.
   close(lun1)
   return

 end subroutine get_tunes
         
      subroutine fft_x_y(time_bin, tbin_max, moments_sum, q_horz, q_vert, data_fft_n, data_fft_x, data_fft_y)
        use precision_def
        use nr
        use bmad
        use compile_mod
        use compile_interface
        implicit none
        TYPE (g2moment_struct), allocatable :: moments_sum(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width
        real(rp) q_horz, q_vert, f_rev, zero/0./, two/2.d0/, f_0
        real(rp) time_bin
        real(rp) mrev_fit, mv_fit, mh_fit
        real(rp) dum1,dum2,dum3
        complex(rp), allocatable ::data_fft_n(:), data_fft_x(:), data_fft_y(:)
        integer isign/1/, n, m, i, peak_bin, m_rev
        integer tbin_max
        integer m_0
        integer low_end/10/
        integer m_top
        
! fourier transform of Number vs time

        n = log(float(tbin_max))/log(two)
        m = two**n
        if(allocated(data_fft_x))deallocate(data_fft_x, data_fft_y, data_fft_n)
        if(.not. allocated(data_fft_x))allocate(data_fft_x(1:m), data_fft_y(1:m), data_fft_n(1:m))
        do i=1,m
           data_fft_x(i) = cmplx(moments_sum(i)%ave(1),zero) * 2*sin((twopi/2*i)/m) * sin((twopi/2.*i)/m)
           data_fft_y(i) = cmplx(moments_sum(i)%ave(3),0.) * 2*sin((twopi/2*i)/m) * sin((twopi/2.*i)/m)
           data_fft_n(i) =cmplx(float(moments_sum(i)%n),zero) * 2*sin((twopi/2*i)/m) * sin((twopi/2.*i)/m)
           
        end do
        call four1(data_fft_x(1:m),isign)
        call four1(data_fft_y(1:m),isign)
        call four1(data_fft_n(1:m),isign)
        
        print *,' tbin_max = ', tbin_max,' log(two) = ', log(two), ' n= ',n,' m = ',m,' size(data_fft_n) = ', size(data_fft_n)

        f_0 = 1./149.2e-3 !guess revolution frequency MHz
        m_0 = f_0 *1.e6*m*time_bin/2
        m_top = m_0
        peak_bin=maxloc(abs(data_fft_n(low_end:2*m_top)),1)+low_end-1
        dum1=abs(data_fft_n(peak_bin-1))
        dum2=abs(data_fft_n(peak_bin))
        dum3=abs(data_fft_n(peak_bin+1))
        call fit_fft_peak(data_fft_n,peak_bin, mrev_fit)               
        m_rev = mrev_fit    !float(peak_bin)
        f_rev = mrev_fit/time_bin*2*1.e-6
        peak_bin=maxloc(abs(data_fft_y(low_end:m_rev/2)),1)+low_end-1
        dum1=abs(data_fft_y(peak_bin-1))
        dum2=abs(data_fft_y(peak_bin))
        dum3=abs(data_fft_y(peak_bin+1))

        call fit_fft_peak(data_fft_y,peak_bin, mv_fit)               
        Q_vert = mv_fit/mrev_fit !float(peak_bin)/m_rev
        peak_bin=maxloc(abs(data_fft_x(low_end:m_rev/2)),1)+low_end-1
        dum1=abs(data_fft_x(peak_bin-1))
        dum2=abs(data_fft_x(peak_bin))
        dum3=abs(data_fft_x(peak_bin+1))
        
        call fit_fft_peak(data_fft_x,peak_bin, mh_fit)                       
        Q_horz = mh_fit/mrev_fit !float(peak_bin)/m_rev
        print '(a,es12.4,a,i10,a,i10)', 'f_rev = ',f_rev,' m_rev = ',m_rev,' peak_bin = ', peak_bin
        print *,' Q_vert = ', Q_vert,' Q_horz = ', Q_horz
        return
   end subroutine fft_x_y
   
           
    subroutine fit_fft_peak(data_fft,peak_bin, mrev_fit)
      use precision_def
      use nr
      implicit none
      complex(rp), allocatable :: data_fft(:)
      real(rp) mrev_fit,a,b
      real(rp) X(3,3), V(3,1),z(-2:2)
      integer i, peak_bin, n/5/, j
      a=(abs(data_fft(peak_bin+1))+abs(data_fft(peak_bin-1)))/2. - abs(data_fft(peak_bin))
      b=(abs(data_fft(peak_bin+1))-abs(data_fft(peak_bin-1)))/2.
      mrev_fit=-b/2./a + peak_bin
      if(abs(b/2./a) > 1)mrev_fit=peak_bin 
!      do i=-2,2
!         z(i) = i
!      end do
!      X(1,1) = n
!      X(1,2) = sum(z(-n/2:n/2))
!      X(2,1)= X(1,2)
!      X(1,3) = sum(z(-n/2:n/2)**2)
!      X(2,2) = X(1,3)
!      X(3,1)=X(1,3)
!      X(2,3) = sum(z(-n/2:n/2)**3)
!      X(3,2) = X(2,3)
!      X(3,3) = sum(z(-n/2:n/2)**4)
!      V(1,1) = sum(abs(data_fft(peak_bin-n/2 : peak_bin+n/2) ))
!      V(2,1) = sum(z(-n/2:n/2)*abs(data_fft(peak_bin-n/2 : peak_bin+n/2)))
!      V(3,1) = sum(z(-n/2:n/2)**2*abs(data_fft(peak_bin-n/2 : peak_bin+n/2)))
!      call gaussj(X,V)
!      mrev_fit=-V(2,1)/2./V(3,1) + peak_bin
      
      return
    end subroutine fit_fft_peak