import unittest

import torch


@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
class Nvfp4NativePackParity(unittest.TestCase):
    def test_h3_projection_widths_match_comfy_layout(self):
        try:
            import comfy_kitchen as ck
            from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
            from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
        except ImportError as error:
            raise unittest.SkipTest("Comfy Kitchen is not installed") from error

        torch.manual_seed(440420)
        extension = _vortex_scale_extension()
        for features in (5376, 7168, 14336):
            with self.subTest(features=features):
                value = torch.randn(129, features, device="cuda", dtype=torch.bfloat16)
                scale = (value.abs().amax() / float(F8_E4M3_MAX * F4_E2M1_MAX)).float()
                expected_qdata, expected_block_scale = ck.quantize_nvfp4(value, scale, pad_16x=True)
                actual_qdata, actual_block_scale = extension.quantize_nvfp4_bf16(value, scale, True, 512)
                self.assertTrue(torch.equal(actual_qdata, expected_qdata))
                self.assertTrue(torch.equal(actual_block_scale.view(torch.uint8), expected_block_scale.view(torch.uint8)))
                into_qdata = torch.empty_like(actual_qdata)
                into_block_scale = torch.empty_like(actual_block_scale)
                extension.quantize_nvfp4_bf16_into(
                    value, scale, into_qdata, into_block_scale, 512,
                )
                self.assertTrue(torch.equal(into_qdata, expected_qdata))
                self.assertTrue(
                    torch.equal(
                        into_block_scale.view(torch.uint8),
                        expected_block_scale.view(torch.uint8),
                    )
                )
                low_first_qdata = torch.empty_like(actual_qdata)
                extension.quantize_nvfp4_bf16_into(
                    value,
                    scale,
                    low_first_qdata,
                    into_block_scale,
                    512,
                    False,
                )
                expected_low_first = (
                    ((expected_qdata & 0x0F) << 4)
                    | ((expected_qdata & 0xF0) >> 4)
                )
                self.assertTrue(torch.equal(low_first_qdata, expected_low_first))

    def test_modulated_producer_matches_materialized_bf16(self):
        try:
            import comfy_kitchen as ck
            from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
            from h3_blackwell_runtime.h3_fusion import fused_modulate_
            from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
        except ImportError as error:
            raise unittest.SkipTest("Comfy Kitchen or Triton is not installed") from error

        torch.manual_seed(440421)
        value = torch.randn(129, 128, device="cuda", dtype=torch.bfloat16)
        row_index = torch.arange(129, device="cuda", dtype=torch.int32) % 3
        divisor = float(F8_E4M3_MAX * F4_E2M1_MAX)
        extension = _vortex_scale_extension()
        for table_dtype in (torch.bfloat16, torch.float32):
            with self.subTest(table_dtype=table_dtype):
                shift = torch.randn(3, 128, device="cuda", dtype=table_dtype)
                scale = torch.randn(3, 128, device="cuda", dtype=table_dtype)
                materialized = fused_modulate_(
                    value.clone(), shift, scale, row_index,
                )
                expected_scale = (materialized.abs().amax() / divisor).to(
                    torch.bfloat16,
                ).float()
                expected_qdata, expected_sfa = ck.quantize_nvfp4(
                    materialized, expected_scale, pad_16x=True,
                )
                actual_scale, actual_qdata, actual_sfa = (
                    extension.quantize_nvfp4_modulated_bf16(
                        value, shift, scale, row_index, divisor, True, 32, 128, 128,
                    )
                )
                self.assertTrue(torch.equal(actual_scale, expected_scale))
                self.assertTrue(torch.equal(actual_qdata, expected_qdata))
                self.assertTrue(
                    torch.equal(actual_sfa.view(torch.uint8), expected_sfa.view(torch.uint8))
                )

    def test_swiglu_producer_matches_materialized_bf16(self):
        try:
            import comfy_kitchen as ck
            from comfy_kitchen.float_utils import F4_E2M1_MAX, F8_E4M3_MAX
            from h3_blackwell_runtime.nvfp4_quant import _vortex_scale_extension
        except ImportError as error:
            raise unittest.SkipTest("Comfy Kitchen is not installed") from error

        torch.manual_seed(440422)
        gate_up = torch.randn(129, 256, device="cuda", dtype=torch.bfloat16)
        gate, up = gate_up.chunk(2, dim=-1)
        materialized = torch.nn.functional.silu(gate).mul_(up)
        divisor = float(F8_E4M3_MAX * F4_E2M1_MAX)
        expected_scale = (materialized.abs().amax() / divisor).to(
            torch.bfloat16,
        ).float()
        expected_qdata, expected_sfa = ck.quantize_nvfp4(
            materialized, expected_scale, pad_16x=True,
        )
        actual_scale, actual_qdata, actual_sfa = (
            _vortex_scale_extension().quantize_nvfp4_swiglu_bf16(
                gate_up, divisor, True, 32, 128, 128,
            )
        )
        self.assertTrue(torch.equal(actual_scale, expected_scale))
        self.assertTrue(torch.equal(actual_qdata, expected_qdata))
        self.assertTrue(
            torch.equal(actual_sfa.view(torch.uint8), expected_sfa.view(torch.uint8))
        )


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