###################################################
### chunk number 4: periodicSeasonByRegion
###################################################
rm(list = ls(all = TRUE))

library(NMMAPSdata)
library(tlnise)
library(tsModelSpec)
source(file.path("data", "model-fitting.R"))
source(file.path("data", "utils.R"))

exclude <- c("anch", "hono")

## Load the results
load(file.path("results", "seasonal.periodic.lag.012.dfSeas.1.rda"))
r <- results[[2]][setdiff(names(results[[2]]), exclude)]  ## Lag 1

## Get region indicators for each city
data(cities)
cityRegion <- with(cities, region[match(names(r), city)])
regionNames <- sort(unique(cityRegion))
regionInd <- 1 * outer(cityRegion, regionNames, "==")
colnames(regionInd) <- regionNames

ndays <- 365
x <- 1:ndays
B <- harmonic(x, 1, ndays, intercept = TRUE)

## Do pooling
bc <- extractBetaCov(r, pollutant = "pm10")
initTLNise()
g <- tlnise(bc$beta, bc$cov, regionInd, prnt = FALSE, labelY = names(r),
            intercept = FALSE, seed = 54321)
pooledRegion <- matrix(g$gamma[,1], byrow = TRUE, ncol = ncol(regionInd),
                       nrow = nrow(g$theta))
pooledRegionSD <- matrix(g$gamma[,2], byrow = TRUE, ncol = ncol(regionInd),
                         nrow = nrow(g$theta))
b <- B %*% pooledRegion
zero <- matrix(0, nrow = ndays, ncol = 7 * 3)
curveRegion <- lapply(1:ncol(pooledRegion), function(i) {    
    B <- zero
    ## This is weird
    B[, seq(i, 7 * 3, 7)] <- harmonic(x, 1, ndays, intercept = TRUE)
    s <- sqrt(diag(B %*% g$Dgamma %*% t(B)))
    cbind(b[,i], b[,i] - 2*s, b[,i] + 2*s)
})

## Overall estimate
initTLNise()
g <- tlnise(bc$beta, bc$cov, rep(1, NROW(bc$beta)), seed = 12345, prnt = FALSE)
V <- diag(B %*% g$Dgamma %*% t(B))
b <- B %*% g$gamma[,1]
curve <- cbind(b, b - 2*sqrt(V), b + 2*sqrt(V))

## Plot
regionFullNames <- c("Industrial Midwest", "Northeast", "Northwest",
                     "Southern California", "Southeast", "Southwest",
                     "Upper Midwest")
Y <- rbind(do.call("rbind", curveRegion), curve) * 1000
conf <- Y[,2:3]
rng <- range(Y) + c(-1, 1) * .05
y <- Y[,1]
nRegions <- length(curveRegion) + 1
x <- rep(1:ndays, nRegions)
g <- ordered(rep(c(regionFullNames, "All Regions"), each = ndays),
             levels = c(regionFullNames, "All Regions"))

library(lattice)
## trellis.device(height = 5, width = 8, color = FALSE)
p <- xyplot(y ~ x | g, as.table = TRUE, subscripts = TRUE, type = "l",
            panel = function(x, y, subscripts, ...) {
                llines(x, Y[subscripts, 2], lty = 3, lwd = 2)
                llines(x, Y[subscripts, 3], lty = 3, lwd = 2)
                panel.xyplot(x, y, ...)
                panel.abline(h = 0, lty = 2)
            },
            strip = strip.custom(bg = 0),
            ylim = rng, layout = c(4, 2),
            ylab = list(expression(paste("% increase in mortality with ",
                "10-", mu, "g/", m^3," increase in ", PM[10])), cex = 0.9),
            xlab = "Day in year",
            scales = list(alternating = 1, tck = c(0.8, 0))
            )
print(p)