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

program fit_to_xavg_vs_time

! 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 bin_min, bin_window !minimum bin and window for fitting to x vs t

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), optional :: ptop_start_time
real(rp) window
real(rp) peak_to_peak

logical maska(4)
real(rp) coef(4),covar(4,4),alpha(4,4)
real(rp) chisq, alamda
real(rp), allocatable :: y(:),x(:),sig(:),xx(:)
!real(rp) ptp_funcs

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
  real(rp) product(1:6,1:6), sum(1:6), third(6)
  integer n
end type
type (running_moment_struct), save, allocatable :: running_moment(:)
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


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

! moments
     allocate(moments(0:n), ptop(0:n))
     ptop(:)=0.
   
  print *,' moments, n= ', n
  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


! 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
   allocate(y(0:bin_window),x(0:bin_window),sig(0:bin_window))
   allocate(xx(0:n))
   y(0:bin_window)= moments(bin_min:bin_min + bin_window)%ave(1) * 1000  !m to mm
   forall(i=0:n) xx(i) = bin_width * i *1.e6  !sec to microseconds
   x(0:bin_window) = xx(bin_min:bin_min+bin_window)
   maska(:)=.false.
   sig(:)=1.
   coef(1)=10.; coef(2)=0.36;coef(3)=0.01;coef(4)=0.01 
   alamda=-1.
   maska(:)=.false.
   maska(3)=.true.
   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))')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda

   do i=1,4
      maska(:)=.false.
      maska(i)=.true.
      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))')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda
   end do
   maska(:)=.true.
   do i=1,4
      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))')'iteration=',i,'  coef[1]=', coef(1),'  coef[2]=', coef(2),'  coef[3]=', coef(3),'  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', 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))')'iteration=',i,'  coef[1]=', coef(1),';  coef[2]=', coef(2),';  coef[3]=', coef(3),';  coef[4]=', coef(4),'  chisq=',chisq,'  alamda=', alamda

   do i = 1,n-bin_min
     write(98,'(2es12.4)')x(i),y(i)
   end do 
   deallocate(y,x, xx)

 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')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

  if(present(ptop_start_time))then
    allocate(mask(0:n))
    mask(0:bin_min)=.false.
    peak_to_peak = maxval(ptop(:), mask(:))
    ptop_start_time = peak_to_peak
    deallocate(mask)
  endif

 if(trim(time_binning)=='position_slice')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)
 deallocate(running_moment)
 deallocate(moments, ptop)
 !  call cbo_decay_time(ptop,bin_width, decay_time)

endif


return
end