library(glmmTMB)
library(dplyr)
library(ggplot2)
library(tidyr)
library(knitr)
library(kableExtra)
library(scales)
library(patchwork)

set.seed(123)

n_area       <- 20
obs_per_area <- 10
n            <- n_area * obs_per_area

# Variabel penjelas
area     <- factor(rep(1:n_area, each = obs_per_area))
x1       <- runif(n, 0, 10)
x2       <- runif(n, 5, 15)
kategori <- factor(sample(c("A", "B", "C"), n, replace = TRUE))

# Random intercept per area
area_effect      <- rnorm(n_area, mean = 0, sd = 0.5)
random_intercept <- area_effect[as.integer(area)]

# True lambda dan respon Poisson
lambda <- exp(1 + 0.3 * x1 - 0.2 * x2 + random_intercept)
y      <- rpois(n, lambda)

data_sae <- data.frame(area, x1, x2, kategori, y)

model_glmm <- glmmTMB(
  y ~ x1 + x2 + kategori + (1 | area),
  data   = data_sae,
  family = poisson
)

smry <- summary(model_glmm)
print(smry)

b    <- fixef(model_glmm)$cond
IRR  <- exp(b)  # Incidence Rate Ratio

irr_df <- data.frame(
  Parameter = names(IRR),
  Estimasi  = round(b, 4),
  IRR       = round(IRR, 4),
  CI_lower  = round(exp(b - 1.96 * coef_df$`Std. Error`), 4),
  CI_upper  = round(exp(b + 1.96 * coef_df$`Std. Error`), 4)
)

kable(irr_df, caption = "Incidence Rate Ratio (IRR) Koefisien Fixed Effect",
      booktabs = TRUE, align = "c",
      col.names = c("Parameter", "Estimasi (log)", "IRR", "CI Lower 95%", "CI Upper 95%")) %>%
  kable_styling(bootstrap_options = c("striped", "hover"), full_width = FALSE)

data_sae$pred_lambda <- predict(model_glmm, type = "response", re.form = NULL)

EB_main <- data_sae %>%
  group_by(area) %>%
  summarise(EB_estimate = mean(pred_lambda), .groups = "drop")

kable(EB_main, caption = "Estimasi Empirical Bayes (EB) per Area",
      booktabs = TRUE, align = "c",
      col.names = c("Area", "EB Estimate (λ̂)"),
      digits = 4) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed"),
                full_width = FALSE)

direct_est <- data_sae %>%
  group_by(area) %>%
  summarise(
    y_direct   = mean(y),
    SE_direct  = sd(y) / sqrt(n()),
    .groups    = "drop"
  ) %>%
  mutate(RSE_direct = 100 * SE_direct / y_direct)

kable(direct_est, caption = "Estimasi Langsung (Direct Estimate) per Area",
      booktabs = TRUE, align = "c",
      col.names = c("Area", "Rata-rata Y", "SE Direct", "RSE Direct (%)"),
      digits = 4) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed"),
                full_width = FALSE) %>%
  column_spec(4, color = ifelse(direct_est$RSE_direct > 25, "red", "black"),
              bold = ifelse(direct_est$RSE_direct > 25, TRUE, FALSE))

compare_df <- left_join(direct_est, EB_main, by = "area") %>%
  mutate(area_num = as.numeric(as.character(area)))

ggplot(compare_df, aes(x = area_num)) +
  geom_line(aes(y = y_direct, color = "Direct"), linewidth = 1) +
  geom_line(aes(y = EB_estimate, color = "Empirical Bayes"), linewidth = 1,
            linetype = "dashed") +
  geom_point(aes(y = y_direct, color = "Direct"), size = 2) +
  geom_point(aes(y = EB_estimate, color = "Empirical Bayes"), size = 2) +
  scale_color_manual(values = c("Direct" = "tomato",
                                "Empirical Bayes" = "steelblue")) +
  labs(title    = "Perbandingan Estimasi: Direct vs Empirical Bayes",
       subtitle = "Per area (n_obs = 10 per area)",
       x = "Area", y = "Estimasi λ",
       color = "Metode") +
  theme_minimal(base_size = 12)

n_area      <- length(unique(data_sae$area))
area_levels <- levels(data_sae$area)
EB_jack     <- matrix(NA, nrow = n_area, ncol = n_area)

for (i in seq_len(n_area)) {
  area_out   <- area_levels[i]
  data_jack  <- filter(data_sae, area != area_out)
  
  model_jack <- glmmTMB(
    y ~ x1 + x2 + kategori + (1 | area),
    data   = data_jack,
    family = poisson
  )
  
  data_jack$pred_lambda <- predict(model_jack, type = "response", re.form = NULL)
  
  pred_jack <- data_jack %>%
    group_by(area) %>%
    summarise(EB_jack = mean(pred_lambda), .groups = "drop")
  
  for (j in seq_len(n_area)) {
    if (area_levels[j] != area_out) {
      idx <- pred_jack$area == area_levels[j]
      if (any(idx)) EB_jack[i, j] <- pred_jack$EB_jack[idx]
    }
  }
}

# Hitung MSE Jackknife
jackknife_mse <- rep(NA_real_, n_area)
for (j in seq_len(n_area)) {
  est_jack    <- EB_jack[, j]
  valid_est   <- est_jack[!is.na(est_jack)]
  theta_bar   <- mean(valid_est)
  jackknife_mse[j] <- ((length(valid_est) - 1) / length(valid_est)) *
    mean((valid_est - theta_bar)^2)
}

EB_main$MSE_jack <- jackknife_mse
EB_main$SE_jack  <- sqrt(jackknife_mse)
EB_main$RSE_jack <- 100 * EB_main$SE_jack / EB_main$EB_estimate

bootstrap_rse_area <- function(data, n_bootstrap = 100) {
  area_list <- unique(data$area)
  results   <- data.frame(area = area_list, RSE_bootstrap = NA_real_)
  
  for (a in area_list) {
    bootstrap_est <- numeric(n_bootstrap)
    
    for (i in seq_len(n_bootstrap)) {
      idx       <- sample(seq_len(nrow(data)), replace = TRUE)
      data_boot <- data[idx, ]
      
      model_boot <- tryCatch(
        glmmTMB(y ~ x1 + x2 + kategori + (1 | area),
                data = data_boot, family = poisson),
        error = function(e) NULL
      )
      if (is.null(model_boot)) {
        bootstrap_est[i] <- NA_real_
        next
      }
      
      data_boot$pred_lambda <- predict(model_boot, type = "response",
                                       re.form = NULL)
      
      eb_boot <- data_boot %>%
        group_by(area) %>%
        summarise(EB_boot = mean(pred_lambda), .groups = "drop")
      
      boot_val           <- eb_boot$EB_boot[eb_boot$area == a]
      bootstrap_est[i]   <- if (length(boot_val) == 0) NA_real_ else boot_val
    }
    
    valid_boot  <- bootstrap_est[!is.na(bootstrap_est)]
    mean_eb     <- mean(valid_boot)
    se_eb       <- sd(valid_boot)
    rse         <- 100 * se_eb / mean_eb
    results$RSE_bootstrap[results$area == a] <- rse
  }
  return(results)
}

bootstrap_result <- bootstrap_rse_area(data_sae, n_bootstrap = 100)

final_result <- left_join(direct_est, EB_main, by = "area") %>%
  left_join(bootstrap_result, by = "area") %>%
  mutate(area_num = as.numeric(as.character(area)))

display_final <- final_result %>%
  select(area, y_direct, EB_estimate, SE_direct, SE_jack,
         RSE_direct, RSE_jack, RSE_bootstrap) %>%
  mutate(across(where(is.numeric), ~ round(.x, 3)))

kable(display_final,
      caption = "Tabel Lengkap: Direct Estimate vs Empirical Bayes per Area",
      booktabs = TRUE, align = "c",
      col.names = c("Area", "Direct Est.", "EB Est.", "SE Direct",
                    "SE Jack", "RSE Direct(%)", "RSE Jack(%)", "RSE Boot(%)")) %>%
  kable_styling(bootstrap_options = c("striped", "hover", "condensed"),
                full_width = TRUE, font_size = 11) %>%
  add_header_above(c(" " = 1, "Estimasi Titik" = 2,
                     "Standar Error" = 2, "RSE (%)" = 3))

rse_summary <- final_result %>%
  summarise(
    `RSE Direct - Mean`   = round(mean(RSE_direct, na.rm = TRUE), 2),
    `RSE Direct - Median` = round(median(RSE_direct, na.rm = TRUE), 2),
    `RSE Direct - Max`    = round(max(RSE_direct, na.rm = TRUE), 2),
    `RSE Jack - Mean`     = round(mean(RSE_jack, na.rm = TRUE), 2),
    `RSE Jack - Median`   = round(median(RSE_jack, na.rm = TRUE), 2),
    `RSE Jack - Max`      = round(max(RSE_jack, na.rm = TRUE), 2),
    `RSE Boot - Mean`     = round(mean(RSE_bootstrap, na.rm = TRUE), 2),
    `RSE Boot - Median`   = round(median(RSE_bootstrap, na.rm = TRUE), 2),
    `RSE Boot - Max`      = round(max(RSE_bootstrap, na.rm = TRUE), 2)
  ) %>%
  pivot_longer(everything(), names_to = "Statistik", values_to = "Nilai") %>%
  mutate(Metode = sub(" - .*", "", Statistik),
         Ukuran = sub(".* - ", "", Statistik)) %>%
  pivot_wider(names_from = Ukuran, values_from = Nilai) %>%
  select(Metode, Mean, Median, Max)

kable(rse_summary, caption = "Ringkasan RSE per Metode (dalam %)",
      booktabs = TRUE, align = "c") %>%
  kable_styling(bootstrap_options = c("striped", "hover"), full_width = FALSE) %>%
  row_spec(0, bold = TRUE, background = "#2c7bb6", color = "white")

rse_long <- final_result %>%
  select(area, area_num, RSE_direct, RSE_jack, RSE_bootstrap) %>%
  pivot_longer(cols = c(RSE_direct, RSE_jack, RSE_bootstrap),
               names_to  = "Method",
               values_to = "RSE") %>%
  mutate(Method = factor(Method,
                         levels = c("RSE_direct", "RSE_jack", "RSE_bootstrap"),
                         labels = c("Direct", "Jackknife", "Bootstrap")))

ggplot(rse_long, aes(x = area_num, y = RSE, color = Method)) +
  geom_line(linewidth = 1.2) +
  geom_point(size = 2.5) +
  geom_hline(yintercept = 25, linetype = "dashed",
             color = "gray40", linewidth = 0.8) +
  annotate("text", x = 18, y = 26.5, label = "RSE = 25%",
           color = "gray40", size = 3.5) +
  scale_color_manual(values = c("Direct"    = "tomato",
                                "Jackknife" = "steelblue",
                                "Bootstrap" = "darkgreen")) +
  scale_x_continuous(breaks = 1:20) +
  labs(
    title    = "Perbandingan RSE: Direct, Jackknife, dan Bootstrap",
    subtitle = "RSE ditampilkan untuk masing-masing metode estimasi per area",
    x        = "Area",
    y        = "Relative Standard Error (%)",
    color    = "Metode Estimasi"
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "top",
        panel.grid.minor = element_blank())

ggplot(rse_long, aes(x = Method, y = RSE, fill = Method)) +
  geom_boxplot(alpha = 0.7, outlier.shape = 21, outlier.size = 3) +
  geom_jitter(width = 0.15, alpha = 0.5, size = 1.5) +
  scale_fill_manual(values = c("Direct"    = "tomato",
                               "Jackknife" = "steelblue",
                               "Bootstrap" = "darkgreen")) +
  geom_hline(yintercept = 25, linetype = "dashed", color = "gray40") +
  labs(
    title    = "Distribusi RSE per Metode Estimasi",
    subtitle = "Titik-titik mewakili nilai RSE per area",
    x        = "Metode",
    y        = "RSE (%)"
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "none")

ggplot(final_result, aes(x = y_direct, y = EB_estimate)) +
  geom_abline(slope = 1, intercept = 0, color = "gray60",
              linetype = "dashed", linewidth = 1) +
  geom_point(aes(color = RSE_direct), size = 3) +
  geom_text(aes(label = area), vjust = -0.7, size = 2.8) +
  scale_color_gradient2(low = "darkgreen", mid = "orange", high = "red",
                        midpoint = 25,
                        name = "RSE Direct (%)") +
  labs(
    title    = "EB Estimate vs Direct Estimate per Area",
    subtitle = "Garis putus-putus = garis y = x (estimasi identik)",
    x        = "Direct Estimate (Rata-rata Y)",
    y        = "Empirical Bayes Estimate"
  ) +
  theme_minimal(base_size = 12)

final_result <- final_result %>%
  mutate(
    gain_jack = RSE_direct - RSE_jack,
    gain_boot = RSE_direct - RSE_bootstrap
  )

gain_long <- final_result %>%
  select(area_num, gain_jack, gain_boot) %>%
  pivot_longer(cols = c(gain_jack, gain_boot),
               names_to  = "method",
               values_to = "gain") %>%
 mutate(method = case_when(
    method == "gain_jack" ~ "Jackknife",
    method == "gain_boot" ~ "Bootstrap",
    TRUE ~ method
  ))

ggplot(gain_long, aes(x = area_num, y = gain, fill = gain > 0)) +
  geom_col(alpha = 0.8) +
  geom_hline(yintercept = 0, color = "black") +
  facet_wrap(~method, nrow = 2) +
  scale_fill_manual(values = c("TRUE" = "steelblue", "FALSE" = "tomato"),
                    labels = c("TRUE" = "EB lebih baik", "FALSE" = "Direct lebih baik"),
                    name = "") +
  scale_x_continuous(breaks = 1:20) +
  labs(
    title    = "Gain RSE: Direct − EB per Area",
    subtitle = "Biru = EB lebih presisi | Merah = Direct lebih presisi",
    x        = "Area",
    y        = "Selisih RSE (Direct − EB)"
  ) +
  theme_minimal(base_size = 11) +
  theme(legend.position = "top",
        panel.grid.minor = element_blank())

final_result <- final_result %>%
  mutate(
    EB_lower = EB_estimate - 1.96 * SE_jack,
    EB_upper = EB_estimate + 1.96 * SE_jack
  )

ggplot(final_result %>% arrange(EB_estimate) %>%
         mutate(area_ord = factor(area, levels = area)),
       aes(y = reorder(area, EB_estimate), x = EB_estimate)) +
  geom_errorbarh(aes(xmin = EB_lower, xmax = EB_upper),
                 height = 0.3, color = "steelblue", linewidth = 0.8) +
  geom_point(size = 2.5, color = "steelblue") +
  geom_point(aes(x = y_direct), shape = 17, color = "tomato", size = 2) +
  labs(
    title    = "Estimasi EB dengan Interval Kepercayaan 95% (Jackknife)",
    subtitle = "Biru (●) = EB | Merah (▲) = Direct Estimate",
    x        = "Estimasi λ",
    y        = "Area"
  ) +
  theme_minimal(base_size = 11)

data_sae$fitted   <- fitted(model_glmm)
data_sae$resid_p  <- (data_sae$y - data_sae$fitted) / sqrt(data_sae$fitted)

p_res1 <- ggplot(data_sae, aes(x = fitted, y = resid_p)) +
  geom_point(alpha = 0.5, color = "steelblue") +
  geom_hline(yintercept = 0, color = "tomato", linetype = "dashed") +
  geom_smooth(method = "loess", se = TRUE, color = "darkred", linewidth = 0.8) +
  labs(title = "Residual Pearson vs Fitted",
       x = "Nilai Fitted (λ̂)", y = "Residual Pearson") +
  theme_minimal(base_size = 11)

p_res2 <- ggplot(data_sae, aes(sample = resid_p)) +
  stat_qq(color = "steelblue", alpha = 0.7) +
  stat_qq_line(color = "tomato", linetype = "dashed") +
  labs(title = "Q-Q Plot Residual Pearson",
       x = "Kuantil Teoretis", y = "Kuantil Sampel") +
  theme_minimal(base_size = 11)

p_res1 | p_res2

# Deviance dan AIC
cat("AIC     :", round(AIC(model_glmm), 2), "\n")

cat("BIC     :", round(BIC(model_glmm), 2), "\n")

cat("LogLik  :", round(logLik(model_glmm)[1], 2), "\n")

# Pearson Chi-Square
pearson_chi2 <- sum((data_sae$y - data_sae$fitted)^2 / data_sae$fitted)
df_resid     <- nrow(data_sae) - length(fixef(model_glmm)$cond) - 1

cat("\nPearson Chi² :", round(pearson_chi2, 2), "\n")

cat("df Residual  :", df_resid, "\n")

cat("Rasio (Chi²/df):", round(pearson_chi2 / df_resid, 3),
    "← (mendekati 1 = tidak overdispersi)\n")