#!/usr/bin/env python
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""π0.5 parity against LeRobot.

Reference: **lerobot 0.6.1**, transformers 5.5.4, torch 2.11.0+cpu. The oracle
is LeRobot's own code, so a later version changing its π0.5 implementation
fails this without anything here regressing — diff ``policies/pi05/`` first.

Compares final action chunks under fixed noise: float32 exactly, bfloat16
against LeRobot's own bfloat16 layout within ``BF16_ATOL``.

Needs a separate ``lerobot[pi]`` venv (its dependencies conflict with
vllm-omni's) with the vllm-omni ``pi05`` package importable::

    python -m pytest tests/diffusion/models/pi05/test_pi05_parity.py -v -s

Skips itself when LeRobot is absent. ``PI05_PARITY_MODEL_PATH`` points at a
local checkpoint to skip the HF download.
"""

from __future__ import annotations

import copy
import importlib.util
import os

import pytest
import torch

# local_model: needs real weights + a lerobot venv, so it runs locally rather
# than in the ready-CI. Additionally gated on lerobot being importable.
_HAS_LEROBOT = importlib.util.find_spec("lerobot") is not None

# cpu: the comparison runs on CPU in float32 so both implementations see
# identical numerics; local_model: it needs a lerobot venv that CI does not have.
pytestmark = [pytest.mark.local_model, pytest.mark.diffusion, pytest.mark.cpu]


# ─── Config (fixed; matches LeRobot defaults for a π0.5 checkpoint) ────
DEVICE = "cpu"
DTYPE_STR = "float32"
ATOL = 1e-4
BF16_ATOL = 5e-2  # measured 2.95e-2, reproducible; margin is for version drift
NUM_STEPS = 10
BATCH_SIZE = 2
ACTION_DIM = 32
STATE_DIM = 32
ACTION_HORIZON = 50
MAX_TOKEN_LEN = 200  # π0 uses 48
NUM_STATE_BINS = 256

MODEL_PATH = os.environ.get("PI05_PARITY_MODEL_PATH", "lerobot/pi05_base")

CAMERAS = ("base_0_rgb", "left_wrist_0_rgb", "right_wrist_0_rgb")


def _resolve_checkpoint_dir() -> str:
    if os.path.isdir(MODEL_PATH):
        return MODEL_PATH
    from huggingface_hub import snapshot_download

    return snapshot_download(repo_id=MODEL_PATH, repo_type="model")


# ─── Dummy dataset stats ──────────────────────────────────────────────
def _dummy_dataset_stats() -> dict:
    """Identity-ish stats. π0.5 normalizes with QUANTILES by default, so q01/q99
    carry the load here where π0's mean/std would.

    Using dummy stats means this test needs no checkpoint with real
    ``norm_stats`` — it verifies numerical equivalence, not calibration.
    """
    return {
        "observation.state": {
            "mean": torch.zeros(STATE_DIM),
            "std": torch.ones(STATE_DIM),
            "q01": -torch.ones(STATE_DIM),
            "q99": torch.ones(STATE_DIM),
        },
        "action": {
            "mean": torch.zeros(ACTION_DIM),
            "std": torch.ones(ACTION_DIM),
            "q01": -torch.ones(ACTION_DIM),
            "q99": torch.ones(ACTION_DIM),
        },
        "images": {
            cam: {
                "mean": torch.zeros(3, 224, 224),
                "std": torch.ones(3, 224, 224),
                "q01": torch.zeros(3, 224, 224),
                "q99": torch.ones(3, 224, 224),
            }
            for cam in CAMERAS
        },
    }


def _create_dummy_batch(batch_size: int = BATCH_SIZE, num_views: int = 3, device: str = DEVICE) -> dict:
    """Reproducible dummy inputs — identical across both implementations."""
    g = torch.Generator(device="cpu").manual_seed(0)
    prompt = "Pick up the red block and place it in the bin"
    batch = {
        "observation.state": torch.randn(batch_size, STATE_DIM, generator=g, dtype=torch.float32).to(device),
        "action": torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device),
        "task": [prompt for _ in range(batch_size)],
    }
    for cam in CAMERAS[:num_views]:
        batch[f"observation.images.{cam}"] = torch.rand(batch_size, 3, 224, 224, generator=g, dtype=torch.float32).to(
            device
        )
    return batch


# ─── Instantiation ────────────────────────────────────────────────────
def _instantiate_lerobot(use_relative_actions: bool = False, action_feature_names=None, dtype: str = DTYPE_STR):
    from lerobot.policies.pi05 import PI05Policy
    from lerobot.policies.pi05.processor_pi05 import make_pi05_pre_post_processors

    policy = PI05Policy.from_pretrained(MODEL_PATH, strict=True, dtype=dtype)
    policy.to(DEVICE)
    policy.config.device = DEVICE
    policy.config.use_relative_actions = use_relative_actions
    if action_feature_names is not None:
        # Must land on the config *before* the processors are built:
        # make_pi05_pre_post_processors captures action_feature_names into
        # RelativeActionsProcessorStep at construction time, and its _build_mask
        # treats ``action_names is None`` as "shift every dim" — including the
        # gripper that relative_exclude_joints is supposed to hold back. Setting
        # it afterwards leaves the step with the all-True mask and silently
        # compares two different transforms.
        policy.config.action_feature_names = action_feature_names
    policy.eval()

    pre, post = make_pi05_pre_post_processors(config=policy.config, dataset_stats=_dummy_dataset_stats())
    return policy, pre, post


def _instantiate_vllm_omni(
    use_relative_actions: bool = False,
    action_feature_names=None,
    dtype: str = DTYPE_STR,
):
    """Build the vllm-omni π0.5 model in isolation (no pipeline, no engine)."""
    from vllm_omni.diffusion.models.pi05 import Pi05Config, Pi05ForActionPrediction
    from vllm_omni.diffusion.models.pi05.pipeline_pi05 import _set_inference_dtype

    cfg = Pi05Config(
        max_action_dim=ACTION_DIM,
        max_state_dim=STATE_DIM,
        chunk_size=ACTION_HORIZON,
        num_inference_steps=NUM_STEPS,
        tokenizer_max_length=MAX_TOKEN_LEN,
        state_num_bins=NUM_STATE_BINS,
        dtype=dtype,
        use_relative_actions=use_relative_actions,
        relative_exclude_joints=["gripper"] if action_feature_names else [],
        action_feature_names=action_feature_names,
    )
    model = Pi05ForActionPrediction(cfg)
    _set_inference_dtype(model, getattr(torch, dtype))
    model.to(device=DEVICE)
    _load_lerobot_weights(model)
    model.eval()
    return model, cfg


def _load_lerobot_weights(model):
    import safetensors.torch

    path = os.path.join(_resolve_checkpoint_dir(), "model.safetensors")
    state = safetensors.torch.load_file(path)
    filled = model.load_weights(list(state.items()))
    assert filled, "no weights were loaded — the remap rules are broken"


def _extract_lerobot_model_inputs(lerobot_policy, processed_batch):
    """Mimic what ``PI05Policy.predict_action_chunk`` feeds into
    ``self.model.sample_actions``. Using these *exact* tensors for vllm-omni
    means any divergence must come from the core model, not preprocessing.

    Note π0.5 has no ``prepare_state`` step: the state is already inside
    ``lang_tokens``.
    """
    from lerobot.utils.constants import (
        OBS_LANGUAGE_ATTENTION_MASK,
        OBS_LANGUAGE_TOKENS,
    )

    images, img_masks = lerobot_policy._preprocess_images(processed_batch)
    return images, img_masks, processed_batch[OBS_LANGUAGE_TOKENS], processed_batch[OBS_LANGUAGE_ATTENTION_MASK]


def _make_fixed_noise(batch_size: int, device: str) -> torch.Tensor:
    g = torch.Generator(device="cpu").manual_seed(42)
    return torch.randn(batch_size, ACTION_HORIZON, ACTION_DIM, generator=g, dtype=torch.float32).to(device)


# ─── Main parity test ─────────────────────────────────────────────────
@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
@pytest.mark.parametrize("num_views", [1, 2, 3])
def test_pi05_vllm_omni_vs_lerobot(num_views):
    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
    omni_model, _ = _instantiate_vllm_omni()

    raw_batch = _create_dummy_batch(num_views=num_views)
    processed = lerobot_pre(copy.deepcopy(raw_batch))
    images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed)
    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)

    # The input contract: 256 tokens per view + a constant 200 text tokens.
    assert lang_tokens.shape[1] == MAX_TOKEN_LEN, "π0.5 must pad text to exactly 200 tokens"

    with torch.no_grad():
        lerobot_actions = lerobot_policy.model.sample_actions(
            images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
        )
        omni_actions = omni_model.sample_actions(
            images=images,
            image_masks=img_masks,
            lang_tokens=lang_tokens,
            lang_masks=lang_masks,
            noise=noise,
            num_steps=NUM_STEPS,
        )

    diff = (lerobot_actions.float() - omni_actions.float()).abs()
    print(f"[parity] views={num_views} |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}")
    if not torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL):
        _diagnose_divergence(lerobot_policy.model, omni_model, images, img_masks, lang_tokens, lang_masks, noise)
    assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=ATOL), (
        f"actions differ beyond atol={ATOL}; max_diff={diff.max().item():.2e}"
    )


@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
def test_pi05_bfloat16_matches_lerobot_bfloat16():
    """Both sides in bfloat16, so this compares implementations rather than
    measuring bfloat16's drift from float32."""
    # LeRobot applies the layout in PaliGemmaWithExpertModel.__init__ from
    # config.dtype, so the policy must be built as bfloat16 rather than cast after.
    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot(dtype="bfloat16")

    omni_model, _ = _instantiate_vllm_omni(dtype="bfloat16")

    raw_batch = _create_dummy_batch(num_views=3)
    processed = lerobot_pre(copy.deepcopy(raw_batch))
    images, img_masks, lang_tokens, lang_masks = _extract_lerobot_model_inputs(lerobot_policy, processed)
    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)

    with torch.no_grad():
        lerobot_actions = lerobot_policy.model.sample_actions(
            images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
        )
        omni_actions = omni_model.sample_actions(
            images=images,
            image_masks=img_masks,
            lang_tokens=lang_tokens,
            lang_masks=lang_masks,
            noise=noise,
            num_steps=NUM_STEPS,
        )

    diff = (lerobot_actions.float() - omni_actions.float()).abs()
    print(f"[parity] bfloat16 |Δ| max={diff.max().item():.2e} mean={diff.mean().item():.2e}")
    assert torch.allclose(lerobot_actions.float(), omni_actions.float(), atol=BF16_ATOL), (
        f"bfloat16 actions differ beyond atol={BF16_ATOL}; max_diff={diff.max().item():.2e}"
    )


@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
def test_pi05_prompt_parity():
    """π0.5's state is carried as prompt tokens, so the tokenized prompt is
    itself part of the contract. A bin-index mismatch here (e.g. from running
    the discretizer before the normalizer) would otherwise only surface as a
    small, plausible-looking action difference.
    """
    from transformers import AutoTokenizer

    from vllm_omni.diffusion.models.pi05.processor_pi05 import (
        apply_norm,
        build_norm_stats,
        build_pi05_prompt,
        tokenize_prompt,
    )

    lerobot_policy, lerobot_pre, _ = _instantiate_lerobot()
    raw_batch = _create_dummy_batch(batch_size=1)
    processed = lerobot_pre(copy.deepcopy(raw_batch))

    from lerobot.utils.constants import OBS_LANGUAGE_TOKENS

    lerobot_tokens = processed[OBS_LANGUAGE_TOKENS][0]

    stats = _dummy_dataset_stats()["observation.state"]
    norm = build_norm_stats({"state": {"q01": stats["q01"], "q99": stats["q99"]}}, "state")
    normalized = apply_norm(raw_batch["observation.state"][0].float(), norm)
    prompt = build_pi05_prompt(
        task=raw_batch["task"][0],
        normalized_state=normalized.numpy(),
        state_num_bins=NUM_STATE_BINS,
    )
    tokenizer = AutoTokenizer.from_pretrained("google/paligemma-3b-pt-224", padding_side="right")
    ids, _ = tokenize_prompt(tokenizer, prompt, MAX_TOKEN_LEN)

    assert len(ids) == MAX_TOKEN_LEN
    assert torch.equal(torch.tensor(ids), lerobot_tokens.cpu()), "π0.5 prompt tokens diverge from LeRobot"


@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
@pytest.mark.parametrize("state_dim", [7, STATE_DIM])
def test_pi05_prompt_parity_across_state_widths(state_dim):
    """``pi05_base`` declares a 32-dim state, which is exactly ``max_state_dim``,
    so the full-policy parity above cannot tell padding apart from the real
    width. This drives LeRobot's tokenizer step directly, at both widths.
    """
    from lerobot.lerobot_types import TransitionKey
    from lerobot.policies.pi05.processor_pi05 import Pi05PrepareStateTokenizerProcessorStep
    from lerobot.utils.constants import OBS_STATE

    from vllm_omni.diffusion.models.pi05.processor_pi05 import build_pi05_prompt

    g = torch.Generator().manual_seed(state_dim)
    state = torch.randn(state_dim, generator=g, dtype=torch.float32)
    task = "pick up the red block"

    step = Pi05PrepareStateTokenizerProcessorStep(max_state_dim=STATE_DIM)
    transition = step(
        {
            TransitionKey.OBSERVATION: {OBS_STATE: state.unsqueeze(0)},
            TransitionKey.COMPLEMENTARY_DATA: {"task": [task]},
        }
    )
    lerobot_prompt = transition[TransitionKey.COMPLEMENTARY_DATA]["task"][0]

    prompt = build_pi05_prompt(task=task, normalized_state=state.numpy())

    assert prompt == lerobot_prompt
    assert len(prompt.split("State: ")[1].split(";")[0].split()) == state_dim


@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
@pytest.mark.parametrize(
    "lerobot_mode,stats,ours,state,state_dim",
    [
        ("MEAN_STD", {"mean": [0.0], "std": [1e-7]}, {"mode": "mean_std", "mean": [0.0], "std": [1e-7]}, [1e-7], 1),
        ("QUANTILES", {"q01": [0.0], "q99": [1e-7]}, {"mode": "quantile", "q01": [0.0], "q99": [1e-7]}, [1e-7], 1),
        ("QUANTILES", {"q01": [2.0], "q99": [2.0]}, {"mode": "quantile", "q01": [2.0], "q99": [2.0]}, [2.0], 1),
        (
            "QUANTILES",
            {"q01": [0.0] * 4, "q99": [10.0] * 4},
            {"mode": "quantile", "q01": [0.0] * 4, "q99": [10.0] * 4},
            [-10.0, 0.0, 5.0, 20.0],
            4,
        ),
        (
            "MEAN_STD",
            {"mean": [1.0, 2.0], "std": [0.5, 4.0]},
            {"mode": "mean_std", "mean": [1.0, 2.0], "std": [0.5, 4.0]},
            [3.0, -6.0],
            2,
        ),
        ("MIN_MAX", {"min": [-1.0], "max": [3.0]}, {"mode": "min_max", "min": [-1.0], "max": [3.0]}, [2.0], 1),
    ],
)
def test_pi05_state_normalization_parity(lerobot_mode, stats, ours, state, state_dim):
    """State normalization against LeRobot's ``NormalizerProcessorStep``.

    The full-policy parity test feeds both sides stats with a wide range, so it
    cannot see the denominator rule. A wrongly normalized state still discretizes
    into a well-formed prompt: for ``q01=0, q99=1e-7`` a threshold substitution
    moves the state token from bin 255 to bin 0.
    """
    from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
    from lerobot.processor.normalize_processor import NormalizerProcessorStep
    from lerobot.utils.constants import OBS_STATE

    from vllm_omni.diffusion.models.pi05.processor_pi05 import apply_norm, build_norm_stats

    step = NormalizerProcessorStep(
        features={OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,))},
        norm_map={FeatureType.STATE: getattr(NormalizationMode, lerobot_mode)},
        stats={OBS_STATE: {k: torch.tensor(v, dtype=torch.float32) for k, v in stats.items()}},
    )
    expected = step._apply_transform(
        torch.tensor(state, dtype=torch.float32), OBS_STATE, FeatureType.STATE, inverse=False
    )

    actual = apply_norm(torch.tensor(state, dtype=torch.float32), build_norm_stats({"state": ours}, "state"))

    assert torch.allclose(actual, expected, atol=1e-6)


@pytest.mark.skipif(not _HAS_LEROBOT, reason="lerobot not installed (run in a lerobot venv).")
def test_pi05_relative_actions_parity():
    """With ``use_relative_actions=True`` the checkpoint's norm_stats live in
    relative space, so the post-pipeline must add the raw state back. Compares
    the full pre → model → post path against LeRobot's."""
    from lerobot.utils.constants import (
        OBS_LANGUAGE_ATTENTION_MASK,
        OBS_LANGUAGE_TOKENS,
    )

    from vllm_omni.diffusion.models.pi05.processor_pi05 import (
        Pi05RelativeActions,
        apply_norm,
        build_norm_stats,
    )

    action_names = [f"joint_{i}" for i in range(ACTION_DIM - 1)] + ["gripper"]
    lerobot_policy, lerobot_pre, lerobot_post = _instantiate_lerobot(
        use_relative_actions=True, action_feature_names=action_names
    )
    omni_model, cfg = _instantiate_vllm_omni(use_relative_actions=True, action_feature_names=action_names)

    raw_batch = _create_dummy_batch()
    processed = lerobot_pre(copy.deepcopy(raw_batch))
    images, img_masks = lerobot_policy._preprocess_images(processed)
    lang_tokens = processed[OBS_LANGUAGE_TOKENS]
    lang_masks = processed[OBS_LANGUAGE_ATTENTION_MASK]
    noise = _make_fixed_noise(raw_batch["observation.state"].shape[0], DEVICE)

    with torch.no_grad():
        lerobot_actions = lerobot_post(
            lerobot_policy.model.sample_actions(
                images, img_masks, lang_tokens, lang_masks, noise=noise, num_steps=NUM_STEPS
            )
        )
        omni_raw = omni_model.sample_actions(
            images=images,
            image_masks=img_masks,
            lang_tokens=lang_tokens,
            lang_masks=lang_masks,
            noise=noise,
            num_steps=NUM_STEPS,
        )

    rel = Pi05RelativeActions(
        enabled=True,
        exclude_joints=cfg.relative_exclude_joints,
        action_names=action_names,
        max_action_dim=ACTION_DIM,
    )
    omni_actions = apply_norm(omni_raw, build_norm_stats(cfg.norm_stats, "action"), inverse=True)
    # Per-sample reference state, matching LeRobot: its
    # RelativeActionsProcessorStep caches the whole (B, D) state on the input
    # side and shifts each sample by its own row. Passing a single row here
    # would shift every sample by sample 0's state, which only agrees for B=1.
    omni_actions = rel.to_absolute(omni_actions, raw_batch["observation.state"])

    lerobot_actions = torch.as_tensor(lerobot_actions).float()
    assert torch.allclose(lerobot_actions, omni_actions.float(), atol=ATOL), (
        "relative-action path diverges from LeRobot; "
        f"max_diff={(lerobot_actions - omni_actions.float()).abs().max().item():.2e}"
    )


# ─── Per-stage divergence diagnostics ─────────────────────────────────
@torch.no_grad()
def _diagnose_divergence(lerobot_flow_model, omni_model, images, img_masks, lang_tokens, lang_masks, noise):
    """Localize a numerical mismatch to a specific stage:
    1. Prefix embeddings (SigLIP / embed_tokens / projector) — image vs lang.
    2. Prefix KV cache layer 0 — PaliGemma LM attention.
    3. The AdaRMS timestep conditioning vector.
    4. A single denoise_step velocity at t=1.0 — action expert forward.
    """
    from vllm_omni.diffusion.models.pi05.modeling_pi05 import (
        make_att_2d_masks,
        prepare_attention_masks_4d,
    )

    lr_embs, lr_pad, lr_att = lerobot_flow_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
    sg_embs, sg_pad, sg_att = omni_model.embed_prefix(images, img_masks, lang_tokens, lang_masks)
    print(f"[diag] prefix_embs max|Δ| = {(lr_embs.float() - sg_embs.float()).abs().max().item():.2e}")
    print(f"[diag] prefix_pad_masks equal: {torch.equal(lr_pad, sg_pad)}")

    img_len = 256 * len(images)
    print(f"[diag]   image slice max|Δ| = {(lr_embs[:, :img_len] - sg_embs[:, :img_len]).abs().max().item():.2e}")
    print(f"[diag]   lang  slice max|Δ| = {(lr_embs[:, img_len:] - sg_embs[:, img_len:]).abs().max().item():.2e}")

    prefix_att_4d = prepare_attention_masks_4d(make_att_2d_masks(sg_pad, sg_att))
    prefix_pos = torch.cumsum(sg_pad, dim=1) - 1
    _, sg_kv = omni_model.paligemma_with_expert.forward(
        attention_mask=prefix_att_4d,
        position_ids=prefix_pos,
        past_key_values=None,
        inputs_embeds=[sg_embs, None],
        use_cache=True,
    )

    t = torch.ones(noise.shape[0], dtype=torch.float32)
    sg_cond = omni_model.embed_timestep(t)
    print(f"[diag] adarms cond: shape={tuple(sg_cond.shape)} mean={sg_cond.mean().item():+.4f}")

    sg_vt = omni_model.denoise_step(sg_pad, sg_kv, noise, t)
    print(f"[diag] denoise_step(t=1) v_t mean={sg_vt.mean().item():+.6f}")


if __name__ == "__main__":
    test_pi05_vllm_omni_vs_lerobot(3)