###############################################################################
###With tree ring, pollen and bore hole
###############################################################################
library(spam)

####range of year is 850-1999
Vol <- read.table("/fs/image/home/boli/Hockey/data/CSM_VolcanicForcing_global.txt",sep='',header=F)
Sol <- read.table("/fs/image/home/boli/Hockey/data/CSM_1.4_Solar.dat",sep='',header=F)
Co2 <- read.table("/fs/image/home/boli/Hockey/data/co2_850-1999.dat",sep='',header=F)

tempt  <- read.table("/fs/image/home/boli/Hockey/data/Instrum-model.dat",sep='', col.names=c('year','temperature'))

tempt.t <- read.table("/fs/image/home/boli/Hockey/data/CSM_b056.82_15TreeRings.dat",sep='',header=F)
tempt.p <- read.table("/fs/image/home/boli/Hockey/data/CSM_b056.82_10Pollen.dat",sep='',header=F)
tempt.b <- read.table("/fs/image/home/boli/Hockey/data/CSM_b056.82_5Borehole.dat",sep='',header=F)

tempt.t <- tempt.t + matrix(rnorm(nrow(tempt.t)*ncol(tempt.t)),nrow(tempt.t), ncol(tempt.t)) %*% diag(2*sd(tempt.t)) 
tempt.p <- tempt.p + matrix(rnorm(nrow(tempt.p)*ncol(tempt.p)),nrow(tempt.p), ncol(tempt.p)) %*% diag(2*sd(tempt.p)) 
tempt.b <- tempt.b + matrix(rnorm(nrow(tempt.b)*ncol(tempt.b)),nrow(tempt.b), ncol(tempt.b)) %*% diag(2*sd(tempt.b))
  
######################################################################
#                             Functions                              #
######################################################################

##Inverse of AR(2) covariance matrix, AR(1) is a special case of AR(2) with phi2=0
Sigma <- function(size, phi1, phi2, sigma2){	
  	gamma0 <- sigma2*(1-phi2)/((1+phi2)*(phi1+phi2-1)*(phi2-phi1-1))
  	rho <- rep(0,size) 
  	rho[1] <- 1
  	rho[2] <- phi1/(1-phi2)
  	for (i in 3:size){
    	rho[i] <- phi1*rho[i-1]+phi2*rho[i-2]
    }
  	Sigma <- matrix(rho[as.matrix(dist(1:size))+1], size, size)
  	#output <- list(gamma0=gamma0, corr=Sigma)
  	Sigma <- gamma0*Sigma
}

SigmaInv <- function(n, phi1, phi2, sigma2){
	SigmaInv <- diag(rep(1+phi1^2+phi2^2,n))
	diag(SigmaInv[-1, -n]) <- diag(SigmaInv[-n, -1]) <- rep(phi1*phi2-phi1, n-1)
	diag(SigmaInv[-c(1,2), -c(n-1,n)]) <- diag(SigmaInv[-c(n-1,n),-c(1,2)]) <- rep(-phi2,n-2)
	SigmaInv[1:2, 1:2] <- SigmaInv[c(n,n-1), c(n,n-1)] <- matrix(c(1,-phi1,-phi1,1+phi1^2))
	return(SigmaInv/sigma2)
}

#########AR(2) errors
error <- function(tt, phi1, phi2, rho){
	error <- rep(NA, tt+500)
	error[1:2] <- rho*rnorm(2)
	for(i in 1:(tt+500-2)){
		error[i+2] <- phi1*error[i+1]+phi2*error[i]+rho*rnorm(1)
	}
	error <- error[501:(tt+500)]
	return(error)
}

#####smooth matrix for tree ring and pollen, freq has to be even number and the actual smooth frequency is freq+1 
mat <- function(tt, freq){
	M <- matrix(0,tt-freq, tt)
	for (i in 1:(tt-freq)){
	    M[i,i:(i+freq)] <- rep(1, freq+1)/(freq+1)
	}
	M1 <- M2 <- matrix(0,freq/2,tt)
	for(i in 1:(freq/2)){
	      M1[i, 1:(2*i-1)] <- rep(1, 2*i-1)/(2*i-1)
	}	
	M2[(freq/2):1, tt:1] <- M1
	M <- rbind(M1, M, M2)
	return(M)	
}



##ind: indicator of proxies. ind=1 means borehole, other values mean other proxies or temperature process
##sigma2P: variance of risiduals
##VPinv: inverse of var-cov matrix 
##n: length of the proxy
##np: number of proxies

phi4ar1 <- function(ind, resid, phi1, xi1, sigma2P, VPinv, flag,n,np){
	phinew <- rnorm(1,phi1,sqrt(xi1))
	while(abs(phinew)>=1)   phinew <- rnorm(1,phi1,sqrt(xi1))

	if (ind==1){
		VPnew <- M.bnn%*%as.spam(Sigma(Ty, phinew, 0, sigma2P))%*%t(M.bnn)
		VPnewInv <- solve(0.5*(VPnew+t(VPnew)))
	}
	
	if (ind!=1){
		VPnew <- as.spam(Sigma(n, phinew, 0, sigma2P))
		VPnewInv <- as.spam(SigmaInv(n, phinew, 0, sigma2P))
	}
	
	tmp <- sum(diag(t(resid)%*%(VPnewInv-VPinv)%*%resid))

	lldif <- -0.5*(np*log(det(as.matrix(VPnew %*% VPinv)))+tmp)

	r <- runif(1) < exp(lldif)
	if (r) {
		phi1 <- phinew
		VPinv   <- VPnewInv
		flag <- flag+1
	}
	output <- list(phi1=phi1, VPinv=VPinv, flag=flag)
	return(output)
}

##ind: indicator of proxies. ind=1 means borehole, other values mean other proxies or temperature process
##sigma2P: variance of risiduals
##VPinv: inverse of var-cov matrix 
##n: length of the proxy
##np: number of proxies

phi4ar2 <- function(ind, resid, phi1, phi2, xi1, xi2, sigma2P, VPinv, flag,n,np){
	phi2new <- rnorm(1,phi2,sqrt(xi2))
	while(abs(phi2new)>=1)   phi2new <- rnorm(1,phi2,sqrt(xi2))
	phi1new <- rnorm(1,phi1, sqrt(xi1))
	while(abs(phi1new)>=(1-phi2new))   phi1new <- rnorm(1,phi1,sqrt(xi1))

	if (ind==1){
		VPnew <- M.bnn%*%as.spam(Sigma(Ty, phi1new, phi2new, sigma2P))%*%t(M.bnn)
		VPnewInv <- solve(0.5*(VPnew+t(VPnew)))
	}
	
	if (ind!=1){
		VPnew <- as.spam(Sigma(n, phi1new, phi2new, sigma2P))
		VPnewInv <- as.spam(SigmaInv(n, phi1new, phi2new, sigma2P))
	}

	tmp <- sum(diag(t(resid)%*%(VPnewInv-VPinv)%*%resid))

	tmpold <-  (1-phi2)*(pnorm(1, phi2new, sqrt(xi2))-pnorm(-1, phi2new, sqrt(xi2)))*(abs(pnorm(1-phi2, phi1new, sqrt(xi1))-pnorm(phi2-1, phi1new, sqrt(xi1))))

	tmpnew <- (1-phi2new)*(pnorm(1, phi2, sqrt(xi2))-pnorm(-1, phi2, sqrt(xi2)))*(abs(pnorm(1-phi2new, phi1, sqrt(xi1))-pnorm(phi2new-1, phi1, sqrt(xi1))))

	phicontrib <- log(tmpnew/tmpold)

	lldif <- -0.5*(np*log(det(as.matrix(VPnew %*% VPinv)))+tmp) + phicontrib

	r <- runif(1) < exp(lldif)
	if (r) {
		phi1 <- phi1new
		phi2 <- phi2new
		#VP <- VPnew
		VPinv   <- VPnewInv
		flag <- flag+1
	}
	output <- list(phi1=phi1, phi2=phi2, VPinv=VPinv, flag=flag)
	return(output)
}

######################################################################
#                         General Parameters                         #
######################################################################
	
Ty  <- 1150
Ty1 <- 1000
Ty2 <- 150

phi1 <- 0.2  
phi2 <- 0.4
rho <- 0.1
uu  <- 3   ##how many eigenvectors are selected
###################################################################
####range of year is 850-1999
tempt.t <- scale(tempt.t)
tempt.p <- scale(tempt.p)
tempt.b <- scale(tempt.b)

Vol <- scale(Vol$V2, center=F)
Sol <- scale(Sol$V2)
Co2 <- scale(Co2$V2)

tempt$temperature <- scale(tempt$temperature)

tempt2 <- tempt$temp[tempt$year>1849]

######################################################################
#                         Generate Data                              #
######################################################################
###############transform matrix for NH mean temperature###############
M.tree <- as.spam(diag(Ty)-mat(Ty, freq.tree))
#M.poll <- mat(Ty, freq.poll[1])-mat(Ty, freq.poll[2])
M.poll <- mat(Ty, freq.poll[1])
M.poll <- as.spam(M.poll[poll.ind,])
np <- length(poll.ind)

#####################Borehole forward algorithm#####################
POM= 0   ##0.5-0.7 degree C
tdiff= 1 #Enter thermal diffusivity (10e-6 m2/s)
tdiff=tdiff*1e-6
tdiff=tdiff*31536000

depth=  seq(0,500,by=5) #Enter file containing depths (m)
nz=length(depth)

#################forward matrix for NH mean temperature ##################
year <- 850:1999
yref=year[Ty]+1   ###The temperature is assumed to be from 1981

nsat=Ty
ybp=yref-year

tmp1 <- matrix(rep(ybp, nz), byrow=TRUE, nrow=nz)
tmp2 <- matrix(rep(depth, nsat), byrow=F, ncol=nsat)
tmp3 <- tmp2/sqrt(4*tdiff*tmp1)
erfc <- 2-2*pnorm(tmp3,sd=1/sqrt(2))
coeffM <- cbind(rep(0, nsat), diag(nsat))+ cbind(-diag(nsat), rep(0, nsat))

M.bore <- erfc %*% coeffM

M.b1  <- M.bore[,1]
M.bnn <- M.bore[,-1]

#############################################################################
### We can define different b0 and b1 for different proxies
proxb0 <- 0
proxb1 <- 1

p.tree <- proxb0 + proxb1*M.tree%*%as.matrix(tempt.t)
p.poll <- proxb0 + proxb1*M.poll%*%as.matrix(tempt.p)
p.bore <- proxb0 + proxb1*(POM*M.b1+M.bnn%*%as.matrix(tempt.b))

nTr <- ncol(p.tree)
nPl <- ncol(p.poll) 
nBr <- ncol(p.bore)

p.tree <- scale(p.tree)
p.poll <- scale(p.poll)
p.bore <- scale(p.bore)
###nz: length of depth
######################################################################
#                         Initial Value                              #
######################################################################

### Ideal proxies without dating or measurement error

P <- p.poll
B <- p.bore
V <- Vol

### Coefficients in data model
beta01 <- as.matrix(rnorm(nTr, mean=mubetap[1], sd=sqrt(tsigmaP0)))
beta11 <- as.matrix(abs(rnorm(nTr, mean=mubetap[2], sd=sqrt(tsigmaP1))))

beta02 <- as.matrix(rnorm(nPl, mean=mubetap[1], sd=sqrt(tsigmaP0)))
beta12 <- as.matrix(abs(rnorm(nPl, mean=mubetap[2], sd=sqrt(tsigmaP1))))

beta03 <- rnorm(nBr, mean=mubetap[1], sd=sqrt(tsigmaP0))
beta13 <- abs(rnorm(nBr, mean=mubetap[2], sd=sqrt(tsigmaP1)))
### Coefficients in Temperauter process
beta0 <- rnorm(1,mean=mubetat[1], sd=sqrt(tsigmaT))
beta123 <- abs(rnorm(3,mean=mubetat[2:4], sd=sqrt(tsigmaT)))
betat <- c(beta0, beta123)
### varince parameters 
sigma2P1 <- 1/rgamma(1,qP,,rP)
sigma2P2 <- 1/rgamma(1,qP,,rP)
sigma2P3 <- 1/rgamma(1,qP,,rP)
sigma2P4 <- 1/rgamma(1,qP,,rP)
sigma2P5 <- 1/rgamma(1,qP,,rP)
sigma2P6 <- 1/rgamma(1,qP,,rP)
 
sigma2T <- 1/rgamma(1,qT,,rT)

### residual correlation parameters
## for tree rings
phi2r <- runif(1, -1, 1)
phi1r <- runif(1, -(1-phi2r), (1-phi2r))

phi2p <- runif(1, -1, 1)
phi1p <- runif(1, -(1-phi2p), (1-phi2p)) # for pollen

phi2b <- runif(1, -1, 1)
phi1b <- runif(1, -(1-phi2b), (1-phi2b))

phi2t <- runif(1, -1, 1)
phi1t <- runif(1, -(1-phi2t), (1-phi2t))

decom <- svd(M.bnn)
M.bnn <- t(decom$u[,1:uu])%*%M.bnn
B <- p.bore <- t(decom$u[,1:uu])%*%p.bore

VPinv.r <- as.spam(SigmaInv(Ty, phi1r, phi2r, sigma2P3))
VPinv.p <- as.spam(SigmaInv(np, phi1p, phi2p, sigma2P4))
#VP.b <- as.spam(Sigma(Ty, phi1b, phi2b, sigma2P5))
VP.b <- sigma2P5*diag(Ty)

cov.b <- solve(M.bnn%*%VP.b%*%t(M.bnn))

VTinv   <- as.spam(SigmaInv(Ty, phi1t, phi2t, sigma2T))