# 係数が時間的に変化するモデル


# ==================================================
# シミュレーションデータの作成
# 収縮期血圧をシミュレーションします
# ==================================================

# ------------------------------
# 薬を投与しなかった時の「状態」
# ------------------------------

N_NotMedicine <- 10              # 薬を投与しなかった日数
N_Medicine <- 100                # 薬を投与した日数
N <- N_NotMedicine + N_Medicine  # 合計日数
muZero <- 160                    # 初期の血圧の「状態」

# 血圧の「状態」のシミュレーション
mu <- numeric()      

set.seed(12)

mu[1] <- rnorm(1, mean=muZero, sd=3)

for(i in 2:N) {
  mu[i] <- rnorm(1, mean=mu[i-1], sd=3)
}

mu

plot(mu, type="b")




# --------------------------------------------------
# 時間によって変化する（慣れる）薬の効果のシミュレーション
# --------------------------------------------------


# 薬を使っても、徐々に血圧は下がらなくなっていく
coefMedicineTrendZero <- 0.005

# 時間的に変化する薬の効果
coefMedicine <- numeric(N_Medicine)
coefMedicineTrend <- numeric(N_Medicine)

# 薬の効果をトレンドモデルで表す
set.seed(1)
coefMedicineTrend[1] <- rnorm(1, mean=coefMedicineTrendZero, sd=0.03)

# トレンドのシミュレーション
for(i in 2:N_Medicine) {
  coefMedicineTrend[i] <- rnorm(1, mean=coefMedicineTrend[i-1], sd=0.03)
}

# シミュレーション結果の確認

plot(coefMedicineTrend, type="b")

sum(coefMedicineTrend)


# 薬の効果のシミュレーション

# 薬の効果の初期値
coefMedicineZero <- -25

coefMedicine[1] <- rnorm(1, mean=coefMedicineTrend[1] + coefMedicineZero, sd=0.5)

for(i in 2:N_Medicine) {
  coefMedicine[i] <- rnorm(1, mean=coefMedicineTrend[i] + coefMedicine[i-1], sd=0.5)
}

# 図示

plot(coefMedicine, type="b")

# 実際の薬の効果は、さらにノイズが加わるとする

coefMedicineReal <- numeric(100)

set.seed(1)
for(i in 1:N_Medicine) {
  coefMedicineReal[i] <- rnorm(1, mean=coefMedicine[i], sd=2)
}


# 薬の効果の時系列変化をプロット

plot(coefMedicineReal, type="b", ylab="薬の効果（マイナスだと血圧が下がる）")
lines(coefMedicine, col=2)
legend(
  "topleft", 
  legend=c("薬の効果", "薬の効果のトレンド"), 
  col=c(1,2), 
  lwd=1
)





# -------------------------------
# 血圧の観測値のシミュレーション
# -------------------------------
# 最初の10日は薬なし
# 70日後に薬を倍にした
# 100日後に薬を3倍にした
medicine <- c(rep(0, N_NotMedicine), rep(1, 60), rep(2, 30), rep(3, 10))
medicine

bloodPressure <- numeric(N)
bloodPressureMean <- numeric(N)

set.seed(1)

# 最初の10日は薬なし
for(i in 1:N_NotMedicine) {
  bloodPressureMean[i] <- mu[i]
  bloodPressure[i] <- rnorm(1, mean=bloodPressureMean[i], sd=10)
}

# 薬を投与した後の血圧のシミュレーション
for(i in (N_NotMedicine+1):N) {
  bloodPressureMean[i] <- mu[i] + coefMedicineReal[i-10] * medicine[i]
  bloodPressure[i] <- rnorm(1, mean=bloodPressureMean[i], sd=10)
}

# 実際の血圧のプロット

plot(bloodPressure, type="b")
abline(v=10)
abline(v=70)
abline(v=100)


# 「状態」の値も重ね合わせたプロット

plot(bloodPressure, type="b")
lines(mu, col=2, lwd=2)
lines(bloodPressureMean, col=3, lty=2, lwd=2)
legend(
  "bottomright",
  legend=c("観測値", "薬がなかった時の「状態」", "薬が入った時の「状態」"),
  lwd=c(1,2,2),
  col=c(1,2,3),
  lty=c(1,1,2)
)


# ==================================================
# Stanによる時変係数モデルの作成
# ==================================================
#---------------------------------
# ステップ1：Stanを使用する準備
#---------------------------------

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

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


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

# モデルの作成
time_variant_coef <- "
  data {
    int N_Medicine;          # 薬を投与した日数
    int N;                   # 合計日数
    vector[N] bloodPressure; # 収縮期血圧値
    vector[N] medicine;      # 投与した薬の量
  }
  parameters {
    real muZero;            # 血圧の「状態」の初期値
    vector[N] mu;           # 血圧の「状態」
    real<lower=0> sigmaV;   # 観測誤差の大きさ
    real<lower=0> sigmaW;   # 過程誤差の大きさ
    real coefMedicineTrendZero;           # 薬の効果のトレンドの初期値
    real coefMedicineZero;                # 薬の効果の初期値
    vector[N_Medicine] coefMedicineTrend; # 薬の効果のトレンド
    vector[N_Medicine] coefMedicine;      # 薬の効果
    vector[N_Medicine] coefMedicineReal;  # ノイズの入った後の薬の効果
    real<lower=0> sigmaWcoef;        # 薬の係数の過程誤差の大きさ
    real<lower=0> sigmaWcoefTrend;   # 薬の係数のトレンドの過程誤差の大きさ
    real<lower=0> sigmaVcoef;        # 薬の係数の観測誤差の大きさ
  }
  model {
    # 状態方程式の部分

    # 血圧の状態の推定
    # 左端から初年度の状態を推定する
    mu[1] ~ normal(muZero, sigmaW/100);  

    # 血圧の状態の遷移
    for(i in 2:N) {
      mu[i] ~ normal(mu[i-1], sigmaW/100);
    }

    # 薬の効果の変化をモデル化する部分
    # 初年度の薬の効果のトレンドを推定する
    coefMedicineTrend[1] ~ normal(coefMedicineTrendZero, sigmaWcoefTrend/100); 

    # 薬の効果のトレンドの遷移
    for(i in 2:N_Medicine) {
      coefMedicineTrend[i] ~ normal(coefMedicineTrend[i-1], sigmaWcoefTrend/100);
    }

    # 初年度の薬の効果を推定する
    coefMedicine[1] ~ normal(coefMedicineTrend[1] + coefMedicineZero, sigmaWcoef/100);  

    # 薬の効果の遷移
    for(i in 2:N_Medicine) {
      coefMedicine[i] ~ normal(coefMedicineTrend[i] + coefMedicine[i-1], sigmaWcoef/100);
    }

    # 薬の効果にノイズが入る
    for(i in 1:N_Medicine) {
      coefMedicineReal[i] ~ normal(coefMedicine[i], sigmaVcoef/100);
    }

    # 観測方程式の部分(薬なし)
    for(i in 1:10) {
      bloodPressure[i] ~ normal(mu[i], sigmaV/100);
    }

    # 観測方程式の部分(薬あり)
    for(i in 11:N) {
      bloodPressure[i] ~ normal(mu[i] + coefMedicineReal[i-10]*medicine[i], sigmaV/100);
    }

  }
"


#---------------------------------
# ステップ3：データの指定
#---------------------------------
# Stanに渡すために整形
stanData <- list(
  N_Medicine = N_Medicine, 
  N = N,
  bloodPressure = bloodPressure,
  medicine = medicine
)

stanData



#---------------------------------
# ステップ4：ベイズ推定の実行
# 注意。相当時間がかかります（2時間以上）
#---------------------------------

# 乱数の種
set.seed(1)

# ローカルレベルモデル
time_variant_coef_Model_1 <- stan(
  model_code = time_variant_coef,
  data=stanData,
  iter=45000,
  warmup=35000,
  thin=10,
  chains=3,
  control=list(
    adapt_delta=0.9,
    max_treedepth=13
  )
)

# 警告メッセージが出てしまいますが、今回はそのまま使いました（本当はあまりよくはありません）。
# 警告メッセージ: 
#1: There were 26 divergent transitions after warmup. Increasing adapt_delta may help. 
#2: There were 70 transitions after warmup that exceeded the maximum treedepth. Increase max_treedepth. 
#3: Examine the pairs() plot to diagnose sampling problems



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

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


# 計算の過程を図示
# あまりきれいな結果ではありません。。。
traceplot(time_variant_coef_Model_1, pars=c("sigmaV", "sigmaW"))
traceplot(time_variant_coef_Model_1, pars=c("sigmaWcoef", "sigmaWcoefTrend", "sigmaVcoef"))

# 推定結果の答え合わせ
print(summary(time_variant_coef_Model_1)$summary[c(1,112:115,416:418),],3)


#---------------------------------------
# 推定結果の図示
#---------------------------------------

# 係数の時間的変化
coefEAP <- summary(time_variant_coef_Model_1)$summary[c(216:315),1]
coefLower95 <- summary(time_variant_coef_Model_1)$summary[c(216:315),4]
coefUpper95 <- summary(time_variant_coef_Model_1)$summary[c(216:315),8]



# 係数の変化の図示
library("ggplot2")

# 描画したいデータをまとめる
d1 <- data.frame(day=1:100, coef=coefMedicine[1:100]) 
d2 <- data.frame(day=1:100, EAP=coefEAP , lower=coefLower95, upper=coefUpper95)
head(d1)
head(d2)

# グラフの外枠の作成
coefGraph <- ggplot(
  data=d1,
  aes(
    x=day,
    y=coef
  )
)

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

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

# 描画
plot(coefGraph)