Alternate implementation

[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 ChanceHashValue
  attr_accessor :count, :key, :parent
  def initialize(parent=nil, key=nil)
    @count = 0
    @prev = nil
    @next = nil
    @parent = parent
    @key = key
  end

  def prev
    @prev = ChanceHash.new(self, nil) unless @prev
    @prev
  end

  def next
    @next = ChanceHash.new(nil, self) unless @next
    @next
  end

end

class ChanceHash

  attr_reader :next_of, :prev_of
  def initialize(prev_of=nil, next_of=nil)
    @hash = Hash.new { |hash, key|
      hash[key] = ChanceHashValue.new(hash, key)
    }
    @picker = []
    @total = 0
    @valid_pick = false
    @prev_of = prev_of
    @next_of = next_of
  end

  def size
    @hash.size
  end

  def [](key)
    @hash[key]
  end

  def increase(el)
    @hash[el].count += 1
    @valid_pick = false
    return @hash[el]
  end

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

  def make_picker
    @picker.clear
    total = 0
    @hash.each { |key, val|
      if val.count > 0
        total += val.count
        @picker << [total, key]
      end
    }
    @total = total
    @valid_pick = true
  end

  def picker
    make_picker unless @valid_pick
    @picker.dup
  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 = ChanceHash.new
  end

  def add_forward(array, base=@mkv)
    return if array.compact.empty?
    sym = array.shift
    s = sym and sym.to_sym
    base.increase(s) if s or base != @mkv
    return if array.empty?
    add_forward(array, base[s].next) if s or base == @mkv
  end

  def add_backwards(array, base=@mkv)
    return if array.compact.empty?
    sym = array.pop
    s = sym and sym.to_sym
    base.increase(s) if s or base != @mkv
    return if array.empty?
    add_backwards(array, base[s].prev) if s or base == @mkv
  end

  def add(array)
    raise "Too many words in new data" if array.size > @max_order + 1
    add_forward(array.dup)
    add_backwards(array.dup)
  end

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

    syms.size.times { |i|
      ord = [@max_order, syms.size-i].min
      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).dup
    ord = ar.size
    base = @mkv
    until ar.empty?
      base = base[ar.shift].next
    end
    base.random
  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).dup
    ord = ar.size
    base = @mkv
    until ar.empty?
      base = base[ar.pop].prev
    end
    base.random
  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.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