"""Tests for FP8/NVFP4 checkpoint dequantization (CPU-only, synthetic).

Covers:
- FP8 block-wise (DeepSeek-style) and per-channel round-trips
- NVFP4 round-trips (direct + reciprocal scale conventions), native and
  manual nibble-unpack paths agreeing with each other
- detect_quant_scheme classification from config.json metadata
- surgery guards rejecting un-dequantized FP8/NVFP4 weights
- end-to-end load_model() on tiny synthetic quantized checkpoints
"""

from __future__ import annotations

import json
import os

import pytest
import torch

from obliteratus.models import quant_dequant as qd

HAS_FP8 = hasattr(torch, "float8_e4m3fn")
requires_fp8 = pytest.mark.skipif(not HAS_FP8, reason="torch build lacks float8 dtypes")

FP8_MAX = 448.0  # e4m3 max
E2M1_MAX = 6.0


# ---------------------------------------------------------------------------
# Reference quantization helpers (test-side)
# ---------------------------------------------------------------------------

def _quantize_fp8_blockwise(w: torch.Tensor, block=(128, 128)):
    """Float weight → (fp8 tensor, scale_inv) using DeepSeek convention."""
    M, N = w.shape
    bm, bn = block
    scale = torch.zeros((M + bm - 1) // bm, (N + bn - 1) // bn)
    q = torch.zeros_like(w)
    for i in range(0, M, bm):
        for j in range(0, N, bn):
            blk = w[i:i + bm, j:j + bn]
            s = blk.abs().max() / FP8_MAX
            s = max(s, 1e-8)
            scale[i // bm, j // bn] = s
            q[i:i + bm, j:j + bn] = blk / s
    return q.to(torch.float8_e4m3fn), scale


def _quantize_fp8_per_channel(w: torch.Tensor, inverse: bool):
    scale = (w.abs().amax(dim=1, keepdim=True) / FP8_MAX).clamp(min=1e-8)
    q = (w / scale).to(torch.float8_e4m3fn)
    return q, (1.0 / scale if inverse else scale)


def _nearest_e2m1(x: torch.Tensor) -> torch.Tensor:
    """Round to nearest E2M1 value, preserving sign. Returns float codes."""
    lut = torch.tensor(qd.E2M1_POSITIVE, dtype=x.dtype)
    sign = torch.sign(x)
    ax = x.abs().clamp(max=E2M1_MAX)
    idx = (ax.unsqueeze(-1) - lut).abs().argmin(-1)
    return sign * lut[idx]


def _pack_nvfp4(w: torch.Tensor, group: int = 16, reciprocal: bool = False):
    """Float (M, N) weight → (packed uint8, e4m3 block scales, fp32 global).

    Follows the ModelOpt convention: w ≈ e2m1 * block_scale * global_scale.
    With reciprocal=True the returned scales are 1/scale (compressed-tensors
    convention).
    """
    M, N = w.shape
    wg = w.reshape(M, N // group, group)
    global_scale = (w.abs().max() / (FP8_MAX * E2M1_MAX)).clamp(min=1e-12)
    bs = (wg.abs().amax(-1) / (E2M1_MAX * global_scale)).clamp(min=1e-8)
    q = _nearest_e2m1(wg / (bs.unsqueeze(-1) * global_scale))
    # values → nibble codes via LUT lookup
    codes = torch.zeros_like(q, dtype=torch.uint8)
    for code, val in enumerate(list(qd.E2M1_POSITIVE) + [-v for v in qd.E2M1_POSITIVE]):
        codes[q == val] = code
    codes = codes.reshape(M, N)
    low = codes[:, 0::2]
    high = codes[:, 1::2]
    packed = (high << 4) | low
    bs_fp8 = bs.to(torch.float8_e4m3fn)
    if reciprocal:
        return packed, (1.0 / bs_fp8.float()).to(torch.float8_e4m3fn), 1.0 / global_scale
    return packed, bs_fp8, global_scale.float()


# ---------------------------------------------------------------------------
# FP8 dequantization
# ---------------------------------------------------------------------------

@requires_fp8
def test_fp8_blockwise_roundtrip():
    torch.manual_seed(0)
    w = torch.randn(256, 256)
    q, scale_inv = _quantize_fp8_blockwise(w)
    out = qd.dequantize_fp8_blockwise(q, scale_inv, (128, 128))
    rel = (out - w).norm() / w.norm()
    assert rel < 0.05, f"blockwise FP8 round-trip error {rel:.4f}"


@requires_fp8
def test_fp8_blockwise_non_divisible_dims():
    torch.manual_seed(1)
    w = torch.randn(200, 300)  # not divisible by 128
    q, scale_inv = _quantize_fp8_blockwise(w)
    out = qd.dequantize_fp8_blockwise(q, scale_inv, (128, 128))
    assert out.shape == w.shape
    rel = (out - w).norm() / w.norm()
    assert rel < 0.05


@requires_fp8
@pytest.mark.parametrize("inverse", [False, True])
def test_fp8_per_channel_roundtrip(inverse):
    torch.manual_seed(2)
    w = torch.randn(64, 128)
    q, scale = _quantize_fp8_per_channel(w, inverse)
    out = qd.dequantize_fp8_per_channel(q, scale, scale_is_inverse=inverse)
    rel = (out - w).norm() / w.norm()
    assert rel < 0.05


@requires_fp8
def test_fp8_per_tensor_scalar_scale():
    torch.manual_seed(3)
    w = torch.randn(32, 64)
    scale = w.abs().max() / FP8_MAX
    q = (w / scale).to(torch.float8_e4m3fn)
    out = qd.dequantize_fp8_per_channel(q, scale.reshape(1))
    rel = (out - w).norm() / w.norm()
    assert rel < 0.05


# ---------------------------------------------------------------------------
# NVFP4 dequantization
# ---------------------------------------------------------------------------

@requires_fp8
@pytest.mark.parametrize("reciprocal", [False, True])
def test_nvfp4_roundtrip(reciprocal):
    torch.manual_seed(4)
    w = torch.randn(64, 128)
    packed, bs, gs = _pack_nvfp4(w, reciprocal=reciprocal)
    out = qd.dequantize_nvfp4(
        packed, bs, gs, scale_is_inverse=reciprocal, force_manual=True,
    )
    assert out.shape == w.shape
    # NVFP4 is coarse — check correlation and relative error, not equality.
    cos = torch.nn.functional.cosine_similarity(
        w.flatten(), out.flatten(), dim=0,
    )
    rel = (out - w).norm() / w.norm()
    # Reciprocal storage double-rounds the inverse scales — slightly coarser.
    assert cos > (0.98 if reciprocal else 0.995), f"NVFP4 cosine {cos:.4f}"
    assert rel < (0.20 if reciprocal else 0.15), f"NVFP4 relative error {rel:.4f}"


@requires_fp8
def test_nvfp4_native_and_manual_unpack_agree():
    if not qd._native_fp4_upcast_works():
        pytest.skip("torch build lacks a working float4_e2m1fn_x2 upcast")
    torch.manual_seed(5)
    packed = torch.randint(0, 256, (16, 64), dtype=torch.uint8)
    native = qd.unpack_e2m1(packed, force_manual=False)
    manual = qd.unpack_e2m1(packed, force_manual=True)
    assert torch.equal(native, manual)


def test_nvfp4_manual_unpack_known_values():
    # 0x1B: low nibble 0xB (-1.5), high nibble 0x1 (+0.5) — low nibble first.
    # 0x84: low nibble 0x4 (+2.0), high nibble 0x8 (-0.0).
    packed = torch.tensor([[0x1B, 0x84]], dtype=torch.uint8)
    out = qd.unpack_e2m1(packed, force_manual=True)
    assert out.shape == (1, 4)
    assert out[0, 0].item() == -1.5
    assert out[0, 1].item() == 0.5
    assert out[0, 2].item() == 2.0
    assert out[0, 3].item() == 0.0


@requires_fp8
def test_nvfp4_no_global_scale():
    torch.manual_seed(6)
    w = torch.randn(32, 64)
    packed, bs, gs = _pack_nvfp4(w)
    out = qd.dequantize_nvfp4(packed, bs, None, force_manual=True)
    assert out.shape == w.shape


def test_nvfp4_bad_group_size_raises():
    packed = torch.zeros((4, 8), dtype=torch.uint8)  # unpacks to 4x16
    bs = torch.ones((4, 1), dtype=torch.float32)
    with pytest.raises(RuntimeError, match="group_size"):
        qd.dequantize_nvfp4(packed, bs, None, group_size=32)


# ---------------------------------------------------------------------------
# Scheme detection
# ---------------------------------------------------------------------------

def _write_config(tmp_path, cfg: dict):
    with open(os.path.join(tmp_path, "config.json"), "w") as fh:
        json.dump(cfg, fh)


def test_detect_none(tmp_path):
    _write_config(tmp_path, {"model_type": "gpt2"})
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.NONE


def test_detect_fp8_blockwise(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {
            "quant_method": "fp8",
            "weight_block_size": [128, 128],
        },
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE
    assert det.block_size == (128, 128)


def test_detect_fp8_blockwise_via_scale_keys(tmp_path):
    _write_config(tmp_path, {"quantization_config": {"quant_method": "fp8"}})
    with open(os.path.join(tmp_path, "model.safetensors.index.json"), "w") as fh:
        json.dump({"weight_map": {"model.layers.0.mlp.weight_scale_inv": "model.safetensors"}}, fh)
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.FP8_BLOCKWISE


def test_detect_fp8_ct(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {
            "quant_method": "compressed-tensors",
            "config_groups": {"group_0": {"weights": {"num_bits": 8, "type": "float"}}},
        },
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT


def test_detect_nvfp4_ct(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {
            "quant_method": "compressed-tensors",
            "config_groups": {"group_0": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}}},
        },
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.NVFP4_CT
    assert det.scale_is_inverse is True
    assert det.group_size == 16


def test_detect_nvfp4_modelopt(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {"quant_method": "modelopt", "quant_algo": "NVFP4"},
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.NVFP4_MODELOPT


def test_detect_modelopt_mixed_precision(tmp_path):
    """Nemotron-3-Nano NVFP4: FP8 mixer + NVFP4 experts in one checkpoint."""
    _write_config(tmp_path, {
        "quantization_config": {
            "quant_method": "modelopt",
            "quant_algo": "MIXED_PRECISION",
            "config_groups": {
                "group_0": {"weights": {"num_bits": 8, "type": "float"}},
                "group_1": {"weights": {"num_bits": 4, "type": "float", "group_size": 16}},
            },
        },
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.NVFP4_MODELOPT


def test_detect_modelopt_mixed_fp8_only(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {
            "quant_method": "modelopt",
            "quant_algo": "MIXED_PRECISION",
            "config_groups": {
                "group_0": {"weights": {"num_bits": 8, "type": "float"}},
            },
        },
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.FP8_PER_CHANNEL_CT


def test_detect_unsupported_fbgemm(tmp_path):
    _write_config(tmp_path, {
        "quantization_config": {"quant_method": "fbgemm_fp8"},
    })
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.UNSUPPORTED
    assert "fbgemm" in det.reason


@pytest.mark.parametrize("method", ["gptq", "awq", "bitsandbytes"])
def test_detect_passthrough_schemes(tmp_path, method):
    _write_config(tmp_path, {"quantization_config": {"quant_method": method}})
    det = qd.detect_quant_scheme(str(tmp_path))
    assert det.scheme is qd.QuantScheme.NONE


# ---------------------------------------------------------------------------
# State-dict dequantization
# ---------------------------------------------------------------------------

@requires_fp8
def test_dequantize_state_dict_fp8_blockwise():
    torch.manual_seed(7)
    w = torch.randn(256, 128)
    q, scale_inv = _quantize_fp8_blockwise(w)
    sd = {
        "layers.0.mlp.weight": q,
        "layers.0.mlp.weight_scale_inv": scale_inv,
        "layers.0.attn.bias": torch.randn(128),
    }
    det = qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE)
    out = qd.dequantize_state_dict(sd, det, out_dtype=torch.bfloat16)
    assert set(out) == {"layers.0.mlp.weight", "layers.0.attn.bias"}
    assert out["layers.0.mlp.weight"].dtype == torch.bfloat16
    rel = (out["layers.0.mlp.weight"].float() - w).norm() / w.norm()
    assert rel < 0.05


@requires_fp8
def test_dequantize_state_dict_nvfp4_modelopt():
    torch.manual_seed(8)
    w = torch.randn(32, 64)
    packed, bs, gs = _pack_nvfp4(w)
    sd = {
        "experts.0.gate_up_proj": packed,
        "experts.0.gate_up_proj_scale": bs,
        "experts.0.gate_up_proj_scale_2": gs,
    }
    # Non-.weight key names (fused MoE style) must also dequantize.
    det = qd.QuantDetection(qd.QuantScheme.NVFP4_MODELOPT)
    out = qd.dequantize_state_dict(sd, det, out_dtype=torch.float32)
    assert set(out) == {"experts.0.gate_up_proj"}
    cos = torch.nn.functional.cosine_similarity(w.flatten(), out["experts.0.gate_up_proj"].flatten(), dim=0)
    assert cos > 0.995


@requires_fp8
def test_dequantize_state_dict_fp8_missing_scale_raises():
    q = torch.randn(16, 16).to(torch.float8_e4m3fn)
    det = qd.QuantDetection(qd.QuantScheme.FP8_BLOCKWISE)
    with pytest.raises(RuntimeError, match="no weight_scale"):
        qd.dequantize_state_dict({"a.weight": q}, det)


# ---------------------------------------------------------------------------
# Surgery guards
# ---------------------------------------------------------------------------

@requires_fp8
def test_is_quantized_param_fp8():
    from obliteratus.abliterate import AbliterationPipeline

    p = torch.nn.Parameter(torch.randn(8, 8).to(torch.float8_e4m3fn))
    assert AbliterationPipeline._is_quantized_param(p) is True
    p2 = torch.nn.Parameter(torch.randn(8, 8))
    assert AbliterationPipeline._is_quantized_param(p2) is False


@requires_fp8
def test_dequantize_weight_guard_fp8():
    from obliteratus.abliterate import AbliterationPipeline

    mod = torch.nn.Linear(8, 8)
    mod.weight = torch.nn.Parameter(torch.randn(8, 8).to(torch.float8_e4m3fn))
    with pytest.raises(RuntimeError, match="without dequantization"):
        AbliterationPipeline._dequantize_weight(mod)


def test_dequantize_weight_guard_packed_uint8_with_scales():
    from obliteratus.abliterate import AbliterationPipeline

    mod = torch.nn.Linear(8, 8)
    mod.weight = torch.nn.Parameter(torch.zeros(8, 4, dtype=torch.uint8), requires_grad=False)
    mod.weight_scale = torch.ones(8, 1, dtype=torch.float32)
    with pytest.raises(RuntimeError, match="without dequantization"):
        AbliterationPipeline._dequantize_weight(mod)


def test_dequantize_weight_plain_uint8_still_converts():
    """Pre-existing behavior for scale-less custom uint8 weights is kept."""
    from obliteratus.abliterate import AbliterationPipeline

    mod = torch.nn.Linear(8, 8)
    mod.weight = torch.nn.Parameter(torch.ones(8, 8, dtype=torch.uint8), requires_grad=False)
    w, is_q = AbliterationPipeline._dequantize_weight(mod)
    assert is_q is True
    assert w.is_floating_point()


# ---------------------------------------------------------------------------
# End-to-end loader on tiny synthetic checkpoints
# ---------------------------------------------------------------------------

@requires_fp8
def test_load_model_fp8_blockwise_checkpoint(tmp_path, monkeypatch):
    from obliteratus.models import loader as loader_mod

    captured = {}

    from safetensors.torch import save_file
    from transformers import GPT2Config, GPT2LMHeadModel

    cfg = GPT2Config(n_layer=1, n_head=2, n_embd=256, n_inner=512, vocab_size=128)
    model = GPT2LMHeadModel(cfg)
    sd = model.state_dict()
    out_sd = {}
    for k, v in sd.items():
        if k.endswith(".weight") and v.dim() == 2 and min(v.shape) >= 32:
            q, s = _quantize_fp8_blockwise(v.float())
            out_sd[k] = q
            out_sd[k[: -len(".weight")] + ".weight_scale_inv"] = s
            captured[k] = v.float()
        else:
            out_sd[k] = v
    save_file(out_sd, str(tmp_path / "model.safetensors"))
    cfg_dict = cfg.to_dict()
    cfg_dict["quantization_config"] = {
        "quant_method": "fp8", "weight_block_size": [128, 128],
    }
    with open(tmp_path / "config.json", "w") as fh:
        json.dump(cfg_dict, fh)

    class _FakeTok:
        pad_token = eos_token = "<|endoftext|>"

    monkeypatch.setattr(
        loader_mod.AutoTokenizer, "from_pretrained", staticmethod(lambda *a, **k: _FakeTok()),
    )
    handle = loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16")
    loaded = handle.model.state_dict()
    assert getattr(handle.model, "_obliteratus_dequantized_scheme", None) == "fp8_blockwise"
    for k, w in captured.items():
        p = loaded[k].detach().float()
        assert p.is_floating_point()
        rel = (p - w).norm() / w.norm()
        assert rel < 0.05, f"{k}: rel error {rel:.4f}"
    # No scale tensors survived into the loaded model
    assert not any("scale" in n for n in loaded)


@requires_fp8
def test_load_model_nvfp4_checkpoint(tmp_path, monkeypatch):
    from obliteratus.models import loader as loader_mod
    from safetensors.torch import save_file
    from transformers import GPT2Config, GPT2LMHeadModel

    cfg = GPT2Config(n_layer=1, n_head=2, n_embd=256, n_inner=512, vocab_size=128)
    model = GPT2LMHeadModel(cfg)
    sd = model.state_dict()
    out_sd = {}
    captured = {}
    for k, v in sd.items():
        if k.endswith(".weight") and v.dim() == 2 and v.shape[1] % 32 == 0 and min(v.shape) >= 32:
            packed, bs, gs = _pack_nvfp4(v.float())
            base = k[: -len(".weight")]
            out_sd[k] = packed
            out_sd[base + ".weight_scale"] = bs
            out_sd[base + ".weight_scale_2"] = gs
            captured[k] = v.float()
        else:
            out_sd[k] = v
    save_file(out_sd, str(tmp_path / "model.safetensors"))
    cfg_dict = cfg.to_dict()
    cfg_dict["quantization_config"] = {
        "quant_method": "modelopt", "quant_algo": "NVFP4",
    }
    with open(tmp_path / "config.json", "w") as fh:
        json.dump(cfg_dict, fh)

    class _FakeTok:
        pad_token = eos_token = "<|endoftext|>"

    monkeypatch.setattr(
        loader_mod.AutoTokenizer, "from_pretrained", staticmethod(lambda *a, **k: _FakeTok()),
    )
    handle = loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16")
    loaded = handle.model.state_dict()
    assert getattr(handle.model, "_obliteratus_dequantized_scheme", None) == "nvfp4_modelopt"
    for k, w in captured.items():
        p = loaded[k].detach().float()
        cos = torch.nn.functional.cosine_similarity(w.flatten(), p.flatten(), dim=0)
        assert cos > 0.99, f"{k}: cosine {cos:.4f}"
    assert not any("scale" in n for n in loaded)


def test_load_model_unsupported_scheme_fails_loudly(tmp_path):
    from obliteratus.models import loader as loader_mod

    _write_config(tmp_path, {
        "model_type": "gpt2",
        "quantization_config": {"quant_method": "fbgemm_fp8"},
    })
    with pytest.raises(RuntimeError, match="Unsupported quantization"):
        loader_mod.load_model(str(tmp_path), task="causal_lm", device="cpu", dtype="bfloat16")