module compute
!!DIR$ NOOPTIMIZE
	use basicmodule
	use polymod
	use iomodule
	use priormod
	implicit none
	
	type tLR
		real, allocatable	:: LR(:)			! for computation of Bayes factors
	end type
	type(tLR)	:: BF_LRs(2,pmax)

	real		:: BFs(pmax),om2
	
	contains
	
	function getstatvar(para) result(val)		! computes unconditional covariance matrix
		use lincg_int
		type(tpara)		:: para
		real			:: val(p,p)
		complex	:: lam(p), Q(p,p), Qinv(p,p)
		complex	:: a(p**2), b(p**2)
		integer	:: i,j
		lam=-getcfromh(para%hAR(1:p))
		do j=1,p
			Q(1,j)=1.0
			do i=2,p
				Q(i,j)=lam(j)*Q(i-1,j)
			enddo
			Q(:,j)=Q(:,j)/norm2(abs(Q(:,j)))
		enddo
		lam=1+lam/Tstar
			
		call lincg(Q,Qinv)
		a=1-[spread(lam,1,p)]*[spread(lam,2,p)]
		b=[spread(Qinv(:,p),dim=2,ncopies=p)*spread(Qinv(:,p),dim=1,ncopies=p)]
		val=reshape([matmul(Q,matmul(reshape(b/a,[p,p]),transpose(Q)))],[p,p])
	end function

	subroutine setstatespace(para,phi,theta,statevar)	! initialize state space recursions using based on para value
		type(tpara)		:: para
		real			:: Phi(p,p),theta(p)
		real			:: poly(p+1)
		real			:: statevar(p,p)
		integer		:: j
		
		Phi=0
		do j=2,p
			Phi(j-1,j)=1
		enddo
		poly=realpolyfromroots(-getcfromh(para%hAR(1:p)))
		Phi(p,:)=-poly(1:p)
		if(p>1) then
			theta=realpolyfromroots(-getcfromh(para%hMA(1:p-1)))
		else
			theta=1
		endif
		Phi=eye(p)+phi/Tstar
		statevar=getstatvar(para)
	end subroutine
	
	function getKalman_ll(para) result(val)	! computes log-likelihood using Kalman filter
		type(tpara)		:: para
		real			:: val
		integer			:: ind
		real			:: Phi(p,p),theta(p)
		real			:: state(p,2), statevar(p,p), vst(p), cv, e(2)
		real			:: ssum(4)
		integer			:: s,j

		call setstatespace(para,phi,theta,statevar)
		if(any(diagonal(statevar)<0)) then
			print *,"statevar negative"
			val=0
			return
		endif
		if(CYrun) statevar=0		! impose zero initial condition for CYrun
		state=0
		ssum=0
		j=1
		do s=1,Tstar
			state=matmul(Phi,state)
			statevar=matmul(Phi,matmul(statevar,transpose(Phi)))
			statevar(p,p)=statevar(p,p)+1.0
			if(s*capN>=j*Tstar) then
				vst=matmul(statevar,theta)
				cv=1.0/sum(vst*theta)
				e=[ydata(j),1.0]-matmul(theta,state)
				ssum(1:3)=ssum(1:3)+cv*[e(1)**2,e(2)**2,e(1)*e(2)]
				ssum(4)=ssum(4)+log(cv)
				state=state+matmul(reshape(vst,[p,1]),reshape(cv*e,[1,2]))
				statevar=statevar-matmul(reshape(vst*cv,[p,1]),reshape(vst,[1,p]))
				j=j+1
			endif
		enddo
		if(CYrun) then
			om2=(ssum(1)-ssum(3)**2/ssum(2))/capN	
			val=.5*ssum(4)-0.5*capN*log(om2)		! return profiled log-likelihood if CYrun
		else
			val=.5*ssum(4)-0.5*log(ssum(2))-0.5*(capN-1)*log(ssum(1)-ssum(3)**2/ssum(2)) ! return integrated-out LL 
		endif	
	end function

	
	function gethalflife(para) result(val)
		type(tpara)		:: para
		real		:: val
		real		:: Phi(p,p), theta(p), statevar(p,p)
		real		:: corr(0:maxhhfac*Tstar)
		integer	:: s
		
		if(hhspecialflag==1) then
			val=product(getcfromh(para%hAR(1:p)))
			if(p>1) val=val/product(getcfromh(para%hMA(1:p-1)))
			val=Tdata*log(2.0)/val
			return
		endif
			
		call setstatespace(para,phi,theta,statevar)
		if(hhspecialflag==2) then
			val=product(getcfromh(para%hAR(1:p)))
			if(p>1) val=val/product(getcfromh(para%hMA(1:p-1)))
			val=Tdata*0.5*log(2.0)/(val**2*(sum(theta*matmul(statevar,theta))/Tstar))
			return
		endif
		
		do s=0,maxhhfac*Tstar
			corr(s)=sum(theta*matmul(statevar,theta))
			statevar=matmul(Phi,statevar)
		enddo

		s=maxhhfac*Tstar
		do
			if(abs(corr(s))>(0.5)**hhf*corr(0)) exit
			s=s-1
		enddo
		val=Tdata*(s/hhf)/Tstar
	end function 
	
	
	subroutine draw_once(para,priorflag) 	! update para in MCMC chain via Metropolis-Hastings
		type(tpara)		:: para
		logical	:: priorflag
		type(tpara)		:: newpara
		real		:: ran(pmax)
		real		:: oldlogp,newlogp,unif(1),hh
	
		newpara=para
		call rnnoa(ran(1:p))
		newpara%hAR(1:p)=newpara%hAR(1:p)+stepsize*ran(1:p)
		if(maxval(newpara%hAR(1:p))>hub .or. minval(newpara%hAR(1:p))<hlb) return		! don't accept new value
		if(p>1) then
			call rnnoa(ran(1:p-1))
			newpara%hMA(1:p-1)=newpara%hMA(1:p-1)+stepsize*ran(1:p-1)
			if(maxval(newpara%hMA(1:p-1))>hub .or. minval(newpara%hMA(1:p-1))<hlb) return
		endif
		if(.not. rawpriorflag) then			! if rawpriorflag=.true.: generate draws from prior baseline prior before half-life adjustment
			if(rootdist(newpara)<0.1) return
			hh=gethalflife(newpara)
			if(hh<hhlb) return
			if(hh>hhub) return
		endif
		oldlogp=getsmoothprior(para)
		newlogp=getsmoothprior(newpara)
		if(.not. rawpriorflag) then
			oldlogp=oldlogp+log(gethhprior(gethalflife(para)))	! half-life adjustmet to prior
			newlogp=newlogp+log(gethhprior(hh))
		endif
		if(.not. priorflag) then
			oldlogp=oldlogp+getKalman_ll(para)					! likelihood
			newlogp=newlogp+getKalman_ll(newpara)
		endif
		call rnun(unif)
		if(exp(newlogp-oldlogp)>unif(1)) then			! accept new draw?
			para=newpara
			acc=acc+1
		endif
	end subroutine
	
	subroutine iterate(priorflag)	! if priorflag=.true.: draw from prior
		logical	:: priorflag
		integer	:: irun,l,leff
		type(tpara)	:: para
		real		:: run(pmax)
		type(tpara)	:: vpara(ns,nruns)
		real		:: vhh(ns,nruns), avaccs(nruns), cmplcount(nruns)
		real, allocatable		:: vhheff(:),unifs(:)
		real		:: hh
		real		:: outtab(3,pmax)=-1
		cmplcount=0
!$omp parallel do private(run,para,l,leff,hh) copyin(p)		! start parallel chains
		do irun=1,nruns
			do 
				call rnun(run(1:p))
				para%hAR(1:p)=hlb+(hub-hlb)*run(1:p)
				if(p>1) then
					call rnun(run(1:p-1))
					para%hMA(1:p-1)=hlb+(hub-hlb)*run(1:p-1)
				endif
				if(rootdist(para)<0.1) cycle
				hh=gethalflife(para)
				if(.not. (hh<hhlb .or. hh>hhub)) exit
			enddo
			acc=0; stepsize=1
			do l=1,5000
				call draw_once(para,priorflag)
				if(mod(l,1000)==0) then
					if(acc<300) then
						stepsize=0.75*stepsize
					else
						stepsize=1.5*stepsize
					endif
					acc=0
				endif
			enddo
!print *,"done with burn in",stepsize
			acc=0
			do l=1,ns*nskip
				call draw_once(para,priorflag)
				if(mod(l,nskip)==0) then
					leff=l/nskip
					vpara(leff,irun)=para
					vhh(leff,irun)=gethalflife(para)
				endif
			enddo
			avaccs(irun)=real(acc)/(ns*nskip)
		enddo
		print *,"average acceptance rates"
		call mdisp(avaccs)
		print *,"quantiles of hh"
		call mdisp(quantile([vhh],[0.01,0.05,0.10,.5,0.9,0.95,.99]))
		if(rawpriorflag) then
			block	
				integer, parameter	:: nout=1000
				integer				:: i
				real				:: paraout(2*p-1,nout)
				do i=1,nout
					paraout(1:p,i)=vpara(i*ns/nout,1)%hAR(1:p)
					paraout(p+1:,i)=vpara(i*ns/nout,1)%hMA(1:p-1)
					call savemat(trim(dir)//"priordraws"//convtos(nint(priorkappa))//".txt",paraout) ! save baseline draws for Figure 3
				enddo
			end block
		endif
		if(priorflag.and.icase==0) then
			block	
				integer, parameter	:: nout=500
				integer				:: i
				real				:: paraout(2*p-1,nout)
				do i=1,nout
					paraout(1:p,i)=vpara(i*ns/nout,1)%hAR(1:p)
					paraout(p+1:,i)=vpara(i*ns/nout,1)%hMA(1:p-1)
					call savemat(trim(dir)//"altweight"//convtos(p)//".txt",paraout) ! save prior draws as weights for power maximizing test
				enddo
			end block
		endif
		block 
			integer	:: i
			do i=1,min(5,nruns)
				call execinML("figure("//convtos(merge(i,i+5,priorflag))//")")
				call plotspecdens(vpara(1:ns:ns/20,i))		! plot implied spectral densities 
			enddo
			vhheff=pack([vhh],[vhh]<hhub)
			unifs=vhheff
			call rnun(unifs)
			vhheff=vhheff+Tdata*(unifs/hhf)/Tstar
			call storeML(vhheff,"hh")
			call execinML("histogram(hh)")
		end block
			
		if(hhpriorset) then			! if half-life prior adjustment is already set, compute Bayes factors
			vhheff=pack([vhh],[vhh]<hhub)
			unifs=vhheff
			call rnun(unifs)
			vhheff=vhheff+Tdata*(unifs/hhf)/Tstar
			call mkpostplot(vhheff)
			outtab(:,p)=quantile_v(vhheff,[0.5,0.05,0.95])
			call mdisp(outtab)
			call setBF(pack([vpara],[vhh]<hhub))
		endif
		if(priorflag .and. (.not. hhpriorset)) then ! set half-life prior adjustemnt
			vhheff=pack([vhh],[vhh]<hhub)
			call sethhprior0(vhheff) 
		endif
		
	end subroutine
	
	subroutine setBF(vpara)		! store ratios of posteriors between models in LR and compute Bayes factors
		type(tpara)	:: vpara(:)

		integer	:: p0,l
		
		if(p>1) then
			p0=p
			if(allocated(BF_LRs(1,p0)%LR)) deallocate(BF_LRs(1,p0)%LR)
			allocate(BF_LRs(1,p0)%LR(size(vpara)))
!$omp parallel do
			do l=1,size(vpara)
				BF_LRs(1,p0)%LR(l)=getmodelLR(p0-1,p0,vpara(l))
			enddo
			p=p0
		endif
		if(p<pmax) then
			p0=p
			if(allocated(BF_LRs(2,p0)%LR)) deallocate(BF_LRs(2,p0)%LR)
			allocate(BF_LRs(2,p0)%LR(size(vpara)))
!$omp parallel do
			do l=1,size(vpara)
				BF_LRs(2,p0)%LR(l)=getmodelLR(p0+1,p0,vpara(l))
			enddo
			p=p0
		endif
		
		if(p>1) then
			BFs(p)=getsingleBF(BF_LRs(1,p)%LR,BF_LRs(2,p-1)%LR)
			print *,"Bayes Factors"
			call mdisp(BFs)
		endif
		print *,product(BFs),1.0/product(BFs)

		
		contains
	
			function getmodelLR(p1,p2,para) result(val)
				integer	:: p1,p2
				type(tpara)	:: para
				real	:: val
				real	:: u(2),hh,log1,log2

				p=p1
				if(p1>p2) then
					call rnun(u)
					para%hAR(p)=hlb+(hub-hlb)*u(1)
					para%hMA(p-1)=hlb+(hub-hlb)*u(2)
					if(rootdist(para)<0.1) then
						val=0
						return
					endif
					hh=gethalflife(para)
				else
					hh=gethalflife(para)
				endif					
				if(hh<hhlb .or. hh>hhub) then
					val=0
					return
				endif
				log1=getsmoothprior(para)+log(gethhprior(hh))+getKalman_ll(para)
				p=p2
				hh=gethalflife(para)
				log2=getsmoothprior(para)+log(gethhprior(hh))+getKalman_ll(para)
				val=exp(log1-log2)
			end function
	end subroutine
	
	function getsingleBF(LR1,LR2) result(val) !LR1 are draws from model 1 and contain q2/q1, LR2 are draws from model 2 and contain q1/q2
		real	:: LR1(:),LR2(:),val
		real	:: r
		integer	:: i
		
		r=1
		do i=1,10
			r=size(LR1)*sum(LR2/(LR2+r))/(size(LR2)*sum(LR1/(1+r*LR1)))
		enddo
		val=r
	end function

end module	