############### R-CODE FOR PROJECT TYCHO CIRCLE DIAGRAM #########################################################
# This code provides the basic structure of a Projec Tycho heatmap. This code has not been generalized          #
# to use as function to various types of data, but is meant to reveal the structure and functions used for      #
# adaptation to other types of data.                                                                            #
#                                                                                                               #
# Please cite the original publication of this graph as the origin of this code or refer to the Project Tycho   #
# website (www.tycho.pitt.edu).                                                                                 #
#                                                                                                               #
# The original study citatiion is: Willem G. van Panhuis, John Grefenstette, Su Yon Jung, Nian Shong Chok,      #
# Anne Cross, Heather Eng, Bruce Y Lee, Vladimir Zadorozhny, Shawn Brown, Derek Cummings, Donald S. Burke.      #
# Contagious Diseases in the United States from 1888 to the present. NEJM 2013; 369(22): 2152-2158              #
#                                                                                                               #
# This code has been created by: Wilbert van Panhuis, University of Pittsburgh Graduate School of Public        #
# Health, Email: wav10@pitt.edu)                                                                                #
#################################################################################################################

#################################################################################################################
# This code includes the following subcomponents:                                                               #
# 1. a modification of the contour function                                                                     #
# 2. function to plot the heatmap including:                                                                    #
#       2.1 top panel with barplot                                                                              #
#       2.2 main color display                                                                                  #
#       2.3 region labels grouping states in epidemiological regions                                            #
#       2.4 side color legend                                                                                   #
#################################################################################################################

#################################################################################################################
# This code uses the following data:                                                                            #
# 1. matrix of values for main color display with rows representing time points and columns representing        #
#   locations                                                                                                   #
# 2. data.frame of values for the top panel barplot with time points in the first column and values in the      #
#   second column                                                                                               #
#################################################################################################################

#   MODIFIED CONTOUR FUNCTION
contour_wvp<-function (x = seq(0, 1, length.out = nrow(z)), y = seq(0, 1, 
    length.out = ncol(z)), z, xlim = range(x, finite = TRUE), 
    ylim = range(y, finite = TRUE), zlim = range(z, finite = TRUE), 
    levels = pretty(zlim, nlevels), nlevels = 20, color.palette = cm.colors, 
    col = color.palette(length(levels) - 1), plot.title, plot.axes, 
    key.title, key.axes, asp = NA, xaxs = "i", yaxs = "i", las = 1, 
    axes = TRUE, frame.plot = axes, ...) 
{
    if (missing(z)) {
        if (!missing(x)) {
            if (is.list(x)) {
                z <- x$z
                y <- x$y
                x <- x$x
            }
            else {
                z <- x
                x <- seq.int(0, 1, length.out = nrow(z))
            }
        }
        else stop("no 'z' matrix specified")
    }
    else if (is.list(x)) {
        y <- x$y
        x <- x$x
    }
    if (any(diff(x) <= 0) || any(diff(y) <= 0)) 
        stop("increasing 'x' and 'y' values expected")
        
    plot.window(xlim, ylim, "", xaxs = xaxs, yaxs = yaxs, asp = asp)
    if (!is.matrix(z) || nrow(z) <= 1L || ncol(z) <= 1L) 
        stop("no proper 'z' matrix specified")
    if (!is.double(z)) 
        storage.mode(z) <- "double"
   .filled.contour(as.double(x), as.double(y), z, as.double(levels), 
                            col = col)
       if (missing(plot.axes)) {
        if (axes) {
            title(main = "", xlab = "", ylab = "")
            Axis(x, side = 1)
            Axis(y, side = 2)
        }
    }
    else plot.axes
    if (frame.plot) 
        box()
    if (missing(plot.title)) 
        title(...)
    else plot.title
    invisible()
}


#   FUNCTION TO PLOT THE HEATMAP, USING THE FOLLOWING INPUT:
#   1. dat: matrix of values for main color display (see data no. 1 above)
#       - it works well to recode NA as -2
#   2. dat2: data.frame of values for the top panel barplot (see data no. 2 above)
#   3. cuts: a vector of cutpoints for the color display
#       - ensure that NA has a separate value, eg. -2
#   4. cutwk: a time point for a vertical red line
#   5. name: name to be used for the plot saving function
#   6. st: a list of location names in the same order as columns in "dat"
#   7. regions: a list of region numbers for each of the locations in "dat"
#   8. ylab: the label to be used for the y-axis
 
heat_time_som<-function(dat, dat2, cuts, cutwk, name, st, regions, ylab){

# derive parameters for cuts/legend
l<-length(cuts)
v=which(substr(rownames(dat),2,7)==cutwk)
n<-ncol(dat)

# define axes
mataxes<-function(dat){
    abline(v=v, lwd=10, col="red")
    axis(side=2, at=seq(from=1.1, to=51.6, by=50/50.3), labels=st, cex.axis=0.9, tick=T, las=1, lwd.ticks=2)
    axis(side=1, at=c(1:(nrow(dat))), labels=substr(rownames(dat),2,5), tick=F, cex.axis=1.2, line=-0.5)
}

# plot
dev.new(width=30, height=30) # this uses the on screen graphics display
par(fig=c(0.08,0.85,0,0.84), mar=c(4,4,0.5,0.5))

# main color display
plot.new()
contour_wvp(x=c(1:(nrow(dat))), y=c(1:ncol(dat)), z=dat, plot.axes=mataxes(dat), levels=cuts, col=c("black",rev(rainbow(l-2, start=0.0, end=0.65))), axes=F, frame.plot=T,
    xlab="")
mtext(side=1, "Week", line=2, cex=1.3)

# legend color bar
colmat<-t(matrix(rep(cuts,2), ncol=2))
par(fig=c(0.85,1,0,0.84), new=T, mar=c(4,0,0.5,5))
contour_wvp(x=c(1:2), y=c(1:l), z=colmat, levels=cuts, col=c("black",rev(rainbow(l-2, start=0.0, end=0.65))), axes=F, frame.plot=T, lwd=2)
axis(side=4, at=c(1.4,2:l), labels=c("NA",0,0,cuts[4:l]), las=1, lwd=0, tick=T, cex.axis=1.2, lwd.ticks=2)
box(lwd=2, col="black")

# region labels on y-axis
regmat<-t(matrix(rep(regions,2), ncol=2))
par(fig=c(0,0.1,0,0.84), new=T, mar=c(4,2.5,0.5,0.5))
contour_wvp(regmat, axes=F, levels=c(1,2,3,4,5,6,7,8,9,9.8,10.1), col=c("white"))
axis(side=4, at=c(1, 0.8864874, 0.8469513, 0.7268786, 0.5658056, 0.4471972, 0.3461605, 0.2670883, 0.1499442, 0.0694077, 0), lwd.ticks=2, tick=T, labels=NA)
text(x=rep(0.5, 10), y=c(0.94359512, 0.86891579, 0.78837926, 0.65073501, 0.51016216, 0.39887532, 0.30808869, 0.21290915, 0.11187241, 0.03426449), 
    labels=c(1:10), cex=1.2)
mtext("US Region/State", side=2, line=0.5, cex=1.3)

# top panel barplot
par(fig=c(0.08,0.85,0.84,1), new=T, mar=c(0,4,2,0.5))
barplot(c(dat2$inc), names.arg="", col="black", border="black", ylim=c(0, max(dat2$inc, na.rm=T)), xaxs="i", ylab="", xlab="", axes=F)
axis(side=2, at=ylab, labels=ylab, cex.axis=1.2, las=1, lwd.ticks=2)
box(lwd=2, col="black")
mtext("Incidence rate", side=2, line=4, cex=1.2)
mtext("(/100,000)", side=2, line=3, cex=1.2)

# save
savePlot(paste("C:\\Vanpanhuis\\Research\\Manuscripts\\In_progress\\Tycho_impactvaccines\\Figures\\Revision\\", name, "_heatsup", sep=""), type="jpeg")
dev.off()
}