import unittest

try:
    import torch
except ImportError:  # the planning host has no torch
    torch = None

if torch is not None:
    from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
                             to_gpu_int4)


@unittest.skipIf(torch is None, "needs torch")
class Quant(unittest.TestCase):
    def setUp(self):
        torch.manual_seed(0)
        self.w = torch.randn(64, 256) * 0.02

    def test_layouts_match_what_quantize_returns(self):
        for quant in ("int8", "int4"):
            got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
            self.assertEqual(got, layout(quant, 64, 256))

    def test_round_trip_error(self):
        for quant, bound in (("int8", 0.01), ("int4", 0.11)):
            err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
            self.assertLess(err / self.w.norm(), bound)

    def test_int4_clipping_beats_plain_round_to_nearest(self):
        g = self.w.unflatten(-1, (-1, GROUP))
        scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
        plain = ((g / scale).round() * scale - g).norm()
        clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
        self.assertLess(clipped, plain)

    def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
        q, _ = quantize(self.w, "int4")
        nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
        self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
        self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))

    def test_cpu_matmul_matches_dequantized_weights(self):
        x = torch.randn(3, 256, dtype=torch.bfloat16)
        for quant in ("int8", "int4"):
            q, scale = quantize(self.w, quant)
            ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
            stored = to_cpu_int4(q) if quant == "int4" else q
            got = cpu_matmul(quant, stored, scale)(x).float()
            self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)