"""Small offline packing tests; requires NumPy, PyTorch and safetensors."""
from pathlib import Path
import sys
import unittest
import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parents[2] / 'tools/community_models'))
import kitten_tts2_weights as weights


def dequantize(payload, bits):
    blocks = np.frombuffer(payload, np.uint8).reshape(-1, 18 if bits == 4 else 34)
    scales = blocks[:, :2].copy().view('<f2').astype(np.float32)
    if bits == 4:
        codes = np.concatenate((blocks[:, 2:] & 15, blocks[:, 2:] >> 4), 1).astype(np.float32) - 8
    else:
        codes = blocks[:, 2:].copy().view(np.int8).astype(np.float32)
    return codes * scales


class PackingTests(unittest.TestCase):
    def test_tq2_1_independent_groups(self):
        values = torch.tensor([-1, 0, 1, 0] * 128).reshape(2, 256).to(torch.bfloat16)
        values = values.reshape(-1, 128) * torch.tensor([0, 0.125, 0.5, 2**-24])[:, None]
        values = values.reshape(2, 256)
        payload = weights.exact_ternary_tq2_1(values)
        blocks = np.frombuffer(payload, np.uint8).reshape(-1, 68)
        scales = blocks[:, 64:].copy().view('<f2').astype(np.float32)
        codes = blocks[:, :64].reshape(-1, 2, 32)
        decoded = np.stack([(codes >> shift) & 3 for shift in (0, 2, 4, 6)], axis=2)
        restored = (decoded.astype(np.float32).reshape(-1, 2, 128) - 1) * scales[..., None]
        self.assertTrue(np.array_equal(restored.reshape(2, 256), values.float().numpy()))
        values[0, 129] = 0.25
        with self.assertRaises(ValueError): weights.exact_ternary_tq2_1(values)
        with self.assertRaises(ValueError): weights.exact_ternary_tq2_1(torch.zeros(1, 128))

    def test_exact_blocks_include_zero_and_small_scales(self):
        values = torch.tensor([-1, 0, 1, 0] * 8).repeat(4, 1).to(torch.bfloat16)
        values *= torch.tensor([0, 0.123046875, 2**-22, 2**-24], dtype=torch.bfloat16)[:, None]
        for bits in (4, 8):
            self.assertTrue(np.array_equal(dequantize(weights.exact_ternary_blocks(values, bits), bits), values.float().numpy()))

    def test_nonternary_and_unrepresentable_scales_rejected(self):
        values = torch.zeros(1, 32, dtype=torch.bfloat16)
        values[0, :2] = torch.tensor([1, 0.5])
        for bits in (4, 8):
            with self.assertRaises(ValueError): weights.exact_ternary_blocks(values, bits)
            with self.assertRaises(ValueError): weights.exact_ternary_blocks(torch.full_like(values, 2**-30), bits)

    def test_base3_digit_order_and_truncation(self):
        # Base-3 byte 102 contains digits 0,1,2,0,1, representing -1,0,1,-1,0.
        codes = torch.full((26,), 102, dtype=torch.uint8)
        expected = torch.tensor(([-1,0,1,-1,0] * 26)[:128], dtype=torch.bfloat16).reshape(1,128)
        actual = weights.unpack_ternary(codes, torch.ones(1,1,dtype=torch.bfloat16), [1,128])
        self.assertTrue(torch.equal(actual, expected))
        with self.assertRaises(ValueError): weights.unpack_ternary(codes[:-1], torch.ones(1,1), [1,128])
        with self.assertRaises(ValueError): weights.unpack_ternary(torch.full_like(codes,255), torch.ones(1,1), [1,128])


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