###
### file: heatmap.R
### author: Ali Tofigh
###
### Functions for displaying heatmaps. Instead of using one giant heatmap
### function that tries to do everything, the common tasks have been divided
### into several functions. It is assumed that the user will use 'layout()' to
### divide the graphic output into segments and will provide space for column-
### and rownames himself. A good example of how to use these functions is
### provided in experiment 01 in project strepi.
###
### The function heatmap.simple() is also provided to support quick full
### featured figures.

library(colorspace)
library(grid)
library(gridBase)

grid.newpage.loop <- function(...) {
    plot.new()
    require(grid)
    grid.clip(...)
}

### heatmap.color.scheme()
###
### Computes break points and color vector for heatmaps. Can also plot
### the HCL space for easier customization of colors.
###
### Arguments:
###     col.palette     either a palette name or index, or a full
###                     specification of an hcl palette. See the
###                     'supported.palettes' definition in the code for
###                     examples.
###     low.breaks
###     high.breaks     breakpoints for the low and high colors. If
###                     there is a gap between these vectors,
###                     'masked.color' will be used for the gap.
###     masked.color    the color to be used for the gap between the
###                     low and high colors.
###     plot            if TRUE, will plot the HCL color space for the
###                     chosen color palette.
###     gp              graphical parameters passed to grid.points and
###                     grid.lines when 'plot' is TRUE.
###
### Returns: a list with three fields:
###     breaks      vector with break points, size is one more that 'col'
###     col         vector with colors, size one less than 'breaks'
###     masked.idx  idx into 'col' and 'breaks' pointing to the color and
###                 the left break point of the masked region (if such a
###                 region exists).
heatmap.color.scheme <- function(low.breaks=seq(-3, 0, length.out=21),
                                 high.breaks=seq(0, 3, length.out=21),
                                 masked.color="black",
                                 col.palette=c("blue.red"),
                                 plot=FALSE, gp=gpar())
{
    library(colorspace)

    stopifnot(length(low.breaks) > 0 || length(high.breaks) > 0)
    stopifnot(length(low.breaks) == 0 || is.numeric(low.breaks) && length(low.breaks) >= 2)
    stopifnot(length(high.breaks) == 0 || is.numeric(high.breaks) && length(high.breaks) >= 2)
    stopifnot(length(low.breaks) == 0 || all(diff(low.breaks) > 0))
    stopifnot(length(high.breaks) == 0 || all(diff(high.breaks) > 0))
    stopifnot(length(low.breaks) == 0 || length(high.breaks) == 0 || tail(low.breaks, 1) <= head(high.breaks, 1))

    supported.palettes <- list("blue.red"=list(h=c(260, 0), c=c(0, 80), l=c(85, 30), power=1.5),
                               "blue.red.2"=list(h=c(260, 0), c=c(0, 100), l=c(90, 50), power=1.5),
                               "blue.orange.1"=list(h=c(240, 60), c=c(0, 77), l=c(98, 72), power=1.5))

    if (is.character(col.palette) || is.numeric(col.palette))
        pal <- supported.palettes[[col.palette]]
    else
        pal <- col.palette

    if (!c("h", "c", "l", "power") %in% names(pal) ||
        !is.numeric(pal$h) || !is.numeric(pal$c) || !is.numeric(pal$l) || !is.numeric(pal$power) ||
        length(pal$h) != 2 || length(pal$c) != 2 || length(pal$l) != 2 || length(pal$power) != 1)
        stop("invalid 'col.palette' specification")

    color.points <- function(b) {
        if (length(b) == 2)
        {
            ret <- 1
        }
        else if (length(b) == 3)
        {
            ret <- c(0, 1)
        }
        else
        {
            mid.b <- b[-c(1, length(b))]
            ret <- c(b[1], mid.b[-length(mid.b)] + diff(mid.b)/2, b[length(b)])
            ret <- ret - ret[1]
            ret <- ret / tail(ret, 1)
        }
        ret
    }

    if (length(low.breaks) > 0)
    {
        low.color.points <- color.points(low.breaks)^pal$power
        low.L <- pal$l[1] + diff(pal$l) * low.color.points
        low.C <- pal$c[1] + diff(pal$c) * low.color.points
        low.colors <- rev(hex(polarLUV(L=low.L, C=low.C, H=pal$h[1]), fixup=TRUE))
    }

    if (length(high.breaks) > 0)
    {
        high.color.points <- color.points(high.breaks)^pal$power
        high.L <- pal$l[1] + diff(pal$l) * high.color.points
        high.C <- pal$c[1] + diff(pal$c) * high.color.points
        high.colors <- hex(polarLUV(L=high.L, C=high.C, H=pal$h[2]), fixup=TRUE)
    }

    scheme <- list()
    if (length(low.breaks) == 0)
    {
        scheme$breaks <- high.breaks
        scheme$col <- c(high.colors)
    }
    else if (length(high.breaks) == 0)
    {
        scheme$breaks <- low.breaks
        scheme$col <- c(low.colors)

    }
    else if (tail(low.breaks, 1) == head(high.breaks, 1))
    {
        scheme$breaks <- c(low.breaks, high.breaks[-1])
        scheme$col <- c(low.colors, high.colors)
    }
    else
    {
        scheme$breaks <- c(low.breaks, high.breaks)
        scheme$col <- c(low.colors, masked.color, high.colors)
        scheme$masked.idx <- length(low.colors) + 1
    }

    if (isTRUE(plot))
    {
        plot.h <- function(h, pixels)
        {
            c <- seq(0,100,length=pixels)[-1];
            l <- seq(0,100,length=pixels)[-1];
            xy <- expand.grid(c=c, l=l);
            rect.col <- hcl(h, xy$c, xy$l, fixup=FALSE)
            grid.rect(xy$c/100, xy$l/100, width=1/length(c), height=1/length(l),
                      gp=gpar(fill=rect.col, col=rect.col, linejoin="mitre", lwd=0.01), just=c("right", "top"),
                      default.unit="native")
            grid.text(h, 0.8, 0.1, default.units="native")
        }

        pushViewport(viewport(layout=grid.layout(3, 2, heights=c(4, 1, 4))))

        pushViewport(viewport(layout.pos.row=1, layout.pos.col=1, xscale=c(1, 0), clip="off"))
        plot.h(pal$h[1], 100)
        grid.lines(pal$c/100, pal$l/100, default.unit="native", gp=gp)
        grid.points(low.C / 100, low.L / 100, default.units="native", gp=gp)
        popViewport()

        pushViewport(viewport(layout.pos.row=1, layout.pos.col=2))
        plot.h(pal$h[2], 100)
        grid.lines(pal$c/100, pal$l/100, default.unit="native", gp=gp)
        grid.points(high.C / 100, high.L / 100, default.units="native", gp=gp)
        popViewport()

        pushViewport(viewport(layout.pos.row=2, layout.pos.col=1:2, xscale=range(scheme$breaks)))
        grid.rect()
        grid.rect(scheme$breaks[-1], 0, width=diff(scheme$breaks), height=1,
                  gp=gpar(fill=scheme$col, col=scheme$col, linejoin="mitre", lwd=0.01),
                  just=c("right", "bottom"),
                  default.units="native")
        popViewport()

        pushViewport(viewport(layout.pos.row=3, layout.pos.col=1:2, xscale=range(scheme$breaks)))
        pushViewport(viewport(layout=grid.layout(4, 9)))

        rc <- expand.grid(c=1:9, r=1:4)
        h <- seq(0, 359, 10)
        for (i in 1:length(h))
        {
            pushViewport(viewport(layout.pos.row=rc[i, "r"], layout.pos.col=rc[i, "c"]))
            plot.h(h[i], 80)
            popViewport()
        }

        popViewport()
        popViewport()

        popViewport()
    }

    return(scheme)
}

### heatmap.is.valid.color.scheme()
###
### internal function for checking color schemes.
heatmap.is.valid.color.scheme <- function(color.scheme)
{
    is.list(color.scheme) &&
    (c("breaks", "col") %in% names(color.scheme)) &&
    is.numeric(color.scheme$breaks) &&
    (length(color.scheme$breaks) == length(color.scheme$col) + 1) &&
    (is.null(color.scheme$masked.idx) || color.scheme$masked.idx %in% 1:length(color.scheme$col))
}

### heatmap.map()
###
### plots a heatmap of 'x' in the current viewport. Returns the color
### scheme used.
###
### Arguments:
###     x               a numeric matrix with the heatmap values
###     color.scheme    see the function heatmap.color.scheme() for details
###     na.color        color used for NA values in 'x'
###
### Returns:
###
###     The color scheme invisibly.
heatmap.map <- function(x, color.scheme=heatmap.color.scheme(), na.color = "grey")
{
    library("colorspace")

    if (length(dim(x)) != 2 || !is.numeric(x))
        stop("'x' must be a numeric matrix")

    bc <- color.scheme
    stopifnot(heatmap.is.valid.color.scheme(bc))

    ## make sure no values are outside the range ['min.val', 'max.val']
    min.val <- head(bc$breaks, 1)
    max.val <- tail(bc$breaks, 1)
    x[x < min.val] <- min.val
    x[x > max.val] <- max.val

    ## convert numbers to colors according to our breaks
    x.colors <- findInterval(x, bc$breaks, all.inside=TRUE)

    ## add na.color at the end of the color scheme
    bc$col <- c(bc$col, na.color)
    x.colors[is.na(x.colors)] <- length(bc$col)

    ## this viewport is set up so that origo is on the top left and the
    ## ranges of the "native" units are number of columns and rows of x.
    pushViewport(viewport(xscale=c(0, ncol(x)), yscale=c(nrow(x),0)))

    ## now draw the heatmap
    grid.rect(col(x), row(x), gp=gpar(col=bc$col[x.colors], fill=bc$col[x.colors], linejoin="mitre", lwd=0.01),
              width=1, height=1, just=c("right", "top"), default.units="native")

    popViewport()

    return(invisible(color.scheme))
}

### heatmap.key()
###
### plots a heatmap key in the current viewport.
###
### Arguments:
###     x               numeric matrix that was used to plot a heatmap
###     color.scheme    the color scheme used to plot the heatmap
###     key.min         should be less than the least break point of the
###                     color scheme
###     key.max         should be greater than the greates break point of the
###                     color scheme
###     key.line.color  the color used to draw the histogram
###
### Returns:
###     invisible(NULL)
heatmap.key <- function(x, color.scheme=heatmap.color.scheme(),
                        key.min = head(color.scheme$breaks, 1) - 1,
                        key.max = tail(color.scheme$breaks, 1) + 1,
                        key.line.color = "cyan")
{
    if (!is.numeric(x))
        stop("'x' must be a numeric matrix")
    stopifnot(heatmap.is.valid.color.scheme(color.scheme))

    bc <- color.scheme

    min.val <- head(bc$breaks, 1)
    max.val <- tail(bc$breaks, 1)
    key.min <- min(key.min, min.val)
    key.max <- max(key.max, max.val)

    ## add breakpoints and colors for the masked region, from min.val to
    ## key.min, and from max.val to key.max. This is so that we can get
    ## a good histogram showing in the key.
    break.width <- min(diff(bc$breaks))

    if (!is.null(bc$masked.idx)) # extend masked region
    {
        masked.color <- bc$col[bc$masked.idx]
        a <- bc$breaks[bc$masked.idx]
        b <- bc$breaks[bc$masked.idx + 1]
        n <- round((b - a) / break.width)
        n <- max(1, n)
        new.breaks <- seq(a, b, length.out = n + 1)[-c(1, n+1)]
        bc$breaks <- c(bc$breaks[1:bc$masked.idx],
                       new.breaks,
                       bc$breaks[(bc$masked.idx + 1):length(bc$breaks)])
        bc$col <- c(bc$col[1:bc$masked.idx],
                    rep(masked.color, n - 1),
                    bc$col[(bc$masked.idx + 1):length(bc$col)])
    }
    if (key.min < min.val) # extend region between min.val and key.min
    {
        n <- round((min.val - key.min) / break.width)
        n <- max(1, n)
        new.breaks <- seq(key.min, min.val, length.out = n + 1)[-(n+1)]
        bc$breaks <- c(new.breaks, bc$breaks)
        bc$col <- c(rep(bc$col[1], n), bc$col)
    }
    if (key.max > max.val) # extend region between max.val and key.max
    {
        n <- round((key.max - max.val) / break.width)
        n <- max(1, n)
        new.breaks <- seq(max.val, key.max, length.out = n + 1)[-1]
        bc$breaks <- c(bc$breaks, new.breaks)
        bc$col <- c(bc$col, rep(bc$col[length(bc$col)], n))
    }

    ## draw the key

    ## we will leave 30% of the area below the key histogram for drawing
    ## an axis, which needs about 2 lines of text.
    axis.text <- do.call(paste, as.list(grid.pretty(c(key.min, key.max))))
    height <- convertHeight(unit(2, "lines"), "points", valueOnly=T)/convertHeight(unit(1, "npc"), "points", valueOnly=T)
    axis.cex <- min((as.double(convertHeight(unit(height, "native"), "lines")) / 2),
                    (1/as.double(convertWidth(unit(1, "strwidth", axis.text), "npc"))))*0.8
    height <- height * axis.cex

    pushViewport(viewport(0.5, 0.9, width=0.8, height=1-height, xscale=range(bc$breaks), yscale=c(1, 0),
                          just=c("centre", "top")))
    grid.rect(bc$breaks[-1], 0,
              width=diff(bc$breaks), height=1, gp=gpar(col=NA, fill=bc$col, lty=0),
              default.units="native", just=c("right", "bottom"))

    if (!all(is.na(x)))
    {
        h <- hist(x, plot = FALSE,
                  breaks = c(min(x, na.rm=TRUE), bc$breaks, max(x, na.rm = TRUE)))
        pushViewport(viewport(xscale=c(key.min, key.max), yscale=c(0, max(h$counts)*1.05)))
        breaks <- h$breaks[-c(1, length(h$breaks))]
        counts <- h$counts[-c(1, length(h$counts))]
        grid.lines(rep(breaks, each=2), c(0, rep(counts, each=2), 0),
                   gp=gpar(col=key.line.color), default.units="native")
        grid.xaxis(gp=gpar(cex=axis.cex, lex=axis.cex))
        popViewport()
    }
    else
    {
        pushViewport(viewport(xscale=c(key.min, key.max)))
        grid.xaxis(gp=gpar(cex=axis.cex, lex=axis.cex))
        popViewport()
    }
    popViewport()

    return(invisible(NULL))
}

### heatmap.clinical()
###
### plots clinical variables as boxes in the current frame.  Assumes 'x'
### is a matrix with color information in each cell. Colors are usually
### given using character strings such as "white", "red", etc. See
### 'colors()' for a list of color names. NAs are plotted with color as
### given by 'na.color'.
###
### Arguments:
###     x           matrix with colors in each cell, or NAs. Will also
###                 accept data frames (with rows representing patients)
###     na.color    color to be used for missing values, i.e., NAs
###
### Returns:
###     NULL
heatmap.clinical <- function(x, na.color = "grey")
{
    if (length(dim(x)) != 2)
        stop("'x' must be a color matrix")

    if (is.data.frame(x))
        x <- t(as.matrix(x))

    x[is.na(x)] <- na.color

    pushViewport(viewport(xscale=c(0, ncol(x)), yscale=c(nrow(x), 0)))

    grid.rect(col(x), row(x), width=1, height=1, gp=gpar(col=x, fill=x, linejoin="mitre", lwd=0.01),
              just=c("right", "top"), default.units="native")
    grid.rect(ncol(x), 1:nrow(x), width=ncol(x), height=1, gp=gpar(col="black", fill=NA),
              just=c("right", "top"), default.units="native")

    popViewport()

    return(invisible(NULL))
}

### heatmap.ddr()
###
### A very flexible function for clustering. It allows multiple ways of
### specifying how the dendrogram is to be computed and it knows about
### the cor() and hclust() functions.
###
### Arguments:
###
###     exprs       a matrix-like object whose columns are to be
###                 clustered. May be omitted if a dendrogram can be
###                 obtained using only the other arguments.
###     dist.elem   specifies how to compute the distance matrix. May be
###                 omitted if a dendrogram can be obtained from
###                 'clust.elem' alone. Possible specifications are: An
###                 object that can be coerced to a distance matrix
###                 (using as.dist()), a function that takes 'exprs' and
###                 produces a distance matrix, or a character string
###                 specifying how a distance matrix is to be computed
###                 from 'exprs'. The acceptable character strings are
###                 "cor" and any method specification that can be
###                 passed on to the dist() function. If 'dist.elem' is
###                 "cor", then 1-correlation (between the columns of
###                 'exprs') is used as distance.
###     clust.elem  specifies how to compute the dendrogram. Possible
###                 specifications are: An object that can be coerced to
###                 "dendrogram" (using as.dendrogram()), a function
###                 that takes a distance matrix and produces an object
###                 that can be coerced to a dendrogram, a character
###                 string specifying how to compute, from a distance
###                 matrix, an object that can be coerced to a
###                 dendrogram, or a character string specifying how to
###                 compute a dendrogram from a distance matrix.
###     p           passed to the dist() function and used only when
###                 'dist.elem' is "minkowski" (see help on the dist()
###                 function for details).
heatmap.ddr <- function(exprs=NULL, dist.elem="cor", clust.elem="ward", p=2)
{
    if (class(clust.elem) != "dendrogram")
    {
        ret <- try(as.dendrogram(clust.elem), silent=TRUE)
        if (class(ret) == "try-error")
        {
            if (mode(dist.elem) == "numeric" && class(dist.elem) == "matrix")
            {
                dist.mat <- as.dist(dist.elem)
            }
            else if (is.null(exprs))
            {
                stop("heatmap.ddr: 'exprs' missing")
            }
            else if (mode(dist.elem) == "function")
            {
                dist.mat <- as.dist(dist.elem(exprs))
            }
            else if (mode(dist.elem) == "character" && length(dist.elem) == 1 && dist.elem == "cor")
            {
                cors <- cor(exprs, use="pairwise.complete.obs")
                cors[is.na(cors)] <- 0
                dist.mat <- as.dist(1 - cors)
            }
            else if (mode(dist.elem) == "character" && length(dist.elem) == 1)
            {
                dist.mat <- dist(t(exprs), method=dist.elem, p=p)
            }
            else
            {
                stop("heatmap.ddr: bad 'dist.elem'")
            }

            if (mode(clust.elem) == "function")
            {
                ret <- as.dendrogram(clust.elem(dist.mat))
            }
            else if (mode(clust.elem) == "character" && length(clust.elem) == 1)
            {
                ret <- as.dendrogram(hclust(dist.mat, method=clust.elem))
            }
            else
            {
                stop("heatmap.ddr: bad 'dist.elem'")
            }
        }
    }
    else
    {
        ret <- clust.elem
    }

    return(ret)
}

### heatmap.labels.cex()
###
### computes a cex value that can be used when plotting text in the
### current viewport. The assumption is that 'labels' is a character
### vector containing text strings that are meant to be drawn in the
### current viewport (either horizontally or vertically). This function
### will compute and return a cex value that can be passed to
### grid.text() such that the strings in 'labels' will fit in the
### current viewport when drawn on separate lines.
###
### Arguments:
###
###     labels
###
###         character vector
###
###     type
###
###         either "row.labels" or "col.labels" indicating whether the
###         labels are intended to be drawn horizontally or vertically.
###
### Returns:
###
###     a cex value to be used when drawing the strings in 'labels' in
###     the current viewport
heatmap.labels.cex <- function(labels, type=c("row.labels", "col.labels"))
{
    type <- match.arg(type)

    rw <- convertWidth(max(do.call(unit.c, lapply(labels, function(label) {unit(1, "strwidth", label)}))), "npc")
    rw <- convertWidth((rw + unit(2, "strwidth", "X")), "npc")
    rh <- convertHeight(unit(1, "strheight", "X"), "npc") * 1.8 * length(labels)
    if (type == "row.labels")
    {
        cex <- min(1/as.double(convertWidth(rh, "npc")), 1/as.double(rw))
    }
    else
    {
        rw <- convertUnit(rw, "npc", "x", "dimension", "y", "dimension")
        rh <- convertUnit(rh, "npc", "y", "dimension", "x", "dimension")
        cex <- min(1/as.double(convertWidth(rh, "npc")), 1/as.double(rw))
    }
    return(cex)
}

### heatmap.labels()
###
### plots each string in 'labels' on a separate line in the current
### viewport. The size of the text and spacing between strings are
### adjusted so that the text will fill the viewport as much as
### possible.
###
### Arguments:
###
###     labels
###
###         character vector with strings to be drawn in the current
###         viewport
###
###     type
###
###         either "row.labels" or "col.lables", determining if the text
###         is to be drawn horizontally or vertically.
###
###     just
###
###         the justification of the text.
###
###     rot
###
###         rotation of the text in angles.
###
### Returns:
###
###     invisible(NULL)
heatmap.labels <- function(labels, type=c("row.labels", "col.labels"), just=c("left", "right", "centre"), rot=0)
{
    type <- match.arg(type)
    just <- match.arg(just)

    if (type == "row.labels")
    {
        pushViewport(viewport(xscale=c(0, 1), yscale=c(1, 0)))
        if (just == "left")
            x <- unit(1, "strwidth", "X")
        else if (just == "right")
            x <- unit(1, "npc") - unit(1, "strwidth", "X")
        else
            x <- 0.5
        y <- ((1:length(labels)) - 0.5) /length(labels)
        grid.text(labels, x, y, default.units="native", just=just, rot = rot,
                  gp=gpar(cex=heatmap.labels.cex(labels, type=type)))
        popViewport()
    }
    else
    {
        pushViewport(viewport(xscale=c(0, 1), yscale=c(0, 1)))
        x <- ((1:length(labels)) - 0.5) /length(labels)
        if (just == "left")
            y <- unit(1, "strwidth", "X")
        else if (just == "right")
            y <- unit(1, "npc") - unit(1, "strwidth", "X")
        else
            y <- 0.5
        grid.text(labels, x, y, default.units="native", just=just, rot=90 + rot,
                  gp=gpar(cex=heatmap.labels.cex(labels, type=type)))
        popViewport()
    }
    invisible(NULL)
}

### heatmap.simple()
###
### Provides an easy way to produce full heatmap figures with the most
### common features such as dendrograms from clustering, row and column
### labels and a matrix of clinical variables. Many of the arguments
### will take on sensible defaults if omitted (as described below).
###
### this function makes extensive use of the grid package and mixes it
### with standard R graphics plots using gridBase. The figure is divided
### into regions (with grid.layout which is similar to base R's layout
### function) as directed by the layout.mat character matrix. The
### strings in layout.mat will determine where each component of the
### figure will be drawn. Other arguments are as described for the
### heatmap functions above.
###
### Arguments:
###     x                   matrix of values for the main heatmap.
###     layout.mat          a character matrix describing the layout of
###                         the figure and its parts. The figure will be
###                         divided into rows and columns corresponding
###                         to the rows and columns of layout.mat. Each
###                         entry should be a string describing what
###                         should be plotted in that region. The list
###                         of recognized entries are:
###
###                         heatmap
###                         key
###                         title
###                         row.labels.ljust
###                         row.labels.rjust
###                         col.labels.ljust
###                         col.labels.rjust
###                         row.ddr.left
###                         row.ddr.right
###                         col.ddr.up
###                         col.ddr.down
###                         clinical
###                         clinical.labels.ljust
###                         clinical.labels.rjust
###
###                         If layout.mat is omitted, the function will
###                         try to provide a sensible one depending on
###                         what other options have been supplied.
###     widths
###     heights             relative widths and heights for the
###                         layout. If 'layout.mat' is omitted, care should be taken
###                         when supplying these since they have to match
###                         the size of the default layout.mat, which can vary
###                         depending on other arguments.
###     color.scheme        a list defining the color scheme to be used for the
###                         heatmap. See the heatmap.color.scheme() function
###                         for details on the requirements for this variable
###     scale               for now only supports "row", "none", or a
###                         user-supplied function that will be applied
###                         to each row.
###     row.labels          If omitted, will default to rownames of x.
###     col.labels          If omitted, will default to colnames of x.
###     row.clust
###     row.clust.dist.fun
###     row.clust.fun
###     col.clust
###     col.clust.dist.fun
###     col.clust.fun       If row.clust is TRUE, then the rows will be
###                         clustered by the heatmap.ddr() function, to
###                         which row.clust.dist.fun and row.clust.fun
###                         are passed. See that function for details on
###                         these arguments. An analogous description
###                         applies to col.clust.*. Note that the expression
###                         matrix will still be reordered if you pass
###                         a dendrogram of your own (using order.dendrogram()).
###                         So make sure the dendrogram was derived from
###                         'exprs' and that 'exprs' has not been reordered.
###     clinical            a matrix of colors or a data frame to be
###                         passed to heatmap.clinical.
###     clinical.labels     If omitted, will default to rownames of
###                         'clinical'.
###     title               text string to be printed as a title
heatmap.simple <- function(exprs,
                           layout.mat=NULL,
                           widths=NULL,
                           heights=NULL,
                           color.scheme=heatmap.color.scheme(),
                           key.min=head(color.scheme$breaks, 1) - 1,
                           key.max=tail(color.scheme$breaks, 1) + 1,
                           na.color="grey",
                           scale="row",
                           key.line.color="cyan",
                           row.labels=NULL,
                           col.labels=NULL,
                           row.clust=TRUE,
                           row.clust.dist.fun="cor",
                           row.clust.fun="ward",
                           col.clust=TRUE,
                           col.clust.dist.fun="cor",
                           col.clust.fun="ward",
                           clinical=NULL,
                           clinical.labels=NULL,
                           clinical.na.color=na.color,
                           title=NULL)
{
    library(grid)
    library(gridBase)

    ## save all standard R graphical parameters
    old.par <- par(no.readonly=TRUE)

    ## if clinical is a vector or data.frame, then convert it to a matrix
    if (!is.null(clinical))
    {
        if (is.vector(clinical))
            clinical <- matrix(clinical, nrow=1)
        else if (is.data.frame(clinical))
            clinical <- t(as.matrix(clinical))
    }

    ## get default values for arguments
    if (is.null(row.labels))
    {
        row.labels <- rownames(exprs)
        if (is.null(row.labels))
            row.labels <- rep("", nrow(exprs))
    }
    if (is.null(col.labels))
    {
        col.labels <- colnames(exprs)
        if (is.null(col.labels))
            col.labels <- rep("", ncol(exprs))
    }
    if (is.null(clinical.labels) && !is.null(clinical))
        clinical.labels <- rownames(clinical)

    ## if layout.mat is omitted, try to come up with a good default layout
    if (is.null(layout.mat))
    {
        layout.mat <- matrix(c("key", "col.ddr.up", "",
                               "row.ddr.left", "heatmap", "row.labels.ljust",
                               "", "col.labels.rjust", ""), ncol=3, byrow=TRUE)
        if (!is.null(clinical))
        {
            layout.mat <- rbind(layout.mat, c("clinical.labels.rjust", "clinical", ""))
            layout.mat <- rbind(layout.mat, c("", "", ""))
        }
        if (!is.null(title))
        {
            layout.mat <- rbind(c("", "title", ""), layout.mat)
        }
        if (is.null(widths))
        {
            widths <- c(0.3, 1, 0.3)
        }
        if (is.null(heights))
        {
            heights <- rep(0.3, nrow(layout.mat))
            heights[which(layout.mat == "heatmap", arr.ind=TRUE)[1]] <- 1
            if (!is.null(clinical))
            {
                heights[length(heights) - 1] <- 0.03*nrow(clinical)
                heights[length(heights)] <- 0.03
            }
            if (!is.null(title))
            {
                heights[1] <- 0.15
            }
        }
        if (!isTRUE(row.clust))
        {
            layout.mat[grep("row.ddr", layout.mat)] <- ""
        }
        if (!isTRUE(col.clust))
        {
            layout.mat[grep("col.ddr", layout.mat)] <- ""
        }

        stopifnot(length(widths) == ncol(layout.mat))
        stopifnot(length(heights) == nrow(layout.mat))
    }

    ## make sure layout is character matrix
    mode(layout.mat) <- "character"
    layout.mat <- as.matrix(layout.mat)

    ## scale the rows (if indicated by caller). When using correlation
    ## as distance, we definitely need to scale the rows *before*
    ## clustering since the mean of the columns don't really mean
    ## anything (expressions of different genes are not comparable in
    ## microarray experiments). Are there situations where we might want
    ## to cluster before scaling?
    if (is.function(scale))
    {
        exprs <- t(apply(exprs, 1, scale))
    }
    else
    {
        stopifnot(length(scale) == 1 && scale %in% c("row", "none"))
        ## scale the matrix if given
        if (scale == "row")
            exprs <- t(scale(t(exprs)))
    }

    ## get dendrograms and reorder the data matrix
    row.ddr <- NULL
    if (isTRUE(row.clust))
    {
        row.ddr <- rev(heatmap.ddr(t(exprs), row.clust.dist.fun, row.clust.fun))
        exprs <- exprs[order.dendrogram(rev(row.ddr)), ]
        row.labels <- row.labels[order.dendrogram(rev(row.ddr))]
    }

    col.ddr <- NULL
    if (isTRUE(col.clust))
    {
        col.ddr <- heatmap.ddr(exprs, col.clust.dist.fun, col.clust.fun)
        exprs <- exprs[ , order.dendrogram(col.ddr)]
        col.labels <- col.labels[order.dendrogram(col.ddr)]
        if (!is.null(clinical))
            clinical <- clinical[, order.dendrogram(col.ddr), drop=FALSE]
    }

    ## create the top viewport
    the.layout <- grid.layout(nrow(layout.mat), ncol(layout.mat), widths=widths, heights=heights)
    top.vp <- viewport(layout=the.layout, name="heatmap.top.vp")
    pushViewport(top.vp)

    ## get all unique strings in 'layout.mat'
    elements <- unique(as.vector(layout.mat))
    elements <- elements[elements != "" & !is.na(elements)]

    ## for the known strings in layout.mat, seek the named viewport and
    ## draw accordingly
    elem <- "heatmap"
    if (elem %in% elements)
    {
        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        heatmap.map(exprs, color.scheme=color.scheme, na.color=na.color)
        upViewport()
    }
    elem <- "key"
    if (elem %in% elements)
    {
        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        heatmap.key(exprs, color.scheme=color.scheme, key.min=key.min, key.max=key.max, key.line.color=key.line.color)
        upViewport()
    }
    elem <- "row.labels.ljust"
    if (elem %in% elements)
    {
        if (length(row.labels) != nrow(exprs))
            warning("the number of row labels is not equal to the number of rows of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        row.labels[is.na(row.labels)] <- ""
        heatmap.labels(row.labels, type="row.labels", just="left")
        upViewport()
    }
    elem <- "row.labels.rjust"
    if (elem %in% elements)
    {
        if (length(row.labels) != nrow(exprs))
            warning("the number of row labels is not equal to the number of rows of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        row.labels[is.na(row.labels)] <- ""
        heatmap.labels(row.labels, type="row.labels", just="right")
        upViewport()
    }
    elem <- "col.labels.ljust"
    if (elem %in% elements)
    {
        if (length(col.labels) != ncol(exprs))
            warning("the number of column labels is not equal to the number of columns of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        col.labels[is.na(col.labels)] <- ""
        heatmap.labels(col.labels, type="col.labels", just="left")
        upViewport()
    }
    elem <- "col.labels.rjust"
    if (elem %in% elements)
    {
        if (length(col.labels) != ncol(exprs))
            warning("the number of column labels is not equal to the number of columns of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        col.labels[is.na(col.labels)] <- ""
        heatmap.labels(col.labels, type="col.labels", just="right")
        upViewport()
    }
    elem <- "row.ddr.left"
    if (elem %in% elements && !is.null(row.ddr))
    {
        if (attr(row.ddr, "members") != nrow(exprs))
            warning("the number of row dendrogram leaves is not equal to the number of rows of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2]), clip="off"))
        par(new=TRUE, fig=gridFIG(), mar=c(0,0,0,0))
        plot(row.ddr, horiz=TRUE, leaflab="none", xaxt="n", yaxt="n", xaxs="r", yaxs="i")
        upViewport()
    }
    elem <- "row.ddr.right"
    if (elem %in% elements && !is.null(row.ddr))
    {
        if (attr(row.ddr, "members") != nrow(exprs))
            warning("the number of row dendrogram leaves is not equal to the number of rows of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        par(new=TRUE, fig=gridFIG(), mar=c(0,0,0,0))
        plot(row.ddr, horiz=TRUE, leaflab="none", xaxt="n", yaxt="n", xaxs="r", yaxs="i",
             xlim=c(0, attr(row.ddr, "height")))
        upViewport()
    }
    elem <- "col.ddr.up"
    if (elem %in% elements && !is.null(col.ddr))
    {
        if (attr(col.ddr, "members") != ncol(exprs))
            warning("the number of column dendrogram leaves is not equal to the number of columns of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        par(new=TRUE, fig=gridFIG(), mar=c(0,0,0,0))
        plot(col.ddr, horiz=FALSE, leaflab="none", xaxt="n", yaxt="n", xaxs="i", yaxs="r")
        upViewport()
    }
    elem <- "col.ddr.down"
    if (elem %in% elements && !is.null(col.ddr))
    {
        if (attr(col.ddr, "members") != ncol(exprs))
            warning("the number of column dendrogram leaves is not equal to the number of columns of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        par(new=TRUE, fig=gridFIG(), mar=c(0,0,0,0))
        plot(col.ddr, horiz=FALSE, leaflab="none", xaxt="n", yaxt="n", xaxs="i", yaxs="r",
             ylim=c(attr(col.ddr, "height"), 0))
        upViewport()
    }
    elem <- "clinical"
    if (elem %in% elements && !is.null(clinical))
    {
        if (ncol(clinical) != ncol(exprs))
            warning("the size of clinical does not correspond to the number of columns of exprs")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))
        heatmap.clinical(clinical, clinical.na.color)
        upViewport()
    }
    elem <- "clinical.labels.ljust"
    if (elem %in% elements && !is.null(clinical.labels))
    {
        if (length(clinical.labels) != nrow(clinical))
            warning("the number of clinical labels is not equal to the number of variables in clinical")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        clinical.labels[is.na(clinical.labels)] <- ""
        heatmap.labels(labels=clinical.labels, type="row.labels", just="left")
        upViewport()
    }
    elem <- "clinical.labels.rjust"
    if (elem %in% elements && !is.null(clinical.labels))
    {
        if (length(clinical.labels) != nrow(clinical))
            warning("the number of clinical labels is not equal to the number of variables in clinical")

        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        clinical.labels[is.na(clinical.labels)] <- ""
        heatmap.labels(labels=clinical.labels, type="row.labels", just="right")
        upViewport()
    }
    elem <- "title"
    if (elem %in% elements && !is.null(title))
    {
        idx <- which(layout.mat == elem, arr.ind=TRUE)
        pushViewport(viewport(name=elem,
                              layout.pos.row=unique(idx[,1]),
                              layout.pos.col=unique(idx[,2])))

        grid.text(title, gp=gpar(cex=0.9 * heatmap.labels.cex(title)))
        upViewport()
    }

    ## restore all standard R graphical parameters
    par(old.par)
    upViewport()

    row.ind <- 1:nrow(exprs)
    col.ind <- 1:ncol(exprs)
    if (!is.null(row.ddr)) row.ind <- order.dendrogram(rev(row.ddr))
    if (!is.null(col.ddr)) col.ind <- order.dendrogram(col.ddr)

    return(invisible(list(layout.mat=layout.mat,
                          widths=widths,
                          heights=heights,
                          color.scheme=color.scheme,
                          key.min=key.min,
                          key.max=key.max,
                          na.color=na.color,
                          scale=scale,
                          key.line.color=key.line.color,
                          row.labels=row.labels,
                          col.labels=col.labels,
                          row.clust=row.clust,
                          row.clust.dist.fun=row.clust.dist.fun,
                          row.clust.fun=row.clust.fun,
                          col.clust=col.clust,
                          col.clust.dist.fun=col.clust.dist.fun,
                          col.clust.fun=col.clust.fun,
                          clinical=clinical,
                          clinical.labels=clinical.labels,
                          clinical.na.color=clinical.na.color,
                          title=title,
                          row.ind=row.ind, col.ind=col.ind,
                          row.ddr=row.ddr, col.ddr=col.ddr)))
}

## creates html making a heatmap image created by heatmap.simple clickable.
## e.g.
## png("heatmap.png", width=800, height=1200)
## heatmap.out <- heatmap.simple(...)
## dev.off()
## writeLines(heatmap.simple.html.map(heatmap.out, "clickable", 800, 1200, "heatmap.png")$html, con="heatmap.html")
heatmap.simple.html.map <- function(heatmap.out, heatmap.name, image.width, image.height, image.filename,
                                    clickable=c("heatmap","col.labels","row.labels","key",
                                      "title","clinical", "clinical.labels")) {
  col.labels <- make.names(heatmap.out$col.labels, unique=T)
  row.labels <- make.names(heatmap.out$row.labels, unique=T)
  clinical.labels <- make.names(heatmap.out$clinical.labels, unique=T)

  names(col.labels) <- heatmap.out$col.labels
  names(row.labels) <- heatmap.out$row.labels
  names(clinical.labels) <- heatmap.out$clinical.labels

  clickable <- intersect(c(clickable, paste0(clickable, ".rjust"), paste0(clickable, ".ljust")), heatmap.out$layout.mat)

  href <- sapply(clickable, function(item) {
    switch(gsub(".[rl]{1}just$", "", item),
           heatmap=sapply(col.labels, function(col)
             sapply(row.labels, function(row)
                    paste(heatmap.name, "-row-", row, "-col-", col, ".html", sep=""))),

           col.labels=matrix(sapply(col.labels, function(col)
             paste(heatmap.name, "-col-", col, ".html", sep="")), nrow=1),

           row.labels=matrix(sapply(row.labels, function(row)
             paste(heatmap.name, "-row-", row, ".html", sep="")), ncol=1),

           key=matrix(paste(heatmap.name, "-key.html", sep=""), nrow=1),

           title=matrix(paste(heatmap.name, "-title.html", sep=""), nrow=1),

           clinical=sapply(col.labels, function(col)
             sapply(clinical.labels, function(clinical)
                    paste(heatmap.name, "-clinical-", clinical, "-col-", col, ".html", sep=""))),

           clinical.labels=matrix(sapply(clinical.labels, function(label)
             paste(heatmap.name, "-clinical-", label, ".html", sep="")), ncol=1))
  }, simplify=F)

  list(html=do.call("layout.mat.html.map",
         c(list(layout.mat=heatmap.out$layout.mat,
                layout.widths=heatmap.out$widths,
                layout.heights=heatmap.out$heights,
                name=heatmap.name,
                image.width=image.width,
                image.height=image.height,
                image.filename=image.filename),
           href)),
       row.labels=row.labels,
       col.labels=col.labels,
       clinical.labels=clinical.labels,
       href=href)
}

## creates an html image map according to the items in the layout.mat
## "..." argument are matrices whose names match items (regions) in layout.mat
## each corresponding region is divided into equal-sized cells corresponding
##  to the dimensions of the matrix
## the matrix contains the link targets
## returns image map code
layout.mat.html.map <- function(layout.mat, layout.widths, layout.heights,
                                name,
                                image.width, image.height,
                                image.filename, ...) {
  stopifnot(nrow(layout.mat) == length(layout.heights))
  stopifnot(ncol(layout.mat) == length(layout.widths))
  stopifnot(all(names(list(...)) %in% layout.mat))

  bounding.box <- function(name) {
    mat.col <- 1:ncol(layout.mat)
    mat.row <- 1:nrow(layout.mat)
    i <- which(name == layout.mat)
    if (length(i) == 0) return(NULL)
    if (length(i) > 1) {
      warning(paste("layout.mat.html:bounding.box: '", name,
                    "' matches multiple items in layout.mat, ",
                    "making only the first one clickable.", sep=""))
      i <- i[1]
    }
    col.ind <- ceiling(i/nrow(layout.mat))
    row.ind <- i - (col.ind-1)*nrow(layout.mat)
    list(x1=sum(layout.widths[which(mat.col < col.ind)])/sum(layout.widths) * image.width,
         x2=sum(layout.widths[which(mat.col <= col.ind)])/sum(layout.widths) * image.width,
         y1=sum(layout.heights[which(mat.row < row.ind)])/sum(layout.heights) * image.height,
         y2=sum(layout.heights[which(mat.row <= row.ind)])/sum(layout.heights) * image.height)
  }
  html.rect <- function(row,col,nrow,ncol,bb,link) {
    cell.width <- (bb$x2 - bb$x1)/ncol
    cell.height <- (bb$y2 - bb$y1)/nrow
    x1 <- cell.width*(col-1) + bb$x1
    x2 <- x1 + cell.width
    y1 <- cell.height*(row-1) + bb$y1
    y2 <- y1 + cell.height
    coords <- paste(floor(x1), ceiling(y1), floor(x2), ceiling(y2), sep=",")
    paste("<area shape=\"rect\" coords=\"", coords, "\" href=\"", link, "\">", sep="")
  }

  out <- paste("<img src=\"",image.filename,"\" usemap=\"#",name,"\", ISMAP>",sep="")
  out <- c(out,paste("<map name=\"",name,"\">",sep=''))

  href <- list(...)
  for (region.name in names(href)) {
    cat("layout.mat.html.map", region.name, "\n")
    region.bb <- bounding.box(region.name)
    if (!is.null(region.bb)) {
      region.nrow <- nrow(href[[region.name]])
      region.ncol <- ncol(href[[region.name]])
      out <- c(out, unlist(lapply(1:region.nrow, function(i)
                                  lapply(1:region.ncol, function(j)
                                         html.rect(i,j,region.nrow,region.ncol,region.bb,
                                                   href[[region.name]][i,j])))))
    }
  }

  c(out,"</map>")
}

## add 1-3 marks to each cell of a simple heatmap
## e.g.
## heatmap.out <- heatmap.simple(exprs, ...)
## marks <- matrix(0, ncol=ncol(exprs), nrow=nrow(exprs))
## marks[pvals < 0.01] <- 1
## marks[pvals < 0.001] <- 2
## heatmap.mark(heatmap.out, marks)
heatmap.mark <- function(heatmap.out, marks, mark="star") {
  stopifnot(ncol(marks) == length(heatmap.out$col.ind))
  stopifnot(nrow(marks) == length(heatmap.out$row.ind))
  stopifnot(all(marks %in% 0:3))

  plot.marks <- function(...) {
      if (mark == "star")
          plot.stars(...)
      else
          plot.box(...)
  }

  plot.stars <- function(x, y, pos) {
    if (length(x) == 0)
      return()

    id <- rep(1:length(x), 3)
    x <- x - (pos - 1)/8
    grid.polygon(c(x, x, x-(1/8)), c(y-1, y-1+(1/4), y-1),
                 gp=gpar(fill="black", lty=0), id = id, default.units = "native")
  }

  plot.box <- function(x,y,pos) {
      if (length(x) == 0) return()

      id <- rep(1:length(x), 4)
      grid.polygon(c(x,x,x-1,x-1), c(y-1,y,y,y-1),
                   gp=gpar(lty=1, col="black"), id=id,
                   default.units="native")
  }

  marks <- marks[heatmap.out$row.ind,]
  marks <- marks[,heatmap.out$col.ind]

  n <- downViewport("heatmap")
  pushViewport(viewport(xscale=c(0, ncol(marks)), yscale=c(nrow(marks), 0)))
  for (i in setdiff(unique(marks), 0)) {
    idx <- which(marks>=i, arr.ind=TRUE)
    plot.marks(idx[, "col"], idx[, "row"], i)
  }
  popViewport()
  upViewport(n)
}

## draw a line under the row given by row.label
## e.g.
## heatmap.out <- heatmap.simple(exprs, ...)
## heatmap.underline.row(heatmap.out, row.label="NR3C1")
heatmap.underline.row <- function(heatmap.out, row.label, below=TRUE, gp=gpar(col="black", lty=1)) {
  y <- 1-(which(heatmap.out$row.labels == row.label) - sign(!below))/length(heatmap.out$row.labels)
  for (element in c("heatmap", "row.labels.rjust", "row.labels.ljust")) {
    if (element %in% heatmap.out$layout.mat) {
      depth <- downViewport(element)
      grid.polyline(c(0,1), c(y,y), gp=gp)
      upViewport(depth)
    }
  }
}

## draw a line next to the column given by col.label
## e.g.
## heatmap.out <- heatmap.simple(exprs, ...)
## heatmap.underline.column(heatmap.out, col.label="curtis.323")
heatmap.underline.column <- function(heatmap.out, col.label, left=TRUE, gp=gpar(col="black", lty=1)) {
  y <- (which(heatmap.out$col.labels == col.label) - sign(left))/length(heatmap.out$col.labels)
  depth <- downViewport("heatmap")
  grid.polyline(c(y,y),c(0,1), gp=gp)
  upViewport(depth)
}


#=================================================================================================#
# harvard_forest_datasets.r
# Datasets and parsing functions for Harvard Forest EMS flux tower (Little Prospect Hill)
# Note: site class is ems (EMS tower)
# Adam Erickson, PhD, Washington State University
# Contact: adam.michael.erickson@gmail.com
# March 14, 2019
# License: Apache 2.0
#
# Copyright 2019 Washington State University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#=================================================================================================#

#-------------------------------------------------------------------------------------------------#
# Datasets specific to EMS tower:
#-------------------------------------------------------------------------------------------------#
# HF004: Canopy-Atmosphere Exchange of Carbon, Water and Energy at Harvard Forest EMS Tower since 1991
# HF006: Soil Respiration, Temperature and Moisture at Harvard Forest EMS Tower since 1995
# HF060: EMS - Automated High-Frequency Methane Data
# HF066: Concentrations and Surface Exchange of Air Pollutants at Harvard Forest EMS Tower since 1990
# HF067: CFCs and Radiatively Important Trace Species at Harvard Forest EMS Tower 1996-2005
# HF068: Soil Respiration Along a Hydrological Gradient at Harvard Forest EMS Tower 2003-2006
# HF069: Biomass Inventories at Harvard Forest EMS Tower since 1993
# HF102: Radiation Measurements at Harvard Forest EMS Tower 1991-2007
# HF145: Hydrocarbon Concentrations at Harvard Forest EMS Tower 1992-2001
# HF206: Microclimate at Harvard Forest HEM, LPH and EMS Towers since 2005
# HF209: Isotopic Composition of Net Ecosystem CO2 Exchange at Harvard Forest EMS Tower since 2011
# HF214: Fluxes of Carbonyl Sulfide at Harvard Forest EMS Tower since 2010
# HF288: Fluxes of Molecular Hydrogen (H2) at Harvard Forest EMS Tower 2010-2012
# HF306: Atmospheric Oxygen and Carbon Dioxide at Harvard Forest EMS Tower since 2006

#-------------------------------------------------------------------------------------------------#
# General data:
#-------------------------------------------------------------------------------------------------#
# "hf069-13-annual-summ.csv"          biomass_ag (woody), ANPP (woody), recruitment, mortality, 1993-2017
# "hf069-15-annual-summ.csv"          biomass_ag (woody), ANPP (woody), recruitment, mortality, 1994-2013 (old)
# "hf069-14-understory-summ.csv"      understory regenertion, biomass, 2004-2006
# "hf001-04-monthly-m.csv"            meteorology, 2001-2014
# "hf004-02-filled.csv"               NEE/GEE, respiration (ecosystem), 1991-2015
# "hf069-09-trees.csv"                species, DBH, 1993-2017
# "hf069-11-trees.csv"                species, DBH, 1998-2013
# "hf008-04-tree.csv"                 species, dbh, 1998-2008
# "hf213-01-hf-inventory.csv"         species, dbh, coordinates, 2009-2011
# "hf154-03-lph-cores.csv"            ANPP, 2005
# "hf133-03-productivity.csv"         biomass, 1996-1997
# "hf137-01-tree.csv"                 biomass, ANPP, species, 1996 & 2006
# "hf008-01-foliage.csv"              canopy chemistry, 1988-2011
# "hf008-02-litter.csv"               canopy chemistry, 1988-2002
# "hf062-01-plot.csv"                 canopy chemistry, 1992
# "hf062-02-tree.csv"                 canopy chemistry, 1992
# "hf069-06-canopy-chem.csv"          canopy chemistry, 1998-2008
# "hf008-07-decomp.csv"               litter chemistry, 1998-2008
# "hf069-05-litter-chem.csv"          litter chemistry, 1998-2000
# "hf047-01-regen.csv"                regeneration, 1990-2000
# "hf005-05-soil-respiration.csv"     r_soil, 1991-2002, C = control
# "hf006-01-soil-respiration.csv"     r_soil, 1995-2010
# "hf006-03-autochamber.csv"          r_soil, 2003
# "hf068-01-soil-resp.csv"            r_soil, 2003-2006
# "hf045-01-soil-resp.csv"            r_soil, 2006-2012
# "hf148-02-lph-soil-respiration.csv" r_soil, 2002-2007
# "hf194-01-soil-resp.csv"            r_soil, 1993-2006
# "hf194-02-synthesis-data"           r_soil, 1991-2008
# "hf243-01-soil-respiration.csv"     r_soil (Finzi), 2009-2010
# "hf040-02-species.csv"              species abundance, 1937-1991
# "hf015-06-species.csv"              species codes
# "hf040-01-species.csv"              species codes
# "hf222-01-vegetation.csv"           species codes
# "hf039-02-tree.csv"                 species DBH-height, 1937
# "hf005-01-trace-gas.csv"            co2, ch4, n20, 1991-2002
# "hf003-01-plant.csv"                tree canopy dimensions, 1991
# "hf069-16-allometries.txt"          allometry equations
#-------------------------------------------------------------------------------------------------#

#-------------------------------------------------------------------------------------------------#
# Metrics and units
#-------------------------------------------------------------------------------------------------#
# Note: kg m-2 (10 Mg ha-1); µmol m-2 sec-1 (1e-6 * 60 * 60 * 24 * 365.24 mol m-2 year-1)
# nee ................. kg C m-2 year-1
# biomass_ag .......... kg m-2
# c_ag ................ kg C m-2
# n_ag ................ kg N m-2
# biomass_bg .......... kg m-2
# c_bg ................ kg C m-2
# n_bg ................ kg N m-2
# c_so ................ kg C m-2
# n_so ................ kg N m-2
# r_soil .............. kg C m-2 year-1
# anpp ................ kg m-2 year-1
# biomass_sp .......... kg m-2
# abundance_sp ........ % of total n-trees

#-------------------------------------------------------------------------------------------------#
# Selected data:
#-------------------------------------------------------------------------------------------------#
# "hf069-13-annual-summ.csv" (annual woody) or...
# "hf069-09-trees.csv"
#   timescale: 1993-2017
#   variables: biomass_ag and modeled c_ag, n_ag, biomass_bg, c_bg, n_bg, anpp, biomass_sp, abundance_sp
# "hf006-01-soil-respiration.csv"
#   timescale: 1995-2010
#   variables: r_soil,
# "hf015-05-soils.csv"
#   timescale: 1937-2013
#   variables: c_so, n_so
# "hf004-02-filled.csv"
#   timescale: 1991-2015
#   variables: nee

#-------------------------------------------------------------------------------------------------#
# Time period:
#-------------------------------------------------------------------------------------------------#
# Intersection: 1996-2009 (14 years)
# Union:        1937-2017 (81 years)
# Climate:      2002-2012 (11 years)

# Example: download observation data
#url  <- "http://harvardforest.fas.harvard.edu/data/p06/hf069/hf069-09-trees.csv"
#dest <- file.path(base, "observations","hf_ems","hf069-09-trees.csv")
#download_file(url=url, destination=dest)

#-------------------------------------------------------------------------------------------------#
# Base directory
#-------------------------------------------------------------------------------------------------#
base_hf_data <- file.path(base, "observations", "hf_ems")

#-------------------------------------------------------------------------------------------------#
# Species look-up table
#-------------------------------------------------------------------------------------------------#
# http://harvardforest.fas.harvard.edu/sites/harvardforest.fas.harvard.edu/files/data/p06/hf069/hf069-17-allometries.txt
# http://harvardforest.fas.harvard.edu:8080/exist/apps/datasets/showData.html?id=hf003
# Using USDA PLANTS database codes: https://plants.usda.gov
species_lut <- data.frame(
  abbr  = c("AB","AC","ash","bb",
            "BB","BC","beech","bo",
            "cherry","chestnut","EH","FPC",
            "gb","haw","hbb","HBB",
            "hem","HT","nwr","NWR",
            "PB","rm","RM","ro",
            "RO","rp","sm","SM",
            "SPI","SS","wb","wh",
            "WH","wo","wp","WP",
            "ws","WS","yb","YB"),
  name  = c("Fagus grandifolia","Castenea dentate","Fraxinus americana","Betula lenta",
            "Betula lenta","Prunus serotina","Fagus grandifolia","Quercus velutina",
            "Prunus serotina","Castenea dentate","Tsuga canadensis","?",
            "Betula populifolia","Crataegus","Vaccinium","Vaccinium",
            "Tsuga canadensis","Crataegus","Viburnum","Viburnum",
            "Betula papyrifera","Acer rubrum","Acer rubrum","Quercus rubra",
            "Quercus rubra","Pinus resinosa","Acer pennsylvanicum","Acer pennsylvanicum",
            "Spiraea","?","Ilex verticillata","Hamamelis virginiania",
            "Hamamelis virginiania","Quercus alba","Pinus strobus","Pinus strobus",
            "Picea glauca","Picea glauca","Betula alleghaniensis","Betula alleghaniensis"),
  genus = c("Fagus","Castenea","Fraxinus","Betula","Betula","Prunus","Fagus","Quercus","Prunus",
            "Castenea","Tsuga","","Betula","Crataegus","Vaccinium","Vaccinium","Tsuga","Crataegus",
            "Viburnum","Viburnum","Betula","Acer","Acer","Quercus","Quercus","Pinus","Acer",
            "Acer","Spiraea","","Ilex","Hamamelis","Hamamelis","Quercus","Pinus","Pinus","Picea",
            "Picea","Betula","Betula"),
  code  = c("FAGR","CADE","FRAM","BELE",
            "BELE","PRSE","FAGR","QUVE",
            "PRSE","CADE","TSCA","",
            "BEPO","CRSP","VACO","VACO",
            "TSCA","CRSP","VICA","VICA",
            "BEPA","ACRU","ACRU","QURU",
            "QURU","PIRE","ACPE","ACPE",
            "","","ILVE","HAVI",
            "HAVI","QUAL","PIST","PIST",
            "PIGL","PIGL","BEAL","BEAL")
)
#message(paste("Species LUT:", jsonlite::toJSON(species_lut, pretty=TRUE)))

#-------------------------------------------------------------------------------------------------#
# Observation file list
#-------------------------------------------------------------------------------------------------#
observation_list <- list(
  climate      = file.path(base_hf_data, "observations", "hf_ems", "hf001-04-monthly-m.csv"),
  fluxes       = file.path(base_hf_data, "observations", "hf_ems", "hf004-02-filled.csv"),
  inventory    = file.path(base_hf_data, "observations", "hf_ems", "hf069-09-trees.csv"),
  root_cn      = file.path(base_hf_data, "observations", "hf_ems", "hf278-04-root-carbon-nitrogen.csv"),
  litter_cn    = file.path(base_hf_data, "observations", "hf_ems", "hf069-06-canopy-chem.csv"),
  soils        = file.path(base_hf_data, "observations", "hf_ems", "hf015-05-soils-cn.csv"),
  r_soil       = file.path(base_hf_data, "observations", "hf_ems", "hf006-01-soil-respiration.csv"),
  anpp         = file.path(base_hf_data, "observations", "hf_ems", "hf069-13-annual-summ.csv")
)

#-------------------------------------------------------------------------------------------------#
# Functions
#-------------------------------------------------------------------------------------------------#

# Tree species aboveground biomass equations; kg; DBH cm; Chojnacky, Heath, Jenkins (2014)
get_tree_biomass_ag <- function(code, dbh) {
  switch(code,                                   #  WSG
    ACPE = { exp(-2.0470 + 2.3852 * log(dbh)) }, # 0.44
    ACRU = { exp(-2.0470 + 2.3852 * log(dbh)) }, # 0.49
    BEAL = { exp(-1.8096 + 2.3480 * log(dbh)) }, # 0.55
    BELE = { exp(-1.8096 + 2.3480 * log(dbh)) }, # 0.60
    BEPA = { exp(-2.2271 + 2.4513 * log(dbh)) }, # 0.48
    BEPO = { exp(-2.2271 + 2.4513 * log(dbh)) }, # 0.45
    CADE = { exp(-2.0705 + 2.4410 * log(dbh)) }, # 0.40
    FAGR = { exp(-2.0705 + 2.4410 * log(dbh)) }, # 0.56
    FRAM = { exp(-1.8384 + 2.3524 * log(dbh)) }, # 0.55
    PIGL = { exp(-2.1364 + 2.3233 * log(dbh)) }, # 0.37
    PIRE = { exp(-2.6177 + 2.4638 * log(dbh)) }, # 0.41
    PIST = { exp(-2.6177 + 2.4638 * log(dbh)) }, # 0.34
    PRSE = { exp(-2.2118 + 2.4133 * log(dbh)) }, # 0.47
    QUAL = { exp(-2.0705 + 2.4410 * log(dbh)) }, # 0.60
    QURU = { exp(-2.0705 + 2.4410 * log(dbh)) }, # 0.56
    QUVE = { exp(-2.0705 + 2.4410 * log(dbh)) }, # 0.56
    TSCA = { exp(-2.3480 + 2.3876 * log(dbh)) }, # 0.38
    CRSP = { NA },                               # shrub
    HAVI = { NA },                               # shrub
    ILVE = { NA },                               # shrub
    VACO = { NA },                               # shrub
    VICA = { NA },                               # shrub
    NA                                           # else
  )
}

# Tree species belowground biomass equations; kg; DBH cm; Chojnacky, Heath, Jenkins (2014)
get_tree_biomass_bg <- function(code, dbh) {
  biomass_ag       = get_tree_biomass_ag(code, dbh)
  root_coarse_frac = exp(-1.4485 - 0.03476 * log(dbh))
  root_fine_frac   = exp(-1.8629 - 0.77534 * log(dbh))
  soil_frac        = 0.68
  biomass_bg       = (biomass_ag * root_coarse_frac) + (biomass_ag * root_fine_frac) + (biomass_ag * soil_frac)
  return(biomass_bg)
}

# Root biomass fraction by DBH class; treated as constants per Chojnacky, Heath, Jenkins (2014)
#plot(1:100, sapply(1:100, function(dbh) exp(-1.4485 - 0.03476 * log(dbh))), type="l")
#plot(1:100, sapply(1:100, function(dbh) exp(-1.8629 - 0.77534 * log(dbh))), type="l")

# Main

# Parent function containing child functions
get_observations <- function(metric="all") {
  funs = list(get_nee(), get_biomass_ag(), get_c_ag(), get_n_ag(), get_biomass_bg(),
              get_c_bg(), get_n_bg(), get_c_so(), get_n_so(), get_r_soil(), get_anpp(),
              get_biomass_sp(), get_abundance_sp())
  switch(metric,
    nee           = { get_nee() },
    biomass_ag    = { get_biomass_ag() },
    c_ag          = { get_c_ag() },
    n_ag          = { get_n_ag() },
    biomass_bg    = { get_biomass_bg() },
    c_bg          = { get_c_bg() },
    n_bg          = { get_n_bg() },
    c_so          = { get_c_so() },
    n_so          = { get_n_so() },
    r_soil        = { get_r_soil() },
    anpp          = { get_anpp() },
    biomass_sp    = { get_biomass_sp() },
    abundance_sp  = { get_abundance_sp() },
    biomass_total = { cbind(get_biomass_ag()[,1], get_biomass_ag()[,2] + get_biomass_bg()[,2]) },
    c_total       = { cbind(get_c_ag()[,1], get_c_ag()[,2] + get_c_bg()[,2]) },
    n_total       = { cbind(get_n_ag()[,1], get_n_ag()[,2] + get_n_bg()[,2]) },
    all           = { setNames(lapply(funs, eval),
                               c("nee","biomass_ag","c_ag","n_ag","biomass_bg",
                                 "c_bg","n_bg","c_so","n_so","r_soil","anpp",
                                 "biomass_sp","abundance_sp")) },
    { stop(paste("Variable", metric, "not found")) }
  )
}

# NEE; kg C m-2 year-1; Converted from µmol CO2 m-2 sec-1 (10^-6 moles)
get_nee <- function() {
  #df = read.csv(self$files$observations$flux, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf004-02-filled.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df$c_kg = df$nee * 1e-6 * 12 / 1000       # convert µmol CO2 to mol CO2 to g C to kg C
  df$c_kg = df$c_kg * 60 * 60 * 24 * 365.24 # convert kg C m-2 second-1 to kg C m-2 year-1
  return(setNames(aggregate(df$c_kg, by=list(df$year), FUN=mean), c("year","nee")))
}

# Aboveground biomass; kg m-2; 34 EMS tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_biomass_ag <- function() {
  #df = read.csv(self$files$observations$inventory, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf069-09-trees.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  #df = df[df$site =="ems" & df$tree.type == "live",]
  df = df[df$site =="ems" & df$tree.type == "live" & !is.na(df$year) & df$year > 1993,]
  df$code = species_lut$code[match(df$species, species_lut$abbr)]
  df = df[df$code %in% c("ACPE","ACRU","BEAL","BELE","BEPO","FAGR","FRAM","PIGL",
                         "PIRE","PIST","PRSE","QURU","QUVE","TSCA"),] # modeled species
  # Mean annual individual tree DBH
  df = setNames(aggregate(df$dbh, by=list(df$year, df$code, df$plot, df$tag), FUN=mean, na.rm=TRUE),
                c("year","code","plot","tag","dbh"))
  df = df[,c("year","code","dbh")]
  df$biomass_ag = apply(df, 1, function(x) get_tree_biomass_ag(x[2], as.numeric(x[3])))
  agb = setNames(aggregate(df$biomass_ag, by=list(df$year), FUN=sum, na.rm=TRUE), c("year","biomass_ag"))
  agb$biomass_ag = agb$biomass_ag / 10681.42 # convert kg to kg m-2
  return(agb)
}

# Aboveground carbon; kg C m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_c_ag <- function() {
  df      = get_biomass_ag()
  df$c_ag = df$biomass_ag * 0.5
  return(df[,c("year","c_ag")])
}

# Aboveground carbon; kg N m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_n_ag <- function() {
  agc = get_c_ag()
  #df = read.csv(self$files$observations$root_cn, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf278-04-root-carbon-nitrogen.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df$cn   = df$c.n / 100
  cn_stem = mean(df$cn, na.rm=TRUE)
  #df = read.csv(self$files$observations$leaf_cn, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf069-06-canopy-chem.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df$cn_leaf = df$c.percent / df$n.percent
  cn_leaf    = mean(df$cn_leaf, na.rm=TRUE)
  agc$n_ag   = agc$c_ag / ((cn_stem * 0.80) + (cn_leaf * 0.20)) # assumes general 80/20 stem/leaf fraction
  return(agc[,c("year","n_ag")])
}

# Belowground biomass; kg m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
# Uses general AGB-to-BGB model of Chojnacky, Heath, Jenkins (2014)
get_biomass_bg <- function() {
  df = read.csv(file.path(base_hf_data, "hf069-09-trees.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df = df[df$site =="ems" & df$tree.type == "live",]
  df$code = species_lut$code[match(df$species, species_lut$abbr)]
  df = df[df$code %in% c("ACPE","ACRU","BEAL","BELE","BEPO","FAGR","FRAM","PIGL",
                         "PIRE","PIST","PRSE","QURU","QUVE","TSCA"),] # modeled species
  # Mean annual individual tree DBH
  df = setNames(aggregate(df$dbh, by=list(df$year, df$code, df$plot, df$tag), FUN=mean, na.rm=TRUE),
                c("year","code","plot","tag","dbh"))
  df = df[,c("year","code","dbh")]
  df$biomass_ag = apply(df, 1, function(x) get_tree_biomass_ag(x[2], as.numeric(x[3])))
  #df$root_coarse = df$biomass_ag * 0.22
  #df$root_fine   = df$biomass_ag * 0.01
  df$root_coarse = df$biomass_ag * exp(-1.4485 - 0.03476 * log(df$dbh))
  df$root_fine   = df$biomass_ag * exp(-1.8629 - 0.77534 * log(df$dbh))
  df$soil        = df$biomass_ag * 0.68
  df$biomass_bg  = df$root_coarse + df$root_fine + df$soil
  bgb = setNames(aggregate(df$biomass_bg, by=list(df$year), FUN=sum, na.rm=TRUE), c("year","biomass_bg"))
  bgb$biomass_bg = bgb$biomass_bg / 10681.42 # convert kg to kg m-2
  return(bgb[,c("year","biomass_bg")])
}

# Belowground C; kg C m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_c_bg <- function () {
  bgb = get_biomass_bg()
  #df = read.csv(self$files$observations$root_cn, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf278-04-root-carbon-nitrogen.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  bgb$c_bg = bgb$biomass_bg * (mean(df$c.per, na.rm=TRUE) / 100)
  return(bgb[,c("year","c_bg")])
}

# Belowground N; kg N m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_n_bg <- function () {
  bgb = get_biomass_bg()
  #df = read.csv(self$files$observations$root_cn, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf278-04-root-carbon-nitrogen.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  bgb$n_bg = bgb$biomass_bg * (mean(df$n.per, na.rm=TRUE) / 100)
  return(bgb[,c("year","n_bg")])
}

# Soil organic C; g C cm-3; 240 (22.5 x 22.5m) 506.25 m2 plots (121500 m2 total) of 269 original plots
get_c_so <- function() {
  #df = read.csv(self$files$observations$soils, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf015-05-soils-cn.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  bd   = df$bdgcc.corr           # bulk density or g soil per cm3
  c_gg = df$gcgsoil              # g C per g soil
  c_so = bd * c_gg / 1000 * 1e+6 # g C per cm3 soil converted to kg per m3
  c_so = mean(c_so, na.rm=TRUE)
  return(data.frame(year=NA, c_so=c_so))
}

# Soil organic N; g N cm-3; 240 (22.5 x 22.5m) 506.25 m2 plots (121500 m2 total) of 269 original plots
get_n_so <- function() {
  #df = read.csv(self$files$observations$soils, header=TRUE, stringsAsFactors=FALSE)
  df  = read.csv(file.path(base_hf_data, "hf015-05-soils-cn.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  bd   = df$bdgcc.corr           # bulk density or g soil per cm3
  n_gg = df$gngsoil              # g N per g soil
  n_so = bd * n_gg / 1000 * 1e+6 # g C per cm3 soil converted to kg per m3
  n_so = mean(n_so, na.rm=TRUE)
  return(data.frame(year=NA, n_so=n_so))
}

# Soil respiration; mg C m-2 hour-1; kg C m-2 year-1; 6 sites near EMS tower
get_r_soil <- function() {
  #df = read.csv(self$files$observations$r_soil, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf006-01-soil-respiration.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df$date = as.POSIXlt(df$date, format="%m/%d/%Y", tz="EST")
  df$year = as.numeric(format(df$date, "%Y"))
  df$flux = df$flux / 1e+6 * 24 * 365.24 # convert mg C m-2 hour-1 to kg C m-2 year-1
  return(setNames(aggregate(df$flux, by=list(df$year), FUN=mean, na.rm=TRUE), c("year","r_soil")))
}

# Annual net primary productivity (ANPP); Mg year-1; kg C m-2 year-1
get_anpp <- function() {
  #df = read.csv(self$files$observations$anpp, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf069-13-annual-summ.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df = df[df$site =="ems",]
  df = df[!is.na(df$anpp),]
  df$anpp = df$anpp * 1000 / 10681.42 * 0.5 # Mg mass year-1 to kg C m-2 year-1
  return(df[,c("year","anpp")])
}

# Total species aboveground biomass; kg m-2; 34 tower plots (10m radius; (pi*10^2)*34 = 10681.42 m2)
get_biomass_sp <- function() {
  #df = read.csv(self$files$observations$inventory, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf069-09-trees.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df = df[df$site =="ems" & df$tree.type == "live",]
  df$code  = species_lut$code[match(df$species, species_lut$abbr)]
  df       = df[df$code %in% c("ACPE","ACRU","BEAL","BELE","BEPO","FAGR","FRAM","PIGL",
                               "PIRE","PIST","PRSE","QURU","QUVE","TSCA"),] # modeled species
  # Mean annual individual tree DBH
  df = setNames(aggregate(df$dbh, by=list(df$year, df$code, df$plot, df$tag), FUN=mean, na.rm=TRUE),
                c("year","code","plot","tag","dbh"))
  df = df[,c("year","code","dbh")]
  df$biomass_ag = apply(df, 1, function(x) get_tree_biomass_ag(x[2], as.numeric(x[3])))
  biomass_sp    = aggregate(df$biomass_ag, by=list(df$year, df$code), FUN=sum, na.rm=TRUE)
  colnames(biomass_sp) = c("year","species","biomass_ag")
  biomass_sp$biomass_ag = biomass_sp$biomass_ag / 10681.42
  return(biomass_sp)
}

# Species relative abundance; percent of n-trees
get_abundance_sp <- function() {
  #df = read.csv(self$files$observations$inventory, header=TRUE, stringsAsFactors=FALSE)
  df = read.csv(file.path(base_hf_data, "hf069-09-trees.csv"), header=TRUE, stringsAsFactors=FALSE)
  colnames(df) = tolower(colnames(df))
  df = df[df$site =="ems" & df$tree.type == "live",]
  df$code        = species_lut$code[match(df$species, species_lut$abbr)]
  df             = df[df$code %in% c("ACPE","ACRU","BEAL","BELE","BEPO","FAGR","FRAM","PIGL",
                                     "PIRE","PIST","PRSE","QURU","QUVE","TSCA"),] # modeled species
  ntrees_yr      = setNames(aggregate(df$code, by=list(df$year), FUN=length), c("year","ntrees"))
  df$ntrees_yr   = ntrees_yr$ntrees[match(df$year, ntrees_yr$year)]
  df$fraction_yr = 1 / df$ntrees_yr # relative abundance of each tree
  return(setNames(aggregate(df$fraction_yr, by=list(df$year, df$code), FUN=sum), c("year","species","abundance")))
}


# ************************************************** #
#              Read exposure and outcome             #
# ************************************************** #
read_gsmr_trait = function(file_con) {
    expo_str = scan(file_con, nlines=1, quiet=TRUE, what="");
    outcome_str = scan(file_con, nlines=1, quiet=TRUE, what="");
    strbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
    return(list(expo_str=expo_str, outcome_str=outcome_str))
}

# ************************************************** #
#                  Read GSMR result                  #
# ************************************************** #
read_gsmr_result = function(file_con) {
    expo_str = outcome_str = bxy = bxy_se = bxy_pval = bxy_m = c()
    while(1) {
        strbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
        if(strbuf[1] == "#gsmr_end") break;
        if(strbuf[1] == "Exposure") next;
        expo_str = c(expo_str, strbuf[1]);
        outcome_str = c(outcome_str, strbuf[2]);
        bxy = c(bxy, as.numeric(strbuf[3]));
        bxy_se = c(bxy_se, as.numeric(strbuf[4]));
        bxy_pval = c(bxy_pval, as.numeric(strbuf[5]));
        bxy_m = c(bxy_m, as.numeric(strbuf[6]));
    }
    return(cbind(expo_str, outcome_str, bxy, bxy_se, bxy_pval, bxy_m))
}

# ************************************************** #
#                  Read SNP effects                  #
# ************************************************** #
read_snp_effect = function(file_con) {
    snp_effect = c()
    while(1) {
        strbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
        if(strbuf[1] == "#effect_end") break;
        snp_effect = rbind(snp_effect, strbuf);
        print(length(strbuf))
        if(length(strbuf)<14) print(strbuf)
    }
    return(snp_effect)
}

# ************************************************** #
#                  Read SNP instruments              #
# ************************************************** #
read_snp_instru = function(file_con, snplist, nexpo, noutcome) {
    nrow = length(snplist); ncol = nexpo+noutcome
    snp_instru = matrix(NA, nrow, ncol)
    while(1) {
        strbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
        if(strbuf[1] == "#marker_end") break;
        expo_indx = as.numeric(strbuf[1]); outcome_indx = as.numeric(strbuf[2]);
        forward_flag = TRUE;
        if(expo_indx < outcome_indx) {
            outcome_indx = outcome_indx - nexpo
        } else {
            expo_indx = expo_indx - nexpo
            forward_flag = FALSE;
        }
        snpbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
        snp_indx = match(snpbuf, snplist)
        posbuf = rep(0, nrow); posbuf[snp_indx] = 1;
        indxbuf = expo_indx
        if(!forward_flag) indxbuf = indxbuf + nexpo
        if(length(which(!is.na(snp_instru[,indxbuf])))==0) {
            snp_instru[,indxbuf] = posbuf;
        } else {
            snp_instru[,indxbuf] = paste(snp_instru[,indxbuf], posbuf, sep="")
        }
    }
    return(snp_instru)
}

# ************************************************** #
#          Read output by GCTA-GSMR for plot         #
# ************************************************** #
read_gsmr_data = function(gsmr_effect_file) {
    trait_flag = gsmr_flag = marker_flag = effect_flag = FALSE;
    file_con = file(gsmr_effect_file, "r")
    while(1) {
        strbuf = scan(file_con, nlines=1, quiet=TRUE, what="");
        if(strbuf == "#trait_begin") {
            # Read the exposures and outcomes
            resbuf = read_gsmr_trait(file_con);
            expo_str = resbuf$expo_str;
            outcome_str = resbuf$outcome_str;
            pheno_str = c(expo_str, outcome_str);
            nexpo = length(expo_str); noutcome = length(outcome_str)
            trait_flag = TRUE;
        } else if(strbuf == "#gsmr_begin") {
            # Read the GSMR result
            bxy_result = read_gsmr_result(file_con);
            colnames(bxy_result) = c("Exposure", "Outcome", "bxy", "se", "p", "n_snps")
            gsmr_flag = TRUE;
        } else if(strbuf == "#effect_begin") {
            # Read the summary statistics
            snp_effect = read_snp_effect(file_con);
            snplist = as.character(snp_effect[,1])
            effect_flag = TRUE;
        } else if(strbuf == "#marker_begin") {
            # Read the SNPs
            snp_instru = read_snp_instru(file_con, snplist, nexpo, noutcome);
            snp_effect = cbind(snp_effect[,1], snp_instru, snp_effect[,-1])
            marker_flag = TRUE;
        }
        if(trait_flag==T & gsmr_flag==T & marker_flag==T & effect_flag==T) break;
    }
    return(list(pheno=c(nexpo, noutcome, pheno_str), bxy_result=bxy_result, snp_effect = snp_effect))
}

# ************************************************** #
#         Display summary of the gsmr data           #
# ************************************************** #
gsmr_summary = function(gsmr_data) {
    message("\n## Exposure and outcome")
    pheno_str = gsmr_data$pheno[c(-1,-2)]
    # exposure
    nexpo = as.numeric(gsmr_data$pheno[1]);
    noutcome = as.numeric(gsmr_data$pheno[2]);
    logger_m = paste(nexpo, "exposure(s):");
    logger_m = paste(logger_m, gsmr_data$pheno[3])
    if(nexpo > 1) {
        for(i in 2 : nexpo)
            logger_m = paste(logger_m, gsmr_data$pheno[i+2], sep=", ")
    }
    message(logger_m)
    # outcome
    logger_m = paste(noutcome, "outcome(s):");
    logger_m = paste(logger_m, gsmr_data$pheno[3+nexpo])
    if(noutcome > 1) {
        for(i in 2 : noutcome)
            logger_m = paste(logger_m, gsmr_data$pheno[i+2+nexpo], sep=", ")
    }
    message(logger_m)

    message("\n## GSMR result")
    m_bxy_rst = data.frame(gsmr_data$bxy_result)
    print(m_bxy_rst)
}

# ************************************************** #
#               Retrieve SNP effects                 #
# ************************************************** #
gsmr_snp_effect = function(gsmr_data, expo_str, outcome_str) {
   # index of SNP instruments
    pheno_str = as.character(gsmr_data$pheno[c(-1,-2)])
    nexpo = as.numeric(gsmr_data$pheno[1])
    noutcome = as.numeric(gsmr_data$pheno[2])
    expo_indx = match(expo_str, pheno_str)
    if(is.na(expo_indx)) stop("\"", expo_str, "\" is not found.")
    outcome_indx = match(outcome_str, pheno_str)
    if(is.na(outcome_indx)) stop("\"", outcome_str, "\" is not found.")
    forward_flag = TRUE;
    if(expo_indx > outcome_indx) forward_flag = FALSE;
    if(forward_flag) {
        outcome_indx = outcome_indx - nexpo;
    } else {
        expo_indx = expo_indx - nexpo;
    }
    indxbuf = expo_indx + 1
    if(!forward_flag) indxbuf = indxbuf + nexpo
    strbuf = as.character(substr(gsmr_data$snp_effect[,indxbuf], outcome_indx, outcome_indx))
    snpindx = which(strbuf=="1")
    if(length(snpindx) < 3) stop("Not enough SNPs retained.")
    # bxy
    indxbuf = which(gsmr_data$bxy_result[,1]==expo_str & gsmr_data$bxy_result[,2]==outcome_str)
    bxy = as.numeric(gsmr_data$bxy_result[indxbuf, 3])
    # SNP effects
    if(forward_flag) {
        indxbuf1 = 1 + nexpo + noutcome + 3 + (expo_indx-1)*2 + 1
        indxbuf2 = 1 + nexpo + noutcome + 3 + nexpo*2 + (outcome_indx-1)*2 + 1
    } else {
        indxbuf1 = 1 + nexpo + noutcome + 3 + nexpo*2 + (expo_indx-1)*2 + 1
        indxbuf2 = 1 + nexpo + noutcome + 3 + (outcome_indx-1)*2 + 1
    }
    snpid = as.character(gsmr_data$snp_effect[snpindx,1])
    bzx = as.numeric(gsmr_data$snp_effect[snpindx,indxbuf1]); indxbuf1 = indxbuf1 + 1;
    bzx_se = as.numeric(gsmr_data$snp_effect[snpindx,indxbuf1]);
    bzx_pval = pchisq((bzx/bzx_se)^2, 1, lower.tail=F);
    bzy = as.numeric(gsmr_data$snp_effect[snpindx,indxbuf2]); indxbuf2 = indxbuf2 + 1;
    bzy_se = as.numeric(gsmr_data$snp_effect[snpindx,indxbuf2]);
    bzy_pval = pchisq((bzy/bzy_se)^2, 1, lower.tail=F);
    return(list(snp=snpid, bxy=bxy, bzx=bzx, bzx_se=bzx_se, bzx_pval=bzx_pval, bzy=bzy, bzy_se=bzy_se, bzy_pval=bzy_pval))
}

# ************************************************** #
#                  Plot bzy vs bzx                   #
# ************************************************** #
plot_snp_effect = function(expo_str, outcome_str, bxy, bzx, bzx_se, bzy, bzy_se, effect_col=colors()[75]) {
    vals = c(bzx-bzx_se, bzx+bzx_se)
    xmin = min(vals); xmax = max(vals)
    vals = c(bzy-bzy_se, bzy+bzy_se)
    ymin = min(vals); ymax = max(vals)
    plot(bzx, bzy, pch=20, cex=0.8, bty="n", cex.axis=1.1, cex.lab=1.2,
         col=effect_col, xlim=c(xmin, xmax), ylim=c(ymin, ymax),
         xlab=substitute(paste(trait, " (", italic(b[zx]), ")", sep=""), list(trait=expo_str)),
         ylab=substitute(paste(trait, " (", italic(b[zy]), ")", sep=""), list(trait=outcome_str)))
    if(!is.na(bxy)) abline(0, bxy, lwd=1.5, lty=2, col="dim grey")
    ## Standard errors
    nsnps = length(bzx)
    for( i in 1:nsnps ) {
        # x axis
        xstart = bzx[i] - bzx_se[i]; xend = bzx[i] + bzx_se[i]
        ystart = bzy[i]; yend = bzy[i]
        segments(xstart, ystart, xend, yend, lwd=1.5, col=effect_col)
        # y axis
        xstart = bzx[i]; xend = bzx[i]
        ystart = bzy[i] - bzy_se[i]; yend = bzy[i] + bzy_se[i]
        segments(xstart, ystart, xend, yend, lwd=1.5, col=effect_col)
    }
}

# ************************************************** #
#             Plot bzy_pval vs bzx_pval              #
# ************************************************** #
plot_snp_pval = function(expo_str, outcome_str, bzx_pval, bzy_pval, gwas_thresh, truncation, effect_col) {
    eps = 1e-300; truncation = -log10(truncation);
    if(truncation > 300) {
        warning("The minimal truncated p-value would be 1e-300.")
        truncation = 300
    }
    bzx_pval = -log10(bzx_pval + eps);
    bzy_pval = -log10(bzy_pval + eps);
    pval = c(bzx_pval, bzy_pval)
    min_val = 0; max_val = max(pval);
    max_val = ifelse(max_val > truncation, truncation, max_val)
    gwas_thresh = -log10(gwas_thresh);
    plot(bzx_pval, bzy_pval, pch=20, cex=0.8, bty="n", cex.axis=1.1, cex.lab=1.2,
         col=effect_col, xlim=c(min_val, max_val), ylim=c(min_val, max_val),
         xlab=substitute(paste(trait, " (", -log[10], italic(P)[zx], ")", sep=""), list(trait=expo_str)),
         ylab=substitute(paste(trait, " (", -log[10], italic(P[zy]), ")", sep=""), list(trait=outcome_str)))
    abline(h=gwas_thresh, lty=2, lwd=1.5, col="maroon")
}

# ************************************************** #
#                Plot bxy vs bzx_pval                #
# ************************************************** #
plot_snp_bxy = function(expo_str, outcome_str, bxy, bzx_pval, effect_col) {
    eps = 1e-300;
    bzx_pval = -log10(bzx_pval + eps);
    xmin = min(bxy, na.rm=T); xmax = max(bxy, na.rm=T)
    ymin = min(bzx_pval); ymax = max(bzx_pval);
    plot(bxy, bzx_pval, pch=20, cex=0.8, bty="n", cex.axis=1.1, cex.lab=1.2,
         col=effect_col, xlim=c(xmin, xmax), ylim=c(ymin, ymax),
         xlab=substitute(paste(italic(hat(b)[xy]), " (", trait1, " -> ", trait2, ")", sep=""), list(trait1=expo_str, trait2=outcome_str)),
         ylab=substitute(paste(trait, " (", -log[10], italic(P[zx]), ")", sep=""), list(trait=expo_str)))
}

# ************************************************** #
#                  Effect size plot                  #
# ************************************************** #
# expo_str, exposure
# outcome_str, outcome
# effect_col, plotting colour
plot_gsmr_effect = function(gsmr_data, expo_str, outcome_str, effect_col=colors()[75]) {
    resbuf = gsmr_snp_effect(gsmr_data, expo_str, outcome_str);
    bxy = resbuf$bxy
    bzx = resbuf$bzx; bzx_se = resbuf$bzx_se;
    bzy = resbuf$bzy; bzy_se = resbuf$bzy_se;
    # plot
    plot_snp_effect(expo_str, outcome_str, bxy, bzx, bzx_se, bzy, bzy_se, effect_col)
}

# ************************************************** #
#                    P-value plot                    #
# ************************************************** #
# expo_str, exposure
# outcome_str, outcome
# effect_col, plotting colour
plot_gsmr_pvalue = function(gsmr_data, expo_str, outcome_str, gwas_thresh=5e-8, truncation=1e-50, effect_col=colors()[75]) {
    resbuf = gsmr_snp_effect(gsmr_data, expo_str, outcome_str);
    bzx_pval = resbuf$bzx_pval; bzy_pval = resbuf$bzy_pval;
    # plot
    plot_snp_pval(expo_str, outcome_str, bzx_pval, bzy_pval, gwas_thresh, truncation, effect_col)
}

# ************************************************** #
#                     bxy distribution plot                         #
# ************************************************** #

# expo_str, exposure
# outcome_str, outcome
# effect_col, plotting colour
plot_bxy_distribution = function(gsmr_data, expo_str, outcome_str, effect_col=colors()[75]) {
    resbuf = gsmr_snp_effect(gsmr_data, expo_str, outcome_str);
    bzx = resbuf$bzx; bzx_pval = resbuf$bzx_pval;
    bzy = resbuf$bzy;
    bxy = bzy/bzx
    # plot
    plot_snp_bxy(expo_str, outcome_str, bxy, bzx_pval, effect_col)
}


## Author: Alan Ruttenberg
## Project: OHD
## Date: 2015-04-24
##
## Survival Analysis: Do it all: resin_restoration_survival_analysis()

library(survival)
library("simPH")

# Create an rsurv "object", mostly so that we can control printing
# Sigh, have to be careful if we want to call by reference - i.e. slot changes persist
# http://www.stat.berkeley.edu/~paciorek/computingTips/Pointers_passing_reference_.html

rsurv <- function()
    { object=new.env(parent=globalenv())
      class(object)="rsurv"
      lastrsurv <<-  object;
      return(object)
  }

# We have a lot of state - don't print it all
print.rsurv <- function(object) { cat("<restoration survival analysis>\n") }

# Main entry point. Query the database twice - once to fetch all
# restorations, once to fetch all the restorations with a known
# failure.
# Returns an rsurv object. Important methods: plot, summary.

resin_restoration_survival_analysis <- function ()
    { s <- rsurv()
      s$all <- collect_all_restorations_and_latest_encounter_after()
      s$fail <- collect_restoration_failures()
      create_rsurv(s)
  }

# Summary, for now, is the correlates table
summary.rsurv <- function (it)
    { summary(it$correlates);
  }

# Take the set of all resin restorations and subtract out those which
# we know failed. The rest will be censored
# Could do this in Sparql, but hadn't learned enough
# SPARQL 1.1 when this was first written.
subset_all_minus_fail <- function(all, fail)
{
  rownames(all)=all[,"proci1"]
  rownames(fail)=fail[,"proci1"]
  keep <- setdiff(rownames(all),rownames(fail));;
  res <- as.matrix(all[keep,]);
}

# Helper functions.

# Takes two date strings(yyyy-mm-yy) and uses seq.date to return the
# time difference in months. Rounds down.

difference_in_months <- function(start,end)
    { length(seq(as.Date(start),as.Date(end),by="month"))-1 }

# Takes two date strings(yyyy-mm-yy) and uses seq.date to return the
# time difference in years. Rounds down.

difference_in_years <- function(start,end)
    { length(seq(as.Date(start),as.Date(end),by="year"))-1 }

# create arrays of correlates aligned with survival data - first the
# the ones we know fail, then the ones that will be censored.  Takes
# as input the rsurv object, where we have already prepared the
# matrices "fail" and "have_last_visit", each having one row per
# "event".
# Sets the "correlates" attribute to
# the created data frame.

prepare_survival_correlates <- function(s)
    { s$correlates <- rbind(data.frame(location=as.factor(s$fail[,"tooth_type"]),
                                       gender=as.factor(s$fail[,"gender"]),
                                       age=s$age_at_fail,
                                       patient=s$fail[,"patienti"],
                                       fail=TRUE),
                            data.frame(location=as.factor(s$have_last_visit[,"tooth_type"]),
                                       gender=as.factor(s$have_last_visit[,"gender"]),
                                       age=s$age_at_last_visit,
                                       patient=s$have_last_visit[,"patienti"],
                                       fail=FALSE))
      # divide the ages into bins, [0,18) , [18-30), [30,60), over 60
      s$correlates[,"age_group"] <- as.factor(cut(s$correlates[,"age"],breaks=c(0,18,30,60,100)));
  }

# having the failures (s$fail) and all events (s$all), create
# s$have_last_visit as the events where we have a restoration and a
# subsequent visit on record. These will be treated as right censored

prepare_censored_events <- function(s)
    {
        # Compute all restorations without a subsequent failure on record
        no_record_of_failure <- NULL;
        no_record_of_failure <- subset_all_minus_fail(s$all,s$fail);

        # Filter out ones without a subsequent visit - we can't use those at all. Call those $have_last_visit.
        s$have_last_visit <- no_record_of_failure[!is.na(no_record_of_failure[,"latest_date2"]),];

        # Calculate the time in months between the restoration and the last visit date on record
        s$no_record_of_failure_in_months <-
            mapply(difference_in_months,s$have_last_visit[,"date1"],s$have_last_visit[,"latest_date2"])

        # Calculate the age at last visit
        s$age_at_last_visit <- mapply(difference_in_years,s$have_last_visit[,"birthdate"],s$have_last_visit[,"date1"])

        s$ncensor = length(s$no_record_of_failure_in_months);
    }

# Calculate time until failure and the age at failure for those events
# where there was a failure. We consider the failure event, for the
# purpose of survival analysis, to be the date the failure is noticed
# or replaced. For future improvement - treat the failure as interval
# censored.

prepare_failure_events <- function(s)
{
    s$fail_difference_in_months <-
        mapply(difference_in_months,s$fail[,"date1"],s$fail[,"soonest_date2"])
    s$age_at_fail <- mapply(difference_in_years,s$fail[,"birthdate"],s$fail[,"date1"])
    s$nfail <- length(s$fail_difference_in_months);
}

# Main analysis
create_rsurv <- function (s)
    {
        prepare_censored_events(s)

        prepare_failure_events(s)

        # Now compute the main inputs to the survival analysis.
        # s$status is 1 if the event was a failure, 0 if right censored.
        # combine the fail time and last visit times to make a vector
        # of time in months, and divide by 12 to get years.
        s$status <- c(rep(1,s$nfail),rep(0,s$ncensor));
        s$months <- c(s$fail_difference_in_months,s$no_record_of_failure_in_months)
        s$years <- c(s$fail_difference_in_months,s$no_record_of_failure_in_months)/12;
        prepare_survival_correlates(s)

        # Do the survival analysis, two argument version.
        surv <- Surv(s$years,s$status) ;
        s$surv <- surv;

        # Now do the univariate fits
        s$fit <- survfit(surv~1,conf.type="none",data=s$correlates);
        s$location_fit <- survfit(surv~location,conf.type="none",data=s$correlates);
        s$gender_fit <- survfit(surv~gender,conf.type="none",data=s$correlates);
        s$age_fit <- survfit(surv~age_group,conf.type="none",data=s$correlates);
        # Make the age groups prettier
        names(s$age_fit$strata)=
            c("age=18 and under", "age=18+ to 30", "age=30+ to 60", "age=60 and over")

        # Compute the proportional hazards
        s$cox<-coxph(surv~location+gender+age, method="breslow",data=s$correlates)

        s
    }

# Draw one of the Kaplan-Meier plots
# We use ggplot, and hunt the web for recipes to show the bits we want
# to show.

plot_rsurv_fit <- function (s,fit=s$location_fit,
                            title="All",
                            xlab="Time in years",
                            ylab="Fraction surviving",
                            ci=F,col=c("dark green","dark blue")
                           )
    { print(
        ggsurv(fit,cens.col='gray',plot.cens=F,back.white=T,
               main=title,xlab=xlab,ylab=ylab,CI=ci,surv.col=col)
        # set up the x scale
       + scale_x_continuous(breaks=1:13,expand=c(0,0),limits=c(0,12.5))
        # set up the y scale
       + scale_y_continuous(expand=c(0,.01))
        # axes font bold (tweak position to look better too)
       + theme(axis.title.y=element_text(vjust=1.5,face="bold"))
       + theme(axis.title.x=element_text(vjust=-.7,face="bold"))
        # So is title
       + theme(title=element_text(vjust=2,face="bold"))
        # Place the legend inside the plot
       + theme(legend.position=c(.9,.5),
               legend.text=element_text(size=12,face="bold"),
               legend.title=element_text(size=14))
        #add guides
        + guides(colour = guide_legend(override.aes = list(size=3))));

  }

# Draw all of the plots. plotWrap prepares the device. bplotf devault
# uses svg device and sends to browser.

plot.rsurv <- function (sa,plotWrap=bplotf)
{ plotWrap(function ()
    {plot_rsurv_fit(sa,title="Resin restoration longevity",
                fit=sa$fit,
                col="dark blue")},
         index="",filebase="overall-longevity")
  plotWrap(function ()
      {plot_rsurv_fit(sa,title="Resin restoration longevity by age at failure",
                  fit=sa$age_fit,
                  col=c("dark gray","dark green","dark blue","dark red"))},
         index="",filebase="longevity-by-age")
  plotWrap(function ()
      {plot_rsurv_fit(sa,title="Resin restoration longevity by gender",
                  fit=sa$gender_fit,
                  col=c("dark blue","dark red"))},
         index="",filebase="longevity-by-gender")
  plotWrap(function ()
      {plot_rsurv_fit(sa,title="Resin restoration longevity by location",
                  fit=sa$location_fit,
                  col=c("dark blue","dark red"))},
         index="",filebase="longevity-by-location")
}

# if factors, number of events for each factor
summary_text <- function(surv,correlate)
    { levels <- levels(as.factor(lastrsurv$correlates[,correlate]))
      c <- surv$correlates;
      factor_summary <- function(value,factor=correlate)
          { patient_count <- length(unique(c[(c[,factor]==value),"patient"]));
            event_count <- nrow(c[(c[,factor]==value),])
            failures_count <- nrow(c[(c[,factor]==value &  c[,"fail"]),]);
            paste0(value,": ",patient_count, " patients", ", ",event_count," events of which ",failures_count," were failures")
        }
      do.call(paste,args=c(lapply(as.list(levels),factor_summary),sep="\n"))
  }

# levels(as.factor(lastrsurv$correlates[,"age_group"]))

## sa$cox$coef["locationposterior"]
## sa$cox$coef["gendermale"]

## bplotf(function(){
## + baseline<-basehaz(coxph(x$gender_fit))
## + plot(baseline$time, baseline$hazard, type='l',main="Hazard rates")
## + lines(baseline$time, exp(0.8245)*baseline$hazard, col="blue") })

## sa$correlates[,"AgeMed"]<-sa$correlates[,"age"]-40
## Sim1 <- coxsimLinear(M1, b = "AgeMed", nsim = 100,Xj = seq(-30, 30, by=1))
## bplotf(function()  {print(simGG(Sim1))})

## baseline<-basehaz(reg_fit)

## #http://www.unc.edu/courses/2010spring/ecol/562/001/docs/lectures/lecture24.htm#coxzph
## #http://www.uni-kiel.de/psychologie/rexrepos/posts/survivalCoxPH.html
## #http://www.statisticsmentor.com/category/r-survival-analysis/