-- Copyright (C) 2010 Tyson Brown
-- 
-- This file is part QuestHelper2, an AddOn for World of Warcraft.
-- 
-- QuestHelper2 is free software: you can redistribute it and/or modify
-- it under the terms of the GNU General Public License as published by
-- the Free Software Foundation, either version 3 of the License, or
-- (at your option) any later version.
-- 
-- QuestHelper2 is distributed in the hope that it will be useful,
-- but WITHOUT ANY WARRANTY; without even the implied warranty of
-- MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-- GNU General Public License for more details.
-- 
-- You should have received a copy of the GNU General Public License
-- along with QuestHelper2. If not, see <http://www.gnu.org/licenses/>.

if LibDebug then LibDebug() end
local shared = select(2, ...)

-- This file contains functions for creating special encoders that predict the symbol
-- being encoded using markov chains.

local wipe = table.wipe

-- creates an object for predicting symbols.
-- Try not to use this with too many symbols, this has the potential to use a lot of memory.
function shared.createMarkovChainEncoder(encodeSymbol, symbols)
  --[===[@debug@
  assert(symbols >= 1)
  --@end-debug@]===]
  
  if symbols == 1 then
    -- Only one symbol? That was easy.
    return function () end
  end
  
  -- Holds information required to encode symbols the chain hasn't ever seen before.
  local fallback_symbols = {[0]=0}
  local fallback_symbols_fwd = {}
  local fallback_symbols_rev = {}
  
  for i=1,symbols do
    fallback_symbols[i] = i
    fallback_symbols_fwd[i] = i
    fallback_symbols_rev[i] = i
  end
  
  -- This holds everything in the first order markov chain, and
  -- acts as the root for the high order markov chains that get
  -- created as necessary.
  local order1 = {[0]=0,1,rev={},next_order={}}
  
  -- When removing symbols from an order temporarly, the temporary result goes in this table.
  local ammended_fwd, ammended_order = {}, {}
  
  -- This returns the new lookup, as it will no longer match what it was.
  local function ammendOrder(order, impossible_symbols, rev, lookup)
    wipe(ammended_fwd)
    wipe(ammended_order)
    
    local original_lookup = lookup
    
    for symbol in pairs(impossible_symbols) do
      local index = rev[symbol]
      
      if index then
        -- Move lookup down, since we'll be removing an element before it.
        if index < original_lookup then lookup = lookup - 1 end
        
        ammended_fwd[index] = true
      end
    end
    
    local sum = 0
    local prev = 0
    local index = 1
    ammended_order[0] = 0
    
    for i, val in ipairs(order) do
      if not ammended_fwd[i] then
        sum = sum + val-prev
        ammended_order[index] = sum
        index = index + 1
      end
      
      prev = val
    end
    
    return lookup
  end
  
  local function encodeSymbolOrder(symbol, order, prev_symbol, ...)
    local impossible_symbols
    
    local rev = order.rev
    -- Lookup is the what this order has symbol stored as.
    local lookup = rev[symbol]
    
    -- If we have context in which to encode symbol, pass that
    -- context to the next order encoder.
    if prev_symbol then
      local next_order = order.next_order[prev_symbol]
      
      if next_order then
        -- Ask the next order to encode the symbol for us.
        impossible_symbols = encodeSymbolOrder(symbol, next_order, ...)
        
        -- if impossible_symbols is not nil, then
        -- the next order was unable to encode the symbol,
        -- and the responsibility falls to us.
        -- Otherwise we're done.
        if not impossible_symbols then
          if lookup then
            -- Increase the probability of the symbol in question.
            for i=lookup,#order do order[i] = order[i]+1 end
          else
            -- Haven't seen this symbol before, remember it.
            lookup = #order
            rev[symbol] = lookup
            --[===[@debug@
            assert(order[lookup]-order[lookup-1] == 1)
            --@end-debug@]===]
            
            if lookup ~= symbols then
              -- Make room to represent an unknown symbol if there might be more.
              order[lookup+1] = order[lookup]+1
            end
          end
          
          return
        end
      else
       -- The next order doesn't exist, so create it for future reference.
        next_order = {[0]=0,1,2,rev={[prev_symbol]=1},next_order={}}
        order.next_order[prev_symbol] = next_order
      end
    end
    
    -- If we've reached this point, then the responsibility of encoding symbol
    -- has fallen to us, as none of the orders above us have seen this context
    -- before.
    local last = #order
    if not lookup then
      lookup = last
    end
      
    if impossible_symbols then
      encodeSymbol(ammended_order, ammendOrder(order, impossible_symbols, rev, lookup))
     
      -- Update probabilities.
      for i=lookup,last do order[i] = order[i]+1 end
    else
      encodeSymbol(order, lookup)
    end
    
    if not rev[symbol] then
      -- Haven't seen this symbol before.
      
      -- Create a table to hold the set of impossible symbols if the next order hasn't already created one.
      impossible_symbols = impossible_symbols or {}
      
      -- Add any symbols we could have handled to the list of symbols that symbol definitely can't be.
      for symbol in pairs(rev) do impossible_symbols[symbol] = true end
      
      -- Remember this symbol for future reference.
      --[[ the decoder puts this part in a closure ]]
        rev[symbol] = lookup
        order[lookup] = order[lookup-1]+1
        
        if lookup ~= symbols then
          -- If there might be a possibility of more symbols, then
          -- make room to remember a new symbol should we encounter it.
          order[lookup+1] = order[lookup]+1
        end
      --[[ the decoder puts this part in a closure ]]
      
      -- Tell the order below us the terrible news.
      return impossible_symbols
    end
  end
  
  -- This is a history of previously encoded symbols.
  local context1, context2 = nil, nil
  
  return function (symbol)
    local impossible_symbols = encodeSymbolOrder(symbol, order1, context1, context2)
    
    if impossible_symbols then
      -- impossible_symbols is naturally only symbols order1 has already encountered,
      -- as such, this table is meaningless except as a flag that the responsibility
      -- for encoding the symbol has fallen to us.
      
      -- Encode the symbol in question.
      local lookup = fallback_symbols_rev[symbol]
      encodeSymbol(fallback_symbols, lookup)
      
      -- To avoid the need to shift everything around, we move the very last symbol
      -- to the position of the symbol we just removed.
      local last = #fallback_symbols
      
      if last ~= lookup then
        local last_symbol = fallback_symbols_fwd[last]
        fallback_symbols_rev[last_symbol] = lookup
        fallback_symbols_fwd[lookup] = last_symbol
      end
      
      -- The current symbol has of course been removed.
      fallback_symbols_rev[symbol] = nil
      fallback_symbols_fwd[last] = nil
      
      -- Since all symbols have the same probability, all we have to do is remove the very last element.
      -- Everything ends up exactly the same.
      fallback_symbols[last] = nil
    end
    
    -- Remember context.
    context1, context2 = symbol, context1
  end
end

-- The decoder is, naturally, largely a copy and paste of the encoder.
function shared.createMarkovChainDecoder(decodeSymbol, symbols)
  --[===[@debug@
  assert(symbols >= 1)
  --@end-debug@]===]
  
  if symbols == 1 then
    -- Only one symbol? That was easy.
    return function () return 1 end
  end
  
  -- Holds information required to encode symbols the chain hasn't ever seen before.
  local fallback_symbols = {[0]=0}
  local fallback_symbols_fwd = {}
  local fallback_symbols_rev = {}
  
  for i=1,symbols do
    fallback_symbols[i] = i
    fallback_symbols_fwd[i] = i
    fallback_symbols_rev[i] = i
  end
  
  -- This holds everything in the first order markov chain, and
  -- acts as the root for the high order markov chains that get
  -- created as necessary.
  local order1 = {[0]=0,1,rev={},fwd={},next_order={}}
  
  -- When removing symbols from an order temporarly, the temporary result goes in this table.
  local ammended_fwd, ammended_order = {}, {}
  
  -- This returns the new lookup, as it will no longer match what it was.
  local function ammendOrder(order, impossible_symbols, rev)
    wipe(ammended_fwd)
    wipe(ammended_order)
    
    for symbol in pairs(impossible_symbols) do
      local index = rev[symbol]
      if index then
        ammended_fwd[index] = true
      end
    end
    
    local sum = 0
    local prev = 0
    local index = 1
    ammended_order[0] = 0
    
    for i, val in ipairs(order) do
      if not ammended_fwd[i] then
        sum = sum + val-prev
        ammended_order[index] = sum
        index = index + 1
      end
      
      prev = val
    end
    
  end
  
  -- We can't update update symbol probabilities until it's known,
  -- so we stick closures in here that once we know what it is
  -- we can go back and finish up.
  local corrections = {}
  
  local function decodeSymbolOrder(order, prev_symbol, ...)
    local symbol, impossible_symbols
    
    local rev = order.rev
    local fwd = order.fwd
    
    -- If we have context in which to encode symbol, pass that
    -- context to the next order encoder.
    if prev_symbol then
      local next_order = order.next_order[prev_symbol]
      
      if next_order then
        -- Ask the next order to encode the symbol for us.
        symbol, impossible_symbols = decodeSymbolOrder(next_order, ...)
        
        -- if impossible_symbols is not nil, then
        -- the next order was unable to encode the symbol,
        -- and the responsibility falls to us.
        -- Otherwise we're done.
        if symbol then
          local lookup = rev[symbol]
          if lookup then
            -- Increase the probability of the symbol in question.
            for i=lookup,#order do order[i] = order[i]+1 end
          else
            -- Haven't seen this symbol before, remember it.
            lookup = #order
            rev[symbol] = lookup
            fwd[lookup] = symbol
            --[===[@debug@
            assert(order[lookup]-order[lookup-1] == 1)
            --@end-debug@]===]
            
            if lookup ~= symbols then
              -- Make room to represent an unknown symbol if there might be more.
              order[lookup+1] = order[lookup]+1
            end
          end
          
          return symbol
        end
      else
       -- The next order doesn't exist, so create it for future reference.
        next_order = {[0]=0,1,2,fwd={[1]=prev_symbol},rev={[prev_symbol]=1},next_order={}}
        order.next_order[prev_symbol] = next_order
      end
    end
    
    -- If we've reached this point, then the responsibility of encoding symbol
    -- has fallen to us, as none of the orders above us have seen this context
    -- before.
    local last = #order
    local lookup
    
    if impossible_symbols then
      ammendOrder(order, impossible_symbols, rev)
      lookup = decodeSymbol(ammended_order)
      
      -- Shift lookup so that it corresponds to the unammended order, i.e., us.
      -- Which turns out to be slightly complicated because we don't know
      -- which of the symbols is before lookup until lookup has its correct
      -- position.
      local need_check = true
      while need_check do
        need_check = false
        for index in pairs(ammended_fwd) do
          if index <= lookup then
            ammended_fwd[index] = nil
            lookup = lookup + 1
            need_check = true
          end
        end
      end
      
      -- Update probabilities.
      for i=lookup,last do order[i] = order[i]+1 end
    else
      lookup = decodeSymbol(order)
    end
    
    local symbol = fwd[lookup]
    if not symbol then
      -- Haven't seen this symbol before.
      
      -- Create a table to hold the set of impossible symbols if the next order hasn't already created one.
      impossible_symbols = impossible_symbols or {}
      
      -- Add any symbols we could have handled to the list of symbols that symbol definitely can't be.
      for symbol in pairs(rev) do impossible_symbols[symbol] = true end
      
      -- Remember this symbol for future reference.
      corrections[function (symbol)
        fwd[lookup] = symbol
        rev[symbol] = lookup
        order[lookup] = order[lookup-1]+1
        
        if lookup ~= symbols then
          -- If there might be a possibility of more symbols, then
          -- make room to remember a new symbol should we encounter it.
          order[lookup+1] = order[lookup]+1
        end
      end] = true
      
      -- Tell the order below us the terrible news.
      return nil, impossible_symbols
    end
    
    return symbol
  end
  
  -- This is a history of previously encoded symbols.
  local context1, context2 = nil, nil
  
  return function (symbol)
    local symbol, impossible_symbols = decodeSymbolOrder(order1, context1, context2)
    
    if impossible_symbols then
      -- impossible_symbols is naturally only symbols order1 has already encountered,
      -- as such, this table is meaningless except as a flag that the responsibility
      -- for encoding the symbol has fallen to us.
      
      -- Decode the symbol in question.
      local lookup = decodeSymbol(fallback_symbols)
      symbol = fallback_symbols_fwd[lookup]
      
      -- To avoid the need to shift everything around, we move the very last symbol
      -- to the position of the symbol we just removed.
      local last = #fallback_symbols
      
      if last ~= lookup then
        local last_symbol = fallback_symbols_fwd[last]
        fallback_symbols_rev[last_symbol] = lookup
        fallback_symbols_fwd[lookup] = last_symbol
      end
      
      -- The current symbol has of course been removed.
      fallback_symbols_rev[symbol] = nil
      fallback_symbols_fwd[last] = nil
      
      -- Since all symbols have the same probability, all we have to do is remove the very last element.
      -- Everything ends up exactly the same.
      fallback_symbols[last] = nil
    end
    
    -- For all those bits that required us to know what symbol was; we finally know.
    for correction in pairs(corrections) do
      correction(symbol)
    end
    
    wipe(corrections)
    
    -- Remember context.
    context1, context2 = symbol, context1
    
    return symbol
  end
end
