# Code primarily to read output of MCMC program and plot deformation 
# using posterior mean deformation.
diagnose <- T

# Input MCMC results --CHANGE PATH TO WHERE YOUR DATA FILEs ARE KEPT
#
directory.path <- "C:\\Documents and Settings\\Paul\\My Documents\\Mesa Air\\NOx\\"

# Input geographic coordinates and covariance matrix
# We are providing both latitude and longitude, and rectangular coordinates
# from a Lambert projection.

latlong.mx <- matrix(scan(paste(directory.path,"ca32.nox.latlong.txt",sep="")),
    byrow=T,ncol=2)
geog.scale <- matrix(scan(paste(directory.path,"ca32.nox.lamb100.txt",sep="")),
    byrow=T,ncol=2)
nsite <- nrow(geog.scale)
covmx <- matrix(scan(paste(directory.path,"ca32.nox.cov.txt",sep="")),ncol=nsite)

region.range <- c(range(pretty(range(latlong.mx[,2],na.rm=T))),
					range(pretty(range(latlong.mx[,1],na.rm=T))))

#state.outline <- map(xlim=region.range[1:2],ylim=region.range[3:4])
county.outline <- map("county",col=3,xlim=region.range[1:2],ylim=region.range[3:4])
points(latlong.mx[,2],latlong.mx[,1])

#THE FOLLOWING ASSUMES THAT THE MCMC OUTPUT FOLDER IS CALLED "Output"
	
newcoord <- read.table(paste(directory.path,"\\Output\\newcoord.txt",sep=""))
W <- read.table(paste(directory.path,"\\Output\\W.new",sep=""))
theta <- scan(paste(directory.path,"\\Output\\theta.new",sep=""))
thetat <- scan(paste(directory.path,"\\Output\\thetat.new",sep=""))
sigma2 <- scan(paste(directory.path,"\\Output\\sigma2.new",sep=""))
nugget <- scan(paste(directory.path,"\\Output\\nugget.new",sep=""))
nu <- read.table(paste(directory.path,"\\Output\\nu.new",sep=""))
mu <- scan(paste(directory.path,"\\Output\\mu.new",sep=""))
c <- read.table(paste(directory.path,"\\Output\\c.new",sep=""))
A <- read.table(paste(directory.path,"\\Output\\A.new",sep=""))

nmcmc <- length(theta)
nburn <- 100 #BURN-IN LENGTH (generally need to reset after examining plots below)

newcoord.array <- array(unlist(newcoord),c(nsite,nmcmc,2))
newcoord.array <- aperm(newcoord.array,c(1,3,2))
newcoord.mean <- apply(newcoord.array[,,nburn:nmcmc],c(1,2),mean)
plot(newcoord.mean)

# Compute posterior means after throwing away a burn-in period
sigma2.mean <- mean(sigma2[nburn:nmcmc])
theta.mean <- mean(theta[nburn:nmcmc])
thetat.mean <- mean(thetat[nburn:nmcmc])
nugget.mean <- mean(nugget[nburn:nmcmc])
nu.mean <- apply(nu[nburn:nmcmc,],2,mean)
mu.mean <- mean(mu[nburn:nmcmc])
W.array <- array(unlist(W),c(30,nmcmc,2))
W.array <- aperm(W.array,c(1,3,2))
W.mean <- apply(W.array[,,nburn:nmcmc],c(1,2),mean)

#When diagnose is set to T the following plots of traces of the MCMC runs 
# are looked at to decide whether to reset burnin, run the MCMC longer, etc

if (diagnose) { 
par(mfrow=c(2,2))
plot(theta,type="l") #,ylim=c(0,.002))
plot(thetat,type="l") #,ylim=c(0,.02))
plot(sigma2,type="l")
plot(nugget,type="l") #,ylim=c(0,.0002))
}

# Selected traces of W coefficients

if (diagnose) {
par(mfcol=c(2,3))
#for (i in seq(1,nsite,2)) {
for (i in 1:nsite) {
	for (j in 1:2) {
		plot(W.array[i,j,],type="l",main=paste("W",i,j))			
    }
	}
}

# Selected traces of coordinates
if (diagnose) {
par(mfcol=c(2,3))
#for (i in seq(1,nsite,2)) {
for (i in 1:nsite) {
	for (j in 1:2) {
		plot(newcoord.array[i,j,],type="l",main=paste("Newcoord",i,j))
		}
	}
}
 
#Now draw the average deformation 

Fdrawtps(t(geog.scale[,1:2]),t(newcoord.mean[,1:2]))


if (diagnose) {
# Plot sample of 4 deformations from posterior
samp4 <- sample(nburn:nmcmc,4)
par(mfrow=c(2,2))
for (i in samp4) {
	newcrdi <-newcoord.array[,,i]
	Fdrawtps(t(geog.scale[,1:2]),t(newcrdi),draw1=F)
		}
	}

# Look at correlations and/or covariances vs geographic and D-plane distance
cormx <- covmx/sqrt(diag(covmx) %*% t(diag(covmx)))

# Draw fitted (posterior mean) correlation function as if constant variance model fitted.

par(mfrow=c(1,2),pty="s")
plot(dist2full(dist(geog.scale)),cormx,
	xlab="G-plane Distance",ylab="Correlation")
plot(dist2full(dist(newcoord.mean)),cormx,
	xlab="D-plane Distance",ylab="Correlation")

dsort <- sort(unique(dist(newcoord.mean)))
nu.mean.mean <- mean(nu.mean)
lines(dsort,(nu.mean.mean*exp(-theta.mean*dsort))/(nu.mean.mean + nugget.mean))

par(mfrow=c(1,2),pty="s")
plot(dist2full(dist(geog.scale)),covmx,
	xlab="G-plane Distance",ylab="Covariance")
plot(dist2full(dist(newcoord.mean)),covmx,
	xlab="D-plane Distance",ylab="Covariance")


# Now play with the tool to link the geographic map with the 
# correlation or covariance vs distance plot


Flink.cor(geog.scale,cormx)

Flink.cor(geog.scale,covmx)

Flink.cor(newcoord.mean,cormx)

Flink.cor(newcoord.mean,covmx)