"""Tests for OrbitQuant loader support in ComfyUI-INT-Crush.

Test strategy: manually construct state dicts with the OrbitQuant format
(packed indices, bfloat16 norms, global RPBH/codebook tensors,
format_version=4 metadata) — no converter dependency needed.
"""

import math
import torch
import pytest

from _intcrush import load

_oq = load("_orbitquant")
_qu = load("_quant_utils")


# ── Helpers ──

def _make_codebook(d: int, bits: int) -> torch.Tensor:
    """Build a simple uniform codebook for testing.

    In production, this comes from Lloyd-Max on the f_d marginal.
    For tests, a uniform grid suffices to verify the dequant path.
    """
    n = 2 ** bits
    centroids = torch.linspace(-1.0, 1.0, n, dtype=torch.float32)
    return centroids


def _make_orbitquant_state_dict(
    d: int = 256,
    out_features: int = 32,
    bits: int = 4,
    seed: int = 42,
    include_adaln: bool = False,
) -> tuple[dict, dict]:
    """Build a synthetic OrbitQuant model's state_dict and metadata.

    Returns (state_dict, metadata) ready for save_file().
    """
    h = _oq.largest_power_of_two_dividing(d)
    perm, signs, h = _oq.make_rpbh_params(d, seed=seed)
    d_padded = perm.shape[0]

    codebook = _make_codebook(d, bits)
    codebook_key = f"_orbitquant_codebook_{d}_{bits}"

    # Build a fake weight: random unit directions * random norms
    W = torch.randn(out_features, d)
    norms = W.norm(dim=1, keepdim=True)
    W_dir = W / (norms + 1e-10)

    # Quantize directions to codebook
    # (simplified: use nearest centroid)
    centroids = codebook
    boundaries = (centroids[:-1] + centroids[1:]) / 2.0
    indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
    indices = indices.clamp(0, len(centroids) - 1).reshape(out_features, d)

    # Offset to signed int8 range
    offset = 1 << (bits - 1)
    indices_int8 = (indices - offset).to(torch.int8)

    if bits == 4:
        packed = _qu.pack_int4(indices_int8).to(torch.uint8)
    else:
        packed = indices_int8

    scale = norms.to(torch.bfloat16)

    state_dict = {
        "layers.0.q_proj.weight": packed,
        "layers.0.q_proj.weight_scale": scale,
        "_rpbh_perm_256": perm,
        "_rpbh_signs_256": signs,
        codebook_key: codebook,
    }

    metadata = {
        "int_crush.format_version": "4",
        "int_crush.orbitquant": "true",
        "int_crush.rpbh_seed": str(seed),
        "int_crush.codebook_bits": str(bits),
        "int_crush.orbitquant_dims": str(d),
        "int_crush.orbitquant_d_padded": f"{d}:{d_padded}",
        "int_crush.layers.0.q_proj.weight.quant_scheme": "orbitquant_codebook",
    }

    return state_dict, metadata


# ── Tests ──

class TestRPBHRotate:
    """Test RPBH rotation matches converter implementation."""

    def test_norm_preservation(self):
        """‖x @ Π^T‖ = ‖x‖ (orthogonal)."""
        d = 256
        perm, signs, h = _oq.make_rpbh_params(d, seed=42)
        x = torch.randn(8, d)
        y = _oq.rpbh_rotate(x, perm, signs, h)
        assert torch.allclose(x.norm(), y.norm(), rtol=1e-3)

    def test_cancellation(self):
        """(W @ Π^T) @ (x @ Π^T)^T = W @ x^T."""
        for d in [256, 512, 1024]:
            perm, signs, h = _oq.make_rpbh_params(d, seed=42)
            W = torch.randn(16, d)
            x = torch.randn(8, d)
            W_rot = _oq.rpbh_rotate(W, perm, signs, h)
            x_rot = _oq.rpbh_rotate(x, perm, signs, h)
            y_rot = W_rot @ x_rot.T
            y_orig = W @ x.T
            rel_err = (y_rot - y_orig).norm() / y_orig.norm()
            assert rel_err < 1e-4

    def test_fwht_matches_sylvester(self):
        """FWHT output matches Sylvester Hadamard @ x."""
        h = 16
        H = torch.tensor([[1.0]])
        while H.shape[0] < h:
            H = torch.kron(torch.tensor([[1., 1.], [1., -1.]]), H)
        H = H / math.sqrt(h)
        x = torch.randn(5, h)
        y_fwht = _oq.fast_walsh_hadamard(x)
        y_matrix = x @ H.T
        assert torch.allclose(y_fwht, y_matrix, atol=1e-4)

    def test_mass_balancing(self):
        """Spiked input → uniform variance after RPBH (CV < 0.3)."""
        d = 1024
        h = _oq.largest_power_of_two_dividing(d)
        k = d // h
        if k < 2:
            pytest.skip("single-block dimension")
        perm, signs, _ = _oq.make_rpbh_params(d, seed=42)
        n_samples = 5000
        x = torch.zeros(n_samples, d)
        x[:, :h] = torch.randn(n_samples, h)
        x = x / x.norm(dim=1, keepdim=True)
        z = _oq.rpbh_rotate(x, perm, signs, h)
        var_per_coord = z.double().pow(2).mean(0)
        cv = var_per_coord.std().item() / var_per_coord.mean().item()
        assert cv < 0.3


class TestCodebookDequant:
    """Test codebook dequantization."""

    def test_int4_roundtrip(self):
        """Quantize → pack → dequant ≈ original (within codebook MSE)."""
        d = 256
        bits = 4
        codebook = _make_codebook(d, bits)
        W = torch.randn(32, d) * 0.1
        norms = W.norm(dim=1, keepdim=True)
        W_dir = W / (norms + 1e-10)
        boundaries = (codebook[:-1] + codebook[1:]) / 2.0
        indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
        indices = indices.clamp(0, 15).reshape(32, d)
        offset = 1 << (bits - 1)
        indices_int8 = (indices - offset).to(torch.int8)
        packed = _qu.pack_int4(indices_int8).to(torch.uint8)

        W_dequant = _oq.codebook_dequant_int4(
            packed, norms.to(torch.bfloat16), codebook, d, d)

        # Should be close to original (within codebook quantization error)
        mse = (W - W_dequant.float()).pow(2).mean().item()
        assert mse < 0.01

    def test_int8_roundtrip(self):
        """INT8 indices → codebook lookup → norms multiply."""
        d = 256
        bits = 8
        codebook = _make_codebook(d, bits)
        W = torch.randn(16, d) * 0.1
        norms = W.norm(dim=1, keepdim=True)
        W_dir = W / (norms + 1e-10)
        boundaries = (codebook[:-1] + codebook[1:]) / 2.0
        indices = torch.searchsorted(boundaries, W_dir.reshape(-1))
        indices = indices.clamp(0, 255).reshape(16, d)
        offset = 1 << (bits - 1)
        indices_int8 = (indices - offset).to(torch.int8)

        W_dequant = _oq.codebook_dequant_int8(
            indices_int8, norms.to(torch.bfloat16), codebook, d, d)

        mse = (W - W_dequant.float()).pow(2).mean().item()
        assert mse < 0.001  # INT8 should be more accurate

    def test_no_padding_trim(self):
        """d_padded > d → output is [out, d_padded] (no trim at inference)."""
        d = 200  # not divisible by most block sizes
        h = 128  # explicit block size
        d_padded = 256  # next multiple of 128
        bits = 4
        codebook = _make_codebook(d_padded, bits)
        norms = torch.ones(8, 1, dtype=torch.bfloat16)
        indices_int8 = torch.zeros(8, d_padded, dtype=torch.int8)

        W_dequant = _oq.codebook_dequant_int4(
            _qu.pack_int4(indices_int8).to(torch.uint8),
            norms, codebook, d, d_padded)

        assert W_dequant.shape == (8, d_padded)  # NOT (8, d)


class TestAdaLNDequant:
    """Test AdaLN RTN dequantization with group scales."""

    def test_group_scale_shape(self):
        """AdaLN scales [out, d/64] → dequant produces [out, d]."""
        d = 256
        group_size = 64
        out = 32
        num_groups = d // group_size
        scales = torch.randn(out, num_groups).abs().clamp(min=1e-4).to(torch.float16)
        packed = _qu.pack_int4(
            torch.randint(-8, 7, (out, d), dtype=torch.int8)
        ).to(torch.uint8)

        W_dequant = _oq.adaln_rtn_dequant(packed, scales, d, bits=4)
        assert W_dequant.shape == (out, d)

    def test_group_scale_correctness(self):
        """Each group of 64 values is multiplied by its scale."""
        d = 128
        group_size = 64
        out = 4
        num_groups = d // group_size
        scales = torch.tensor([[2.0, 3.0]] * out, dtype=torch.float16)
        # All indices = 0 → dequant value = 0 * scale = 0
        # Use index 8 (offset 0) → value = -8 * scale
        values = torch.full((out, d), -8, dtype=torch.int8)
        packed = _qu.pack_int4(values).to(torch.uint8)

        W_dequant = _oq.adaln_rtn_dequant(packed, scales, d, bits=4)
        # Group 0: -8 * 2.0 = -16.0
        # Group 1: -8 * 3.0 = -24.0
        assert torch.allclose(W_dequant[:, :64], torch.full((out, 64), -16.0), atol=0.1)
        assert torch.allclose(W_dequant[:, 64:], torch.full((out, 64), -24.0), atol=0.1)


class TestOrbitQuantLayoutDequantize:
    """Test OrbitQuantLayout.dequantize() works for both paths."""

    def test_codebook_path(self):
        """dequantize with codebook produces correct output."""
        _ql = load("quant_layout")
        OrbitQuantLayout = _ql.OrbitQuantLayout
        d = 256
        bits = 4
        codebook = _make_codebook(d, bits)
        norms = torch.ones(8, 1, dtype=torch.bfloat16)
        packed = _qu.pack_int4(
            torch.zeros(8, d, dtype=torch.int8)
        ).to(torch.uint8)

        params = OrbitQuantLayout.Params(
            scale=norms, orig_dtype=torch.float16, orig_shape=(8, d),
            quant_scheme="orbitquant_codebook", d=d, d_padded=d,
            bits=bits, codebook=codebook,
        )
        W = OrbitQuantLayout.dequantize(packed, params)
        assert W.shape == (8, d)

    def test_adaln_path(self):
        """dequantize with adaln_rtn uses group scales."""
        _ql = load("quant_layout")
        OrbitQuantLayout = _ql.OrbitQuantLayout
        d = 256
        out = 8
        scales = torch.ones(out, d // 64, dtype=torch.float16)
        packed = _qu.pack_int4(
            torch.zeros(out, d, dtype=torch.int8)
        ).to(torch.uint8)

        params = OrbitQuantLayout.Params(
            scale=scales, orig_dtype=torch.float16, orig_shape=(out, d),
            quant_scheme="adaln_rtn", d=d, d_padded=d, bits=4,
        )
        W = OrbitQuantLayout.dequantize(packed, params)
        assert W.shape == (out, d)


class TestMetadataParsing:
    """Test OrbitQuant metadata parsing."""

    def test_parse_orbitquant_metadata(self):
        """Valid metadata is parsed correctly."""
        metadata = {
            "int_crush.format_version": "4",
            "int_crush.orbitquant": "true",
            "int_crush.rpbh_seed": "42",
            "int_crush.codebook_bits": "4",
            "int_crush.orbitquant_dims": "256,512",
            "int_crush.orbitquant_d_padded": "256:256,512:512",
        }
        result = _oq.parse_orbitquant_metadata(metadata)
        assert result is not None
        assert result["rpbh_seed"] == 42
        assert result["codebook_bits"] == 4
        assert result["dims"] == [256, 512]
        assert result["d_padded_map"] == {256: 256, 512: 512}

    def test_parse_non_orbitquant(self):
        """Non-OrbitQuant metadata returns None."""
        assert _oq.parse_orbitquant_metadata({}) is None
        assert _oq.parse_orbitquant_metadata(None) is None
        assert _oq.parse_orbitquant_metadata({
            "int_crush.format_version": "3"
        }) is None

    def test_get_quant_scheme_from_metadata(self):
        """Per-layer quant_scheme is read correctly."""
        metadata = {
            "int_crush.layers.0.q_proj.weight.quant_scheme": "orbitquant_codebook",
            "int_crush.layers.0.k_proj.weight.quant_scheme": "adaln_rtn",
        }
        assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.q_proj.") == "orbitquant_codebook"
        assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.k_proj.") == "adaln_rtn"
        assert _oq.get_quant_scheme_from_metadata(metadata, "layers.0.v_proj.") is None

    def test_get_quant_scheme_strips_model_prefix(self):
        """Metadata lookup strips model-level prefix from ComfyUI."""
        metadata = {
            "int_crush.single_blocks.0.linear1.weight.quant_scheme": "orbitquant_codebook",
            "int_crush.double_blocks.0.img_attn.q_proj.weight.quant_scheme": "adaln_rtn",
        }
        # ComfyUI passes prefix with model.diffusion_model. prefix
        assert _oq.get_quant_scheme_from_metadata(
            metadata, "model.diffusion_model.single_blocks.0.linear1.") == "orbitquant_codebook"
        assert _oq.get_quant_scheme_from_metadata(
            metadata, "model.diffusion_model.double_blocks.0.img_attn.q_proj.") == "adaln_rtn"
        # Non-existent layer still returns None
        assert _oq.get_quant_scheme_from_metadata(
            metadata, "model.diffusion_model.single_blocks.0.linear2.") is None

    def test_infer_metadata_from_state_dict(self):
        """Metadata is inferred from global tensor keys in state_dict."""
        state_dict = {
            "_rpbh_perm_256": torch.randperm(256),
            "_rpbh_signs_256": torch.ones(256),
            "_orbitquant_codebook_256_4": torch.linspace(-1, 1, 16),
            "layers.0.weight": torch.zeros(1),
        }
        result = _oq.infer_orbitquant_metadata_from_state_dict(state_dict)
        assert result is not None
        assert result["int_crush.format_version"] == "4"
        assert result["int_crush.orbitquant"] == "true"
        assert "256" in result["int_crush.orbitquant_dims"]
        assert result["int_crush.codebook_bits"] == "4"

    def test_infer_metadata_returns_none_without_globals(self):
        """Returns None when no OrbitQuant globals in state_dict."""
        state_dict = {"layers.0.weight": torch.zeros(1)}
        assert _oq.infer_orbitquant_metadata_from_state_dict(state_dict) is None

    def test_infer_quant_scheme_from_scale(self):
        """Scale shape inference works correctly."""
        # [out, 1] → orbitquant_codebook
        scale_cb = torch.ones(8, 1)
        assert _oq.infer_quant_scheme_from_scale(scale_cb) == "orbitquant_codebook"

        # [out, 4] → adaln_rtn (d=256, group_size=64 → 4 groups)
        scale_adaln = torch.ones(8, 4)
        assert _oq.infer_quant_scheme_from_scale(scale_adaln) == "adaln_rtn"


class TestGlobalTensorCache:
    """Test OrbitQuant global tensor caching."""

    def test_load_and_cache_globals(self):
        """Global tensors are loaded and cached correctly."""
        state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
        model_path = "test_model"

        globals_cache = _oq.load_orbitquant_globals(state_dict, metadata, model_path)

        assert 256 in globals_cache["rpbh"]
        assert "perm" in globals_cache["rpbh"][256]
        assert "signs" in globals_cache["rpbh"][256]
        assert (256, 4) in globals_cache["codebooks"]

        # Verify global tensors were popped from state_dict
        assert "_rpbh_perm_256" not in state_dict
        assert "_rpbh_signs_256" not in state_dict
        assert "_orbitquant_codebook_256_4" not in state_dict

        # Verify per-layer tensors remain
        assert "layers.0.q_proj.weight" in state_dict

        # Cleanup
        _oq.clear_orbitquant_globals(model_path)

    def test_cache_idempotent(self):
        """Second call returns cached result without popping again."""
        state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
        model_path = "test_model_idempotent"

        cache1 = _oq.load_orbitquant_globals(state_dict, metadata, model_path)
        # State dict now has globals popped — reconstruct for second call
        # But the function should return cached result without touching state_dict
        cache2 = _oq.load_orbitquant_globals(state_dict, metadata, model_path)
        assert cache1 is cache2

        _oq.clear_orbitquant_globals(model_path)

    def test_clear_globals(self):
        """clear_orbitquant_globals removes cached data."""
        state_dict, metadata = _make_orbitquant_state_dict(d=256, bits=4)
        model_path = "test_model_clear"

        _oq.load_orbitquant_globals(state_dict, metadata, model_path)
        assert model_path in _oq._orbitquant_globals

        _oq.clear_orbitquant_globals(model_path)
        assert model_path not in _oq._orbitquant_globals

    def test_clear_all_globals(self):
        """clear_orbitquant_globals(None) clears everything."""
        sd1, meta1 = _make_orbitquant_state_dict(d=256, bits=4)
        sd2, meta2 = _make_orbitquant_state_dict(d=256, bits=4)

        _oq.load_orbitquant_globals(sd1, meta1, "model_a")
        _oq.load_orbitquant_globals(sd2, meta2, "model_b")

        _oq.clear_orbitquant_globals()
        assert len(_oq._orbitquant_globals) == 0


class TestLoadersDetection:
    """Test OrbitQuant detection helpers in _loaders."""

    def test_detect_orbitquant(self):
        """_detect_orbitquant returns True for valid metadata."""
        _ld = load("_loaders")
        _detect_orbitquant = _ld._detect_orbitquant

        assert _detect_orbitquant({
            "int_crush.format_version": "4",
            "int_crush.orbitquant": "true",
        }) is True

    def test_detect_non_orbitquant(self):
        """_detect_orbitquant returns False for non-OrbitQuant metadata."""
        _ld = load("_loaders")
        _detect_orbitquant = _ld._detect_orbitquant

        assert _detect_orbitquant({}) is False
        assert _detect_orbitquant(None) is False
        assert _detect_orbitquant({
            "int_crush.format_version": "3",
        }) is False
        assert _detect_orbitquant({
            "int_crush.format_version": "4",
            # missing orbitquant=true
        }) is False