#################################################################
# Carl Schmertmann
# 15 Nov 2013
#
# draw and save a standardardized set of plots 
#  from a forecast output file
#
#   this program must be run in a directory that ALREADY has a 
#   subdirectory called "images"
#################################################################

rm(list=ls())
graphics.off()

require(RColorBrewer)
require(denstrip)

windows(record=T)

fcast    = dget("Forecast Fri 15 Nov 2013 1204.dput")
imageDir = "./images/"                                    ## (relative) location of image directory

## alternative if using simulated 1985 forecast
# fcast    = dget("1985 Forecast Fri 08 Nov 2013 1719.dput")
# imageDir = "./1985images/"                             ## (relative) location of image directory


full.name = c(
  AUT         = "Austria",
  BGR         = "Bulgaria",
  CAN         = "Canada",
  CZE         = "Czech Rep",
  EST         = "Estonia",
  FIN         = "Finland",
  FRATNP      = "France",
  DEUTNP      = "Germany",
  DEUTW       = "Western Germany",
  DEUTE       = "Eastern Germany",
  HUN         = "Hungary",
  LTU         = "Lithuania",
  NLD         = "Netherlands",
  PRT         = "Portugal",
  RUS         = "Russia",
  SVK         = "Slovakia",        
  SVN         = "Slovenia",
  SWE         = "Sweden",
  CHE         = "Switzerland",
  GBR_NP      = "Great Britain",
  GBRTENW     = "England and Wales",
  GBR_SCO     = "Scotland",
  GBR_NIR     = "Northern Ireland",
  USA         = "USA",   
  Australia   = "Australia",
  Belgium     = "Belgium",
  Denmark     = "Denmark",
  Greece      = "Greece",
  Iceland     = "Iceland",
  Italy       = "Italy",
  Japan       = "Japan",
  Korea       = "Korea",
  Luxembourg  = "Luxembourg",
  New_Zealand = "New Zealand",
  Romania     = "Romania",
  Singapore   = "Singapore",
  Brazil      = "Brazil"
)

#############################################################
# save the current screen plot to a png file
#############################################################

plotsave = function(fname) {
  png( file=fname, bg="white", width=8, height=8, units="in", res=250)
  dev.set(dev.prev())                       # swap to previous
  dev.copy(which=dev.next())                # copy previous -> new
  dev.off()                                 # close new
}

#############################################################


Lexis = function( X ,ti="") {
  hues = colorRampPalette( brewer.pal(9,"YlOrRd"))(41)
  brks = c( seq(0,.20,len=41), 1)
  image( coh, age, t(X), 
         col=hues,breaks=brks,
         xlab="Cohort Year of Birth", ylab="Age", main=ti)
  abline(v=seq(1950,2000,5), h=seq(15,50,5),lty=2,col=grey(.30))
} #Lexis()

LexisSD = function( X ,ti="") {
  hues = colorRampPalette( brewer.pal(9,"Greys"))(26)
  brks = c( seq(0,.025,len=26), 1)
  image( coh, age, t(X), 
         col=hues,breaks=brks,
         xlab="Cohort Year of Birth", ylab="Age", main=ti)
  abline(v=seq(1950,2000,5), h=seq(15,50,5),lty=2,col=grey(.30))
  contour( coh, age, t(X), breaks=brks, col="orangered", lwd=2, add=TRUE)
} #Lexis()


penalty.plot = function(pen, ti="") {
  plot(seq(pen), pen ,  type="n",  ylim=c(0,max(100,pen,na.rm=TRUE)), 
       main=ti, axes=F, xlab="Penalty", ylab="")
  
  lines( 1:30   , pen[1:30]   , lty=1, col="dodgerblue", lwd=3)
  lines( 30+1:30, pen[30+1:30], lty=1, col="dodgerblue", lwd=3)
  lines( 60+1:30, pen[60+1:30], lty=1, col="dodgerblue", lwd=3)
  
  axis(1,pos=0,
       at=c(seq(5,25,10), seq(31,56,5), seq(61,86,5)), 
       labels=c(seq(1970,1990,10), seq(15,40,5), seq(15,40,5)), tck=0)
  axis(2)
  
  yy = mean( c(0,max(100,pen,na.rm=TRUE)))
  text( 16, yy, "Shape", cex=1.50 , col='darkblue')
  text( 46, yy, "Freeze\nSlope", cex=1.50 , col='darkblue')
  text( 76, yy, "Freeze\nRate", cex=1.50 , col='darkblue')
  abline( v=c(30.5, 60.5), col='darkblue', lty=2)
  
  segments( 1,27,30,27, col='darkblue', lty=2, lwd=1)
  segments(31,30,60,30, col='darkblue', lty=2, lwd=1)
  segments(61,30,90,30, col='darkblue', lty=2, lwd=1)
  
}


###############################################################################
for (k in names(fcast)) {

  tmp  = fcast[[k]]
  
  if (!is.null(tmp)) {

    coh  = as.numeric( colnames(tmp$mu))
    age  = as.numeric( rownames(tmp$mu))
    vis  = tmp$vis
    surf = tmp$surf
    forecast.year = tmp$forecast.year
    last.coh      = tmp$last.coh
    
    Lexis( ifelse(vis,1,NA) * surf, paste(full.name[k],"Data\nOfficial Statistics"))
    plotsave( fname = paste(imageDir,k," Lexis data.png",sep=""))
    
    Lexis( tmp$mu, paste(full.name[k],"Forecast\nPosterior Mean"))
    plotsave( fname = paste(imageDir,k," Lexis mean.png",sep=""))
    
    LexisSD( tmp$sd, paste(full.name[k],"Uncertainty\nPosterior SD"))
    plotsave( fname = paste(imageDir,k," Lexis sd.png",sep=""))

    penalty.plot( tmp$pi,paste(full.name[k],"Penalties")) 
    plotsave( fname = paste(imageDir,k," penalties.png",sep=""))
    
    ## cohort schedules
    sel = last.coh + seq(-35,-10,5)
    nsel = length(sel)
    matplot( age, tmp$mu[,paste(sel)], type="l",
             lty=2,,lwd=3, col=rainbow(nsel), 
             main=paste(full.name[k],"Cohort Schedules\nPosterior Means"),
             xlab="Age",ylab="rate",bty="l",ylim=c(0,.20))
    matlines( age, (ifelse(vis,1,NA)*tmp$mu)[,paste(sel)], type="l",
            lty=1, col=rainbow(nsel), lwd=6)

    legend(35, .20, paste(sel), lwd=6, lty=1, col=rainbow(nsel),bty="n")

    plotsave( fname = paste(imageDir,k," cohorts.png",sep=""))
    
    
    ## time series
    sel = seq(20,40,5)
    ix  = match( sel, age)
    nsel = length(sel)
    
    matplot( coh, t(tmp$mu), type="l",
             lty=2,,lwd=2, col=rainbow(30), 
             main=paste(full.name[k],"Time Series\nPosterior Means"),
             xlab="Cohort Year of Birth",ylab="rate",bty="l",ylim=c(0,.25))
    
    matlines( coh, t((ifelse(vis,1,NA)*tmp$mu)), type="l",
              lty=1, col=rainbow(30), lwd=4)
    
    legend(1980, .25, paste(sel), lwd=6, lty=1, col=rainbow(30)[ix],bty="n")
    
    plotsave( fname = paste(imageDir,k," time series.png",sep=""))
    
    ## CFR forecast
    
    xval = seq(min(coh), max(coh)-5, len=10*length(coh))
    
    xx = spline( coh[1:35], tmp$mu.cfr[1:35], xout=xval)$x
    mm = spline( coh[1:35], tmp$mu.cfr[1:35], xout=xval)$y
    ss = spline( coh[1:35], tmp$sd.cfr[1:35], xout=xval)$y
    
    delta = xx[2]-xx[1]
    
    hue = "dodgerblue"
    nbands = 256
    plot(coh[1:35], tmp$mu.cfr[1:35], type="n", main=paste(full.name[k],forecast.year,"forecast"),
         ylim=c(0.8,max(3, 1.10*tmp$mu.cfr )), bty="l", xlim=range(coh,coh-3),
         xlab="COHORT YEAR OF BIRTH", ylab="CFR", cex.lab=1.40, cex.axis=1.40)
    
    for (cc in seq(xx)) {
      denstrip.normal( mm[cc], ss[cc], 
                       at = xx[cc], horiz=FALSE, colmax=hue, width=delta, n=nbands)
    }
    
    abline(h=seq(0,6,.25),lty=2,col="grey")   
    points( max(coh)-5+0.5, tmp$mu.cfr[35], pch=16, col=hue, cex=1.50)
    segments( max(coh)-5+0.5, qnorm(.75, tmp$mu.cfr[35], tmp$sd.cfr[35]), 
              max(coh)-5+0.5, qnorm(.95, tmp$mu.cfr[35], tmp$sd.cfr[35]), 
              col=hue, lwd=5)
    segments( max(coh)-5+0.5, qnorm(.25, tmp$mu.cfr[35], tmp$sd.cfr[35]), 
              max(coh)-5+0.5, qnorm(.05, tmp$mu.cfr[35], tmp$sd.cfr[35]), 
            col=hue, lwd=5)  
    text( max(coh)-3, qnorm(c(.05,.25,.50,.75,.95), tmp$mu.cfr[35], tmp$sd.cfr[35]),
          c("5%ile", "25%", "50%","75%","95%"), cex=0.80, adj=0,
          col=hue)
    
    text( last.coh - seq(5,30,5), 0.9, seq(20,45,5), col=hue, cex=1.30)
    text( 1+min(coh), 1.0, "Age at Forecast", adj=0, col=hue, cex=1.30)
    abline(h=2)
    
    solid = apply(surf, 2, function(x) all(is.finite(x)))
    
    cfr   = apply(surf,2,sum)
    points(coh[solid], cfr[solid], type="p", pch=16, col="navy", cex=1.5)
    points(coh[!solid], cfr[!solid], type="p", pch=16, col="navy", cex=0.8)
    points(coh[!solid], cfr[!solid], type="p", pch=1, col="navy", cex=1.5)
    
    plotsave( fname = paste(imageDir,k," cfr.png",sep=""))
    
    } # if not is.null
} # for k