# ========================================================
# POISSON AND NEGATIVE BINOMIAL REGRESSION IN R
# PHAA outcomes data (the Lesson 4 dataset), 795 adults
# ========================================================
# Question: are age and smoking associated with the number
# of urgent care or emergency department visits in the
# past 12 months?
# ========================================================


# ========================================================
# SECTION 1: PACKAGES AND DATA
# ========================================================

library(MASS)    # glm.nb() fits the negative binomial model

# The data file is stored on the course website
phaa_url <- "https://www.sfu-epi.ca/r-activities/data/phaa_outcomes.csv"
phaa <- read.csv(phaa_url)
phaa <- na.omit(phaa)      # keep people with no missing values
nrow(phaa)

phaa$smoker <- factor(phaa$smoker, levels = c("No", "Yes"))   # "No" = reference


# ========================================================
# SECTION 2: LOOK AT THE COUNTS
# ========================================================

# --- 2.1 How many people had 0, 1, 2, ... visits?
table(phaa$urgent_visits)

barplot(table(phaa$urgent_visits),
        main = "Urgent care visits in the past 12 months",
        xlab = "Number of visits",
        ylab = "Number of people",
        col = "grey80")

# --- 2.2 The mean and the variance of the counts
mean(phaa$urgent_visits)
var(phaa$urgent_visits)

# --- 2.3 The average number of visits by smoking status
tapply(phaa$urgent_visits, phaa$smoker, mean)


# ========================================================
# SECTION 3: POISSON REGRESSION
# ========================================================

pois <- glm(urgent_visits ~ age + smoker, data = phaa, family = poisson)
summary(pois)

# Incidence rate ratios (IRR) with 95% confidence intervals
exp(cbind(IRR = coef(pois), confint(pois)))

# IRR for a 10-year difference in age
exp(10 * coef(pois)["age"])


# ========================================================
# SECTION 4: CHECK THE POISSON MODEL FOR OVERDISPERSION
# ========================================================

# --- 4.1 The dispersion ratio (close to 1 if the Poisson model fits)
sum(residuals(pois, type = "pearson")^2) / df.residual(pois)

# --- 4.2 Zeros: observed and expected by the Poisson model
sum(phaa$urgent_visits == 0)
sum(dpois(0, fitted(pois)))

# --- 4.3 Every count from 0 to 6: observed and expected
visits <- 0:6
observed <- sapply(visits, function(k) sum(phaa$urgent_visits == k))
expected_pois <- sapply(visits, function(k) sum(dpois(k, fitted(pois))))
round(cbind(visits, observed, expected_pois))

mids <- barplot(observed, names.arg = visits,
                main = "Observed and expected numbers of people",
                xlab = "Urgent care visits in the past 12 months",
                ylab = "Number of people",
                col = "grey85", ylim = c(0, 500))
lines(mids, expected_pois, type = "b", pch = 16, lwd = 2, col = "firebrick")


# ========================================================
# SECTION 5: NEGATIVE BINOMIAL REGRESSION
# ========================================================

nb <- glm.nb(urgent_visits ~ age + smoker, data = phaa)
summary(nb)

# Incidence rate ratios with 95% confidence intervals
exp(cbind(IRR = coef(nb), confint(nb)))


# ========================================================
# SECTION 6: COMPARE THE TWO MODELS
# ========================================================

# --- 6.1 AIC: lower is better
AIC(pois, nb)

# --- 6.2 Dispersion ratio for the negative binomial model
sum(residuals(nb, type = "pearson")^2) / df.residual(nb)

# --- 6.3 Zeros and other counts expected by the negative binomial model
sum(dnbinom(0, mu = fitted(nb), size = nb$theta))
expected_nb <- sapply(visits, function(k)
  sum(dnbinom(k, mu = fitted(nb), size = nb$theta)))
round(cbind(visits, observed, expected_pois, expected_nb))

lines(mids, expected_nb, type = "b", pch = 17, lwd = 2, col = "steelblue")
legend("topright", legend = c("Poisson", "Negative binomial"),
       col = c("firebrick", "steelblue"), pch = c(16, 17), lwd = 2)


# ========================================================
# SECTION 7: IS AGE LINEAR ON THE LOG SCALE?
# ========================================================

# Add a squared age term and test whether it improves the model
nb_sq <- glm.nb(urgent_visits ~ age + I(age^2) + smoker, data = phaa)
anova(nb, nb_sq)


# ========================================================
# SECTION 8: DIFFERENT FOLLOW-UP TIMES AND THE OFFSET
# ========================================================

# GP visits counted over follow-up periods of different lengths
fu_url <- "https://www.sfu-epi.ca/r-activities/data/phaa_followup.csv"
fu <- read.csv(fu_url)
summary(fu$fu_years)
fu$smoker <- factor(fu$smoker, levels = c("No", "Yes"))

# offset(log(fu_years)) turns the counts into rates per year
gp <- glm(gp_visits ~ age + smoker + offset(log(fu_years)),
          data = fu, family = poisson)
round(exp(cbind(IRR = coef(gp), confint(gp))), 3)

# The same overdispersion check applies to this model
sum(residuals(gp, type = "pearson")^2) / df.residual(gp)