Bidirectional completion

[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 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
  # Maximum depth
  MAX_ORDER = 5

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

  def initialize
    # mkv[i] holds the chains of order i
    @mkv = Array.new

    # Each chain is in the form
    # [:array, :of, :symbols] => {
    #   :prev => ChanceHash,
    #   :next => ChanceHash
    # }
    # except for order 0, which is a simple ChanceHash
    # itself
    @mkv[0] = ChanceHash.new
    MAX_ORDER.times { |i|
      @mkv[i+1] = Hash.new { |hash, key|
        hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
      }
    }

  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)
    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))
      end
    end
  end

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

  def raw_prev(syms)
    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_next(ar.first(ord-1))
      end
    end
  end

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

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

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

  def complete(text)
    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)
    nxt = raw_next(syms)
    while nxt or prev do
      syms.unshift(prev)
      syms.push(nxt)
      # Keep adding only on the side where we
      # didn't come across a nil already
      prev = raw_prev(syms) if prev
      nxt = raw_next(syms) if nxt
    end
    return syms.to_s
  end

end