module compute
!$dir nooptimize	
	use rpmod
	use Sigmod
	use scpcmod
	implicit none

	integer, parameter	:: nstates=48
	integer, parameter		:: ndata=217,nseries=1000
	
	contains
	
	function getstatelocs(istate,lightpower,n,gdens) result(val)
		integer	:: istate,n
		real	:: lightpower,val(2,n)
		real, optional	:: gdens(n)
		
		logical, save	:: initflag=.true.		
		real, allocatable, save	:: mdata(:,:)
		character(len=:),allocatable,save	:: names(:)		
		real, allocatable	:: sdata(:,:)
		real, allocatable	:: dG(:),G(:)
		
		real	:: u(3),s0(2),s(2)
		integer	:: i,j,l
		
		if(initflag) then
			mdata=loadcsv("c:/dropbox/SHAR/data/US_Light_ext.csv")
			names=loadstrings("c:/dropbox/SHAR/data/US_Light_extnames.txt")
			initflag=.false.
		endif
		if(istate>nstates) then
			print *,"getstatelocs encountered state with no data"
			stop
		endif
!		print *,"generating random data from state",istate, names(istate:istate)
		sdata=selectifr(mdata,mdata(:,1)==istate)

		dG=sdata(:,6)
		dG=merge(dG,0.0,isfinite(dG))+1E-100
		dG=dG**lightpower
		G=dG
		do i=2,size(G)
			G(i)=G(i-1)+G(i)
		enddo
		G=G/(0.99999999*G(size(G)))
		do i=1,n
			call rnun(u)
			j=count(u(3)>G)+1
			s0=sdata(j,2:3)+u(1:2)*(sdata(j,4:5)-sdata(j,2:3))
truncnorm:	do
				call rnnoa(s)
				s=s0+0.2*s*(sdata(j,4:5)-sdata(j,2:3))			! now check if s is in some rectangle
				if(all(sdata(j,4:5)-s>0) .and. all(s-sdata(j,2:3)>0)) exit truncnorm
				do l=1,size(G)			
					if(all(sdata(l,4:5)-s>0) .and. all(s-sdata(l,2:3)>0)) exit truncnorm
				enddo
			enddo truncnorm
			val(:,i)=s
			if(present(gdens)) gdens(i)=dG(i)
		enddo
		disttype=0
	end function
	
	subroutine setWBpanel(m,k,xind,feflag0)
		use rnsri_int
		use rnper_int
		use rlse_int
		type(mtype)	:: m
		integer		:: k
		integer, optional	:: xind
		logical, optional	:: feflag0
		logical				:: feflag
		logical, save	:: init=.true.
		real, allocatable, save	:: adata(:,:,:),mdata(:,:)

		real		:: sse,sst
		real, allocatable	:: b(:),bx(:)
		integer		:: i,j,jc,t,inds(k+1),rper(k+1),l
		integer,allocatable	:: oinds(:),nco(:)

		if(present(feflag0)) then
			feflag=feflag0
		else
			feflag=.false.
		endif
		call erset(0,0,0)
		if(init) then
			mdata=loadcsv("c:/dropbox/SpatialRegression/WorldBank_Data/Series_Data_Panel.csv",.false.)
			mdata=mdata(:,2:)	! get rid of country code
			do i=1,capT
				if(i==1) then
					if(allocated(adata)) deallocate(adata)
					allocate(adata(ndata,size(mdata,2),capT))
				endif
				adata(:,:,i)=mdata((i-1)*(ndata+1)+2:i*(ndata+1),:)
			enddo
			init=.false.
		endif
		if(present(xind)) then
			if(xind>size(adata,2)-2) then
				m%n=-1
				return
			endif
		endif
		allocate(oinds(size(adata,1)),nco(size(adata,1)))
		l=0
		do 		
			do 
				l=l+1
				if(present(xind) .and. (l>10)) then
					m%n=-1
					return
				endif
				call rnsri(size(adata,2)-2,inds)
				call rnper(rper)
				inds=inds(rper)
				if(present(xind)) inds(1:2)=xind
				inds=inds+2
				m%n=0; m%m=0
				do i=1,size(adata,1)
					j=0
					do t=1,capT
						if(.not.any(adata(i,inds,t).ne.adata(i,inds,t))) j=j+1
					enddo
					if(j==0) cycle
					m%n=m%n+1
					oinds(m%n)=i
					m%m=m%m+j
					nco(m%n)=j		! record number of obs in cluster
				enddo
				if(m%n>99) exit
			enddo
			if(allocated(m%y)) deallocate(m%y,m%x,m%Z,m%s,m%nclust,m%clustid)
			allocate(m%y(m%m),m%x(m%m),m%Z(m%m,merge(capT-1,0,feflag)+k-1+m%n),m%s(2,m%n),m%nclust(m%n),m%clustid(m%n))
			if(allocated(b)) deallocate(b,bx)
			allocate(b(size(m%z,2)),bx(size(m%z,2)+1))
			m%z=0
			j=0
			do i=1,m%n
				m%nclust(i)=nco(i)
				m%clustid(i)=j+1
				do t=1,capT
					if(any(adata(oinds(i),inds,t).ne.adata(oinds(i),inds,t))) cycle
					j=j+1
					m%y(j)=adata(oinds(i),inds(1),t)
					m%x(j)=adata(oinds(i),inds(2),t)
					if(k>1) m%z(j,1:k-1)=adata(oinds(i),inds(3:2+k-1),t)
					m%z(j,k-1+i)=1
					if(feflag.and.t<capT) m%z(j,k-1+m%n+t)=1
				enddo
				m%s(:,i)=adata(oinds(i),1:2,1)
			enddo
			m%x=m%x/stddev(m%x)
			m%y=m%y/stddev(m%y)
			do i=1,k-1
				m%z(:,i)=m%z(:,i)/stddev(m%z(:,i))
			enddo
			call rlse(m%x,m%z,b,intcep=0,sse=sse,sst=sst)
			if(iercd()>0) cycle
			if(sse/sst>.01) exit
		enddo
		m%x=m%x-matmul(m%z,b)
		m%xZ=m%x.clr.m%Z
		call rlse(m%y,m%xZ,bx,intcep=0)
		if(iercd()>0) print *,"issue in pop y xZ reg",iercd()
		m%y=m%y-matmul(m%xZ,bx)
		if(allocated(m%dxsi)) deallocate(m%dxsi)
		allocate(m%dxsi(m%n))
		do j=1,m%n
			m%dxsi(j)=sqrt(sum(m%x(m%clustid(j):m%clustid(j)+m%nclust(j)-1)**2))
			if(m%dxsi(j)==0) m%dxsi(j)=1.0
		enddo
		disttype=1
	end subroutine

	subroutine setWB(m,k,xind)
		use rnsri_int
		use rnper_int
		use rlse_int
		integer		:: k
		integer, optional	:: xind
		type(mtype)	:: m
		logical, save	:: init=.true.
		real, allocatable, save	:: adata(:,:)

		integer, allocatable	:: cid(:)
		real, allocatable	:: s0(:,:)
		real		:: b(k),bx(k+1),sse,sst
		integer		:: i,j,inds(k+1),rper(k+1),l
		integer,allocatable	:: oinds(:)

		call erset(0,0,0)
		if(init) then
			adata=loadcsv("c:/dropbox/SpatialRegression/WorldBank_Data/Series_Data_2015_2005.csv",.true.)
			init=.false.
		endif
		if(present(xind)) then
			if(xind>size(adata,2)-2) then
				m%n=-1
				return
			endif
		endif
		allocate(oinds(size(adata,1)))
		l=0
		do			
			do
				l=l+1
				if(present(xind) .and. l>10 ) then
					m%n=-1
					return
				endif
				call rnsri(size(adata,2)-2,inds)
				call rnper(rper)
				inds=inds(rper)
				if(present(xind)) inds(1:2)=xind
				inds=inds+2
				m%m=0
				do i=1,size(adata,1)
					if(any(adata(i,inds).ne.adata(i,inds))) cycle
					m%m=m%m+1
					oinds(m%m)=i
				enddo
				if(m%m>99) exit
			enddo
			if(allocated(m%y)) deallocate(m%x,m%y,m%Z)
			allocate(m%y(m%m),m%x(m%m),m%Z(m%m,k))
			m%y=adata(oinds(1:m%m),inds(1))
			if(sum(m%y**2)==0) cycle
			m%y=m%y/stddev(m%y)
			m%x=adata(oinds(1:m%m),inds(2))
			if(sum(m%x**2)==0) cycle
			m%x=m%x/stddev(m%x)
			m%Z(:,1)=1
			do i=2,k
				m%Z(:,i)=adata(oinds(1:m%m),inds(i+1))
				if(sum(m%Z(:,i)**2)==0) cycle
				m%Z(:,i)=m%Z(:,i)/stddev(m%Z(:,i))
			enddo
			call rlse(m%x,m%z,b,intcep=0,sse=sse,sst=sst)
			if(iercd()>0) cycle
			if(sse/sst>0.01) exit
		enddo
		m%x=m%x-matmul(m%z,b)
		m%x=m%x/sqrt(sum(m%x**2)/m%m)
		m%xZ=m%x.clr.m%Z
		call rlse(m%y,m%xZ,bx,intcep=0)
		m%y=m%y-matmul(m%xZ,bx)
		m%s=transpose(adata(oinds(1:m%m),1:2))
		m%n=m%m
		if(allocated(m%nclust)) deallocate(m%nclust)
		allocate(m%nclust(m%n))
		m%nclust=1
		m%clustid=[(i,i=1,m%n)]
		if(allocated(m%dxsi)) deallocate(m%dxsi)
		allocate(m%dxsi(m%n))
		do j=1,m%n
			m%dxsi(j)=abs(m%x(j))
		enddo
		disttype=1
	end subroutine

	
	function getmaxcv(W,s,c0) result(val)
		real	:: W(:,:),s(:,:),c0,val
		integer, parameter	:: nc=30
		integer				:: i
		real, parameter	:: cfacs(nc)=exp([(5.0*(i-1.0)/(nc-1),i=1,nc)])
		real, allocatable	:: Sig(:,:),Om(:,:)
		real				:: cvlist(nc)
		integer	:: n,q
		
		n=size(s,2)
		q=size(W,2)
		val=getStudentcv(q,level)
		do i=1,nc
			Sig=getcSig(c0*cfacs(i),s)
			Om=matmul(transpose(W),matmul(Sig,W))
			cvlist(i)=getcv(Om,level)
		enddo
		val=max(val,maxval(cvlist))
	end function		
	
	function getetacv(m,c0) result(val)
		real	:: c0,val
		type(mtype)	:: m
		
		val=getcvfromOms(getOms(m%Weta,m%s,c0),size(m%Weta,2)-1,level)
	end function
	
	subroutine setmfromls(m,m0,ls)
		use rlse_int
		use iercd_int
		type(mtype)	:: m0,m
		integer		:: ls(:)
		integer		:: j,i,cl,mm,l
		real		:: beta(size(m0%Z,2)),sse,sst
		real		:: zss(size(m0%Z,2))
		integer, allocatable	:: izss(:)

		zss=0; mm=0
		do j=1,size(ls)
			cl=ls(j)
			zss=zss+sum(m0%z(m0%clustid(cl):m0%clustid(cl)+m0%nclust(cl)-1,:)**2,dim=1)
			mm=mm+m0%nclust(cl)
		enddo
		izss=pack([(i,i=1,size(m0%Z,2))],zss>0)

		if(.not. allocated(m%x)) then
			m%m=mm;m%n=size(ls)
			allocate(m%x(m%m),m%y(m%m),m%Z(m%m,size(izss)),m%xZ(m%m,size(izss)+1),m%s(2,m%n),m%dxsi(m%n),m%nclust(m%n),m%clustid(m%n))
		elseif (.not.(m%n==size(ls) .and. mm==m%m .and. size(m%Z,2)==size(izss))) then
			deallocate(m%x,m%Z,m%xZ,m%y,m%s,m%dxsi,m%nclust,m%clustid)
			m%m=mm;m%n=size(ls)
			allocate(m%x(m%m),m%y(m%m),m%Z(m%m,size(izss)),m%xZ(m%m,size(izss)+1),m%s(2,m%n),m%dxsi(m%n),m%nclust(m%n),m%clustid(m%n))
		endif
		l=1
		do j=1,size(ls)
			m%clustid(j)=l
			cl=ls(j)
			m%nclust(j)=m0%nclust(cl)
			m%s(:,j)=m0%s(:,cl)
			m%x(m%clustid(j):m%clustid(j)+m%nclust(j)-1)=m0%x(m0%clustid(cl):m0%clustid(cl)+m0%nclust(cl)-1)
			m%y(m%clustid(j):m%clustid(j)+m%nclust(j)-1)=m0%y(m0%clustid(cl):m0%clustid(cl)+m0%nclust(cl)-1)
			m%Z(m%clustid(j):m%clustid(j)+m%nclust(j)-1,:)=m0%Z(m0%clustid(cl):m0%clustid(cl)+m0%nclust(cl)-1,izss)
			l=l+m%nclust(j)
		enddo
		
		call rlse(m%x,m%Z,beta(1:size(izss)),intcep=0,sse=sse,sst=sst)
		if(sse<.01*sst .or. iercd()>0) then
			m%n=-1
			return
		endif
		m%x(:)=m%x-matmul(m%z,beta(1:size(izss)))
		m%xZ(:,:)=m%x.clr.m%Z
		do j=1,m%n
			m%dxsi(j)=sqrt(sum(m%x(m%clustid(j):m%clustid(j)+m%nclust(j)-1)**2))
			if(m%dxsi(j)==0) m%dxsi(j)=1
		enddo
	end subroutine
	
	function robust_avcx(m,avc0,etaflag) result(val)
		type(mtype)	:: m
		logical	:: etaflag
		real	:: avc0,val
		real, dimension(size(m%W,2),size(m%W,2)):: Om
		integer, parameter	:: navcgrid=20
		real	:: rpgrid(5,navcgrid),avcgrid(navcgrid)
		real	:: avc
		integer	:: p,j
	
		if(disttype==1) disttype=2
		avc=avc0
		do j=1,navcgrid
			avcgrid(j)=avc
			avc=avc/1.2
		enddo
!$omp parallel do collapse(2) private(Om,j)		
		do p=-1,3
			do j=1,navcgrid
				call setom(getavc_c(avcgrid(j),p,m%s),p,Om)
				rpgrid(p+2,j)=getrp(Om,1.0)
			enddo
		enddo
		val=maxval(rpgrid)
		if(disttype==2) disttype=1
	contains
	
		subroutine setom(c0,p,Om)
			real	:: c0,Om(:,:)
			integer	:: p
			if(etaflag) then
				Om=getOmfromW(m%Weta,getcpSig(c0,p,m%s))
			else
				if(m%m>m%n) then
					Om=getOmfromW(m%Wx,geteSig(m,getcpSig(c0,p,m%s)))
				else
					Om=getOmfromW(m%Wx,getcpSig(c0,p,m%s))
				endif
			endif
		end subroutine
	end function

end module