#' ---
#' title: "Regression and Other Stories: Elections Economy"
#' author: "Andrew Gelman, Jennifer Hill, Aki Vehtari"
#' date: "`r format(Sys.Date())`"
#' output:
#'   html_document:
#'     theme: readable
#'     toc: true
#'     toc_depth: 2
#'     toc_float: true
#'     code_download: true
#' ---

#' Present uncertainty in parameter estimates. See Chapter 7 in
#' Regression and Other Stories.
#' 
#' -------------
#' 

#+ setup, include=FALSE
knitr::opts_chunk$set(message=FALSE, error=FALSE, warning=FALSE, comment=NA)
# switch this to TRUE to save figures in separate files
savefigs <- FALSE

#' #### Load packages
library("rprojroot")
root<-has_file(".ROS-Examples-root")$make_fix_file()
library("arm")
library("rstanarm")

#' #### Load data
hibbs <- read.table(root("ElectionsEconomy/data","hibbs.dat"), header=TRUE)
head(hibbs)

#' ## Likelihood for 2 parameters
M1 <- lm(vote ~ growth, data = hibbs)
display(M1)
summ <- summary(M1)

#' #### Plot likelihood (a, b| y)
#+ eval=FALSE, include=FALSE
if (savefigs) pdf(root("ElectionsEconomy/figs","hill_2a.pdf"), height=4, width=5)
#+
# Contour plots etc of simple likelihoods
trans3d <- function(x,y,z, pmat) {
       tr <- cbind(x,y,z,1) %*% pmat
       list(x = tr[,1]/tr[,4], y= tr[,2]/tr[,4])
     }
dmvnorm <- function (y, mu, Sigma, log=FALSE){
  # multivariate normal density
  n <- nrow(Sigma)
  logdens <- -(n/2)*log(2*pi*det(Sigma)) - t(y-mu)%*%solve(Sigma)%*%(y-mu)/2
  return (logdens)
#  return (ifelse (log, logdens, exp(logdens)))
}
#
rng.x <- summ$coef[1,1] + summ$coef[1,2]*c(-4,4)
rng.y <- summ$coef[2,1] + summ$coef[2,2]*c(-4,4)
x <- seq(rng.x[1], rng.x[2], length=30)
y <- seq(rng.y[1], rng.y[2], length=30)
z <- array(NA, c(length(x),length(y)))
for (i.x in 1:length(x))
  for (i.y in 1:length(y))
    z[i.x,i.y] <- dmvnorm(c(x[i.x],y[i.y]), summ$coef[,1], summ$cov.unscaled*summ$sigma^2, log=TRUE)
z <- exp(z-max(z))
par(mar=c(0, 0, 0, 0))
persp(x, y, z,
  xlim=c(rng.x[1]-.15*(rng.x[2]-rng.x[1]), rng.x[2]), ylim=c(rng.y[1]-.15*(rng.y[2]-rng.y[1]), rng.y[2]),
  xlab="a", ylab="b", zlab="likelihood", d=2, box=FALSE, axes=TRUE, expand=.6) -> res
text(trans3d(mean(rng.x), rng.y[1]-.12*(rng.y[2]-rng.y[1]), 0, pm = res), expression(beta[0]))
text(trans3d(rng.x[1]-.08*(rng.x[2]-rng.x[1]), mean(rng.y), 0, pm = res), expression(beta[1]))
mtext("likelihood, p(a, b |y)", side=3, line=-1.5)
#+ eval=FALSE, include=FALSE
if (savefigs) dev.off()

#' #### Plot maximum likelihood estimate and std errs
#+ eval=FALSE, include=FALSE
if (savefigs) pdf(root("ElectionsEconomy/figs","hill_2b.pdf"), height=5, width=5)
#+
par(mar=c(3, 3, 3, 1), mgp=c(1.7, .5, 0), tck=-.01)
plot(rng.x, rng.y, xlab="a", ylab="b", main=expression(paste("(", hat(a) %+-% 1, " std err,  ", hat(b) %+-% 1, " std err)")), type="n")
lines(rep(summ$coef[1,1], 2), summ$coef[2,1] + c(-1,1)*summ$coef[2,2], col="gray20")
lines(summ$coef[1,1] + c(-1,1)*summ$coef[1,2], rep(summ$coef[2,1], 2), col="gray20")
points(summ$coef[1,1], summ$coef[2,1], pch=19)
#+ eval=FALSE, include=FALSE
if (savefigs) dev.off()

#' #### Plot maximum likelihood estimate and covariance
#+ eval=FALSE, include=FALSE
if (savefigs) pdf(root("ElectionsEconomy/figs","hill_2c.pdf"), height=5, width=5)
#+
par(mar=c(3, 3, 3, 1), mgp=c(1.7, .5, 0), tck=-.01)
plot(rng.x, rng.y, xlab="a", ylab="b", main=expression(paste("(", hat(a), ", ", hat(b), ") and covariance matrix")), type="n")
points(summ$coef[1,1], summ$coef[2,1], pch=19)
rho <- summ$cov.unscaled[1,2]/sqrt(summ$cov.unscaled[1,1]*summ$cov.unscaled[2,2])
aa <- seq(-1,1,length=500)
bb <- sqrt(1-aa^2)
xx <- summ$coef[1,1] + summ$coef[1,2]*(aa*sqrt(1+rho)-bb*sqrt(1-rho))
yy <- summ$coef[2,1] + summ$coef[2,2]*(aa*sqrt(1+rho)+bb*sqrt(1-rho))
lines (xx, yy)
xx <- summ$coef[1,1] + summ$coef[1,2]*(aa*sqrt(1+rho)+bb*sqrt(1-rho))
yy <- summ$coef[2,1] + summ$coef[2,2]*(aa*sqrt(1+rho)-bb*sqrt(1-rho))
lines (xx, yy)
#+ eval=FALSE, include=FALSE
if (savefigs) dev.off()

#' ## Bayesian model with flat prior
M3 <- stan_glm(vote ~ growth, data = hibbs, 
               prior_intercept=NULL, prior=NULL, prior_aux=NULL,
               refresh = 0)
sims <- as.data.frame(M3)
a <- sims[,1]
b <- sims[,2]

#' #### Plot posterior draws
#+ eval=FALSE, include=FALSE
if (savefigs) pdf(root("ElectionsEconomy/figs","hill_3c.pdf"), height=5, width=5)
#+
par(mar=c(3, 3, 3, 1), mgp=c(1.7, .5, 0), tck=-.01)
plot(c(39.8, 52.5), c(.3, 5.8), xlab="a", ylab="b", main="4000 posterior draws of (a, b)", type="n", cex.main=1.5, cex.lab=1.5, cex.axis=1.5)
points(a, b, pch=20, cex=.2)
#+ eval=FALSE, include=FALSE
if (savefigs) dev.off()
