# SPDX-License-Identifier: Apache-2.0
"""
Parity test for Z-Image Qwen3 text encoder support in FastVideo.

This compares:
1) direct transformers AutoModel output, and
2) FastVideo's production ``TextEncoderLoader`` passthrough, and
3) FastVideo's shared native Qwen3 encoder (``Qwen3ForCausalLM``, added for Flux2
   Klein and reused here — Z-Image-Turbo's ``Qwen3Model`` checkpoint routes
   to it via the model registry),

using identical local checkpoint, tokenization, and inputs.

Usage:
    pytest tests/local_tests/zimage/test_zimage_encoder_parity.py -v -s
"""

from __future__ import annotations

from pathlib import Path
from types import SimpleNamespace
import gc
import json
import os

import pytest
import torch


def _reclaim_vram() -> None:
    """Actually reclaim VRAM after the caller has ``del``'d its model refs.

    HF transformer models hold reference cycles, so ``del`` alone does not free
    them — a ``gc.collect()`` is required before the CUDA caching allocator will
    release the blocks. Without this the parity test keeps two full encoder
    copies resident and OOMs on a 44 GB GPU.
    """
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()


def _gpu_mem_mb() -> float:
    if torch.cuda.is_available():
        return torch.cuda.memory_allocated() / (1024**2)
    return 0.0


from torch.testing import assert_close
from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
from safetensors.torch import safe_open

from fastvideo.configs.models.encoders.qwen3 import Qwen3TextConfig
from fastvideo.distributed.parallel_state import cleanup_dist_env_and_memory, maybe_init_distributed_environment_and_model_parallel
from fastvideo.layers.quantization.absmax_fp8 import AbsMaxFP8MergedParameter
from fastvideo.models.encoders.qwen3 import Qwen3ForCausalLM
from fastvideo.models.loader.component_loader import TextEncoderLoader
from fastvideo.models.loader.utils import set_default_torch_dtype
from fastvideo.models.registry import ModelRegistry

PARITY_SCOPE = "both"

# Strict-load contract: the shared Qwen3 encoder is body-only (embed_tokens +
# layers + norm, no lm_head). Z-Image-Turbo ships a full Qwen3 checkpoint, so
# `lm_head.weight` is the only key that may go unmatched; anything else means a
# real silent drop. Enforced here in the test since the shared encoder's loader
# is intentionally lenient (it's used by multiple models).
_ALLOWED_UNEXPECTED_KEYS = {"lm_head.weight"}

REPO_ROOT = Path(__file__).resolve().parents[3]
ZIMAGE_TEXT_ENCODER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "text_encoder"
ZIMAGE_TOKENIZER_DIR = REPO_ROOT / "official_weights" / "Z-Image" / "tokenizer"


@pytest.fixture(scope="module", autouse=True)
def _init_dist_and_tp_groups():
    os.environ.setdefault("MASTER_ADDR", "localhost")
    os.environ.setdefault("MASTER_PORT", "29531")
    os.environ.setdefault("RANK", "0")
    os.environ.setdefault("WORLD_SIZE", "1")
    os.environ.setdefault("LOCAL_RANK", "0")

    maybe_init_distributed_environment_and_model_parallel(1, 1)
    yield
    cleanup_dist_env_and_memory()


def _load_json(path: Path) -> dict:
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _iter_safetensors(path: Path):
    with safe_open(str(path), framework="pt", device="cpu") as f:
        for k in f.keys():
            yield k, f.get_tensor(k)


def _iter_pretrained_safetensors(model_dir: Path):
    single = model_dir / "model.safetensors"
    if single.exists():
        yield from _iter_safetensors(single)
        return

    index = model_dir / "model.safetensors.index.json"
    if index.exists():
        idx = _load_json(index)
        shard_names = sorted(set(idx["weight_map"].values()))
        for shard in shard_names:
            yield from _iter_safetensors(model_dir / shard)
        return

    raise FileNotFoundError(
        f"Missing safetensors checkpoint in {model_dir} (expected model.safetensors or model.safetensors.index.json)")


def _pretrained_safetensor_keys(model_dir: Path) -> set[str]:
    return {name for name, _ in _iter_pretrained_safetensors(model_dir)}


def _load_qwen3_config() -> Qwen3TextConfig:
    cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
    for key in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
        cfg_raw.pop(key, None)
    config = Qwen3TextConfig()
    config.update_model_arch(cfg_raw)
    assert config.num_hidden_layers == 36
    assert config.hidden_size == 2560
    assert config.intermediate_size == 9728
    assert config.num_attention_heads == 32
    assert config.num_key_value_heads == 8
    return config


def _loader_args(cpu_offload: bool) -> SimpleNamespace:
    return SimpleNamespace(
        text_encoder_cpu_offload=cpu_offload,
        override_text_encoder_quant=None,
        override_text_encoder_safetensors=None,
        pin_cpu_memory=False,
    )


def _precision_name(dtype: torch.dtype) -> str:
    if dtype == torch.float32:
        return "fp32"
    if dtype == torch.bfloat16:
        return "bf16"
    raise ValueError(f"Unsupported parity dtype: {dtype}")


def _tokenize_official_prompts(tokenizer, prompts: list[str]):
    formatted = [
        tokenizer.apply_chat_template(
            [{
                "role": "user",
                "content": prompt
            }],
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=True,
        ) for prompt in prompts
    ]
    return tokenizer(
        formatted,
        padding="max_length",
        truncation=True,
        max_length=512,
        return_tensors="pt",
    )


def _assert_native_load_surface(
    model,
    checkpoint_keys: set[str],
    loaded_params: set[str],
) -> None:
    param_names = {name for name, _ in model.named_parameters()}
    missing = param_names - loaded_params
    assert not missing, f"FastVideo parameters not loaded: {sorted(missing)}"

    normalized_keys = {name[len("model."):] if name.startswith("model.") else name for name in checkpoint_keys}
    mapped_keys: set[str] = set()
    for checkpoint_name in normalized_keys:
        for param_name, weight_name, _ in model.config.arch_config.stacked_params_mapping:
            if weight_name in checkpoint_name:
                mapped_keys.add(checkpoint_name.replace(weight_name, param_name))
                break
        else:
            mapped_keys.add(checkpoint_name)

    unexpected = mapped_keys - param_names - {
        "rotary_emb.inv_freq",
        "rotary_emb.cos_cached",
        "rotary_emb.sin_cached",
    }
    assert unexpected <= _ALLOWED_UNEXPECTED_KEYS, ("Unexpected checkpoint keys not in allowlist: "
                                                    f"{sorted(unexpected - _ALLOWED_UNEXPECTED_KEYS)}")


def _fake_qwen3_weight_target(include_quant_scales: bool = False):
    stacked = Qwen3TextConfig().arch_config.stacked_params_mapping
    params = {
        "layers.0.self_attn.qkv_proj.weight": torch.nn.Parameter(torch.empty(1)),
        "layers.0.mlp.gate_up_proj.weight": torch.nn.Parameter(torch.empty(1)),
    }
    if include_quant_scales:
        qkv_scale = AbsMaxFP8MergedParameter(torch.zeros(3), requires_grad=False)
        qkv_scale.output_partition_sizes = [1, 1, 1]
        gate_up_scale = AbsMaxFP8MergedParameter(torch.zeros(2), requires_grad=False)
        gate_up_scale.output_partition_sizes = [1, 1]
        params.update({
            "layers.0.self_attn.qkv_proj.scale_weight": qkv_scale,
            "layers.0.mlp.gate_up_proj.scale_weight": gate_up_scale,
        })
    for param in params.values():
        if not hasattr(param, "weight_loader"):
            param.weight_loader = (lambda target, loaded, *args: target.data.copy_(loaded))

    model = SimpleNamespace(
        config=SimpleNamespace(arch_config=SimpleNamespace(stacked_params_mapping=stacked)),
        named_parameters=lambda: params.items(),
    )
    return model, params, stacked


# The provisional bf16 thresholds below come from the historical unpinned
# 35-layer run. Per add-model-02-parity's calibration block, bf16 max-based
# asserts are meaningless for deep encoders — fused-vs-unfused linear order
# alone produces a long tail, so the test uses mean + median instead. The
# pinned snapshot has 36 blocks and must be rerun before these thresholds are
# treated as current evidence.
#
# Historical empirical numbers (unrecorded snapshot, 35 layers, CUDA):
#   last_hidden_state (post-norm):  mean=0.0152, median=0.0117
#   hidden_states[-2] (pre-norm):   mean=0.0754, median=0.0625
# Thresholds below add ~1.6x headroom over observed.
_BF16_MEAN_DRIFT_POST_NORM = 0.025
_BF16_MEAN_DRIFT_PRE_NORM = 0.120
_BF16_MEDIAN_DRIFT_POST_NORM = 0.020
_BF16_MEDIAN_DRIFT_PRE_NORM = 0.100
_FP32_ATOL = 2e-4
_FP32_RTOL = 1e-4


def _print_diag(label: str, ref: torch.Tensor, fv: torch.Tensor) -> torch.Tensor:
    diff = (ref.float() - fv.float()).abs()
    flat = diff.flatten()
    p99 = flat.kthvalue(max(1, int(0.99 * flat.numel()))).values
    print(
        f"[{label}] max_diff={diff.max():.4f}  mean_diff={diff.mean():.4f}  "
        f"median_diff={diff.median():.4f}  p99_diff={p99:.4f}",
        flush=True,
    )
    return diff


@pytest.mark.parametrize(
    ("cpu_offload", "target_device", "expected_device"),
    [
        (True, torch.device("cuda"), torch.device("cpu")),
        (False, torch.device("meta"), torch.device("meta")),
    ],
)
def test_qwen3_production_loader_avoids_device_context_and_honors_placement(
    monkeypatch: pytest.MonkeyPatch,
    cpu_offload: bool,
    target_device: torch.device,
    expected_device: torch.device,
):
    """HF loading stays ambient while final placement honors target/offload."""
    import fastvideo.platforms as platforms

    placements: list[torch.device] = []
    auto_model_calls: list[dict] = []

    class FakeHFModel:

        def eval(self):
            return self

        def to(self, device: torch.device):
            placements.append(torch.device(device))
            return self

    fake_model = FakeHFModel()

    ambient_device = torch.empty(0).device
    default_devices: list[torch.device] = []

    def fake_auto_model_from_pretrained(*args, **kwargs):
        auto_model_calls.append(kwargs)
        default_devices.append(torch.empty(0).device)
        return fake_model

    def fail_causal_lm_from_pretrained(*args, **kwargs):
        raise AssertionError("Qwen text encoding must not load AutoModelForCausalLM")

    monkeypatch.setattr(
        AutoModel,
        "from_pretrained",
        staticmethod(fake_auto_model_from_pretrained),
    )
    monkeypatch.setattr(
        AutoModelForCausalLM,
        "from_pretrained",
        staticmethod(fail_causal_lm_from_pretrained),
    )
    monkeypatch.setattr(
        platforms,
        "_current_platform",
        SimpleNamespace(
            is_mps=lambda: False,
            verify_model_arch=lambda _arch: None,
        ),
    )
    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
    monkeypatch.setattr(
        "fastvideo.distributed.get_local_torch_device",
        lambda: torch.device("cuda"),
    )

    config = Qwen3TextConfig()
    # The pinned checkpoint names the causal wrapper even though the official
    # Z-Image loader intentionally asks Transformers for the body-only model.
    config.arch_config.architectures = ["Qwen3ForCausalLM"]
    loaded = TextEncoderLoader().load_model(
        "unused-by-mock",
        config,
        target_device,
        _loader_args(cpu_offload=cpu_offload),
        dtype="fp32",
        use_text_encoder_override=True,
    )

    assert loaded is fake_model
    assert auto_model_calls == [{
        "local_files_only": True,
        "dtype": torch.float32,
        "low_cpu_mem_usage": True,
    }]
    assert default_devices == [ambient_device]
    assert placements == [expected_device]
    assert loaded._fastvideo_input_device == expected_device


def test_set_default_torch_dtype_restores_after_exception():
    original = torch.get_default_dtype()

    with pytest.raises(RuntimeError, match="expected failure"):
        with set_default_torch_dtype(torch.bfloat16):
            raise RuntimeError("expected failure")

    assert torch.get_default_dtype() == original


@pytest.mark.parametrize("checkpoint_prefix", ["", "model."])
def test_qwen3_native_load_accepts_already_fused_weights_and_scales(checkpoint_prefix: str, ):
    """Direct FastVideo/state-dict keys satisfy each fused target atomically."""
    fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
    expected = {name: torch.arange(1, param.numel() + 1, dtype=param.dtype) for name, param in params.items()}
    weights = [(f"{checkpoint_prefix}{name}", expected[name]) for name in params]

    loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)

    assert loaded == set(params)
    for name, param in params.items():
        assert torch.equal(param.data, expected[name])
    _assert_native_load_surface(
        fake_model,
        {name
         for name, _ in weights},
        loaded,
    )


def test_qwen3_native_load_rejects_incompatible_fused_quant_scale_shape():
    fake_model, params, _ = _fake_qwen3_weight_target(include_quant_scales=True)
    weights = [(name, torch.ones_like(param)) for name, param in params.items()
               if name != "layers.0.self_attn.qkv_proj.scale_weight"]
    weights.append(("layers.0.self_attn.qkv_proj.scale_weight", torch.ones(1)))

    with pytest.raises(AssertionError, match="Attempted to load weight"):
        Qwen3ForCausalLM.load_weights(fake_model, weights)


def test_qwen3_native_load_uses_real_shard_loader_for_split_quant_scales():
    fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
    weights = [(name, torch.ones_like(param)) for name, param in params.items() if name.endswith(".weight")]
    weights.extend((f"model.layers.0.self_attn{weight_name}.scale_weight", torch.tensor(float(index + 1)))
                   for index, (_, weight_name, _) in enumerate(stacked[:3]))
    weights.extend((f"model.layers.0.mlp{weight_name}.scale_weight", torch.tensor(float(index + 4)))
                   for index, (_, weight_name, _) in enumerate(stacked[3:]))

    loaded = Qwen3ForCausalLM.load_weights(fake_model, weights)

    assert loaded == set(params)
    assert torch.equal(
        params["layers.0.self_attn.qkv_proj.scale_weight"].data,
        torch.tensor([1.0, 2.0, 3.0]),
    )
    assert torch.equal(
        params["layers.0.mlp.gate_up_proj.scale_weight"].data,
        torch.tensor([4.0, 5.0]),
    )


@pytest.mark.parametrize(
    ("missing_weight_name", "missing_shard_id"),
    [
        (".q_proj", "q"),
        (".k_proj", "k"),
        (".v_proj", "v"),
        (".gate_proj", 0),
        (".up_proj", 1),
    ],
)
def test_qwen3_native_load_rejects_missing_fused_weight_shard(
    missing_weight_name: str,
    missing_shard_id: str | int,
):
    fake_model, _, stacked = _fake_qwen3_weight_target()
    weights = [(f"model.layers.0.self_attn{weight_name}.weight", torch.empty(1)) for _, weight_name, _ in stacked[:3]
               if weight_name != missing_weight_name]
    weights.extend((f"model.layers.0.mlp{weight_name}.weight", torch.empty(1)) for _, weight_name, _ in stacked[3:]
                   if weight_name != missing_weight_name)

    with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
        Qwen3ForCausalLM.load_weights(fake_model, weights)

    destination = ("layers.0.self_attn.qkv_proj.weight"
                   if isinstance(missing_shard_id, str) else "layers.0.mlp.gate_up_proj.weight")
    assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)


@pytest.mark.parametrize(
    ("missing_weight_name", "missing_shard_id"),
    [
        (".q_proj", "q"),
        (".k_proj", "k"),
        (".v_proj", "v"),
        (".gate_proj", 0),
        (".up_proj", 1),
    ],
)
def test_qwen3_native_load_rejects_missing_fused_quant_scale_shard(
    missing_weight_name: str,
    missing_shard_id: str | int,
):
    fake_model, params, stacked = _fake_qwen3_weight_target(include_quant_scales=True)
    weights = [(name, torch.empty(1)) for name in params if name.endswith(".weight")]
    weights.extend((f"model.layers.0.self_attn{weight_name}.scale_weight", torch.empty(1))
                   for _, weight_name, _ in stacked[:3] if weight_name != missing_weight_name)
    weights.extend((f"model.layers.0.mlp{weight_name}.scale_weight", torch.empty(1))
                   for _, weight_name, _ in stacked[3:] if weight_name != missing_weight_name)

    with pytest.raises(ValueError, match="Missing required stacked checkpoint shards") as exc_info:
        Qwen3ForCausalLM.load_weights(fake_model, weights)

    destination = ("layers.0.self_attn.qkv_proj.scale_weight"
                   if isinstance(missing_shard_id, str) else "layers.0.mlp.gate_up_proj.scale_weight")
    assert f"{destination}[{missing_shard_id}]" in str(exc_info.value)


@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
@pytest.mark.parametrize(
    "dtype",
    [
        pytest.param(torch.float32, id="fp32"),
        pytest.param(
            torch.bfloat16,
            id="bf16",
            marks=pytest.mark.skipif(
                not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
                reason="bf16 parity requires a bf16-capable CUDA device",
            ),
        ),
    ],
)
def test_zimage_qwen3_encoder_parity_forward(dtype: torch.dtype):
    if not ZIMAGE_TOKENIZER_DIR.exists():
        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.manual_seed(11)

    tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)

    prompts = [
        "A cinematic shot of a rainy neon street at night.",
        "A watercolor illustration of a fox in autumn leaves.",
    ]
    toks = _tokenize_official_prompts(tokenizer, prompts)
    input_ids = toks["input_ids"].to(device)
    attention_mask = toks["attention_mask"].to(device)
    config = _load_qwen3_config()

    ref = AutoModel.from_pretrained(
        str(ZIMAGE_TEXT_ENCODER_DIR),
        local_files_only=True,
        trust_remote_code=True,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).eval().to(device=device)

    with torch.no_grad():
        ref_out = ref(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
        ref_last = ref_out.last_hidden_state.detach().float().cpu()
        assert ref_out.hidden_states is not None
        ref_hs_m2 = ref_out.hidden_states[-2].detach().float().cpu()

    print(f"[mem] after HF ref forward: {_gpu_mem_mb():.0f} MiB", flush=True)
    # Keep only CPU outputs so each full encoder is resident one at a time.
    del ref, ref_out
    _reclaim_vram()
    print(f"[mem] after freeing HF ref:  {_gpu_mem_mb():.0f} MiB", flush=True)

    production = TextEncoderLoader().load_model(
        str(ZIMAGE_TEXT_ENCODER_DIR),
        config,
        device,
        _loader_args(cpu_offload=False),
        dtype=_precision_name(dtype),
        use_text_encoder_override=True,
    )
    assert production.__class__.__name__ == "Qwen3Model"
    assert not hasattr(production, "lm_head")
    assert production._fastvideo_input_device == device

    with torch.no_grad():
        production_out = production(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
        production_last = production_out.last_hidden_state.detach().float().cpu()
        assert production_out.hidden_states is not None
        production_hs_m2 = production_out.hidden_states[-2].detach().float().cpu()

    # Production loading must preserve the independent Transformers oracle
    # exactly; native-kernel drift is assessed separately below.
    assert_close(ref_last, production_last, atol=0.0, rtol=0.0)
    assert_close(ref_hs_m2, production_hs_m2, atol=0.0, rtol=0.0)
    del production, production_out, production_last, production_hs_m2
    _reclaim_vram()

    fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
    fv = fv_cls(config).eval()
    loaded = fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
    assert loaded, "No Qwen3 weights were loaded into FastVideo model"
    fv = fv.to(device=device, dtype=dtype)
    print(f"[mem] after FastVideo load:  {_gpu_mem_mb():.0f} MiB", flush=True)

    _assert_native_load_surface(
        fv,
        _pretrained_safetensor_keys(ZIMAGE_TEXT_ENCODER_DIR),
        loaded,
    )

    with torch.no_grad():
        fv_out = fv(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
        )
        fv_last = fv_out.last_hidden_state.detach().float().cpu()
        assert fv_out.hidden_states is not None
        fv_hs_m2 = fv_out.hidden_states[-2].detach().float().cpu()

    # Z-Image uses only valid token embeddings selected by attention mask.
    # Compare parity on that exact path to avoid padded-token artifacts.
    mask_cpu = attention_mask.detach().bool().cpu()
    is_bf16 = dtype == torch.bfloat16

    for i in range(mask_cpu.shape[0]):
        valid = mask_cpu[i]
        ref_last_v = ref_last[i][valid]
        fv_last_v = fv_last[i][valid]
        ref_hs_m2_v = ref_hs_m2[i][valid]
        fv_hs_m2_v = fv_hs_m2[i][valid]

        last_diff = _print_diag(f"batch{i} last_hidden_state {dtype}", ref_last_v, fv_last_v)
        hs_m2_diff = _print_diag(f"batch{i} hidden_states[-2] {dtype}", ref_hs_m2_v, fv_hs_m2_v)

        if is_bf16:
            # Distribution checks instead of element-wise assert_close: cross-
            # kernel bf16 produces a long max tail (fused vs unfused linears),
            # but a real op bug pushes mean *and* median high. Mean catches
            # systematic drift; median catches "wrong weights / swapped
            # layers" where most elements diverge. Pre-norm tensors get a
            # looser bound because the final RMSNorm compresses drift ~5x.
            assert last_diff.mean().item() < _BF16_MEAN_DRIFT_POST_NORM, (
                f"batch{i} last_hidden_state bf16 mean drift "
                f"{last_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_POST_NORM}")
            assert last_diff.median().item() < _BF16_MEDIAN_DRIFT_POST_NORM, (
                f"batch{i} last_hidden_state bf16 median drift "
                f"{last_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_POST_NORM}")
            assert hs_m2_diff.mean().item() < _BF16_MEAN_DRIFT_PRE_NORM, (
                f"batch{i} hidden_states[-2] bf16 mean drift "
                f"{hs_m2_diff.mean():.4f} >= {_BF16_MEAN_DRIFT_PRE_NORM}")
            assert hs_m2_diff.median().item() < _BF16_MEDIAN_DRIFT_PRE_NORM, (
                f"batch{i} hidden_states[-2] bf16 median drift "
                f"{hs_m2_diff.median():.4f} >= {_BF16_MEDIAN_DRIFT_PRE_NORM}")
        else:
            assert_close(ref_last_v, fv_last_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)
            assert_close(ref_hs_m2_v, fv_hs_m2_v, atol=_FP32_ATOL, rtol=_FP32_RTOL)


@pytest.mark.skipif(not ZIMAGE_TEXT_ENCODER_DIR.exists(), reason="Z-Image text encoder checkpoint required")
@pytest.mark.skipif(
    not (torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
    reason="bf16 per-layer diagnostic requires a bf16-capable CUDA device",
)
def test_zimage_qwen3_encoder_per_layer_bf16_diagnostic():
    """Diagnostic-only: prints per-layer FastVideo-vs-HF drift for the bf16
    forward to distinguish monotonic accumulation (bf16 tail) from a
    single-layer spike (real op mismatch).

    This test always passes; review the captured stdout to interpret.
    Expected pattern for healthy bf16 accumulation: max/mean/median grow
    smoothly with layer depth. Suspect pattern: one layer's diff is
    >>2x the previous, suggesting a real bug at that block.
    """
    if not ZIMAGE_TOKENIZER_DIR.exists():
        pytest.skip(f"Z-Image tokenizer dir not found: {ZIMAGE_TOKENIZER_DIR}")

    device = torch.device("cuda")
    dtype = torch.bfloat16
    torch.manual_seed(11)

    tokenizer = AutoTokenizer.from_pretrained(str(ZIMAGE_TOKENIZER_DIR), local_files_only=True)
    toks = _tokenize_official_prompts(
        tokenizer,
        ["A cinematic shot of a rainy neon street at night."],
    )
    input_ids = toks["input_ids"].to(device)
    attention_mask = toks["attention_mask"].to(device)

    ref = AutoModel.from_pretrained(
        str(ZIMAGE_TEXT_ENCODER_DIR),
        local_files_only=True,
        trust_remote_code=True,
        dtype=dtype,
        low_cpu_mem_usage=True,
    ).eval().to(device=device)

    with torch.no_grad():
        ref_out = ref(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
            return_dict=True,
        )
        # Cache the reference hidden states on CPU, then free the HF model so
        # only one encoder is resident when FastVideo runs (avoids OOM).
        ref_hidden_cpu = [h.detach().float().cpu() for h in ref_out.hidden_states]
    del ref, ref_out
    _reclaim_vram()

    fv_cls, _ = ModelRegistry.resolve_model_cls("Qwen3Model")
    cfg_raw = _load_json(ZIMAGE_TEXT_ENCODER_DIR / "config.json")
    for k in ("_name_or_path", "transformers_version", "model_type", "torch_dtype"):
        cfg_raw.pop(k, None)
    cfg = Qwen3TextConfig()
    cfg.update_model_arch(cfg_raw)
    fv = fv_cls(cfg).eval()
    fv.load_weights(_iter_pretrained_safetensors(ZIMAGE_TEXT_ENCODER_DIR))
    fv = fv.to(device=device, dtype=dtype)

    with torch.no_grad():
        fv_out = fv(
            input_ids=input_ids,
            attention_mask=attention_mask,
            output_hidden_states=True,
        )
        fv_hidden_cpu = [h.detach().float().cpu() for h in fv_out.hidden_states]

    assert fv_out.hidden_states is not None
    assert len(ref_hidden_cpu) == 37, (f"expected embeddings plus 36 block outputs, got {len(ref_hidden_cpu)} entries")
    assert len(ref_hidden_cpu) == len(fv_hidden_cpu), (
        f"len mismatch: HF={len(ref_hidden_cpu)} FV={len(fv_hidden_cpu)}")

    mask_cpu = attention_mask.detach().bool().cpu()[0]
    print(
        f"\n[per-layer bf16 diagnostic] num_hidden_states={len(ref_hidden_cpu)}  "
        f"valid_tokens={mask_cpu.sum().item()}/{mask_cpu.numel()}",
        flush=True,
    )
    print(
        f"{'idx':>3}  {'kind':<10}  {'max':>8}  {'mean':>8}  {'median':>8}  {'p99':>8}",
        flush=True,
    )
    for idx, (ref_h, fv_h) in enumerate(zip(ref_hidden_cpu, fv_hidden_cpu)):
        ref_v = ref_h[0][mask_cpu]
        fv_v = fv_h[0][mask_cpu]
        diff = (ref_v - fv_v).abs()
        flat = diff.flatten()
        p99 = torch.quantile(flat, 0.99).item()
        if idx == 0:
            kind = "embedding"
        elif idx == len(ref_hidden_cpu) - 1:
            kind = "post-norm"
        else:
            kind = f"layer{idx - 1}-out"
        print(
            f"{idx:>3}  {kind:<10}  "
            f"{diff.max().item():>8.4f}  {diff.mean().item():>8.4f}  "
            f"{diff.median().item():>8.4f}  {p99:>8.4f}",
            flush=True,
        )