"""DinoVolumeSceneV4Pi0 — v4 head + pi0's PaliGemma backbone.

Mirrors :mod:`lib.model_volumetric_v4_paligemma` but the PaliGemma weights
are NOT google's web-pretrained version. They come from Physical
Intelligence's **pi0** model, ported to PyTorch by HuggingFace LeRobot
(``lerobot/pi0_base``). pi0 = PaliGemma vision-language stack + a
flow-matching action expert, jointly pretrained on a large robot data
mixture. We discard the action expert and use ONLY the PaliGemma
sub-weights — the scientific question is whether robot-data-pretrained
vision/language transfers better to YAM than internet-pretrained
PaliGemma.

We bypass lerobot's runtime entirely. lerobot/pi0 requires a forked
``transformers`` (custom ``siglip.check`` submodule) at init; vanilla
transformers fails. Instead we:

1. Load a blank ``PaliGemmaForConditionalGeneration`` from
   ``google/paligemma-3b-pt-448`` (vanilla transformers, same
   architecture).
2. Open lerobot/pi0_base's ``model.safetensors`` with the safetensors
   reader.
3. Filter tensors whose names start with
   ``model.paligemma_with_expert.paligemma.`` and strip that prefix —
   matching transformers' own ``PaliGemmaForConditionalGeneration``
   state-dict key naming.
4. ``self.vlm.load_state_dict(filtered, strict=False)`` — strict=False
   tolerates pi0's extra expert keys staying ignored, and any minor
   buffer-shape differences. We print a sanity report of matched /
   missing / unexpected key counts so regressions surface loudly.

Everything else (vision-token attention, multi-modal projector training,
LLM frozen-by-default, head sizing to Gemma hidden_dim) inherits
unchanged from :class:`DinoVolumeSceneV4PaliGemma`.
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional

import torch

from .model_volumetric_v4_paligemma import DinoVolumeSceneV4PaliGemma


# Prefix lerobot/pi0_base safetensors uses for the PaliGemma sub-state.
# Verified empirically (777 total keys; 602 under this prefix). pi0's
# PyTorch port wraps everything under
# ``paligemma_with_expert.paligemma.model.*`` for the VLM half (the
# action-expert side is under ``paligemma_with_expert.gemma_expert.model.*``
# and gets dropped here). The trailing ``.model.`` is a lerobot-specific
# wrap on top of transformers' PaliGemmaForConditionalGeneration, whose
# OWN state_dict starts with ``vision_tower.*`` / ``language_model.*``.
_PI0_PALIGEMMA_PREFIX = "paligemma_with_expert.paligemma.model."


class DinoVolumeSceneV4Pi0(DinoVolumeSceneV4PaliGemma):
    """v4 + pi0's pretrained PaliGemma backbone (LeRobot port).

    Architectural note: pi0 was trained on the **224²** variant of
    PaliGemma, not the 448². The vision-tower position-embedding table
    has 256 (=16²) entries, so the base model MUST be
    ``google/paligemma-3b-pt-224``. Using the 448 variant raises a
    size-mismatch during ``load_state_dict``.
    """

    def __init__(self, *args,
                 paligemma_model_name: str = "google/paligemma-3b-pt-224",
                 pi0_safetensors_path: Optional[str] = None,
                 pi0_repo_id: str = "lerobot/pi0_base",
                 **kwargs):
        # Standard PaliGemma init (downloads google weights → builds
        # the architecture + tokenizer + head sizes everything to Gemma
        # hidden_dim). The vlm state-dict gets overwritten right after.
        super().__init__(*args,
                         paligemma_model_name=paligemma_model_name,
                         **kwargs)

        # Resolve safetensors path (default = pull from HF cache).
        if pi0_safetensors_path is None:
            from huggingface_hub import hf_hub_download
            pi0_safetensors_path = hf_hub_download(
                repo_id=pi0_repo_id, filename="model.safetensors")
        pi0_path = Path(pi0_safetensors_path)
        if not pi0_path.is_file():
            raise FileNotFoundError(
                f"pi0 safetensors not found: {pi0_path}")
        print(f"  pi0 weights ← {pi0_path}")

        # safetensors header read is mmap'd — no full-file load until we
        # explicitly get_tensor each one. Filter to the paligemma subset
        # AND remap language_model.* → language_model.model.* (lerobot
        # flattened the GemmaForCausalLM .model.* wrap when re-serialising
        # pi0, but vanilla transformers PaliGemma keeps it).
        from safetensors import safe_open
        filtered: dict[str, torch.Tensor] = {}
        skipped = 0
        with safe_open(str(pi0_path), framework="pt") as f:
            for k in f.keys():
                if not k.startswith(_PI0_PALIGEMMA_PREFIX):
                    skipped += 1
                    continue
                new_k = k[len(_PI0_PALIGEMMA_PREFIX):]
                if (new_k.startswith("language_model.")
                        and not new_k.startswith("language_model.model.")):
                    new_k = "language_model.model." + new_k[len("language_model."):]
                t = f.get_tensor(k)
                filtered[new_k] = t.to(torch.bfloat16)

        print(f"  pi0 paligemma keys: kept={len(filtered)} "
              f"skipped(non-paligemma)={skipped}")
        if not filtered:
            raise RuntimeError(
                f"No keys matched prefix {_PI0_PALIGEMMA_PREFIX!r} in "
                f"{pi0_path}. Verify pi0 safetensors structure.")

        # Overwrite the vanilla google paligemma weights with pi0's.
        report = self.vlm.load_state_dict(filtered, strict=False)
        n_match = len(filtered) - len(report.unexpected_keys)
        print(f"  vlm.load_state_dict: matched={n_match}/"
              f"{len(filtered)} pi0 keys; "
              f"missing(in-model-not-in-pi0)={len(report.missing_keys)} "
              f"unexpected(in-pi0-not-in-model)="
              f"{len(report.unexpected_keys)}")
        if report.missing_keys:
            example = ", ".join(report.missing_keys[:3])
            print(f"    missing first 3: {example}")
        if report.unexpected_keys:
            example = ", ".join(report.unexpected_keys[:3])
            print(f"    unexpected first 3: {example}")
        # Hard-fail if no overlap — usually means the prefix changed.
        if n_match == 0:
            raise RuntimeError(
                "pi0 weights had 0 matching keys after prefix-strip. "
                "Update _PI0_PALIGEMMA_PREFIX after inspecting the "
                "safetensors header.")