-- 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 implements functions for encoding information into binary strings using arithmetic encoding.
-- Due to the complexity involved; there is a program to test this file in dev/serialize_test.

local char, byte = string.char, string.byte
local concat, wipe, insert = table.concat, table.wipe, table.insert
local floor, max = math.floor, math.max

-- The encoder has limits on the amout of information it can encode.
-- This function checks if the list exceeds the legal amount, and if so,
-- reduces the values to remain within legal ranges.
local function normalizeList(list)
  --[===[@debug@
  assert(list[0] == 0)
  --@end-debug@]===]
  
  local elements = #list
  local range = list[elements]
  
  if range >= 4194304 then
    local scale = 2097152/range
    local prev = 0
    for i = 1,elements do
      local val = range[i]
      local prob = val-prev
      
      if prob > 0 then
        range[i] = range[i-1] + max(1, floor(prob*scale+0.5))
      else
        range[i] = range[i-1]
      end
      
      prev = val
    end
    
    -- Could be false if the list contains over 4 million elements;
    -- but if it does, you've doing something wrong.
    assert(list[elements] < 4194304)
  end
end

-- This increases the probability of symbol in list. We do this for the caller on their behalf,
-- so they don't need to.
local function amendListProbability(list, symbol)
  for i=symbol,#list do
    list[i] = list[i]+1
  end
end

-- These decoders all take lists, and each index in the list represents a symbol.
-- Element 0 of the list is special and must be zero.
-- list[x] must be equal or greater than list[x-1].
-- (list[x-1]-list[x])/list[#list] is the probability of symbol x appearing.
-- You must not attempt to encode a symbol with a probability of zero.
-- list[#list] should be less than 16384; 2^14 is the maximum precision the encoder is guarenteed to be able to encode.

function shared.createArithmeticEncoder()
  local encoder = {}
  local data = {}
  
  local encoded_low, encoded_high = 0, 1073741823
  local outstanding = 0
  
  local bits = 0
  local byte = 0
  
  local function bit0()
    byte = byte*2
    
    if bits == 7 then
      insert(data, char(byte))
      byte = 0
      bits = 0
    else
      bits = bits+1
    end
  end
  
  local function bit1()
    byte = byte*2+1
    
    if bits == 7 then
      insert(data, char(byte))
      byte = 0
      bits = 0
    else
      bits = bits+1
    end
  end
  
  -- Returns the encoded data, and the number of bits encoded.
  -- If you were to continue encoding, you would create another stream concatenated onto the previous.
  -- The decoder can retrieve the individual streams by passing the bit offset.
  function encoder.flush()
    if encoded_low < 268435456 then
      bit0()
      for i=0,outstanding do bit1() end
    else
      bit1()
      for i=0,outstanding do bit0() end
    end
    
    encoded_low, encoded_high, outstanding = 0, 1073741823, 0
    
    local raw_data = concat(data)
    wipe(data)
    data[1] =raw_data
    
    if bits == 0 then
      return raw_data, raw_data:len()*8
    else
      return raw_data .. char(byte * 2^(8-bits)), raw_data:len()*8+bits
    end
  end
  
  local function normalize()
    while true do
      --[===[@debug@
      assert(encoded_low < encoded_high)
      --@end-debug@]===]
      
      if encoded_high < 536870912 then
        bit0() for i=1,outstanding do bit1() end
        outstanding = 0
        encoded_low = encoded_low*2
        encoded_high = encoded_high*2+1
      elseif encoded_low >= 536870912 then
        bit1() for i=1,outstanding do bit0() end
        outstanding = 0
        encoded_low = encoded_low*2-1073741824
        encoded_high = encoded_high*2-1073741823
      elseif encoded_low >= 268435456 and encoded_high < 805306368 then
        outstanding = outstanding + 1
        encoded_low = encoded_low*2-536870912
        encoded_high = encoded_high*2-536870911
      else
        break
      end
    end
  end
  
  function encoder.encodeSymbol(list, symbol)
    --[===[@debug@
    assert(list[symbol] > list[symbol-1])
    --@end-debug@]===]
    
    normalizeList(list)
    
    local scaling = (encoded_high-encoded_low+1)/list[#list]
    encoded_high = encoded_low+floor(list[symbol]*scaling-1)
    encoded_low = encoded_low+floor(list[symbol-1]*scaling)
    
    amendListProbability(list, symbol)
    
    normalize()
  end
  
  -- Encodes a boolean value.
  function encoder.encodeBoolean(false_prob, true_prob, state)
    --[===[@debug@
    assert(floor(false_prob) == false_prob)
    assert(floor(true_prob) == true_prob)
    assert(false_prob >= 0)
    assert(true_prob >= 0)
    assert(state == true or state == false)
    --@end-debug@]===]
    
    local prob_total = false_prob + true_prob
    
    if prob_total >= 4194304 then
      local scale = 2097152/total
      false_prob = max(1, floor(false_prob*scale+0.5))
      true_prob = max(1, floor(true_prob*scale+0.5))
      prob_total = false_prob+true_prob
    end
    
    local scaling = (encoded_high-encoded_low+1)/(prob_total)
    
    if state then
      --[===[@debug@
      assert(true_prob > 0)
      --@end-debug@]===]
      encoded_high  = encoded_low+floor(false_prob*scaling)
    else
      --[===[@debug@
      assert(false_prob > 0)
      --@end-debug@]===]
      encoded_low = encoded_low+floor(false_prob*scaling-1)
    end
    
    normalize()
  end
  
  -- Encodes an integer in the range [low,high]
  local function encodeInteger(low, high, val)
    --[===[@debug@
    assert(floor(low) == low)
    assert(floor(high) == high)
    assert(floor(val) == val)
    assert(high >= low)
    assert(val >= low)
    assert(val <= high)
    --@end-debug@]===]
    
    val = val - low
    local range = high-low+1
    
    if range >= 4194304 then
      -- Shave 11 bits off val and try again.
      
      local end_range = floor(range/2048)
      local a, b = floor(val/2048), val%2048
      
      encodeInteger(0, end_range, a)
      
      if a == end_range then
        encodeInteger(0, range-end_range*2048-1, b)
      else
        encodeInteger(0, 2047, b)
      end
      
      return
    end
    
    local scaling = (encoded_high-encoded_low+1)/range
    
    encoded_high = encoded_low+floor((val+1)*scaling-1)
    encoded_low  = encoded_low+floor(val*scaling)
    
    normalize()
  end
  
  encoder.encodeInteger = encodeInteger
  
  return encoder
end

-- TODO: accept a bit offset so multiple streams can appear in the same raw string.
function shared.createArithmeticDecoder(raw_string)
  local decoder = {}
  
  local byte_offset = 1
  local encoded_low, encoded_high = 0, 1073741823
  local encoded_value = (byte(raw_string, byte_offset  ) or 0)*4194304
                       +(byte(raw_string, byte_offset+1) or 0)*16384
                       +(byte(raw_string, byte_offset+2) or 0)*64
                       +floor((byte(raw_string, byte_offset+3) or 0)*0.25)
  
  local next_byte_bits = 6
  local next_byte = ((byte(raw_string, byte_offset+3) or 0)*64)%256
  byte_offset = byte_offset+4
  
  local function nextBit()
    if next_byte < 128 then
      if next_byte_bits == 7 then
        next_byte = byte(raw_string, byte_offset) or 0
        byte_offset = byte_offset+1
        next_byte_bits = 0
      else
        next_byte = next_byte*2
        next_byte_bits = next_byte_bits + 1
      end
      
      return 0
    else
      if next_byte_bits == 7 then
        next_byte = byte(raw_string, byte_offset) or 0
        byte_offset = byte_offset+1
        next_byte_bits = 0
      else
        next_byte = (next_byte-128)*2
        next_byte_bits = next_byte_bits + 1
      end
      
      return 1
    end
  end
  
  local function normalize()
    while true do
      --[===[@debug@
      assert(encoded_low < encoded_high)
      assert(encoded_value >= encoded_low and encoded_value <= encoded_high)
      --@end-debug@]===]
      if encoded_high < 536870912 then
        encoded_low = encoded_low*2
        encoded_high = encoded_high*2+1
        encoded_value = encoded_value*2+nextBit()
      elseif encoded_low >= 536870912 then
        encoded_low = encoded_low*2-1073741824
        encoded_high = encoded_high*2-1073741823
        encoded_value = encoded_value*2-1073741824+nextBit()
      elseif encoded_low >= 268435456 and encoded_high < 805306368 then
        encoded_low = encoded_low*2-536870912
        encoded_high = encoded_high*2-536870911
        encoded_value = encoded_value*2-536870912+nextBit()
      else
        break
      end
    end
  end
  
  function decoder.decodeSymbol(list)
    normalizeList(list)
    
    local symbol_min, symbol_max = 1,#list
    local scaling = (encoded_high-encoded_low+1)/list[symbol_max]
    local code = floor(((encoded_value-encoded_low+1)*list[symbol_max]-1)/(encoded_high-encoded_low+1))
    
    while true do
      local symbol_mid = floor((symbol_min+symbol_max)*0.5)
      
      if code < list[symbol_mid-1] then
        symbol_max = symbol_mid
      elseif code >= list[symbol_mid] then
        symbol_min = symbol_mid+1
      else
        encoded_high = encoded_low+floor(list[symbol_mid]*scaling-1)
        encoded_low = encoded_low+floor(list[symbol_mid-1]*scaling)
        
        amendListProbability(list, symbol_mid)
        
        normalize()
        return symbol_mid
      end
    end
  end
  
  -- Decodes a boolean value.
  function decoder.decodeBoolean(false_prob, true_prob)
    --[===[@debug@
    assert(floor(false_prob) == false_prob)
    assert(floor(true_prob) == true_prob)
    assert(false_prob >= 0)
    assert(true_prob >= 0)
    --@end-debug@]===]
    
    local prob_total = false_prob + true_prob
    
    if prob_total >= 4194304 then
      local scale = 2097152/total
      false_prob = max(1, floor(false_prob*scale+0.5))
      true_prob = max(1, floor(true_prob*scale+0.5))
      prob_total = false_prob+true_prob
    end
    
    local scaling = (encoded_high-encoded_low+1)/(prob_total)
    local code = floor(((encoded_value-encoded_low+1)*prob_total-1)/(encoded_high-encoded_low+1))
    local state
    
    if code < false_prob then
      state = true
      encoded_high  = encoded_low+floor(false_prob*scaling)
    else
      state = false
      encoded_low = encoded_low+floor(false_prob*scaling-1)
    end
    
    normalize()
    
    return state
  end
  
  -- Decodes an integer in the range [low,high]
  local function decodeInteger(low, high)
    --[===[@debug@
    assert(floor(low) == low)
    assert(floor(high) == high)
    assert(high >= low)
    --@end-debug@]===]
    
    local range = high-low+1
    
    if range >= 4194304 then
      local end_range = floor(range/2048)
      local a = decodeInteger(0, end_range)
      
      if a == end_range then
        return low + a*2048 + decodeInteger(0, range-end_range*2048-1)
      else
        return low + a*2048 + decodeInteger(0, 2047)
      end
    end
    
    local scaling = (encoded_high-encoded_low+1)/range
    
    local val = floor(((encoded_value-encoded_low+1)*range-1)/(encoded_high-encoded_low+1))
    
    encoded_high = encoded_low+floor((val+1)*scaling-1)
    encoded_low  = encoded_low+floor(val*scaling)
    
    normalize()
    
    return low+val
  end
  
  decoder.decodeInteger = decodeInteger
  
  return decoder
end
