#
# This file is part of mymc+, based on mymc by Ross Ridge.
#
# mymc+ 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.
#
# mymc+ 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 mymc+.  If not, see <http://www.gnu.org/licenses/>.
#

"""
Implementation of Haruhiko Okumura's LZARI data compression algorithm
in Python.  Largely based on LZARI.C, one key difference is the use of
a two level dicitionary look up during compression rather than
LZARI.C's binary search tree.
"""

import sys
import array
import binascii
import string
import time
from bisect import bisect_right
from math import log

try:
    import ctypes
    import mymcsup
except ImportError:
    mymcsup = None

hexlify = binascii.hexlify

__ALL__ = ['lzari_codec', 'string_to_bit_array', 'bit_array_to_string']

#
# Fundamental constants of the LZARI compression alogorithm.
#
# Changing any of these values will create an incompatible implementation.
#

HIST_LEN = 4096
MIN_MATCH_LEN = 3
MAX_MATCH_LEN = 60

ARITH_BITS = 15
QUADRANT1 = 1 << ARITH_BITS
QUADRANT2 = QUADRANT1 * 2
QUADRANT3 = QUADRANT1 * 3
QUADRANT4 = QUADRANT1 * 4
MAX_CUM = QUADRANT1 - 1
MAX_CHAR = (256 + MAX_MATCH_LEN - MIN_MATCH_LEN + 1)

#
# Other constants specific to this implementation
#

MAX_SUFFIX_CHAIN = 50    # limit on how many identical suffixes to try to match

#def debug(value, msg):
#    print "@@@ %s %04x" % (msg, value)
debug = lambda value, msg: None

_tr_16 = bytes.maketrans(b"0123456789abcdef",
              b"\x00\x01\x02\x03"
              b"\x10\x11\x12\x13"
              b"\x20\x21\x22\x23"
              b"\x30\x31\x32\x33")
_tr_4 = bytes.maketrans(b"0123",
             b"\x00\x01"
             b"\x10\x11")
_tr_2 = bytes.maketrans(b"01", b"\x00\x01")

def string_to_bit_array(s):
    """Convert a string to an array containing a sequence of bits."""
    s = binascii.hexlify(s).translate(_tr_16)
    s = binascii.hexlify(s).translate(_tr_4)
    s = binascii.hexlify(s).translate(_tr_2)
    a = array.array('B', s)
    return a

_tr_rev_2 = bytes.maketrans(b"\x00\x01", b"01")
_tr_rev_4 = bytes.maketrans(b"\x00\x01\x10\x11", b"0123")
_tr_rev_16 = bytes.maketrans(b"\x00\x01\x02\x03"
                  b"\x10\x11\x12\x13"
                  b"\x20\x21\x22\x23"
                  b"\x30\x31\x32\x33",
                  b"0123456789abcdef")
def bit_array_to_string(a):
    """Convert an array containing a sequence of bits to a string."""
    remainder = len(a) % 8
    if remainder != 0:
        a.fromlist([0] * (8 - remainder))
    s = a.tobytes()
    s = binascii.unhexlify(s.translate(_tr_rev_2))
    s = binascii.unhexlify(s.translate(_tr_rev_4))
    return binascii.unhexlify(s.translate(_tr_rev_16))

def _match(src, pos, hpos, mlen, end):
    mlen += 1
    if not src.startswith(src[hpos : hpos + mlen], pos):
        return None
    for i in range(mlen, end):
        if src[pos + i] != src[hpos + i]:
            return i
    return end

def _rehash_table2(src, chars, head, next, next2, hist_invalid):
    p = head
    table2 = {}
    l = []
    while p > hist_invalid:
        l.append(p)
        p = next[p % HIST_LEN]
    l.reverse()
    for p in l:
        p2 = p + MIN_MATCH_LEN
        key2 = src[p2 : p2 + chars]
        head2 = table2.get(key2, hist_invalid)
        next2[p % HIST_LEN] = head2
        table2[key2] = p
    return table2

class lzari_codec(object):
    # despite the name this does not implement a codec compatible
    # with Python's codec system
    
    def init(self, decode):
        self.high = QUADRANT4
        self.low = 0
        if decode:
            self.code = 0
            # reverse the order of sym_cum so bisect_right() can
            # be used for faster searching
            self.sym_cum = list(range(0, MAX_CHAR + 1))
        else:
            self.shifts = 0
            self.char_to_symbol = list(range(1, MAX_CHAR + 1))
            self.sym_cum = list(range(MAX_CHAR, -1, -1))
            self.next_table = [None] * HIST_LEN
            self.next2_table = [None] * HIST_LEN
            self.suffix_table = {}

        self.symbol_to_char = [0] + list(range(MAX_CHAR))
        self.sym_freq = [0] + [1] * MAX_CHAR
        self.position_cum = [0] * (HIST_LEN + 1)
        a = 0
        for i in range(HIST_LEN, 0, -1):
            a =  a + 10000 // (200 + i)
            self.position_cum[i - 1] = a
        
    def search(self, table, x):
        c = 1
        s = len(table) - 1
        while True:
            a = (s + c) // 2
            if table[a] <= x:
                s = a
            else:
                c = a + 1
            if c >= s:
                break
        return c

    def update_model_decode(self, symbol):
        # A compatible implemention to the one used while compressing.
        
        sym_freq = self.sym_freq
        sym_cum = self.sym_cum
        
        if self.sym_cum[MAX_CHAR] >= MAX_CUM:
            c = 0
            for i in range(MAX_CHAR, 0, -1):
                self.sym_cum[MAX_CHAR - i] = c
                a = (self.sym_freq[i] + 1) // 2
                self.sym_freq[i] = a
                c += a
            self.sym_cum[MAX_CHAR] = c
        freq = sym_freq[symbol]
        new_symbol = symbol
        while self.sym_freq[new_symbol - 1] == freq:
            new_symbol -= 1
        # new_symbol = sym_freq.index(freq)
        if new_symbol != symbol:
            symbol_to_char = self.symbol_to_char
            swap_char = symbol_to_char[new_symbol]
            char = symbol_to_char[symbol]
            symbol_to_char[new_symbol] = char
            symbol_to_char[symbol] = swap_char
        sym_freq[new_symbol] = freq + 1
        for i in range(MAX_CHAR - new_symbol + 1, MAX_CHAR + 1):
            sym_cum[i] += 1
            
    def update_model_encode(self, symbol):
        sym_freq = self.sym_freq
        sym_cum = self.sym_cum
        
        if sym_cum[0] >= MAX_CUM:
            c = 0
            for i in range(MAX_CHAR, 0, -1):
                sym_cum[i] = c
                a = (sym_freq[i] + 1) // 2
                sym_freq[i] = a
                c += a
            sym_cum[0] = c
        freq = sym_freq[symbol]
        new_symbol = symbol
        while sym_freq[new_symbol - 1] == freq:
            new_symbol -= 1
        if new_symbol != symbol:
            debug(new_symbol, "a")
            swap_char = self.symbol_to_char[new_symbol]
            char = self.symbol_to_char[symbol]
            self.symbol_to_char[new_symbol] = char
            self.symbol_to_char[symbol] = swap_char
            self.char_to_symbol[char] = new_symbol
            self.char_to_symbol[swap_char] = symbol
        sym_freq[new_symbol] += 1
        for i in range(new_symbol):
            sym_cum[i] += 1

    def decode_char(self):
        high = self.high
        low = self.low
        code = self.code
        sym_cum = self.sym_cum
        
        _range = high - low
        max_cum_freq = sym_cum[MAX_CHAR]
        n = ((code - low + 1) * max_cum_freq - 1) // _range
        i = bisect_right(sym_cum, n, 1)
        high = low + sym_cum[i] * _range // max_cum_freq
        low += sym_cum[i - 1] * _range // max_cum_freq
        symbol = MAX_CHAR + 1 - i

        while True:
            if low < QUADRANT2:
                if low < QUADRANT1 or high > QUADRANT3:
                    if high > QUADRANT2:
                        break
                else:
                    low -= QUADRANT1
                    code -= QUADRANT1
                    high -= QUADRANT1
            else:
                low -= QUADRANT2
                code -= QUADRANT2
                high -= QUADRANT2
            low *= 2
            high *= 2
            code = code * 2 + self.in_iter()

        ret = self.symbol_to_char[symbol]
        self.high = high
        self.low = low
        self.code = code
        self.update_model_decode(symbol)
        return ret
    
    def decode_position(self):
        _range = self.high - self.low
        max_cum = self.position_cum[0]
        pos = self.search(self.position_cum, ((self.code - self.low + 1) * max_cum - 1) // _range) - 1
        self.high = (self.low + self.position_cum[pos] * _range // max_cum)
        self.low += self.position_cum[pos + 1] * _range // max_cum
        while True:
            if self.low < QUADRANT2:
                if (self.low < QUADRANT1
                    or self.high > QUADRANT3):
                    if self.high > QUADRANT2:
                        return pos
                else:
                    self.low -= QUADRANT1
                    self.code -= QUADRANT1
                    self.high -= QUADRANT1
            else:
                self.low -= QUADRANT2
                self.code -= QUADRANT2
                self.high -= QUADRANT2
            self.low *= 2
            self.high *= 2
            self.code = self.in_iter() + self.code * 2

    def add_suffix_1(self, pos, find):
        # naive implemention used for testing
        
        if not find:
            return (None, 0)
        src = self.src
        mlen = min(1000, self.max_match, len(src) - pos)
        hist_start = max(pos - HIST_LEN, 0)
        while mlen >= MIN_MATCH_LEN:
            i = src.rfind(src[pos : pos + mlen], hist_start, pos)
            if i != -1:
                assert (src[pos : pos + mlen]
                    == src[i: i + mlen])
                return (i, mlen)
            mlen -= 1
        return (None, -1)
            
    def add_suffix_2(self, pos, find):
        # a two level dictionary look up that leverages Python's
        # built-in dicts to get something that's hopefully faster
        # than implementing binary trees in completely in Python.
        
        src = self.src
        suffix_table = self.suffix_table
        max_match = min(self.max_match, len(src) - pos)

        mlen = -1
        mpos = None
        
        hist_invalid = pos - HIST_LEN - 1
        modpos = pos % HIST_LEN
        pos2 = pos + MIN_MATCH_LEN
        
        key = src[pos : pos2]
        a = suffix_table.get(key)
        if a != None:
            next = self.next_table
            next2 = self.next2_table
            
            [count, head, table2, chars] = a
            
            pos3 = pos2 + chars
            key2 = src[pos2 : pos3]
            min_match2 = MIN_MATCH_LEN + chars
            if find:
                p = table2.get(key2, hist_invalid)
                maxmlen = max_match - min_match2
                while p > hist_invalid and mlen != maxmlen:
                    p3 = p + min_match2
                    if mpos == None and p3 <= pos:
                        mpos = p
                        mlen = 0
                    if p3 >= pos:
                        p = next2[p % HIST_LEN]
                        continue
                    rlen = _match(src, pos3, p3, mlen,
                              min(maxmlen, pos - p3))
                    if rlen != None:
                        mpos = p
                        mlen = rlen
                    p = next2[p % HIST_LEN]
            if mpos != None:
                mlen += min_match2
            elif find:
                p = head
                maxmlen = min(chars, max_match - MIN_MATCH_LEN)
                i = 0
                while (p > hist_invalid and i < 50000
                       and mlen < maxmlen):
                    assert i < count
                    i += 1
                    p2 = p + MIN_MATCH_LEN
                    l2 = pos - p2
                    if mpos == None and l2 >= 0:
                        mpos = p
                        mlen = 0
                    if l2 <= 0:
                        p = next[p % HIST_LEN]
                        continue
                    if l2 > maxmlen:
                        l2 = maxmlen
                    m = mlen + 1
                    if src.startswith(src[p2 : p2 + m],
                              pos2):
                        mpos = p
                        for j in range(m, l2):
                            if (src[pos2 + j]
                                != src[p2 + j]):
                                mlen = j
                                break
                        else:
                            mlen = l2
                    #rlen = _match(src, pos2, p2, mlen, l2)
                    #if rlen != None:
                    #    mpos = p
                    #    mlen = rlen
                    p = next[p % HIST_LEN]
                    
                if mpos != None:
                    mlen += MIN_MATCH_LEN
                    
            count += 1
            new_chars = int(log(count, 2))
            # new_chars = 50
            new_chars = min(new_chars, max_match - MIN_MATCH_LEN)
            if new_chars > chars:
                chars = new_chars
                table2 = _rehash_table2(src, chars, head,
                            next, next2,
                            hist_invalid)

            next[modpos] = head
            head = pos
            
            key2 = src[pos2 : pos2 + chars]
            head2 = table2.get(key2, hist_invalid)
            next2[modpos] = head2
            table2[key2] = pos

            a[0] = count
            a[1] = head
            a[2] = table2
            a[3] = chars
        else:
            self.next_table[modpos] = hist_invalid
            self.next2_table[modpos] = hist_invalid
            key2 = b""
            # key2 = src[pos2 : pos2 + 1]
            suffix_table[key] = [1, pos, {key2: pos}, len(key2)]

        p = pos - HIST_LEN
        if p >= 0:
            p2 = p + MIN_MATCH_LEN
            key = src[p : p2]
            a = suffix_table[key]
            (count, head, table2, chars) = a
            count -= 1
            if count == 0:
                assert head == p
                del suffix_table[key]
            else:
                key2 = src[p2 : p2 + chars]
                if table2[key2] == p:
                    del table2[key2]
                a[0] = count
        assert (mpos == None
            or src[pos : pos + mlen] == src[mpos : mpos + mlen])
        return (mpos, mlen)

    def _add_suffix(self, pos, find):
        r = self.add_suffix_2(pos, find)
        start_pos = self.start_pos
        if find and r[0] != None:
            print ("%4d %02x %4d %2d"
                   % (pos - start_pos, ord(self.src[pos]),
                  r[0] - start_pos, r[1]))
        else:
            print ("%4d %02x"
                       % (pos - start_pos, ord(self.src[pos])))
        return r
    
    add_suffix = add_suffix_2
    
    def output_bit(self, bit):
        self.append_bit(bit)
        bit ^= 1
        for i in range(self.shifts):
            self.append_bit(bit)
        self.shifts = 0
        
    def encode_char(self, char):
        low = self.low
        high = self.high
        sym_cum = self.sym_cum
        
        symbol = self.char_to_symbol[char]
        range = high - low
    
        high = low + range * sym_cum[symbol - 1] // sym_cum[0]
        low += range * sym_cum[symbol] // sym_cum[0]
        debug(high, "high");
        debug(low, "low");
        while True:
            if high <= QUADRANT2:
                self.output_bit(0)
            elif low >= QUADRANT2:
                self.output_bit(1)
                low -= QUADRANT2
                high -= QUADRANT2
            elif low >= QUADRANT1 and high <= QUADRANT3:
                self.shifts += 1
                low -= QUADRANT1
                high -= QUADRANT1
            else:
                break
            low *= 2
            high *= 2
        self.low = low
        self.high = high
        self.update_model_encode(symbol)

    def encode_position(self, position):
        position_cum = self.position_cum
        low = self.low
        high = self.high

        range = high - low
        high = low + range * position_cum[position] // position_cum[0]
        low += range * position_cum[position + 1] // position_cum[0]

        debug(high, "high");
        debug(low, "low");
        while True:
            if high <= QUADRANT2:
                self.output_bit(0)
            elif low >= QUADRANT2:
                self.output_bit(1)
                low -= QUADRANT2
                high -= QUADRANT2
            elif low >= QUADRANT1 and high <= QUADRANT3:
                self.shifts += 1
                low -= QUADRANT1
                high -= QUADRANT1
            else:
                break
            low *= 2
            high *= 2
            
        self.low = low
        self.high = high
            
    def encode(self, src, progress = None):
        """Compress a string."""
        
        length = len(src)
        if length == 0:
            return b""

        out_array = array.array('B')
        self.out_array = out_array
        self.append_bit = out_array.append
        
        self.init(False)

        max_match = min(MAX_MATCH_LEN, length)
        self.max_match = max_match
        self.src = src = b"\x20" * max_match + src
            
        in_length = len(src)
        
        self.start_pos = max_match
        
        for in_pos in range(max_match):
            self.add_suffix(in_pos, False)
        in_pos += 1
        last_percent = -1
        while in_pos < in_length:
            if progress:
                percent = (in_pos - max_match) * 100 // length
                if percent != last_percent:
                    sys.stderr.write("%s%3d%%\r"
                             % (progress, percent))
                    last_percent = percent
            debug(src[in_pos], "src")
            (match_pos, match_len) = self.add_suffix(in_pos, True)
            if match_len < MIN_MATCH_LEN:
                self.encode_char(src[in_pos])
            else:
                debug(in_pos - match_pos - 1, "match_pos")
                debug(match_len, "match_len")
                self.encode_char(256 - MIN_MATCH_LEN + match_len)
                self.encode_position(in_pos - match_pos - 1)
                for i in range(match_len - 1):
                    in_pos += 1
                    self.add_suffix(in_pos, False)
            in_pos += 1
                
        self.shifts += 1
        if self.low < QUADRANT1:
            self.output_bit(0)
        else:
            self.output_bit(1)

        #for k, v in sorted(self.suffix_table.items()):
        #    count, head, table2, chars = v
        #    print hexlify(k), count, head, len(table2), chars
            
        if progress:
            sys.stderr.write("%s100%%\n" % progress)
        
        return bit_array_to_string(out_array)
        
    def decode(self, src, out_length, progress = None):
        """Decompress a string."""
        
        a = string_to_bit_array(src)
        a.fromlist([0] * 32)     # add some extra bits 
        self.in_iter = iter(a).__next__

        out = array.array('B', b"\0") * out_length
        outpos = 0
        
        self.init(True)

        self.code = 0
        for i in range(ARITH_BITS + 2):
            self.code += self.code + self.in_iter()

        hist_pos = HIST_LEN - MAX_MATCH_LEN
        history = [0x20] * hist_pos + [0] * MAX_MATCH_LEN

        decode_char = self.decode_char
        last_percent = -1
        last_time = time.time()
        while outpos < out_length:
            if progress:
                percent = outpos * 100 // out_length
                if percent != last_percent:
                    now = time.time()
                    if now - last_time >= 1:
                        sys.stderr.write("%s%3d%%\r"
                            % (progress, percent))
                        last_percent = percent
                        last_time = now
            char = decode_char()
            if char >= 0x100:
                pos = self.decode_position()
                length = char - 0x100 + MIN_MATCH_LEN
                base = (hist_pos - pos - 1) % HIST_LEN
                for off in range(length):
                    a = history[(base + off) % HIST_LEN]
                    out[outpos] = a
                    outpos += 1
                    history[hist_pos] = a
                    hist_pos = (hist_pos + 1) % HIST_LEN
            else:
                out[outpos] = char
                outpos += 1
                history[hist_pos] = char
                hist_pos = (hist_pos + 1) % HIST_LEN
        
        self.in_iter = None
        if progress:
            sys.stderr.write("%s100%%\n" % progress)
        return out.tobytes()

if mymcsup == None:
    def decode(src, out_length, progress = None):
        return lzari_codec().decode(src, out_length, progress)
    
    def encode(src, progress = None):
        return lzari_codec().encode(src, progress)
else:
    mylzari_decode = mymcsup.mylzari_decode
    mylzari_encode = mymcsup.mylzari_encode
    mylzari_free_encoded = mymcsup.mylzari_free_encoded
    
    def decode(src, out_length, progress = None):
        out = ctypes.create_string_buffer(out_length)
        if (mylzari_decode(src, len(src), out, out_length, progress)
            == -1):
            raise ValueError("compressed input is corrupt")
        return ctypes.string_at(out, out_length)

    def encode(src, progress = None):
        (r, compressed, comp_len) = mylzari_encode(src, len(src),
                               progress)
        # print r, compressed.value, comp_len
        if r == -1:
            raise MemoryError("out of memory during compression")
        if compressed.value == None:
            return b""
        ret = ctypes.string_at(compressed.value, comp_len.value)
        mylzari_free_encoded(compressed)
        return ret;

def main2(args):
    import struct
    import os
    
    src = open(args[2], "rb").read()
    lzari = lzari_codec()
    out = open(args[3], "wb")
    start = os.times()
    if args[1] == "c":
        dest = lzari.encode(src)
        now = os.times()
        out.write(struct.pack("L", len(src)))
    else:
        dest = lzari.decode(src[4:],
                    struct.unpack("L", src[:4])[0])
        now = os.times()
    out.write(dest)
    out.close()
    print("time:", now[0] - start[0], now[1] - start[1], now[4] - start[4])


def _get_hotshot_lineinfo(filename):
    import hotshot.log
    log = hotshot.log.LogReader(filename)
    timings = {}
    for what, loc, tdelta in log:
        if what == hotshot.log.LINE:
            a = timings.get(loc)
            if a == None:
                timings[loc] = [1, tdelta]
            else:
                a[0] += 1
                a[1] += tdelta
    return list(timings.items())

def _dump_hotshot_lineinfo(log):
    a = sorted(_get_hotshot_lineinfo(log))
    total_count = sum((time[0]
               for (loc, time) in a))
    total_time = sum((time[1]
              for (loc, time) in a))
    for (loc, [count, time]) in a:
        print(("%8d %6.3f%%  %8d %6.3f%%"
               % (time, time * 100.0 / total_time,
              count, count * 100.0 / total_count)), end=' ')
        print("%s:%d(%s)" % loc)

def _dump_hotshot_lineinfo2(log):
    cur = None
    a = sorted(_get_hotshot_lineinfo(log))
    total_count = sum((time[0]
               for (loc, time) in a))
    total_time = sum((time[1]
              for (loc, time) in a))
    for ((filename, lineno, fn), [count, time]) in a:
        if cur != filename:
            if cur != None and f != None:
                for line in f:
                    print(line[:-1])
                f.close()
            try:
                f = open(filename, "r")
            except OSError:
                f = None
            cur = filename
            l = 0
            print("#", filename)
        if f != None:
            while l < lineno:
                print(f.readline()[:-1])
                l += 1
        print ("# %8d %6.3f%%  %8d %6.3f%%"
               % (time, time * 100.0 / total_time,
              count, count * 100.0 / total_count))
    if cur != None and f != None:
        for line in f:
            print(line[:-1])
        f.close()
    
def main(args):
    import os
    
    if args[1] == "pc":
        import profile
        pr = profile.Profile()
        for i in range(5):
            print(pr.calibrate(100000))
        return
    elif args[1] == "p":
        import profile
        ret = 0
        # profile.Profile.bias = 5.26e-6
        profile.runctx("ret = main2(args[1:])",
                   globals(), locals())
        return ret
    elif args[1].startswith("h"):
        import hotshot, hotshot.stats
        import warnings

        warnings.filterwarnings("ignore")
        tmp = os.tempnam()
        try:
            l = args[1].startswith("hl")
            p = hotshot.Profile(tmp, l)
            ret = p.runcall(main2, args[1:])
            p.close()
            p = None
            if l:
                if args[1] == "hl2":
                    _dump_hotshot_lineinfo2(tmp)
                else:
                    _dump_hotshot_lineinfo(tmp)
            else:
                hotshot.stats.load(tmp).print_stats()
        finally:
            try:
                os.remove(tmp)
            except OSError:
                pass
        return ret
            
    return main2(args)

if __name__ == '__main__':
    sys.exit(main(sys.argv))
    
