#---------------------------------
# ステップ1：Stanを使用する準備
#---------------------------------

# ライブラリの読み込み
require(rstan)

# 計算の並列化
rstan_options(auto_write = TRUE)
options(mc.cores = parallel::detectCores())


#---------------------------------
# ステップ2：統計モデルの記述
#---------------------------------

# モデルの作成
localLevelModel_2 <- "
  data {
    int n;                  # サンプルサイズ
    vector[n] Nile;         # ナイル川の流量データ
  }
  parameters {
    real muZero;            # 左端
    vector[n] mu;           # 確率的レベル
    real<lower=0> sigmaV;   # 観測誤差の大きさ
    real<lower=0> sigmaW;   # 過程誤差の大きさ
  }
  model {
    # 状態方程式の部分
    # 左端から初年度の状態を推定する
    mu[1] ~ normal(muZero, sqrt(sigmaW));  

    # 状態の遷移
    for(i in 2:n) {
      mu[i] ~ normal(mu[i-1], sqrt(sigmaW));
    }

    # 観測方程式の部分
    for(i in 1:n) {
      Nile[i] ~ normal(mu[i], sqrt(sigmaV));
    }
  }
"


#---------------------------------
# ステップ3：データの指定
#---------------------------------

# Nile川流量データ
Nile

# Stanに渡すために整形
NileData <- list(Nile = as.numeric(Nile), n=length(Nile))
NileData



#---------------------------------
# ステップ4：ベイズ推定の実行
#---------------------------------

# 乱数の種
set.seed(1)

# ローカルレベルモデル
Nile_LocalLevelModel_1 <- stan(
  model_code = localLevelModel_2,
  data=NileData,
  iter=2500,
  warmup=1000,
  thin=1,
  chains=3
)

#---------------------------------------
# ステップ5：うまく推定されているかをチェック〜修正
#---------------------------------------

# 結果の確認
# 推定されたパラメタ
Nile_LocalLevelModel_1


# 計算の過程を図示
traceplot(Nile_LocalLevelModel_1, pars=c("sigmaV", "sigmaW"))


# 乱数の種
set.seed(1)

# ローカルレベルモデル
Nile_LocalLevelModel_2 <- stan(
  model_code = localLevelModel_2,
  data=NileData,
  iter=30000,
  warmup=1000,
  thin=10,
  chains=3
)


# 結果の確認
# 推定されたパラメタ
Nile_LocalLevelModel_2

# 計算の過程を図示
traceplot(Nile_LocalLevelModel_2, pars=c("sigmaV", "sigmaW"))


#---------------------------------------
# 推定結果の取得と利用方法
#---------------------------------------

# 計算結果をより細かく取得する
# summary(Nile_LocalLevelModel_2)
# summary(Nile_LocalLevelModel_2)$summary
head(summary(Nile_LocalLevelModel_2)$summary)

# EAP推定量一覧
summary(Nile_LocalLevelModel_2)$summary[,1]


# muだけ取り出す
summary(Nile_LocalLevelModel_2)$summary[2:101,1]


# EAP推定量と95%信用区間を取得
muEAP <- summary(Nile_LocalLevelModel_2)$summary[2:101,1]
muLower95 <- summary(Nile_LocalLevelModel_2)$summary[2:101,4]
muUpper95 <- summary(Nile_LocalLevelModel_2)$summary[2:101,8]


# 時系列データ(ts)形式に修正
muEAP_ts <- ts(as.numeric(muEAP), start=1871, end=1970)
muLower95_ts <- ts(as.numeric(muLower95), start=1871, end=1970)
muUpper95_ts <- ts(as.numeric(muUpper95), start=1871, end=1970)


# プロット
plot(Nile, type="p", pch=16, main="ナイル川流量のローカルレベルモデル")
lines(muEAP_ts, col=2, lwd=2)
lines(muLower95_ts, col=4, lwd=2)
lines(muUpper95_ts, col=4, lwd=2)



#-----------------------------------------------
# おまけ1．ggplot2を使ったきれいなグラフを作る
#-----------------------------------------------

# ggplot2を使ったもっときれいなグラフ
library("ggplot2")

# 描画したいデータをまとめる
d1 <- data.frame(year=1871:1970, as.numeric(Nile))
d2 <- data.frame(year=1871:1970, EAP=muEAP, lower=muLower95, upper=muUpper95)
head(d1)
head(d2)

# グラフの外枠の作成
NileGraph <- ggplot(
  data=d1,
  aes(
    x=year,
    y=Nile
  )
)

# データ点を追加
NileGraph <- NileGraph + geom_point()

# EAP推定量と95%信用区間を追加
NileGraph <- NileGraph + geom_smooth(
  aes(
    ymin=lower, ymax=upper, y=EAP
  ),
  lwd=1.2,
  color=1,
  data=d2,
  stat="identity"
)

# 描画
plot(NileGraph)


#---------------------------------------
# おまけ2．乱数の一覧を取得する方法
#---------------------------------------

# 乱数の一覧を取得
resultMCMC <- extract(Nile_LocalLevelModel_2)

# シミュレーション結果は何個ある？
length(resultMCMC$muZero)

# ((iter - warmup) /thin)*chain
((30000 - 1000)/10)*3

# muは1年ごとに8700個のシミュレーション結果がある
length(resultMCMC$mu[,1])

# シミュレーション結果の期待値をとると、EAP推定量が求まる
mean(resultMCMC$mu[,1])
muEAP[1]

# 95%区間も同様
quantile(resultMCMC$mu[,1], probs=c(0.025, 0.975))
muLower95[1]
muUpper95[1]


#---------------------------------------
# おまけ3．結果の保存
#---------------------------------------

# 作業スペースを読み込んだ後は、以下を実行するだけでOK

library(rstan)
library(ggplot2)

# 推定されたパラメタ
Nile_LocalLevelModel_2

# 計算の過程を図示
traceplot(Nile_LocalLevelModel_2, pars=c("sigmaV", "sigmaW"))

# ggplot2による図示
plot(NileGraph)