"""DinoVolumeSceneV4Clip — v4 head with OpenAI CLIP ViT-L/14-336 backbone.

Same v4 head (per-voxel MLP fusion + sin/cos height PE + adaptive per-view
z range), with the DINOv3 backbone replaced by **CLIP ViT-L/14-336** from
OpenAI (``openai/clip-vit-large-patch14-336``). CLIP's vision tower is
the canonical web-image-text-contrastive representation; used here as a
robotics-transfer baseline against DINOv3, PaliGemma, and pi0.

──────────────────────────────────────────────────────────────────────
Architecture deltas vs v4
──────────────────────────────────────────────────────────────────────

1. **Backbone**: ``openai/clip-vit-large-patch14-336`` loaded via
   ``CLIPVisionModel.from_pretrained(...)``. Vision-only tower is used;
   no text encoder is loaded. Output is
   ``.last_hidden_state`` of shape (B, N+1, D) with the CLS token at
   index 0 and (H/14)*(W/14) patches. At 336² input, that's 24×24 = 576
   patch tokens; embed_dim = 1024.

2. **Per-view loop, no backbone-level cross-view attention**. Same as
   the default DINOv3 path.

3. **Per-pixel features**: 1024-d patch tokens → 1×1 conv (feat_head) →
   bilinear-up to ``pred_size`` → 5-layer refine CNN (the standard v4
   pipeline; no DPT, no extra projection).

4. **Global descriptor**: CLIP's CLS token is 1024-d (= patch embed_dim).
   Flows into ``input_proj`` as the per-view CLS.

5. **Image size + normalisation**: CLIP's processor advertises 336×336
   native input and its own OpenAI-CLIP mean/std. Inputs from the PARA
   pipeline arrive ImageNet-normalised at 504×504; we denormalise,
   resize to 336×336, then renormalise with CLIP's stats.

Everything else (per-voxel MLP fusion, sin/cos PE, adaptive per-view z,
per-voxel grip/rot, voxel sampling, loss) is unchanged from v4.
"""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from .model_volumetric_v4 import DinoVolumeSceneV4


# Standard ImageNet stats — the PARA pipeline normalises with these
# before handing tensors to the model.
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)


class DinoVolumeSceneV4Clip(DinoVolumeSceneV4):
    """v4 head + OpenAI CLIP ViT-L/14-336 vision backbone (single-view loop)."""

    def __init__(self, *args,
                 clip_model_name: str = "openai/clip-vit-large-patch14-336",
                 **kwargs):
        super().__init__(*args, **kwargs)
        # v4 already built the DINOv3 backbone; drop it before loading.
        del self.backbone

        from transformers import CLIPVisionModel, CLIPImageProcessor
        print(f"Loading CLIP vision backbone: {clip_model_name}")
        self.clip = CLIPVisionModel.from_pretrained(clip_model_name)
        processor = CLIPImageProcessor.from_pretrained(clip_model_name)

        # Native input size + normalisation from the processor.
        clip_size = 336
        clip_mean = (0.48145466, 0.4578275, 0.40821073)   # OpenAI CLIP defaults
        clip_std = (0.26862954, 0.26130258, 0.27577711)
        try:
            if hasattr(processor, "size") and processor.size:
                sz = processor.size
                if isinstance(sz, dict):
                    clip_size = int(sz.get("height",
                                            sz.get("shortest_edge", clip_size)))
                elif isinstance(sz, int):
                    clip_size = int(sz)
            if getattr(processor, "image_mean", None) is not None:
                clip_mean = tuple(float(x) for x in processor.image_mean)
            if getattr(processor, "image_std", None) is not None:
                clip_std = tuple(float(x) for x in processor.image_std)
        except Exception:
            pass
        self.clip_img_size = int(clip_size)
        self.register_buffer(
            "_clip_mean", torch.tensor(clip_mean).view(1, 3, 1, 1))
        self.register_buffer(
            "_clip_std", torch.tensor(clip_std).view(1, 3, 1, 1))
        self.register_buffer(
            "_imn_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1))
        self.register_buffer(
            "_imn_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1))

        # Resolve embed dim + patch count via a tiny dummy forward.
        with torch.no_grad():
            dummy = torch.zeros(
                1, 3, self.clip_img_size, self.clip_img_size,
                dtype=next(self.clip.parameters()).dtype,
                device=next(self.clip.parameters()).device,
            )
            out = self.clip(pixel_values=dummy)
            hidden = out.last_hidden_state        # (1, 1+N, D)
            self._clip_embed_dim = int(hidden.shape[-1])
            n_tokens = int(hidden.shape[1]) - 1   # drop CLS
            side = int(round(math.sqrt(n_tokens)))
            assert side * side == n_tokens, (
                f"CLIP patch count {n_tokens} is not a perfect square; "
                f"clip_size={self.clip_img_size} may not be a clean multiple "
                f"of the ViT patch size.")
            self._clip_side = side

        print(f"  CLIP: input {self.clip_img_size}², "
              f"embed_dim={self._clip_embed_dim}, "
              f"patch_grid={self._clip_side}², "
              f"mean={clip_mean[:1]}…, std={clip_std[:1]}…")

        # Rebuild head modules sized for CLIP's dims.
        self.embed_dim = self._clip_embed_dim
        feat_dim = self.feat_dim
        d_model = self.d_model
        self.feat_head = nn.Conv2d(self._clip_embed_dim, feat_dim, 1)
        cls_fan = (self.n_views * self.embed_dim
                    if self.cls_fusion == "concat" else self.embed_dim)
        past_fan = self.past_enc_dim if self.past_n > 0 else 0
        self.input_proj = nn.Linear(feat_dim + cls_fan + past_fan, d_model)

        # Rebuild the optional cross-view attention stack at the CLIP
        # patch-token embed_dim.
        if self.cross_view_attn is not None:
            from .model_volumetric_v4 import CrossViewAttnStack
            n_heads = self.cross_view_attn.blocks[0].attn.num_heads
            self.cross_view_attn = CrossViewAttnStack(
                n_layers=self.cross_view_layers,
                embed_dim=self._clip_embed_dim,
                n_heads=n_heads,
            )

    # ────────────────────────────────────────────────────────────────────
    # Per-view backbone forward
    # ────────────────────────────────────────────────────────────────────

    def _to_clip_norm(self, x: torch.Tensor) -> torch.Tensor:
        """Convert (B, 3, S, S) from ImageNet stats → CLIP stats, resized
        to CLIP's native 336² input."""
        x01 = x * self._imn_std + self._imn_mean
        if x01.shape[-1] != self.clip_img_size or x01.shape[-2] != self.clip_img_size:
            x01 = F.interpolate(
                x01, size=(self.clip_img_size, self.clip_img_size),
                mode="bilinear", align_corners=False)
        return (x01 - self._clip_mean) / self._clip_std

    def _per_view_patch_features(self, rgb: list[torch.Tensor]):
        """Per-view CLIP forward.

        Returns:
          view_patches: list of N tensors of shape (B, D, ph, pw)
          view_global:  list of N tensors of shape (B, D) — CLS token
        """
        assert len(rgb) == self.n_views
        view_patches: list[torch.Tensor] = []
        view_global: list[torch.Tensor] = []
        side = self._clip_side
        for k in range(self.n_views):
            x = self._to_clip_norm(rgb[k])
            out = self.clip(pixel_values=x)
            hidden = out.last_hidden_state             # (B, 1+N, D)
            cls_tok = hidden[:, 0, :]                   # (B, D)
            patches = hidden[:, 1:, :]                   # (B, N, D)
            B = patches.shape[0]
            patch_grid = patches.transpose(1, 2).reshape(
                B, self._clip_embed_dim, side, side).contiguous()
            view_patches.append(patch_grid)
            view_global.append(cls_tok)
        return view_patches, view_global

    def patch_features(self, rgb: torch.Tensor):
        raise NotImplementedError(
            "CLIP backbone uses the multi-view entry point. "
            "Use ``_per_view_patch_features(rgb_list)`` instead.")