import random
import unittest
import numpy as np

from codec import (
    ArithmeticEncoder, ArithmeticDecoder, base91_encode, base91_decode,
    logits_to_cdf, pack, unpack, verify_text, profile_id, MAX_TOKENS, MAX_BYTES, MAX_CODE,
)


class CodecTests(unittest.TestCase):
    def test_base91_binary_roundtrip(self):
        rng = random.Random(721)
        for size in [0, 1, 2, 3, 16, 100, 511, 10000]:
            raw = rng.randbytes(size)
            code = base91_encode(raw)
            self.assertTrue(all(33 <= ord(c) <= 126 for c in code))
            self.assertEqual(base91_decode(code), raw)
        self.assertEqual(base91_encode(b'Hello World!'), '>OwJh>Io0Tv!8PE')

    def test_arithmetic_changing_distributions(self):
        rng = np.random.default_rng(52)
        for length, vocab in [(0, 3), (1, 7), (1000, 256), (15, 262144)]:
            distributions = [logits_to_cdf(rng.normal(0, 8, vocab)) for _ in range(length)]
            symbols = rng.integers(0, vocab, length).tolist()
            enc = ArithmeticEncoder()
            for symbol, cdf in zip(symbols, distributions):
                enc.write(symbol, cdf)
            dec = ArithmeticDecoder(enc.finish())
            self.assertEqual([dec.read(cdf) for cdf in distributions], symbols)

    def test_distribution_has_full_support(self):
        cdf = logits_to_cdf(np.array([-10000.0, 0.0, 10000.0]))
        self.assertTrue(np.all(np.diff(cdf) > 0))
        self.assertLessEqual(int(cdf[-1]), 1 << 24)
        with self.assertRaises(ValueError):
            logits_to_cdf([float('nan')])

    def test_packet_roundtrip_and_damage(self):
        profile = profile_id({'model': 'test'})
        for text in ['', 'hello', '  Hello 🌍\n你好\t café  ']:
            code = pack(b'\x40', text, 0, profile)
            packet = unpack(' \n' + code + '\t')
            self.assertEqual(packet.profile, profile)
            verify_text(text, packet)
            with self.assertRaises(ValueError):
                verify_text(text + '!', packet)
            replacement = 'A' if code[5] != 'A' else 'B'
            with self.assertRaises(ValueError):
                unpack(code[:5] + replacement + code[6:])
            with self.assertRaises(ValueError):
                unpack(code[:-2])

    def test_limits_and_profile(self):
        p = profile_id({'revision': 'a'})
        self.assertNotEqual(p, profile_id({'revision': 'b'}))
        with self.assertRaises(ValueError):
            pack(b'\x40', 'x', MAX_TOKENS + 1, p)
        with self.assertRaises(ValueError):
            unpack('A' * (MAX_CODE + 1))
        text = 'a' * MAX_BYTES
        packet = unpack(pack(b'\x40', text, MAX_TOKENS, p))
        self.assertEqual(packet.token_count, MAX_TOKENS)
        verify_text(text, packet)
        with self.assertRaises(ValueError):
            pack(b'\x40', text + 'a', 1, p)


if __name__ == '__main__':
    unittest.main()