subroutine ptp_funcs(x,coef,yfit,dydcoef)
  use precision_def
  use nrtype
  implicit none
  real(rp), dimension(:), intent(in) :: x, coef
  real(rp), dimension(:), intent(out) :: yfit
  real(rp), dimension(:,:), intent(out) :: dydcoef
      
  yfit(:)=0
  yfit(:)  = coef(1) * sin(twopi*coef(2)*x(:)+coef(3))+coef(4)
  dydcoef(:,1) = sin(twopi*coef(2)*x(:)+coef(3))
  dydcoef(:,2) = coef(1) * twopi*x(:)*cos(twopi*coef(2)*x(:)+coef(3))
  dydcoef(:,3) = coef(1) * cos(twopi*coef(2)*x(:)+coef(3))
  dydcoef(:,4) = 1
     
end subroutine ptp_funcs

subroutine compute_moments_vs_time(nturns, turn,nmuons, muons, string, coef_save, chisq_save, ndf_save, test)

! In BMAD tracking, the coordinate along the reference trajectory is the dependent variable.
!The tracking routines, propagate the coordinates of the particles in a distribution from 
!the end of one element to the end of the next.
!The phase space coordinates, and the time, are given with respect to the reference particle.
!Detectors in the g-2 ring will measure the properties of the distribution (centroid, width, etc.)
!at a particular location and within some window of time. It is therefore convenient to compute
!moments of the distribution within that window.
! In this routine the phase space coordinates at the end of an element are sorted by arrival time.
! Running moments are accumulated.

! peak_to_peak is the peak to peak average x for all times beyond to ptop_start_time
! ptop_start_time is optional. If present than peak_to_peak is evaluated and returned as ptop_start_time

use muon_mod
use parameters_bmad
use muon_interface, dummy => compute_moments_vs_time
use nr

implicit none
integer  nmuons
integer i, tot
integer lun
integer, save :: n
integer j,k
integer time_bin
integer m
integer unit
integer nturns
integer turn
integer ndf
integer, optional :: ndf_save

real(rp) num
real(rp) x_max, x_min
real(rp) fturn
real(rp), allocatable :: ptop(:)
real(rp), allocatable :: dphi_dp(:),dphi_dp_err(:)
real(rp) siga, sigb, a,b, chi2,q
real(rp) peak_to_peak
real(rp) ptop_start_time
real(rp) coef(4)
real(rp) chisq
real(rp), optional :: coef_save(4), chisq_save

integer, optional :: test

character*(*) string

logical first/.true./
logical, allocatable :: mask(:)

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


if(turn > 0) then
   if(first .or. .not. allocated(running_moment))then
      n = nturns*149.e-9/bin_width + 100
      allocate(running_moment(0:n))
   do i=1,n
    running_moment(i)%product(:,:)=0
    running_moment(i)%sum(:)=0
    running_moment(i)%n=0
    running_moment(i)%third(:) = 0
   end do
   first=.false.
 endif

!Sort into time bins
   do i=1,nmuons
     if (muons(i)%coord%state /= alive$)cycle ! .and. muons(i)%decay == .false.)cycle
     time_bin = muons(i)%coord%t/bin_width +0.5 ! associate integer with each bin_width time slice
     if(time_bin > n .or. time_bin <0)cycle
! running sums
     do j=1,6
      running_moment(time_bin)%product(1:6,j) = running_moment(time_bin)%product(1:6,j) + muons(i)%coord%vec(1:6)*muons(i)%coord%vec(j)
!      print '(3i12,6es12.4)',i,j,time_bin,running_moment(time_bin)%product(1:6,j)
     end do
     running_moment(time_bin)%sum(1:6) = running_moment(time_bin)%sum(1:6) + muons(i)%coord%vec(1:6)
     running_moment(time_bin)%n = running_moment(time_bin)%n+1
     running_moment(time_bin)%third(1:6) = running_moment(time_bin)%third(1:6) + muons(i)%coord%vec(1:6)**3
     
  end do
  fturn = float(turn)
  unit=1
  if(pbin_width > 0 .and. trim(time_binning) == 'position_slice')call bin_spin_angle_by_time(fturn,unit, nmuons,muons)

  elseif(turn == -1)then !compute moments and write. <turn> is the flag that tracking is complete

 ! extra so can run independently

     n = nturns*149.e-9/bin_width + 100     
! moments
     allocate(moments(0:n), ptop(0:n))
     ptop(:)=0.
   
  print *,' moments, n= ', n
 if(.not. present(test))then  !!!!!
     do i=1,n
        moments(i)%ave(1:6)=0
!       time_bin = muons(i)%coord%t/bin_width +0.5 ! associate integer with each bin_width time slice
        if(running_moment(i)%n == 0)cycle
        forall(j=1:6)moments(i)%ave(j) = running_moment(i)%sum(j)/float(running_moment(i)%n)
        num = float(running_moment(i)%n)
        do j=1,6
           do k=1,6
              moments(i)%sigma(j,k) = running_moment(i)%product(j,k)/num  -  moments(i)%ave(j)*moments(i)%ave(k)
           end do
           moments(i)%third(j) = running_moment(i)%third(j)/num -3*running_moment(i)%product(j,j)*moments(i)%ave(j)/num + 2*moments(i)%ave(j)**3
        end do
     end do
  else !!!! test is present read data from file
     allocate(running_moment(0:n))     
     if(test == 0)call get_time_dep_moments(n,moments, running_moment)
     if(test >=0)call get_xavg(n,moments, running_moment, test) 
  endif
  
  call fit_average_x(n, moments, running_moment, coef, chisq, ndf)
  if(present(coef_save))coef_save=coef
  if(present(chisq_save))chisq_save=chisq
  if(present(ndf_save))ndf_save=ndf

 if(.not. present(test))then 
    lun=lunget()
    open(unit=lun,file=trim(directory)//'/'//'Time_Dep_Moments_at_'//trim(string)//'.dat')
    print *,'Write: Time_Dep_Moments_at_'//trim(string)//'.dat'
    write(lun,'(4(a12,es12.4))')'  coef[1] = ',coef(1),'; coef[2] = ',coef(2) ,'; coef[3] = ',coef(3) ,'; coef[4] = ',coef(4) 
    write(lun,'(a10,a18,a10,24a12)')'Time bin','time','number','<x>','<px>','<y>','<py>','<z>','<pz>','<xx>','<x-px>','<px-px>','<yy>','<y-py>','<py-py>', &
         '<zz>','<z-pz>','<pz-pz>','<xxx>','<xpxpxp>','<yyy>','<ypypyp>','<zzz>','<zpzpzp>','<ptop>','<dphi_dp>','<dphi_dp_err>'
    fturn = float(turn)
    if(trim(time_binning)=='position_slice' .and. pbin_width>0)call bin_spin_angle_by_time(fturn,unit, nmuons,muons, dphi_dp, dphi_dp_err)
    

    do i=1,n
       if(running_moment(i)%n == 0)cycle
       if(i>15 .and. nmuons >1)then
          x_max = maxval(moments(i-15:i)%ave(1))
          x_min = minval(moments(i-15:i)%ave(1))
          ptop(i) = x_max-x_min
       endif

!    allocate(mask(0:n))
!    mask(0:bin_min)=.false.
!    peak_to_peak = maxval(ptop(:), mask(:))
!    ptop_start_time = peak_to_peak
!    deallocate(mask)

       if(trim(time_binning)=='position_slice' .and. allocated(dphi_dp))then
          write(lun,'(i10,es18.8,i10,21es12.4,es12.4,2es12.4)')i, i*bin_width,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), ptop(i), dphi_dp(i), dphi_dp_err(i)
       else
          write(lun,'(i10,es18.8,i10,21es12.4,es12.4)')i, i*bin_width,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), ptop(i)
       endif
       
    end do

    close(unit=lun)
 endif
 
 deallocate(running_moment)
 deallocate(moments, ptop)
 !  call cbo_decay_time(ptop,bin_width, decay_time)

endif


return
end

subroutine get_time_dep_moments(n,moments, running_moment)
  use precision_def
  use muon_mod
  use muon_interface, dummy => get_time_dep_moments
  implicit none
  
  TYPE (g2moment_struct), allocatable :: moments(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width
  type (running_moment_struct), allocatable :: running_moment(:)

  integer n,i,j,idum
  integer lun
  real(rp)x1, x(6)
  character*120 word
  character*3 fitn
  
  lun=lunget()
  open(lun,file='Time_Dep_Moments_at_END.dat')
  read(lun,'(a)')word
  read(lun,'(a)')word  
  write(111,'(3a12)')'j','time','x'      
     do i=1,n
!    write(lun,'(i10,es18.8,i10,21es12.4,es12.4,2es12.4)')i, i*bin_width,running_moment(i)%n, moments(i)%ave(:),&        
       read(lun,'(i10,es18.8,i10,6es12.4)',end=10)j, x1,idum,x(1),x(2),x(3),x(4),x(5),x(6)
       
       moments(j)%ave(1:6) = x(1:6)
       running_moment(j)%n=idum
       write(111,'(i12,2es12.4)')j,x1,x(1)
    end do
10    return
  end subroutine get_time_dep_moments
  
  subroutine fit_average_x(n,moments, running_moment, coef, chisq, ndf)
  use precision_def
  use muon_mod
  use muon_interface, dummy => fit_average_x
  use parameters_bmad
  use nr
  implicit none
  
  TYPE (g2moment_struct), allocatable :: moments(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width
  type (running_moment_struct), allocatable :: running_moment(:)

  integer n,i,j,idum
  integer lun
  real(rp), allocatable :: y(:),x(:),sig(:)
  logical maska(4)
  real(rp) coef(4),covar(4,4),alpha(4,4)
  real(rp) chisq, alamda
  real(rp) alamda_last
  character*7 warning   
  real(rp) ptop_start_time
  real(rp) window
  integer bin_min, bin_window !minimum bin and window for fitting to x vs t
  integer ndf
  integer, save :: fit_number
  character*3 fitn
  logical first/.true./

  interface
     subroutine ptp_funcs(x,coef,yfit,dydcoef)
       use precision_def
       use nrtype
       implicit none
       real(rp), dimension(:), intent(in) :: x, coef
       real(rp), dimension(:), intent(out) :: yfit
       real(rp), dimension(:,:), intent(out) :: dydcoef
     end subroutine ptp_funcs
  end interface
  ! fit sin function to average x vs time

   ptop_start_time = 30.e-6
   window = 25.e-6      !fit from ptop_start_time to ptop_start_time + window
   bin_min = ptop_start_time/bin_width  !start fit at ptop_start_time
   bin_window = window/bin_width

   j=0
   do i=1,n
      if(running_moment(i)%n == 0 .or. i*bin_width < ptop_start_time .or. i*bin_width > ptop_start_time+window)cycle
      j=j+1
   end do
   allocate(y(1:j),x(1:j),sig(1:j))
   j=0
   do i=1,n
      if(running_moment(i)%n == 0 .or. i*bin_width < ptop_start_time .or. i*bin_width > ptop_start_time+window)cycle
      j=j+1
      y(j)= moments(i)%ave(1) * 1000  !m to mm
      x(j) = bin_width*i*1.e6  !sec to microseconds
   enddo
   ndf=j

   if(ndf < 100)then
      coef=0.; chisq=0.
      return
   endif
   

   if(first)then
      fit_number=0
      first=.false.
   endif 
   fit_number = fit_number+1
   if(fit_number < 10) write(fitn,'(i1)')fit_number
   if(fit_number >= 10) write(fitn,'(i2)')fit_number
   open(98,file=trim(directory)//'/'//'xavg_vs_t_'//trim(fitn)//'.dat')
   maska(:)=.false.
   sig(:)=1.
   coef(2)=0.36; coef(3)=twopi/4.
   coef(4) = sum(y)/ndf
   coef(1) = (maxval(y) - minval(y))/2.
   alamda=-1.
   maska(:)=.true.
!   call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda) !first time fit phase
!   alamda=0.
!   i=0
!   call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda) !first time fit phase
!   write(98,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=', chisq,' alamda=', alamda, &
!        ' sqrt(chisq/ndf)=',sqrt(chisq/j)
!   write(6,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
!        ' sqrt(chisq/ndf)=',sqrt(chisq/j)   
   do i=3,2,-1
      maska(:)=.false.
      maska(i)=.true.
      alamda=-1.      
      call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda)  !fit with each of coefficients inorder
      alamda=0.
      call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda)  !fit with each of coefficients inorder      
      write(98,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
        ' sqrt(chisq/ndf)=',sqrt(chisq/j)
      write(6,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
        ' sqrt(chisq/ndf)=',sqrt(chisq/j)      
   end do
   maska(:)=.true.
   alamda=-1.
   alamda_last= 1.
   do i=1,20
      call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda)  !fit with all four coefficients
      write(98,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
        ' sqrt(chisq/ndf)=',sqrt(chisq/j)
      write(6,'(a10,i10,6(a10,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
           ' sqrt(chisq/ndf)=',sqrt(chisq/j)
      if(alamda > alamda_last)exit
      alamda_last=alamda
   end do   
   alamda=0.
   call mrqmin(x, y, sig, coef, maska,covar,alpha,chisq, ptp_funcs, alamda)  !fit to get covariance
   write(98,'(a10,i10,6(a11,es12.4),a17,es12.4)')'iteration=',i,'  coef[1]=', coef(1),';  coef[2]=', coef(2),';  coef[3]=', coef(3),';  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda, &
        ' sqrt(chisq/ndf)=',sqrt(chisq/j)

   do i = 1,j
      warning=''
      if(y(i) == 0)warning = 'WARNING'
      write(98,'(i10,2es12.4,1x,a)')i,x(i),y(i),warning
   end do
   close(98)
   deallocate(y,x,sig)
   return
 end subroutine fit_average_x
 
 subroutine get_xavg( n,moments, running_moment, fit_number)
     use precision_def
  use muon_mod
  use muon_interface, dummy => get_xavg
  use parameters_bmad
  implicit none
  
  TYPE (g2moment_struct), allocatable :: moments(:) !moment structure for each time bin, 1000 turns => an array with 1000*149e-9/bin_width
  type (running_moment_struct), allocatable :: running_moment(:)

  integer n,i,j,idum
  integer lun
  integer fit_number
  
  real(rp)x,y
  character*120 word
  character*7 warning
  character*3 fitn
  character*200 string

  integer ios
  
  lun=lunget()
  if(fit_number < 10) write(fitn,'(i1)')fit_number
  if(fit_number >= 10) write(fitn,'(i2)')fit_number
  open(lun,file='xavg_vs_t_'//trim(fitn)//'.dat', STATUS='old',IOSTAT=ios)
  do while (.true.)
     read(lun,'(a)')string
     if(index(string,'iteration') == 0)exit
  end do

  read(string,'(i10,2es12.4,1x,a)')j,x,y,warning   !read the first line of x,y data
  i = x/1.e6/bin_width + 0.1 !the extra 0.1 to be sure we get the right integer i
  moments(i)%ave(1) = y/1000.   !mm to m
  running_moment(i)%n =  200
  running_moment(1: i-1)%n=0     !no particles for times less than x =i*bin_width*1.e6 

  do while(Ios >= 0)
     read(lun,'(i10,2es12.4,1x,a)',IOSTAT=ios)j,x,y,warning
      i = x/1.e6/bin_width + 0.1 !the extra 0.1 to be sure we get the right integer i
      moments(i)%ave(1) = y/1000. !mm to m
      running_moment(i)%n =  200   
   end do
   n=i
   
    close(lun)
    return
  end subroutine get_xavg