# param_infer_unsupervised.rb
def main
  # πi: それぞれのサイコロが箱の中にどのような割合で混ざっているか
  pi = [0.1, 0.4, 0.5]

  # θi: それぞれのサイコロを投げた時に表が出る確率
  theta = [[0.8, 0.2],
           [0.6, 0.4],
           [0.3, 0.7]]

  # 試行回数
  n = 10000
  r = make_trial(n, pi, theta)

  # 推定する含有率πの初期値
  initial_infer_pi = [0.3, 0.5, 0.2]

  pis, log_likelihoods = infer_params_unsupervised(r, theta, initial_infer_pi)

  # 結果表示
  pis.size.times {|i| puts "#{i}: #{pis[i].inspect}, #{log_likelihoods[i]}"}
end

# 教師なし学習アルゴリズムでパラメータの推定を行う
# _r_ :: サイコロの各目が出た回数 r_k
# _theta_ :: 各サイコロの、各目の出る確率
# _initial_pi_ :: 推定する、サイコロの含有率の初期値
def infer_params_unsupervised(r, theta, initial_pi)
  # 試行で出た目だけを使って（どのサイコロだったかはわからない状態で）
  # サイコロのそれぞれの含有率piを推定する
  # 各サイコロを振った時に出る目の割合は既知(theta)とする

  c = theta.size  # サイコロの種類数
  m = theta[0].size  # サイコロの出る目の種類数
  n = r.inject(:+)  # 試行回数：各目が出た回数の合計

  # Step 1: πi, θikの初期値を与える
  pi = initial_pi

  pis = []
  log_likelihoods = []
  51.times do
    if i > 0
      # Step 2: ベイズの定理により、P(ωi|vk)を計算する
      p_wi_vk = (0...c).map do |i|
        (0...m).map do |k|
          prob = sigma(0...c) do |j|
            pi[j] * theta[j][k]
          end
          pi[i] * theta[i][k] / prob
        end
      end
      # Step 3-1: この値を用いてπiを更新し、以下の^πiを求める
      pi_hat = (0...c).map do |i|
        sigma(0...m) do |k|
          r[k] * p_wi_vk[i][k] / n
        end
      end

      # Step 4: πi = ^πi, θik = ^θikと設定する
      pi = pi_hat
    end

    # 式(5.8)、(5.51)より対数尤度logP(x)を求め、その増分が予め定めた
    # 閾値以下なら終了し、さもなければ、Step 2に戻って同じ処理を繰り返す。
    p_vk = (0...m).map do |k|  # 式(5.8): サイコロの目vkが出る確率
      sigma(0...c) do |i|
        pi[i] * theta[i][k]
      end
    end
    log_px = sigma(0...m) do |k|  # 式(5.51): 対数尤度logP(x)
      r[k] * Math::log(p_vk[k])
    end

    pis.push(pi)
    log_likelihoods.push(log_px)
  end
  return pis, log_likelihoods
end

# 箱からサイコロを取り出して投げる、という試行をn回行い、
# その結果サイコロの各目が出た回数を配列で返す
def make_trial(n, pi, theta)
  m = theta[0].size  # 出る目の種類の数
  r = Array.new(m) {0}
  n.times do
    st = pick_dice(pi)
    xt = roll_dice(theta[st])
    r[xt] += 1
  end
  return r
end

# 箱からサイコロを無作為に取り出す
def pick_dice(pi)
  random_choise(pi)
end

# サイコロを投げてどの目が出たかを返す
def roll_dice(theta)
  random_choise(theta)
end

# 確率でランダムに選び、選んだインデクスを返す
def random_choise(probs)
  r = rand
  probs.each_with_index do |p, i|
    return i if r < p
    r -= p
  end
  return probs.size - 1  # 誤差対策
end

# Rangeに対してブロックを呼び出し、その結果の和を返す：∑
def sigma(range, &block)
  range.map do |i|
    block.call(i)
  end.inject(:+)
end

main