# IWLS for example in class

link.func <- function(mu) {-1/mu}
inv.link.func <- function(eta) {-1/eta}
deriv.link <- function(mu) {1/mu^2}

variance.func <- function(mu) {mu^2}
phi <- 1

# function to do one step of IWLS
new.beta <- function(beta, X, y){
    eta <- X %*% beta
    mu <- inv.link.func(eta)
    u <- (y - mu)/(variance.func(mu) * phi * deriv.link(mu))
    W <- diag(as.numeric(1/(variance.func(mu) * phi * deriv.link(mu)^2)))
    diff <- solve(t(X) %*% W %*% X) %*% t(X) %*% u
    new.beta <- beta + diff
    list(beta = new.beta, diff = diff)
}

# some fake data
n <- 500
x1 <- 1:n
x2 <- rep(1:5, length.out = n)
y <- rexp(n, rate = 5 + 1*x1 + 2*x2) # exponential uses 1/lambda in R
X <- cbind(int = rep(1, n), x1, x2)
p <- ncol(X)

# initial guess for beta
mu <- y
eta <- link.func(mu)
W <- diag(1/(variance.func(mu) * phi * deriv.link(mu)^2))
beta_0 <- solve(t(X) %*% W %*% X) %*% t(X) %*% W %*% eta

# one iteration
new.beta(beta_0, X, y)$beta

# iterate 50 times - see convergence after only about 4 iterations

beta_old <- beta_0
for (iteration in 1:50) {
    cat("iteration", iteration, "")
    beta_new <- new.beta(beta_old, X, y)
    beta_old <- beta_new$beta
    cat("beta is:", beta_old, "\n")
}

# or alternatively keep iterating until diff is small
diff <- 1
beta_old <- beta_0
iteration <- 1
while (abs(diff) > 1e-10){
    cat("iteration", iteration, "")
    beta_new <- new.beta(beta_old, X, y)
    diff <- sum(beta_new$diff)
    beta_old <- beta_new$beta
    cat("beta is:", beta_old, "\n")
    iteration <- iteration + 1
}
beta_old