Don't add useless chains

[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 [](key)
    @hash[key]
  end

  def keys
    @hash.keys
  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)
    @max_order = ord
    @mkv = Hash.new { |hash, key|
      hash[key] = {:prev => ChanceHash.new, :next => ChanceHash.new}
    }
    @mkv[nil] = ChanceHash.new
  end

  def words
    @mkv[nil].keys
  end

  def add_one(sym)
    # Don't add nil to order 0
    return unless sym
    @mkv[nil].increase(sym.to_sym)
  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
    # Don't add prev to chains whose first element is nil
    return unless array.first
    h = @mkv[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
    # Don't add next to chains whose last element is nil
    return unless array.last
    h = @mkv[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 => false}.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)
    if @mkv.key?(ar)
      @mkv[ar][:next].random
    else
      raw_next(ar.butfirst, o)
    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)
    if @mkv.key?(ar)
      @mkv[ar][:prev].random
    else
      raw_prev(ar.butlast, o)
    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[nil].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