Display amount of known words

[rbot-mark] / mark2.rb

# vim: set sw=2 et:
# Author: Giuseppe Bilotta <giuseppe.bilotta@gmail.com>
# New markov chain plugin

class Array
  def butlast
    first(self.size-1)
  end

  def butfirst
    last(self.size-1)
  end

  def pick_one
    self[rand(self.size)]
  end

  def pick_some(n)
    return nil if self.empty?
    i = rand(self.size)
    if n < 0
      count = rand(self.size-i)
    else
      count = rand([n, self.size-i].min)
    end
    self[i, count]
  end
end

class ChanceHash

  def initialize
    @hash = Hash.new(0)
    @picker = []
    @total = 0
    @valid_pick = false
  end

  def size
    @hash.size
  end

  def increase(el)
    if @hash.key?(el)
      @hash[el] += 1
    else
      @hash[el] = 1
    end
    @valid_pick = false
    return @hash[el]
  end

  def decrease(el)
    if @hash.key?(el)
      @hash[el] -= 1
      @hash.delete(el) if @hash[el] == 0
    end
    @valid_pick = false
    return @hash[el]
  end

  def make_picker
    @picker.clear
    total = 0
    @hash.each { |el, ch|
      total += ch
      @picker << [total, el]
    }
    @total = total
    @valid_pick = true
  end

  def random
    case @hash.size
    when 0
      return nil
    when 1
      return @hash.keys.first
    else
      make_picker unless @valid_pick
      pick = rand(@total)
      @picker.each { |ch, el|
        return el if pick < ch
      }
    end
  end
end

class MarkovChainer
  # Word or nonword regexp:
  # can be used to scan a string splitting it into
  # words and nonwords.
  WNW = /\w+|\W/u

  attr_reader :max_order
  def initialize(ord=5)
    # mkv[i] holds the chains of order i
    @max_order = ord
    @mkv = Array.new(@max_order+1) { |i|
      if i==0
        ChanceHash.new
      else
        Hash.new { |hash, key|
          hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
        }
      end
    }
  end

  def add_one(sym)
    s = sym.to_sym rescue nil
    @mkv[0].increase(s)
  end

  def add_before(array, prev)
    raise "Not enough words in new data" if array.empty?
    raise "Too many words in new data" if array.size > @max_order
    size = array.size
    h = @mkv[size][array.dup]
    h[:prev].increase(prev)
  end

  def add_after(array, nxt)
    raise "Not enough words in new data" if array.empty?
    raise "Too many words in new data" if array.size > @max_order
    size = array.size
    h = @mkv[size][array.dup]
    h[:next].increase(nxt)
  end

  def add_multi(array)
    raise "Too many words in new data" if array.size > @max_order + 1
    add_before(array.butfirst, array.first)
    add_after(array.butlast, array.last)
  end

  def add(*data)
    if data.size == 1
      add_one(data.first)
    else
      add_multi(data)
    end
  end

  def simple_learn(text)
    syms = text.scan(WNW).map { |w| w.intern } 
    syms.unshift(nil)
    syms.push(nil)

    syms.size.times { |i|
      ([@max_order, syms.size-i].min+1).times { |ord|
        v = syms[i, ord+1]
        # puts "Learning #{v.inspect}"
        add(*v)
        # pp @mkv
      }
    }
  end

  def learn(text, o={})
    opts = {:lowercase => true}.merge o

    lc = opts[:lowercase]

    simple_learn(text)
    if lc
      simple_learn(text.downcase)
    end

    pp @mkv if defined? pp
  end

  def raw_next(syms, o={})
    max_order = o.fetch(:max_order, @max_order)
    ar = syms.last([max_order, syms.size].min)
    ord = ar.size
    if ord == 0
      @mkv[0].random
    else
      if @mkv[ord].key?(ar)
        @mkv[ord][ar][:next].random
      else
        raw_next(ar.last(ord-1), o)
      end
    end
  end

  def next(text, o={})
    syms = text.scan(WNW).map { |w| w.intern }
    raw_next(syms, o)
  end

  def raw_prev(syms, o={})
    max_order = o.fetch(:max_order, @max_order)
    ar = syms.first([max_order, syms.size].min)
    ord = ar.size
    if ord == 0
      @mkv[0].random
    else
      if @mkv[ord].key?(ar)
        @mkv[ord][ar][:prev].random
      else
        raw_prev(ar.first(ord-1), o)
      end
    end
  end

  def prev(text, o={})
    syms = text.scan(WNW).map { |w| w.intern }
    raw_prev(syms, o)
  end

  def complete_prev(text, o={})
    syms = text.scan(WNW).map { |w| w.intern }
    prev = raw_prev(syms, o)
    while prev do
      syms.unshift(prev)
      prev = raw_prev(syms, o)
    end
    return syms.to_s
  end

  def complete_next(text, o={})
    syms = text.scan(WNW).map { |w| w.intern }
    nxt = raw_next(syms, o)
    while nxt do
      syms.push(nxt)
      nxt = raw_next(syms, o)
    end
    return syms.to_s
  end

  def complete(text, o={})
    txt = text
    while txt.empty? do
      txt = @mkv[0].random.to_s
    end
    syms = txt.scan(WNW).map { |w| w.intern }
    prev = raw_prev(syms, o)
    nxt = raw_next(syms, o)
    while nxt or prev do
      # Keep adding only on the side where we
      # didn't come across a nil already
      if prev
        syms.unshift(prev)
        prev = raw_prev(syms, o)
      end
      if nxt
        syms.push(nxt)
        nxt = raw_next(syms, o)
      end
    end
    return syms.to_s
  end

end