###############################################################################
## Fit seasonal models
## Copyright (C) 2004, Roger D. Peng <rpeng@jhsph.edu>
##     
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 2 of the License, or
## (at your option) any later version.
## 
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
## 
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
###############################################################################

## beta: vector of length N, or N x p matrix.
## cov: a vector of length N or a p x p x N array.
## Return value????
poolCoef <- function(b, cov = NULL, w = NULL, 
                     method = c("tlnise", "lme", "fixed"),
                     extractors = NULL, ...) {
    method <- match.arg(method)

    if(is.null(cov) && !is.list(b))
        stop(sQuote("b"), " must be list if ", sQuote("cov"),
             " is ", sQuote("NULL"))
    if(is.list(b)) {
        cov <- b$cov
        b <- b$beta
    }
    if(is.null(w))
        w <- rep(1, NROW(b))
    if(method == "tlnise") {
        stopifnot(exists("tlnise"))
        fit <- tlnise(b, cov, w, prnt = FALSE, ...)
        if(is.null(extractors))
            pooled <- drop(fit$gamma[, c("est", "se")])
    }
    else if(method == "fixed") {
        stopifnot(NCOL(b) == 1)
        fit <- lm(b ~ w, weights = 1 / cov)
        if(is.null(extractors))
            pooled <- summary(fit)$coefficients[1, 1:2]
    }   
    else if(method == "lme") 
        stop("Not implemented yet")
    else
        stop("Invalid ", sQuote("method"), " specified")
    if(!is.null(extractors)) 
        pooled <- lapply(extractors, function(f) f(fit))
    pooled
}

coefSeasonal <- function(results, pollutant, method = "factor2",
                         Seasons = c("Winter", "Spring", "Summer", "Fall")) {
    if(method == "factor2") {
        betacov <- extractBetaCov(results, pollutant)
        coefmat <- poolCoef(betacov$beta, betacov$cov)
        rownames(coefmat) <- Seasons
        colnames(coefmat) <- c("Estimate", "Std. Error")
    }
    else
        stop("Invalid ", sQuote("method"), " specified")
    coefmat
}


## Return a n x 2 matrix where n is the number of seasons.  The first
## column is the pooled beta coefficient and the second column is the
## standard error.  `results' is a list of length equal to the number
## of cities.
extractBetaCov <- function(results, pollutant) {
    use <- !sapply(results, inherits, what = "try-error")
    results <- results[use]
    covlist <- lapply(results, function(x) {
        stopifnot(!is.null(x$cov))
        i <- grep(pollutant, rownames(x$cov))
        if(length(i) == 0)
            stop(sQuote("grep"), " failed: ", sQuote(pollutant),
                 " did not match any row names")
        x$cov[i, i]
    })
    cov <- array(unlist(covlist), c(dim(covlist[[1]]), length(covlist)))

    beta <- t(sapply(results, function(x) {
        stopifnot(!is.null(x$coefficients))
        i <- grep(pollutant, rownames(x$coefficients))
        if(length(i) == 0)
            stop(sQuote("grep"), " failed: ", sQuote(pollutant),
                 " did not match any coefficient names")
        x$coefficients[i, 1]
    }))
    if(NCOL(cov) == 1)
        cov <- as.vector(cov)
    list(beta = drop(beta), cov = cov, lcov = covlist)
}


LouisFormat <- function(x, type = c("stderr", "confint"), digits = 2) {
    type <- match.arg(type)

    if(!is.matrix(x))
        x <- as.matrix(x)

    switch(type,
           stderr = {
               stopifnot(ncol(x) == 2)
               minus <- ifelse(x[,1] < 0, "-", "")
               
               r <- paste("$",
                          formatC(abs(x[,1]), digits = digits, format = "f"),
                          "_", "{(",
                          formatC(x[,2], digits = digits, format = "f"),
                          ")}", "$", sep = "")
               paste(minus, r, sep = "")
           },
           confint = {
               stopifnot(ncol(x) == 3)  ## est, lo, hi
               minus <- ifelse(x[,1] < 0, "-", "")
               r <- paste("$", formatC(abs(x[,1]), digits = digits,
                                       format = "f"),
                          "_", "{(", formatC(x[,2], digits = digits,
                                             format = "f"),
                          ",\\,", formatC(x[,3], digits = digits,
                                          format = "f"),
                          ")}", "$", sep = "")
               paste(minus, r, sep = "")
           })
}