require 'open-uri'
require 'open_uri_redirections'
require 'nokogiri'
require 'natto'
require 'mecab'
require 'pp'
require 'json'

class OneGramModel
  
  N = 1000000
  
  def initialize
    @word_cnt_for_p_ml = {}
    @word_cnt_for_lambda = {}
    @lambda = 0
    get_stop_words
  end
  
  def get_stop_words
    @stop_words = []
    File.foreach('stop_words.txt') do |line|
      @stop_words << line.strip unless line.strip == ''
    end
  end
  
  def countup_word_cnt_for_p_ml(word)
    @word_cnt_for_p_ml[word] = 0 unless @word_cnt_for_p_ml.has_key?(word)
    @word_cnt_for_p_ml[word] += 1
  end
  
  def countup_word_cnt_for_lambda(word)
    @word_cnt_for_lambda[word] = 0 unless @word_cnt_for_lambda.has_key?(word)
    @word_cnt_for_lambda[word] += 1
  end
  
  def in_docs(word)
    if @word_cnt_for_p_ml.has_key?(word)
      @word_cnt_for_p_ml[word]
    else
      0
    end
  end
  
  def get_p_ml(word)
    in_docs(word).to_f/@word_cnt_for_p_ml.values.inject(0){|sum,e| sum += e }.to_f
  end
  
  def filter(surface, feature)
    ['名詞','形容詞','形容動詞'].include?(feature.split(',')[0]) && !@stop_words.include?(surface)
  end
  
  def get_prob(txt)
    log_prob = 0
    nm = Natto::MeCab.new
    nm.enum_parse(txt).each do |n|
      if !n.is_eos? && filter(n.surface, n.feature)
        log_prob += Math.log(@lambda * get_p_ml(n.surface) + (1-@lambda) / N)
      end
    end
    return Math.exp(log_prob)
  end
    
  
  def optimize_lambda
    learning_rate = 1e-6
    before_lambda = after_lambda = 0.84
    epoch = 0
    while true
      grad = 0
      before_lambda = after_lambda
      printf("epoch = %d, lambda = %.5f\n",epoch,after_lambda)
      @word_cnt_for_lambda.each do |word,cnt|
        p_ml = get_p_ml(word)
        grad += cnt * (N * p_ml - 1)/(N * before_lambda * p_ml + 1 - before_lambda)
      end
      after_lambda = grad * learning_rate + before_lambda
      epoch += 1
      if (after_lambda - before_lambda).abs < 1e-4
        @lambda = after_lambda
        return
      end
    end
  end
  
  def save
    File.open("p_ml_wc.json", 'w') do |file|
      JSON.dump(@word_cnt_for_p_ml, file)
    end
    File.open("lambda_wc.json", 'w') do |file|
      JSON.dump(@word_cnt_for_lambda, file)
    end
  end
  
  def load
    @word_cnt_for_p_ml = open('p_ml_wc.json', 'r') do |io|
      JSON.load(io)
    end
    @word_cnt_for_lambda = open('lambda_wc.json', 'r') do |io|
      JSON.load(io)
    end
  end
  
end


url = 'https://www.aozora.gr.jp/cards/000148/files/789_14547.html'
model = OneGramModel.new

puts "1.補間前確率と補間係数最適化用のコーパスを取得"
open(url) do |res|
  charset = res.charset
  html = res.read
  doc = Nokogiri::HTML.parse(html)
  doc.css('.midashi_anchor').remove
  txt = doc.css('.main_text').text
  txt.gsub!(/（.*?）/, '')
  nm = Natto::MeCab.new
  txt.split("\n").each_with_index do |line,idx|
    nm.enum_parse(line).each do |n|
      if !n.is_eos? && model.filter(n.surface, n.feature)
        case idx % 2
        when 0
          model.countup_word_cnt_for_p_ml(n.surface)
        when 1
          model.countup_word_cnt_for_lambda(n.surface)
        end
      end
    end
  end
end

puts "2.コーパスを保存"
model.save

puts "3.コーパスを読込"
model.load

puts "4.補間係数を最適化"
model.optimize_lambda

puts "\n5.テストデータから確率を推定\n"

puts test_text =  "吾輩は猫である"
printf("確率:%.5e\%\n",model.get_prob(test_text)*100)
puts test_text =  "恋は宇宙的な活力である"
printf("確率:%.5e\%\n",model.get_prob(test_text)*100)

puts test_text =  "吾輩は犬である"
printf("確率:%.5e\%\n",model.get_prob(test_text)*100)

4.補間係数を最適化
epoch = 0, lambda = 0.84000
epoch = 1, lambda = 0.83945
epoch = 2, lambda = 0.83900
epoch = 3, lambda = 0.83861
epoch = 4, lambda = 0.83829
epoch = 5, lambda = 0.83802
epoch = 6, lambda = 0.83780
epoch = 7, lambda = 0.83761
epoch = 8, lambda = 0.83745
epoch = 9, lambda = 0.83731
epoch = 10, lambda = 0.83720

5.テストデータから確率を推定
吾輩は猫である
確率:3.94505e-03%
恋は宇宙的な活力である
確率:5.15099e-10%
吾輩は犬である
確率:3.08339e-04%