#Markov-chain in #Lua

markov-chain.lua

-- PiL Exercise 10.4: Generalize the Markov-chain algorithm so
-- that it can use any size for the sequence of previous
-- words used in the choice of the next word.


local N = tonumber(arg[1])
local MAXGEN = tonumber(arg[2])

local function allwords()
    local line = io.read()
    local pos = 1
    return function()
        while line do
            local s, e = string.find(line, "%w+", pos)
            if s then
                pos = e + 1
                return string.sub(line, s, e)
            else
                line = io.read()
                pos = 1
            end
        end
        return nil
    end
end

local function prefix(words)
    local prefix_sequence = ''
    for _,w in pairs(words) do
      prefix_sequence =  prefix_sequence .. w .. " "
    end
    return prefix_sequence
end

local statetab = {}

local function insert(index, value)
    local list = statetab[index]
    if list == nil then
        statetab[index] = {value}
    else
        list[#list+1] = value
    end
end

local NOWORD = "\n"

-- build table
local function init_words ()
  local prefix_words = {}
  local i = 1
  while i <= N do
    prefix_words[i] = NOWORD
    i = i + 1
  end
  return prefix_words
end

local prefix_words = init_words()

for w in allwords() do
    insert(prefix(prefix_words), w)
    table.remove(prefix_words, 1)
    table.insert(prefix_words, w)
end
insert(prefix(prefix_words), NOWORD)

prefix_words = init_words()


-- generate text

for i = 1, MAXGEN do
    local list = statetab[prefix(prefix_words)]

    -- choose random item from list
    local r = math.random(#list)
    local nextword = list[r]
    if nextword == NOWORD then return end
    io.write(nextword, " ")
    table.remove(prefix_words, 1)
    table.insert(prefix_words, nextword)
end