"""
Tests for bitarray.util module
"""
from __future__ import absolute_import

import os
import sys
import shutil
import tempfile
import unittest
from string import hexdigits
from random import choice, randint, random
from collections import Counter

from bitarray import (bitarray, frozenbitarray, bits2bytes, decodetree,
                      get_default_endian, _set_default_endian)
from bitarray.test_bitarray import Util

from bitarray.util import (
    zeros, urandom, pprint, make_endian, rindex, strip, count_n,
    parity, count_and, count_or, count_xor, subset,
    serialize, deserialize, ba2hex, hex2ba,
    ba2int, int2ba, huffman_code
)

if sys.version_info[0] == 3:
    from io import StringIO
else:
    from io import BytesIO as StringIO

tests = []

# ---------------------------------------------------------------------------

class TestsZeros(unittest.TestCase):

    def test_1(self):
        for default_endian in 'big', 'little':
            _set_default_endian(default_endian)
            a = zeros(0)
            self.assertEqual(a, bitarray())
            self.assertEqual(a.endian(), default_endian)

            b  = zeros(0, endian=None)
            self.assertEqual(b.endian(), default_endian)

            for n in range(100):
                a = zeros(n)
                self.assertEqual(a, bitarray(n * '0'))

            for endian in 'big', 'little':
                a = zeros(3, endian)
                self.assertEqual(a, bitarray('000'))
                self.assertEqual(a.endian(), endian)

    def test_wrong_args(self):
        self.assertRaises(TypeError, zeros) # no argument
        self.assertRaises(TypeError, zeros, '')
        self.assertRaises(TypeError, zeros, bitarray())
        self.assertRaises(TypeError, zeros, [])
        self.assertRaises(TypeError, zeros, 1.0)
        self.assertRaises(ValueError, zeros, -1)

        # endian not string
        for x in 0, 1, {}, [], False, True:
            self.assertRaises(TypeError, zeros, 0, x)
        # endian wrong string
        self.assertRaises(ValueError, zeros, 0, 'foo')

tests.append(TestsZeros)

# ---------------------------------------------------------------------------

class TestsRandom(unittest.TestCase):

    def test_1(self):
        for default_endian in 'big', 'little':
            _set_default_endian(default_endian)
            a = urandom(0)
            self.assertEqual(a, bitarray())
            self.assertEqual(a.endian(), default_endian)

            b  = urandom(0, endian=None)
            self.assertEqual(b.endian(), default_endian)

            for n in range(100):
                a = urandom(n)
                self.assertEqual(len(a), n)
                self.assertEqual(b.endian(), default_endian)

            a = urandom(1000)
            b = urandom(1000)
            self.assertNotEqual(a, b)
            self.assertTrue(400 < a.count() < 600)
            self.assertTrue(400 < b.count() < 600)

    def test_wrong_args(self):
        self.assertRaises(TypeError, urandom)
        self.assertRaises(TypeError, urandom, '')
        self.assertRaises(TypeError, urandom, bitarray())
        self.assertRaises(TypeError, urandom, [])
        self.assertRaises(TypeError, urandom, 1.0)
        self.assertRaises(ValueError, urandom, -1)

        self.assertRaises(TypeError, urandom, 0, 1)
        self.assertRaises(ValueError, urandom, 0, 'foo')

tests.append(TestsRandom)

# ---------------------------------------------------------------------------

class TestsPPrint(unittest.TestCase):

    @staticmethod
    def get_code_string(a):
        f = StringIO()
        pprint(a, stream=f)
        return f.getvalue()

    def round_trip(self, a):
        b = eval(self.get_code_string(a))
        self.assertEqual(b, a)
        self.assertEqual(type(b), type(a))

    def test_bitarray(self):
        a = bitarray('110')
        self.assertEqual(self.get_code_string(a), "bitarray('110')\n")
        self.round_trip(a)

    def test_frozenbitarray(self):
        a = frozenbitarray('01')
        self.assertEqual(self.get_code_string(a), "frozenbitarray('01')\n")
        self.round_trip(a)

    def test_formatting(self):
        a = bitarray(200)
        for width in range(40, 130, 10):
            for n in range(1, 10):
                f = StringIO()
                pprint(a, stream=f, group=n, width=width)
                r = f.getvalue()
                self.assertEqual(eval(r), a)
                s = r.strip("bitary(')\n")
                for group in s.split()[:-1]:
                    self.assertEqual(len(group), n)
                for line in s.split('\n'):
                    self.assertTrue(len(line) < width)

    def test_fallback(self):
        for a in None, 'asd', [1, 2], bitarray(), frozenbitarray('1'):
            self.round_trip(a)

    def test_subclass(self):
        class Foo(bitarray):
            pass

        a = Foo()
        code = self.get_code_string(a)
        self.assertEqual(code, "Foo()\n")
        b = eval(code)
        self.assertEqual(b, a)
        self.assertEqual(type(b), type(a))

    def test_random(self):
        for n in range(150):
            self.round_trip(urandom(n))

    def test_file(self):
        tmpdir = tempfile.mkdtemp()
        tmpfile = os.path.join(tmpdir, 'testfile')
        a = bitarray(1000)
        try:
            with open(tmpfile, 'w') as fo:
                pprint(a, fo)
            with open(tmpfile, 'r') as fi:
                b = eval(fi.read())
            self.assertEqual(a, b)
        finally:
            shutil.rmtree(tmpdir)

tests.append(TestsPPrint)

# ---------------------------------------------------------------------------

class TestsMakeEndian(unittest.TestCase, Util):

    def test_simple(self):
        a = bitarray('1110001', endian='big')
        b = make_endian(a, 'big')
        self.assertTrue(b is a)
        c = make_endian(a, 'little')
        self.assertTrue(c == a)
        self.assertEqual(c.endian(), 'little')

        # wrong arguments
        self.assertRaises(TypeError, make_endian, '', 'big')
        self.assertRaises(TypeError, make_endian, bitarray(), 1)
        self.assertRaises(ValueError, make_endian, bitarray(), 'foo')

    def test_empty(self):
        a = bitarray(endian='little')
        b = make_endian(a, 'big')
        self.assertTrue(b == a)
        self.assertEqual(len(b), 0)
        self.assertEqual(b.endian(), 'big')

    def test_from_frozen(self):
        a = frozenbitarray('1101111', 'big')
        b = make_endian(a, 'big')
        self.assertTrue(b is a)
        c = make_endian(a, 'little')
        self.assertTrue(c == a)
        self.assertEqual(c.endian(), 'little')

    def test_random(self):
        for a in self.randombitarrays():
            aa = a.copy()
            for endian in 'big', 'little':
                b = make_endian(a, endian)
                self.assertEqual(a, b)
                self.assertEqual(b.endian(), endian)
                if a.endian() == endian:
                    self.assertTrue(b is a)
            self.assertEQUAL(a, aa)

tests.append(TestsMakeEndian)

# ---------------------------------------------------------------------------

class TestsRIndex(unittest.TestCase, Util):

    def test_simple(self):
        self.assertRaises(TypeError, rindex)
        self.assertRaises(TypeError, rindex, None)
        self.assertRaises(TypeError, rindex, bitarray(), 1, 2)
        for endian in 'big', 'little':
            a = bitarray('00010110000', endian)
            self.assertEqual(rindex(a), 6)
            self.assertEqual(rindex(a, 1), 6)
            self.assertEqual(rindex(a, 'A'), 6)
            self.assertEqual(rindex(a, True), 6)

            a = bitarray('00010110111', endian)
            self.assertEqual(rindex(a, 0), 7)
            self.assertEqual(rindex(a, None), 7)
            self.assertEqual(rindex(a, False), 7)

            a = frozenbitarray('00010110111', endian)
            self.assertEqual(rindex(a, 0), 7)
            self.assertEqual(rindex(a, None), 7)
            self.assertEqual(rindex(a, False), 7)

            for v in 0, 1:
                self.assertRaises(ValueError, rindex,
                                  bitarray(0, endian), v)
            self.assertRaises(ValueError, rindex,
                              bitarray('000', endian), 1)
            self.assertRaises(ValueError, rindex,
                              bitarray('11111', endian), 0)

    def test_random(self):
        for a in self.randombitarrays():
            v = randint(0, 1)
            try:
                i = rindex(a, v)
            except ValueError:
                i = None
            s = a.to01()
            try:
                j = s.rindex(str(v))
            except ValueError:
                j = None
            self.assertEqual(i, j)

    def test_3(self):
        for _ in range(100):
            n = randint(1, 100000)
            v = randint(0, 1)
            a = bitarray(n)
            a.setall(1 - v)
            lst = [randint(0, n - 1) for _ in range(100)]
            for i in lst:
                a[i] = v
            self.assertEqual(rindex(a, v), max(lst))

    def test_one_set(self):
        for _ in range(100):
            N = randint(1, 10000)
            a = bitarray(N)
            a.setall(0)
            a[randint(0, N - 1)] = 1
            self.assertEqual(rindex(a), a.index(1))

tests.append(TestsRIndex)

# ---------------------------------------------------------------------------

class TestsStrip(unittest.TestCase, Util):

    def test_simple(self):
        self.assertRaises(TypeError, strip, '0110')
        self.assertRaises(TypeError, strip, bitarray(), 123)
        self.assertRaises(ValueError, strip, bitarray(), 'up')
        for default_endian in 'big', 'little':
            _set_default_endian(default_endian)
            a = bitarray('00010110000')
            self.assertEQUAL(strip(a), bitarray('0001011'))
            self.assertEQUAL(strip(a, 'left'), bitarray('10110000'))
            self.assertEQUAL(strip(a, 'both'), bitarray('1011'))
            b = frozenbitarray('00010110000')
            self.assertEqual(strip(b, 'both'), bitarray('1011'))

        for mode in 'left', 'right', 'both':
            self.assertEqual(strip(bitarray('000'), mode), bitarray())
            self.assertEqual(strip(bitarray(), mode), bitarray())

    def test_random(self):
        for a in self.randombitarrays():
            b = a.copy()
            s = a.to01()
            self.assertEqual(strip(a, 'left'), bitarray(s.lstrip('0')))
            self.assertEqual(strip(a, 'right'), bitarray(s.rstrip('0')))
            self.assertEqual(strip(a, 'both'), bitarray(s.strip('0')))
            self.assertEQUAL(a, b)

    def test_one_set(self):
        for _ in range(100):
            N = randint(1, 10000)
            a = bitarray(N)
            a.setall(0)
            a[randint(0, N - 1)] = 1
            self.assertEqual(strip(a, 'both'), bitarray('1'))

tests.append(TestsStrip)

# ---------------------------------------------------------------------------

class TestsCount_N(unittest.TestCase, Util):

    @staticmethod
    def count_n(a, n):
        "return the index i for which a[:i].count() == n"
        i, j = n, a.count(1, 0, n)
        while j < n:
            if a[i]:
                j += 1
            i += 1
        return i

    def check_result(self, a, n, i):
        self.assertEqual(a.count(1, 0, i), n)
        if i > 0:
            self.assertTrue(a[i - 1])

    def test_simple(self):
        a = bitarray('111110111110111110111110011110111110111110111000')
        b = a.copy()
        self.assertEqual(len(a), 48)
        self.assertEqual(a.count(), 37)
        self.assertRaises(TypeError, count_n, '', 0)
        self.assertEqual(count_n(a, 0), 0)
        self.assertEqual(count_n(a, 20), 23)
        self.assertEqual(count_n(a, 37), 45)
        self.assertRaisesMessage(ValueError, "non-negative integer expected",
                                 count_n, a, -1) # n < 0
        self.assertRaisesMessage(ValueError, "n larger than bitarray size",
                                 count_n, a, 49) # n > len(a)
        self.assertRaisesMessage(ValueError, "n exceeds total count",
                                 count_n, a, 38) # n > a.count()
        self.assertRaises(TypeError, count_n, a, "7")
        for n in range(0, 37):
            i = count_n(a, n)
            self.check_result(a, n, i)
            self.assertEqual(a[:i].count(), n)
            self.assertEqual(i, self.count_n(a, n))
        self.assertEQUAL(a, b)

    def test_frozen(self):
        a = frozenbitarray('001111101111101111101111100111100')
        self.assertEqual(len(a), 33)
        self.assertEqual(a.count(), 24)
        self.assertRaises(TypeError, count_n, '', 0)
        self.assertEqual(count_n(a, 0), 0)
        self.assertEqual(count_n(a, 10), 13)
        self.assertEqual(count_n(a, 24), 31)
        self.assertRaises(ValueError, count_n, a, -1) # n < 0
        self.assertRaises(ValueError, count_n, a, 25) # n > a.count()
        self.assertRaises(ValueError, count_n, a, 34) # n > len(a)
        self.assertRaises(TypeError, count_n, a, "7")

    def test_large(self):
        for N in list(range(100)) + [1000, 10000, 100000]:
            a = bitarray(N)
            v = randint(0, 1)
            a.setall(not v)
            for _ in range(randint(0, min(N, 100))):
                a[randint(0, N - 1)] = v
            n = randint(0, a.count())
            self.check_result(a, n, count_n(a, n))
            # check for total count
            tc = a.count()
            self.assertTrue(count_n(a, tc) <= N)
            self.assertRaises(ValueError, count_n, a, tc + 1)

    def test_one_set(self):
        N = 100000
        for _ in range(10):
            a = bitarray(N)
            a.setall(0)
            self.assertEqual(count_n(a, 0), 0)
            self.assertRaises(ValueError, count_n, a, 1)
            i = randint(0, N - 1)
            a[i] = 1
            self.assertEqual(count_n(a, 1), i + 1)
            self.assertRaises(ValueError, count_n, a, 2)

    def test_random(self):
        for a in self.randombitarrays():
            n = a.count() // 2
            i = count_n(a, n)
            self.check_result(a, n, i)

tests.append(TestsCount_N)

# ---------------------------------------------------------------------------

class TestsBitwiseCount(unittest.TestCase, Util):

    def test_count_byte(self):
        ones = bitarray(8)
        ones.setall(1)
        zeros = bitarray(8)
        zeros.setall(0)
        for i in range(0, 256):
            a = bitarray()
            a.frombytes(bytes(bytearray([i])))
            cnt = a.count()
            self.assertEqual(count_and(a, zeros), 0)
            self.assertEqual(count_and(a, ones), cnt)
            self.assertEqual(count_and(a, a), cnt)
            self.assertEqual(count_or(a, zeros), cnt)
            self.assertEqual(count_or(a, ones), 8)
            self.assertEqual(count_or(a, a), cnt)
            self.assertEqual(count_xor(a, zeros), cnt)
            self.assertEqual(count_xor(a, ones), 8 - cnt)
            self.assertEqual(count_xor(a, a), 0)

    def test_bit_count1(self):
        a = bitarray('001111')
        aa = a.copy()
        b = bitarray('010011')
        bb = b.copy()
        self.assertEqual(count_and(a, b), 2)
        self.assertEqual(count_or(a, b), 5)
        self.assertEqual(count_xor(a, b), 3)
        for f in count_and, count_or, count_xor:
            # not two arguments
            self.assertRaises(TypeError, f)
            self.assertRaises(TypeError, f, a)
            self.assertRaises(TypeError, f, a, b, 3)
            # wrong argument types
            self.assertRaises(TypeError, f, a, '')
            self.assertRaises(TypeError, f, '1', b)
            self.assertRaises(TypeError, f, a, 4)
        self.assertEQUAL(a, aa)
        self.assertEQUAL(b, bb)

        b.append(1)
        for f in count_and, count_or, count_xor:
            self.assertRaises(ValueError, f, a, b)
            self.assertRaises(ValueError, f,
                              bitarray('110', 'big'),
                              bitarray('101', 'little'))

    def test_bit_count_frozen(self):
        a = frozenbitarray('001111')
        b = frozenbitarray('010011')
        self.assertEqual(count_and(a, b), 2)
        self.assertEqual(count_or(a, b), 5)
        self.assertEqual(count_xor(a, b), 3)

    def test_bit_count_random(self):
        for n in list(range(50)) + [randint(1000, 2000)]:
            a = urandom(n)
            b = urandom(n)
            self.assertEqual(count_and(a, b), (a & b).count())
            self.assertEqual(count_or(a, b),  (a | b).count())
            self.assertEqual(count_xor(a, b), (a ^ b).count())

tests.append(TestsBitwiseCount)

# ---------------------------------------------------------------------------

class TestsSubset(unittest.TestCase, Util):

    def test_basic(self):
        a = frozenbitarray('0101')
        b = bitarray('0111')
        self.assertTrue(subset(a, b))
        self.assertFalse(subset(b, a))
        self.assertRaises(TypeError, subset)
        self.assertRaises(TypeError, subset, a, '')
        self.assertRaises(TypeError, subset, '1', b)
        self.assertRaises(TypeError, subset, a, 4)
        b.append(1)
        self.assertRaises(ValueError, subset, a, b)

    def subset_simple(self, a, b):
        return (a & b).count() == a.count()

    def test_True(self):
        for a, b in [('', ''), ('0', '1'), ('0', '0'), ('1', '1'),
                     ('000', '111'), ('0101', '0111'),
                     ('000010111', '010011111')]:
            a, b = bitarray(a), bitarray(b)
            self.assertTrue(subset(a, b) is True)
            self.assertTrue(self.subset_simple(a, b) is True)

    def test_False(self):
        for a, b in [('1', '0'), ('1101', '0111'),
                     ('0000101111', '0100111011')]:
            a, b = bitarray(a), bitarray(b)
            self.assertTrue(subset(a, b) is False)
            self.assertTrue(self.subset_simple(a, b) is False)

    def test_random(self):
        for a in self.randombitarrays(start=1):
            b = a.copy()
            # we set one random bit in b to 1, so a is always a subset of b
            b[randint(0, len(a) - 1)] = 1
            self.assertTrue(subset(a, b))
            # but b in not always a subset of a
            self.assertEqual(subset(b, a), self.subset_simple(b, a))
            # we set all bits in a, which ensures that b is a subset of a
            a.setall(1)
            self.assertTrue(subset(b, a))

tests.append(TestsSubset)

# ---------------------------------------------------------------------------

class TestsParity(unittest.TestCase, Util):

    def test_bitarray(self):
        a = bitarray()
        self.assertTrue(parity(a) is False)
        par = False
        for _ in range(100):
            self.assertTrue(parity(a) is par)
            a.append(1)
            par = not par

    def test_unused_bits(self):
        a = bitarray(1)
        a.setall(1)
        self.assertTrue(parity(a))

    def test_return_type(self):
        p = parity(bitarray('101'))
        self.assertIsInstance(p, bool)

    def test_frozenbitarray(self):
        self.assertTrue(parity(frozenbitarray()) is False)
        self.assertTrue(parity(frozenbitarray('0010011')) is True)
        self.assertTrue(parity(frozenbitarray('10100110')) is False)

    def test_wrong_args(self):
        self.assertRaises(TypeError, parity, '')
        self.assertRaises(TypeError, bitarray(), 1)

    def test_byte(self):
        for i in range(256):
            a = bitarray()
            a.frombytes(bytes(bytearray([i])))
            self.assertEqual(parity(a), a.count() % 2)

    def test_random(self):
        for a in self.randombitarrays():
            self.assertEqual(parity(a), a.count() % 2)

tests.append(TestsParity)

# ---------------------------------------------------------------------------

CODEDICT = {'little': {}, 'big': {
    '0': bitarray('0000'),    '1': bitarray('0001'),
    '2': bitarray('0010'),    '3': bitarray('0011'),
    '4': bitarray('0100'),    '5': bitarray('0101'),
    '6': bitarray('0110'),    '7': bitarray('0111'),
    '8': bitarray('1000'),    '9': bitarray('1001'),
    'a': bitarray('1010'),    'b': bitarray('1011'),
    'c': bitarray('1100'),    'd': bitarray('1101'),
    'e': bitarray('1110'),    'f': bitarray('1111'),
}}
for k, v in CODEDICT['big'].items():
    CODEDICT['little'][k] = v[::-1]


class TestsHexlify(unittest.TestCase, Util):

    def test_ba2hex(self):
        self.assertEqual(ba2hex(bitarray(0, 'big')), '')
        self.assertEqual(ba2hex(bitarray('1110', 'big')), 'e')
        self.assertEqual(ba2hex(bitarray('1110', 'little')), '7')
        self.assertEqual(ba2hex(bitarray('0000 0001', 'big')), '01')
        self.assertEqual(ba2hex(bitarray('1000 0000', 'big')), '80')
        self.assertEqual(ba2hex(bitarray('0000 0001', 'little')), '08')
        self.assertEqual(ba2hex(bitarray('1000 0000', 'little')), '10')
        self.assertEqual(ba2hex(frozenbitarray('1100 0111', 'big')), 'c7')
        # length not multiple of 4
        self.assertRaises(ValueError, ba2hex, bitarray('10'))
        self.assertRaises(TypeError, ba2hex, '101')

        c = ba2hex(bitarray('1101', 'big'))
        self.assertIsInstance(c, str)

        for n in range(7):
            a = bitarray(n * '1111', 'big')
            b = a.copy()
            self.assertEqual(ba2hex(a), n * 'f')
            # ensure original object wasn't altered
            self.assertEQUAL(a, b)

    def test_hex2ba(self):
        _set_default_endian('big')
        self.assertEqual(hex2ba(''), bitarray())
        for c in 'e', 'E', b'e', b'E', u'e', u'E':
            a = hex2ba(c)
            self.assertEqual(a.to01(), '1110')
            self.assertEqual(a.endian(), 'big')
        self.assertEQUAL(hex2ba('01'), bitarray('0000 0001', 'big'))
        self.assertEQUAL(hex2ba('08', 'little'),
                         bitarray('0000 0001', 'little'))
        self.assertEQUAL(hex2ba('aD'), bitarray('1010 1101', 'big'))
        self.assertEQUAL(hex2ba(b'10aF'),
                         bitarray('0001 0000 1010 1111', 'big'))
        self.assertEQUAL(hex2ba(b'10aF', 'little'),
                         bitarray('1000 0000 0101 1111', 'little'))

    def test_hex2ba_errors(self):
        self.assertRaises(TypeError, hex2ba, 0)

        for endian in 'little', 'big':
            _set_default_endian(endian)
            self.assertRaises(ValueError, hex2ba, '01a7g89')
            self.assertRaises(UnicodeEncodeError, hex2ba, u'10\u20ac')
            # check for NUL bytes
            for b in b'\0', b'\0f', b'f\0', b'\0ff', b'f\0f', b'ff\0':
                self.assertRaises(ValueError, hex2ba, b)

    @staticmethod
    def hex2ba(s, endian=None):
        a = bitarray(0, endian or get_default_endian())
        a.encode(CODEDICT[a.endian()], s)
        return a

    @staticmethod
    def ba2hex(a):
        return ''.join(a.iterdecode(CODEDICT[a.endian()]))

    def test_explicit(self):
        data = [ #                       little   big
            ('',                         '',      ''),
            ('1000',                     '1',     '8'),
            ('1000 1100',                '13',    '8c'),
            ('1000 1100 1110',           '137',   '8ce'),
            ('1000 1100 1110 1111' ,     '137f',  '8cef'),
            ('1000 1100 1110 1111 0100', '137f2', '8cef4'),
        ]
        for bs, hex_le, hex_be in data:
            a_be = bitarray(bs, 'big')
            a_le = bitarray(bs, 'little')
            self.assertEQUAL(hex2ba(hex_be, 'big'), a_be)
            self.assertEQUAL(hex2ba(hex_le, 'little'), a_le)
            self.assertEqual(ba2hex(a_be), hex_be)
            self.assertEqual(ba2hex(a_le), hex_le)
            # test simple encode / decode implementation
            self.assertEQUAL(self.hex2ba(hex_be, 'big'), a_be)
            self.assertEQUAL(self.hex2ba(hex_le, 'little'), a_le)
            self.assertEqual(self.ba2hex(a_be), hex_be)
            self.assertEqual(self.ba2hex(a_le), hex_le)

    def test_round_trip(self):
        s = ''.join(choice(hexdigits) for _ in range(randint(10, 100)))
        for default_endian in 'big', 'little':
            _set_default_endian(default_endian)
            a = hex2ba(s)
            self.check_obj(a)
            self.assertEqual(len(a) % 4, 0)
            self.assertEqual(a.endian(), default_endian)
            t = ba2hex(a)
            self.assertEqual(t, s.lower())
            b = hex2ba(t, default_endian)
            self.assertEQUAL(a, b)
            # test simple encode / decode implementation
            self.assertEQUAL(a, self.hex2ba(t))
            self.assertEqual(t, self.ba2hex(a))

tests.append(TestsHexlify)

# ---------------------------------------------------------------------------

class TestsIntegerization(unittest.TestCase, Util):

    def test_ba2int(self):
        self.assertEqual(ba2int(bitarray('0')), 0)
        self.assertEqual(ba2int(bitarray('1')), 1)
        self.assertEqual(ba2int(bitarray('00101', 'big')), 5)
        self.assertEqual(ba2int(bitarray('00101', 'little')), 20)
        self.assertEqual(ba2int(frozenbitarray('11')), 3)
        self.assertRaises(ValueError, ba2int, bitarray())
        self.assertRaises(ValueError, ba2int, frozenbitarray())
        self.assertRaises(TypeError, ba2int, '101')
        a = bitarray('111')
        b = a.copy()
        self.assertEqual(ba2int(a), 7)
        # ensure original object wasn't altered
        self.assertEQUAL(a, b)

    def test_int2ba(self):
        self.assertEqual(int2ba(0), bitarray('0'))
        self.assertEqual(int2ba(1), bitarray('1'))
        self.assertEqual(int2ba(5), bitarray('101'))
        self.assertEQUAL(int2ba(6, endian='big'), bitarray('110', 'big'))
        self.assertEQUAL(int2ba(6, endian='little'),
                         bitarray('011', 'little'))
        self.assertRaises(TypeError, int2ba, 1.0)
        self.assertRaises(TypeError, int2ba, 1, 3.0)
        self.assertRaises(ValueError, int2ba, 1, 0)
        self.assertRaises(TypeError, int2ba, 1, 10, 123)
        self.assertRaises(ValueError, int2ba, 1, 10, 'asd')
        # signed integer requires length
        self.assertRaises(TypeError, int2ba, 100, signed=True)

    def test_signed(self):
        for s, i in [
                ('0',  0),
                ('1', -1),
                ('00',  0),
                ('10',  1),
                ('01', -2),
                ('11', -1),
                ('000',  0),
                ('100',  1),
                ('010',  2),
                ('110',  3),
                ('001', -4),
                ('101', -3),
                ('011', -2),
                ('111', -1),
                ('00000',   0),
                ('11110',  15),
                ('00001', -16),
                ('11111',  -1),
                ('000000000',    0),
                ('111111110',  255),
                ('000000001', -256),
                ('111111111',   -1),
                ('0000000000000000000000', 0),
                ('1001000011000000100010', 9 + 3 * 256 + 17 * 2 ** 16),
                ('1111111111111111111110', 2 ** 21 - 1),
                ('0000000000000000000001', -2 ** 21),
                ('1001000011000000100011', -2 ** 21
                                           + (9 + 3 * 256 + 17 * 2 ** 16)),
                ('1111111111111111111111', -1),
        ]:
            self.assertEqual(ba2int(bitarray(s, 'little'), signed=1), i)
            self.assertEqual(ba2int(bitarray(s[::-1], 'big'), signed=1), i)

            self.assertEQUAL(int2ba(i, len(s), 'little', signed=1),
                             bitarray(s, 'little'))
            self.assertEQUAL(int2ba(i, len(s), 'big', signed=1),
                             bitarray(s[::-1], 'big'))

    def test_int2ba_overflow(self):
        self.assertRaises(OverflowError, int2ba, -1)
        self.assertRaises(OverflowError, int2ba, -1, 4)

        self.assertRaises(OverflowError, int2ba, 128, 7)
        self.assertRaises(OverflowError, int2ba, 64, 7, signed=1)
        self.assertRaises(OverflowError, int2ba, -65, 7, signed=1)

        for n in range(1, 20):
            self.assertRaises(OverflowError, int2ba, 2 ** n, n)
            self.assertRaises(OverflowError, int2ba, 2 ** (n - 1), n,
                              signed=1)
            self.assertRaises(OverflowError, int2ba, -2 ** (n - 1) - 1, n,
                              signed=1)

    def test_int2ba_length(self):
        self.assertRaises(TypeError, int2ba, 0, 1.0)
        self.assertRaises(ValueError, int2ba, 0, 0)
        self.assertEqual(int2ba(5, length=6, endian='big'),
                         bitarray('000101'))
        for n in range(1, 100):
            ab = int2ba(1, n, 'big')
            al = int2ba(1, n, 'little')
            self.assertEqual(ab.endian(), 'big')
            self.assertEqual(al.endian(), 'little')
            self.assertEqual(len(ab), n),
            self.assertEqual(len(al), n)
            self.assertEqual(ab, bitarray((n - 1) * '0') + bitarray('1'))
            self.assertEqual(al, bitarray('1') + bitarray((n - 1) * '0'))

            ab = int2ba(0, n, 'big')
            al = int2ba(0, n, 'little')
            self.assertEqual(len(ab), n)
            self.assertEqual(len(al), n)
            self.assertEqual(ab, bitarray(n * '0', 'big'))
            self.assertEqual(al, bitarray(n * '0', 'little'))

            self.assertEqual(int2ba(2 ** n - 1), bitarray(n * '1'))
            self.assertEqual(int2ba(2 ** n - 1, endian='little'),
                             bitarray(n * '1'))
            for endian in 'big', 'little':
                self.assertEqual(int2ba(-1, n, endian, signed=True),
                                 bitarray(n * '1'))

    def test_explicit(self):
        _set_default_endian('big')
        for i, sa in [( 0,     '0'),    (1,         '1'),
                      ( 2,    '10'),    (3,        '11'),
                      (25, '11001'),  (265, '100001001'),
                      (3691038, '1110000101001000011110')]:
            ab = bitarray(sa, 'big')
            al = bitarray(sa[::-1], 'little')
            self.assertEQUAL(int2ba(i), ab)
            self.assertEQUAL(int2ba(i, endian='big'), ab)
            self.assertEQUAL(int2ba(i, endian='little'), al)
            self.assertEqual(ba2int(ab), ba2int(al), i)

    def check_round_trip(self, i):
        for endian in 'big', 'little':
            a = int2ba(i, endian=endian)
            self.check_obj(a)
            self.assertEqual(a.endian(), endian)
            self.assertTrue(len(a) > 0)
            # ensure we have no leading zeros
            if a.endian == 'big':
                self.assertTrue(len(a) == 1 or a.index(1) == 0)
            self.assertEqual(ba2int(a), i)
            if i > 0:
                self.assertEqual(i.bit_length(), len(a))
            # add a few trailing / leading zeros to bitarray
            if endian == 'big':
                a = zeros(randint(0, 3), endian) + a
            else:
                a = a + zeros(randint(0, 3), endian)
            self.assertEqual(a.endian(), endian)
            self.assertEqual(ba2int(a), i)

    def test_many(self):
        for i in range(20):
            self.check_round_trip(i)
            self.check_round_trip(randint(0, 10 ** randint(3, 300)))

    @staticmethod
    def twos_complement(i, num_bits):
        # https://en.wikipedia.org/wiki/Two%27s_complement
        mask = 2 ** (num_bits - 1)
        return -(i & mask) + (i & ~mask)

    def test_random_signed(self):
        for a in self.randombitarrays(start=1):
            i = ba2int(a, signed=True)
            b = int2ba(i, len(a), a.endian(), signed=True)
            self.assertEQUAL(a, b)

            j = ba2int(a, signed=False)  # unsigned
            if i >= 0:
                self.assertEqual(i, j)

            self.assertEqual(i, self.twos_complement(j, len(a)))

tests.append(TestsIntegerization)

# ---------------------------------------------------------------------------

class MixedTests(unittest.TestCase, Util):

    def test_bin(self):
        for i in range(100):
            s = bin(i)
            self.assertEqual(s[:2], '0b')
            a = bitarray(s[2:], 'big')
            self.assertEqual(ba2int(a), i)
            t = '0b%s' % a.to01()
            self.assertEqual(t, s)
            self.assertEqual(eval(t), i)

    def test_hex(self):
        for i in range(1000):
            s = hex(i)
            self.assertEqual(s[:2], '0x')
            a = hex2ba(s[2:], 'big')
            self.assertEqual(ba2int(a), i)
            t = '0x%s' % ba2hex(a)
            self.assertEqual(t, s)
            self.assertEqual(eval(t), i)

tests.append(MixedTests)

# ---------------------------------------------------------------------------

class TestsSerialization(unittest.TestCase, Util):

    def test_explicit(self):
        for a, b in [
                (bitarray(0, 'little'),   b'\x00'),
                (bitarray(0, 'big'),      b'\x10'),
                (bitarray('1', 'little'), b'\x07\x01'),
                (bitarray('1', 'big'),    b'\x17\x80'),
                (bitarray('11110000', 'little'), b'\x00\x0f'),
                (bitarray('11110000', 'big'),    b'\x10\xf0'),
        ]:
            self.assertEqual(serialize(a), b)
            self.assertEQUAL(deserialize(b), a)

    def test_zeros_and_ones(self):
        for endian in 'little', 'big':
            for n in range(100):
                a = zeros(n, endian)
                s = serialize(a)
                self.assertEqual(s[1:], b'\0' * bits2bytes(n))
                self.assertEQUAL(a, deserialize(s))
                a.setall(1)
                self.assertEQUAL(a, deserialize(serialize(a)))

    def test_wrong_args(self):
        self.assertRaises(TypeError, serialize, '0')
        self.assertRaises(TypeError, serialize, bitarray(), 1)
        self.assertRaises(TypeError, deserialize, u'01')
        self.assertRaises(TypeError, deserialize, b'\x00', 1)

    def test_invalid_bytes(self):
        self.assertRaises(ValueError, deserialize, b'')

        for i in range(256):
            b = bytes(bytearray([i]))
            if i == 0 or i == 16:
                self.assertEqual(deserialize(b), bitarray())
            else:
                self.assertRaises(ValueError, deserialize, b)

            b += b'\0'
            if i < 32 and i % 16 < 8:
                self.assertEqual(deserialize(b), zeros(8 - i % 8))
            else:
                self.assertRaises(ValueError, deserialize, b)

    def test_bits_ignored(self):
        # the unused bits (with the last bytes) are ignored
        for b, a in [
                (b'\x07\x01', bitarray('1', 'little')),
                (b'\x07\x03', bitarray('1', 'little')),
                (b'\x07\xff', bitarray('1', 'little')),
                (b'\x17\x80', bitarray('1', 'big')),
                (b'\x17\xc0', bitarray('1', 'big')),
                (b'\x17\xff', bitarray('1', 'big')),
        ]:
            self.assertEQUAL(deserialize(b), a)

    def test_random(self):
        for a in self.randombitarrays():
            b = deserialize(serialize(a))
            self.assertEQUAL(a, b)
            self.check_obj(b)

tests.append(TestsSerialization)

# ---------------------------------------------------------------------------

class TestsHuffman(unittest.TestCase):

    def test_simple(self):
        freq = {0: 10, 'as': 2, None: 1.6}
        code = huffman_code(freq)
        self.assertEqual(len(code), 3)
        self.assertEqual(len(code[0]), 1)
        self.assertEqual(len(code['as']), 2)
        self.assertEqual(len(code[None]), 2)

    def test_tiny(self):
        code = huffman_code({0: 0})
        self.assertEqual(len(code), 1)
        self.assertEqual(code, {0: bitarray()})

        code = huffman_code({0: 0, 1: 0})
        self.assertEqual(len(code), 2)
        for i in range(2):
            self.assertEqual(len(code[i]), 1)

    def test_endianness(self):
        freq = {'A': 10, 'B': 2, 'C': 5}
        for endian in 'big', 'little':
            code = huffman_code(freq, endian)
            self.assertEqual(len(code), 3)
            for v in code.values():
                self.assertEqual(v.endian(), endian)

    def test_wrong_arg(self):
        self.assertRaises(TypeError, huffman_code, [('a', 1)])
        self.assertRaises(TypeError, huffman_code, 123)
        self.assertRaises(TypeError, huffman_code, None)
        # cannot compare 'a' with 1
        self.assertRaises(TypeError, huffman_code, {'A': 'a', 'B': 1})
        self.assertRaises(ValueError, huffman_code, {})

    def check_tree(self, code):
        n = len(code)
        tree = decodetree(code)
        self.assertEqual(tree.todict(), code)
        # ensure tree has 2n-1 nodes (n symbol nodes and n-1 internal nodes)
        self.assertEqual(tree.nodes(), 2 * n - 1)

    def test_balanced(self):
        n = 6
        freq = {}
        for i in range(2 ** n):
            freq[i] = 1
        code = huffman_code(freq)
        self.assertEqual(len(code), 2 ** n)
        self.assertTrue(all(len(v) == n for v in code.values()))
        self.check_tree(code)

    def test_unbalanced(self):
        N = 27
        freq = {}
        for i in range(N):
            freq[i] = 2 ** i
        code = huffman_code(freq)
        self.assertEqual(len(code), N)
        for i in range(N):
            self.assertEqual(len(code[i]), N - (1 if i <= 1 else i))
        self.check_tree(code)

    def test_counter(self):
        message = 'the quick brown fox jumps over the lazy dog.'
        code = huffman_code(Counter(message))
        a = bitarray()
        a.encode(code, message)
        self.assertEqual(''.join(a.decode(code)), message)
        self.check_tree(code)

    def test_random_list(self):
        plain = [randint(0, 100) for _ in range(500)]
        code = huffman_code(Counter(plain))
        a = bitarray()
        a.encode(code, plain)
        self.assertEqual(a.decode(code), plain)
        self.check_tree(code)

    def test_random_freq(self):
        N = randint(2, 1000)
        # create Huffman code for N symbols
        code = huffman_code({i: random() for i in range(N)})
        self.check_tree(code)

tests.append(TestsHuffman)

# ---------------------------------------------------------------------------

def run(verbosity=1):
    import os
    import bitarray

    print('bitarray is installed in: %s' % os.path.dirname(bitarray.__file__))
    print('bitarray version: %s' % bitarray.__version__)
    print('Python version: %s' % sys.version)

    suite = unittest.TestSuite()
    for cls in tests:
        suite.addTest(unittest.makeSuite(cls))

    runner = unittest.TextTestRunner(verbosity=verbosity)
    return runner.run(suite)


if __name__ == '__main__':
    run()
