# ====================================================================
# 階層ベイズモデルによる店舗の実力（平均パラメータ）のMCMC推定
# ====================================================================
rm(list = ls())
set.seed(123)

# 1. 架空の売上データの作成（5つの店舗、各店舗10日間のデータ）
# 店舗ごとに実力（真の平均）が少し異なる（全体平均は10、店舗間のばらつきは2）
true_mu_shop <- c(8.5, 11.0, 9.2, 12.5, 9.8) 
n_shops <- length(true_mu_shop)
n_days  <- 10

# データの生成（売上のばらつき SD=1.5）
y <- matrix(NA, nrow = n_days, ncol = n_shops)
for(j in 1:n_shops) {
  y[, j] <- rnorm(n_days, mean = true_mu_shop[j], sd = 1.5)
}

# 2. MCMC（メトロポリス法）の設定
N <- 50000  # 歩数（サンプリング数）

# サンプル保存用の入れ物
# 推定したいもの：①各店舗の実力(mu_shop), ②全店共通の平均(mu_all)
trace_mu_shop <- matrix(NA, nrow = N, ncol = n_shops)
trace_mu_all  <- numeric(N)

# 初期値の設定
current_mu_shop <- colMeans(y) # 初期値はとりあえず各店舗のデータ平均
current_mu_all  <- mean(y)     # 初期値は全体のデータ平均

# 固定するパラメータ（簡単のため、売上のばらつきと店舗間のばらつきのSDは1.5に固定）
sd_y <- 1.5
sd_shop <- 1.5

# 3. 対数尤度＋対数事前分布を計算する関数（アンダーフロー防止のため対数で計算）
calc_log_posterior <- function(mu_shop, mu_all, y, sd_y, sd_shop) {
  # ① 事前分布：各店舗の実力は、全店共通の平均(mu_all)の周りに分布する
  log_prior_shop <- sum(dnorm(mu_shop, mean = mu_all, sd = sd_shop, log = TRUE))
  # ② 事前分布：全店共通の平均自体は、広めの分布（無情報に近い）に従う
  log_prior_all  <- dnorm(mu_all, mean = 10, sd = 100, log = TRUE)
  
  # ③ 尤度：実際のデータが、各店舗の実力(mu_shop)からどれくらい発生しやすいか
  log_likelihood <- 0
  for(j in 1:ncol(y)) {
    log_likelihood <- log_likelihood + sum(dnorm(y[, j], mean = mu_shop[j], sd = sd_y, log = TRUE))
  }
  
  return(log_prior_shop + log_prior_all + log_likelihood)
}

# 4. MCMCサンプリング開始
for (i in 1:N) {
  
  # --- (A) 各店舗の mu_shop を更新 ---
  # 現在の「もっとらしさ（対数事後確率）」
  current_log_post <- calc_log_posterior(current_mu_shop, current_mu_all, y, sd_y, sd_shop)
  
  # 次の候補地を提案（5店舗分まとめて少しずらす）
  proposed_mu_shop <- current_mu_shop + rnorm(n_shops, mean = 0, sd = 0.1)
  proposed_log_post <- calc_log_posterior(proposed_mu_shop, current_mu_all, y, sd_y, sd_shop)
  
  # 受容確率 α の計算（対数なので引き算が割り算の意味になる）
  alpha <- exp(proposed_log_post - current_log_post)
  if (runif(1) < alpha) {
    current_mu_shop <- proposed_mu_shop
  }
  
  # --- (B) 全店共通の mu_all を更新 ---
  current_log_post <- calc_log_posterior(current_mu_shop, current_mu_all, y, sd_y, sd_shop)
  
  proposed_mu_all <- current_mu_all + rnorm(1, mean = 0, sd = 0.1)
  proposed_log_post <- calc_log_posterior(current_mu_shop, proposed_mu_all, y, sd_y, sd_shop)
  
  alpha <- exp(proposed_log_post - current_log_post)
  if (runif(1) < alpha) {
    current_mu_all <- proposed_mu_all
  }
  
  # 記録
  trace_mu_shop[i, ] <- current_mu_shop
  trace_mu_all[i]    <- current_mu_all
}

# 5. 結果の可視化（最初の5000歩をバーンインとして捨てる）
burn_in <- 5001:N
clean_mu_shop <- trace_mu_shop[burn_in, ]

# 各店舗の「真の実力」の推定結果を箱ひげ図で表示
boxplot(clean_mu_shop, names = paste("店舗", 1:5),
        main = "MCMCによる各店舗の実力（平均売上）の事後分布",
        col = "lightgreen", ylab = "推定された平均売上")

# 本物の正解データを赤点でプロットして答え合わせ
points(1:5, true_mu_shop, col = "red", pch = 16, cex = 2)
legend("topleft", legend = "真の設定値（正解）", col = "red", pch = 16)