program create_correlated_distribution
use nr
use muon_mod
use muon_interface
use parameters_bmad
use sim_utils

implicit none
integer unit, nmuons/-1/
integer nmuon_first/0/
character*100 string/'correlation'/
character*120 muon_file
character*10 nmuons_cha
type (muon_struct), allocatable :: muons(:)

real(rp), allocatable :: p(:), Sigma(:,:), temp(:,:), ave(:)
real(rp), allocatable ::  L(:,:), LT(:,:), Diag(:,:), M(:,:)
integer lun, ndim
integer i,j
integer tot/0/, turn/2/
integer nargs


ndim=8

allocate(p(1:ndim))
allocate(Sigma(1:ndim, 1:ndim))
allocate(L(1:ndim, 1:ndim),LT(1:ndim,1:ndim))
allocate(M(1:ndim,1:ndim), Diag(1:ndim,1:ndim))
allocate(temp(1:ndim,1:ndim))
allocate(ave(1:ndim))

nargs = command_argument_count()
call cesr_getarg(1,muon_file)
if(nargs ==2)then
   call cesr_getarg(2,nmuons_cha)
   read(nmuons_cha,*)nmuons
endif

  
lun= lunget()
open(unit = lun, file = trim(muon_file))

if(nmuons<0)then
   call get_number_events(lun,  tot, nmuon_first)
   nmuons = tot-4
endif


allocate(muons(1:nmuons))

call read_phase_space(lun, nmuons, muons, string, tot, nmuon_first)

call compute_correlation(nmuons, muons, Sigma, ave)

print *,' average and Covariance matrix'
print '(8es12.4,/)',ave(1:ndim)
do j=1,ndim
   print '(8es12.4)',(Sigma(j,i),i=1,ndim)
end do

temp = Sigma

p=0

call choldc(temp,p) !cholesky_decomposition

L=0

do i=1,ndim
   do j=i,ndim
      L(j,i) = temp(j,i)
   end do
   L(i,i) = p(i)
end do

temp=0
forall(i=1:ndim)temp(i,i)=1

print '(/,a)',' Choleksy L'
do j=1,ndim
   print '(8es12.4)',(L(j,i),i=1,ndim)
end do

LT=transpose(L)

print '(/,a)',' Choleksy LT'
do j=1,ndim
   print '(8es12.4)',(LT(j,i),i=1,ndim)
end do


M = matmul(L,LT)

print '(/,a)',' LLT'
do j=1,ndim
   print '(8es12.4)',(M(i,j),i=1,ndim)
end do

call create_dist_with_correlation(L, muons, ave)

call compute_correlation(nmuons, muons, Sigma)

!directory='junk'
!muons(:)%coord%state =alive$
!call write_phase_space(nmuons, muons, 'END',tot,turn)

print '(/,a)',' average and Covariance matrix of new distribution'
print '(8es12.4)',ave(1:ndim)
print *
do j=1,ndim
   print '(8es12.4)',(Sigma(j,i),i=1,ndim)
end do


end program create_correlated_distribution