### data and setup
#rm(list = ls())
source("datasetup.r")
###

### prior 1 prior distribution 
alpha<-1
mu0<-eY0
V0<-  cor(Y)*sqrt( outer(diag(CY0),diag(CY0)) )
k0<- n0/(alpha+1)
nu0<-n0
S0<- nu0*V0
theta0.p1<-list(mu0=mu0,k0=k0,S0=S0,nu0=nu0)

spri<-2500

PSI<-NULL
cat("1a: sampling from the prior pi0 \n")
for(s in 1:spri)
{
  tmp<-rpsi.0(1,theta0.p1,alpha)   # generate functional from DPM
  PSI<-rbind(PSI,tmp)
  if(s%%25==0) { cat(s/spri,"\n") }
}

vpsi<-apply(PSI,2,var)
b<-apply(1/exp(PSI[,3:4]),2,mean)/apply(1/exp(PSI[,3:4]),2,var)
a<-apply(1/exp(PSI[,3:4]),2,mean)*b
theta1<-list(em=eYN,Vm=diag(vpsi[1:2]), s20=b/a, nu0=2*a)

### prior 3 (parametric) prior distribution 
PSIPRI<-cbind(rnorm(spri,mu0[1],sqrt(vpsi[1])),
           rnorm(spri,mu0[2],sqrt(vpsi[2])),
           log(1/rgamma(spri,a[1],b[1])),
           log(1/rgamma(spri,a[2],b[2]))  )


### for comparing p1 and p0
alpha<-1
theta0<-list(mu0=eYn,k0=1/10,S0=CYn,nu0=p+2)  # from prior2
PSI0<- NULL
cat("1b: sampling from the prior pi1 \n")
for(s in 1:spri)
{
  tmp<-rpsi.0(1,theta0,alpha)   # generate functional from DPM
  PSI0<-rbind(PSI0,tmp)
  if(s%%25==0) { cat(s/spri,"\n") }
}

p0pars<-p0fit(PSI0)


### starting values
n<-dim(Y)[1]
G<-1:n
ypsi<-rypsi.YG(Y,G,alpha,theta0) ; psi<-ypsi$psi ; ypp<-ypsi$y
fast.approx<-FALSE ; use.p1<-TRUE
###


### MCMC
cat("2: sampling from the posterior \n")
S<-25000 ; ALPHA<-LPY0<-YPP<-PSIPOST<-NULL ; acs<-0 ; plotres<-TRUE
for(s in 1:S)
{

  ## update 	
  for(i in sample(1:n))
  {  
    G1<-G
    G1[i]<-NA
    G1[(1:n)[-i]]<-match(G1[-i],unique(G1[-i]))
    K<-max(G1,na.rm=TRUE)
    G1[i]<-K+1
    lpG<-rep(NA,K+1)
    for(k in 1:K)
    { 
      Y0<-Y[G1==k,,drop=FALSE] ; Y1<-rbind(Y0,Y[i,])
      lpG[k]<-log(sum(G1==k))+lpy.marg(Y1,theta0)-lpy.marg(Y0,theta0)
    }
    lpG[K+1]<-log(alpha) + lpy.marg(Y[i,,drop=FALSE],theta0)
    G1[i]<-sample(1:(K+1),1,prob=exp(lpG-max(lpG)))
    ypsi<-rypsi.YG(Y,G1,alpha,theta0) ; psi1<-ypsi$psi ; ypp1<-ypsi$y  
    lhr<-0
    if(use.p1)
    {
      lhr<-lhr+(lppsi.p1(psi1,theta1) -  lppsi.p1(psi,theta1))  
      if(!fast.approx){lhr<-lhr + (lppsi.p0(psi1,p0pars)-lppsi.p0(psi,p0pars))}
    }     
    if(log(runif(1))<lhr) { G<-G1 ; psi<-psi1 ; ypp<-ypp1 ; acs<-acs+1 }
  }  
  G<-match(G,unique(G)) 
  ## 

  ## output
  if(s%%10==0)
  { 
    LPY0<-c(LPY0,lpY.G(Y,G,theta0))
    YPP<-rbind(YPP,ypp) 
    PSIPOST<-rbind(PSIPOST,psi) 
    cat(s,acs/(s*n)," ",sort(table(G),decreasing=TRUE),"\n")  
  
    if(plotres) 
    {
      par(mfrow=c(3,2),mar=c(3,3,1,1),mgp=c(1.5,.5,0)) 
      matplot(PSIPOST) ; abline(h=apply(PSIPOST,2,median),col=1:ncol(PSIPOST))
      plot(YPP[,1],YPP[,2],pch=16,cex=.5,col="gray",
            xlim=range(c(YPP[,1],Y[,1])),ylim=range(c(YPP[,2],Y[,2])) )
      if(dim(YPP)[1]>5) {plot2dd(YPP,lwd=2,col="gray",add=TRUE) }
      points(Y[,1],Y[,2],col=rainbow(max(G))[G],pch=16)  
      for(j in 1:length(psi))
      { 
        hist(PSIPOST[,j] ,main="") 
        abline(v=psi.true[j],col="green") 
        abline(v=mean(PSIPOST[,j]),col="blue")
      }
    }
  }
}  
###

dput(list(LPY0=LPY0,YPP=YPP,PSIPOST=PSIPOST,PSIPRI=PSIPRI),"results.prior3")