# param_infer_supervised.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 = 100
  nss = make_trial(n, pi, theta)
  p nss

  # パラメータを最尤推定
  ps, tss = infer_params_supervised(nss)

  # 結果表示
  ps.size.times do |i|
    printf("%2d: %.2f: %s\n", i, ps[i], tss[i].map {|t| sprintf('%.2f', t)}.join(', '))
  end
end

# 教師付き学習の最尤推定でパラメータの推定を行う
# _ns_ :: サイコロの各種類・各目が出た回数 n_ik
def infer_params_supervised(nss)
  n = nss.flatten.sum  # 総試行回数
  nis = nss.map(&:sum)  # 各サイコロの種類が取り出された回数

  pi = nis.map do |ni|  # 式(5.44): 最尤推定による各サイコロの含有率の推定値
    ni.to_f / n  # $\pi_i = \frac{n_i}{n}$
  end

  theta = nss.each_with_index.map do |ns, i|  # 式(5.47): 最尤推定による各サイコロの出目の推定値
    ns.map do |nik|
      nik.to_f / nis[i]  # $\theta_{ik} = \frac{n_{ik}}{n_i}$
    end
  end

  return pi, theta
end

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

class Array
  def sum
    self.inject(:+)
  end
end

main

# 結果：
#[[9, 3], [23, 14], [13, 38]]
# 0: 0.12: 0.75, 0.25
# 1: 0.37: 0.62, 0.38
# 2: 0.51: 0.25, 0.75