source('https://myweb.uiowa.edu/pbreheny/7110/f20/notes/fun.R')

# Download data
britdoc <- read.delim('https://s3.amazonaws.com/pbreheny-data-sets/britdoc.txt')

# Format
Time <- britdoc$PersonYears / 1000
X <- model.matrix( ~ Age + Smoking, britdoc)
y <- britdoc$Deaths

# MLE (simple)
beta <- rep(0, ncol(X))
for (i in 1:30) {
  eta <- drop(log(Time) + X %*% beta)
  mu <- exp(eta)
  W <- diag(mu)
  Info <- crossprod(X, W) %*% X
  beta <- beta + solve(Info) %*% crossprod(X, y-mu)
}
fit <- glm(Deaths ~ offset(log(PersonYears/1000)) + Age + Smoking, britdoc, family=poisson)
coef(fit)
drop(beta)

# Improvement 1 (monitor convergence)
converged <- function(old, new, eps=1e-6) {all(abs(old-new) < eps)}
beta <- rep(0, ncol(X))
iter <- 0
repeat {
  old <- beta
  iter <- iter + 1
  eta <- drop(log(Time) + X %*% beta)
  mu <- exp(eta)
  W <- diag(mu)
  Info <- crossprod(X, W) %*% X
  beta <- old + solve(Info) %*% crossprod(X, y-mu)
  if (converged(old, beta) | iter > 50) break
}
iter

# Final version
pois_fit <- function(X, y, Time, eps=1e-6) {
  beta <- c(log(sum(y) / sum(Time)), rep(0, ncol(X)-1))
  iter <- 0
  repeat {
    old <- beta
    iter <- iter + 1
    eta <- drop(log(Time) + X %*% beta)
    mu <- exp(eta)
    W <- diag(mu)
    Info <- crossprod(X, W) %*% X
    beta <- old + solve(Info) %*% crossprod(X, y-mu)
    if (converged(old, beta, eps=eps)) break
  }
  eta <- drop(log(Time) + X %*% beta)
  loglik <- drop(crossprod(y, eta) - sum(exp(eta)))
  return(list(loglik=loglik, beta=drop(beta), iter=iter, r=y-exp(eta), eta=eta, Info=Info))
}
res <- pois_fit(X, y, Time)

# Newton animation: From https://myweb.uiowa.edu/pbreheny/7210/f19/notes/10-03.R

# Table
wald_test <- function(res, j) {
  z <- res$beta / sqrt(diag(solve(res$Info)))
  2*pnorm(-abs(z[j]))
}
wald_ci <- function(res, j) {
  SE <- sqrt(diag(solve(res$Info)))[j]
  lwr <- res$beta[j] + qnorm(0.025) * SE
  upr <- res$beta[j] + qnorm(0.975) * SE
  cbind(lwr, upr)
}
lr_test <- function(res1, res0) {
  x <- 2*(res1$loglik-res0$loglik)
  pchisq(x, 1, lower.tail=FALSE)
}
lr_ci <- function(res1, X, y, Time, j) {
  lr_chisq <- function(b, X, x, y, Time, ll1) {
    out <- double(length(b))
    for (j in length(b)) {
      res <- pois_fit(X, y, exp(log(Time) + x*b[j]))
      out[j] <- 2*(ll1-res$loglik)
    }
    out
  }
  f <- function(b) lr_chisq(b, X[,-j], X[,j], y, Time, res1$loglik) - qchisq(0.95, 1)
  lwr <- uniroot(f, interval=c(-5, res$beta[6]))$root
  upr <- uniroot(f, interval=c(res$beta[6], 5))$root
  c(lwr, upr)
}
score_test <- function(X, y, Time, j) {
  x <- X[, j]
  X0 <- X[, -j]
  res0 <- pois_fit(X0, y, Time)
  W <- diag(exp(res0$eta))
  v <- solve(crossprod(X, W) %*% X)[6,6] # Same as
  v <- 1/(crossprod(x, W) %*% x - crossprod(x, W) %*% X0 %*% solve(res0$Info) %*% crossprod(X0, W) %*% x)
  z <- drop(crossprod(x, res0$r)) * drop(sqrt(v))
  2*pnorm(-abs(z))
}
score_ci <- function(X, y, Time, j) {
  score_chisq <- function(b, X, x) {
    out <- double(length(b))
    for (j in length(b)) {
      res <- pois_fit(X, y, exp(log(Time) + x*b))
      W <- diag(exp(res$eta))
      v <- 1/(crossprod(x, W) %*% x - crossprod(x, W) %*% X %*% solve(res$Info) %*% crossprod(X, W) %*% x)
      out[j] <- drop(crossprod(x, res$r)) * sqrt(v)
    }
    out^2
  }
  f <- function(b) score_chisq(b, X[,-j], X[,j]) - qchisq(0.95, 1)
  lwr <- uniroot(f, interval=c(-5, res$beta[6]))$root
  upr <- uniroot(f, interval=c(res$beta[6], 5))$root
  c(lwr, upr)
}
res <- pois_fit(X, y, Time)
res0 <- pois_fit(model.matrix( ~ Age, britdoc), y, Time)
wald_ci(res, 6)
wald_test(res, 6)
lr_ci(res, X, y, Time, 6)
lr_test(res, res0)
score_ci(X, y, Time, 6)
score_test(X, y, Time, 6)

# Compare with
fit0 <- glm(Deaths ~ offset(log(PersonYears/1000)) + Age + Smoking, britdoc, family=poisson)
fit <- glm(Deaths ~ offset(log(PersonYears/1000)) + Age, britdoc, family=poisson)