"""
OpenVLA-7B language grounding module.

In this stack OpenVLA plays two roles:

  1. Task embedding (encode):
       instruction + current frame  →  (embedding_dim,) task vector
       This conditions the Cosmos Policy on what the robot should do.

  2. Direct action prediction (predict_action):
       instruction + frame  →  (action_dim,) immediate action
       Used as a high-level fallback or for language-specified primitives.

OpenVLA (openvla/openvla-7b) is a 7B VLA model built on a
Prismatic VLM (SigLIP + DINOv2 vision, Llama-2 language backbone).
It was trained on Open X-Embodiment with 7-DoF manipulator actions;
you will need to fine-tune the action head for a 21-DoF humanoid.

If language.enabled is false the module returns a constant embedding
so the policy operates task-agnostic.
"""
from __future__ import annotations

import logging

import numpy as np
from PIL import Image

from models.base import BaseLanguageModel

log = logging.getLogger(__name__)


# ─────────────────────────────────────────────── instruction grounding (shared)

_VERBS = {
    "pick":  ("pick", "grab", "grasp", "take", "lift", "get", "retrieve", "hold"),
    "place": ("place", "put", "drop", "pack", "insert", "into the box", "in the box"),
    "wave":  ("wave", "hello", "hi", "greet"),
    "raise": ("raise", "arms up", "hands up", "overhead"),
    "look":  ("look", "head", "face"),
    "rest":  ("rest", "stand", "idle", "ready", "stop", "wait", "home"),
}
_COLORS = ("red", "blue", "green", "yellow")


def ground_instruction(text: str) -> dict:
    """
    Parse a natural-language instruction into a structured grounding:
        {"verb": <skill>, "target": <colour|None>, "raw": <text>}

    This is the language→symbol step of the VLA loop. The policy then binds
    `target` to a PERCEIVED object (by colour) to get a live 3-D goal — so
    nothing about the object's location is hardcoded. (The real OpenVLA model
    overrides this with model-based grounding; the stub uses keyword rules,
    matching the rx1_brain reference's classify_action / match_object_by_text.)
    """
    t = (text or "").lower()
    verb = "rest"
    for v, kws in _VERBS.items():
        if any(k in t for k in kws):
            verb = v
            break
    target = next((col for col in _COLORS if col in t), None)
    # "pick"/"place" need an object; if a colour is named, that's the target.
    return {"verb": verb, "target": target, "raw": text}


# ─────────────────────────────────────────────────────────────── stub

class _StubLanguageModel(BaseLanguageModel):
    def __init__(self, embedding_dim: int, rng: np.random.Generator):
        self._dim = embedding_dim
        self._rng = rng
        self._cache: dict[str, np.ndarray] = {}

    def encode(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        # Cache so the same instruction returns the same embedding.
        if instruction not in self._cache:
            emb = self._rng.normal(0, 1, size=self._dim).astype(np.float32)
            emb /= np.linalg.norm(emb) + 1e-8
            self._cache[instruction] = emb
            log.debug("Stub task embedding for: '%s'", instruction)
        return self._cache[instruction]

    def predict_action(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        return np.zeros(21, dtype=np.float32)

    def ground(self, instruction: str) -> dict:
        return ground_instruction(instruction)


class _DisabledLanguageModel(BaseLanguageModel):
    """Used when language.enabled is false."""

    def __init__(self, embedding_dim: int):
        self._emb = np.zeros(embedding_dim, dtype=np.float32)

    def encode(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        return self._emb

    def predict_action(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        return np.zeros(21, dtype=np.float32)

    def ground(self, instruction: str) -> dict:
        return ground_instruction(instruction)


# ─────────────────────────────────────────────────────────────── real

class _RealLanguageModel(BaseLanguageModel):
    """
    OpenVLA-7B via HuggingFace transformers.

    Usage note — action-head adaptation:
      OpenVLA was trained for 7-DoF manipulator actions.
      For a 21-DoF humanoid you need to either:
        a) use the model only for task embeddings (not direct actions), OR
        b) fine-tune the output projection layer on humanoid demonstrations.

      The `predict_action` method here adapts the 7-DoF output to 21-DoF
      by zero-padding; replace with a properly fine-tuned head for real use.
    """

    ACTION_DIM_OPENVLA = 7

    def __init__(self, model_id: str, device: str, dtype_str: str,
                 embedding_dim: int, robot_action_dim: int = 21):
        import torch
        from transformers import AutoModelForVision2Seq, AutoProcessor

        self._device = device
        self._embedding_dim = embedding_dim
        self._robot_action_dim = robot_action_dim

        # 4-bit quantization path (dtype="int4") — fits OpenVLA-7B in ~4 GB VRAM.
        # Requires: pip install bitsandbytes
        use_4bit = (dtype_str == "int4")
        if use_4bit:
            from transformers import BitsAndBytesConfig
            quant_cfg = BitsAndBytesConfig(load_in_4bit=True,
                                           bnb_4bit_compute_dtype=torch.bfloat16)
            self._dtype = torch.bfloat16
        else:
            quant_cfg = None
            self._dtype = getattr(torch, dtype_str, torch.bfloat16)

        log.info("Loading language model %s%s ...",
                 model_id, " (4-bit)" if use_4bit else "")
        self._processor = AutoProcessor.from_pretrained(
            model_id, trust_remote_code=True
        )
        self._model = AutoModelForVision2Seq.from_pretrained(
            model_id,
            attn_implementation="flash_attention_2",
            torch_dtype=self._dtype,
            quantization_config=quant_cfg,
            low_cpu_mem_usage=True,
            trust_remote_code=True,
        )
        if not use_4bit:
            self._model = self._model.to(device)
        self._model.eval()
        log.info("Language model loaded.")

    def encode(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        """Extract the final decoder hidden state as a task embedding."""
        import torch

        prompt = f"In: What robot action should be performed to {instruction}?\nOut:"
        pil = Image.fromarray(frame)
        inputs = self._processor(prompt, pil, return_tensors="pt").to(
            self._device, dtype=self._dtype
        )
        with torch.no_grad():
            out = self._model(**inputs, output_hidden_states=True)
        # Last token of the last hidden layer → project to embedding_dim
        hidden = out.hidden_states[-1][:, -1, :].float().squeeze(0).cpu().numpy()
        emb = hidden[: self._embedding_dim]
        if emb.shape[0] < self._embedding_dim:
            emb = np.pad(emb, (0, self._embedding_dim - emb.shape[0]))
        emb /= np.linalg.norm(emb) + 1e-8
        return emb.astype(np.float32)

    def predict_action(self, instruction: str, frame: np.ndarray) -> np.ndarray:
        """Run the VLA's action head; zero-pads to 21-DoF."""
        import torch

        prompt = f"In: What action should the robot take to {instruction}?\nOut:"
        pil = Image.fromarray(frame)
        inputs = self._processor(prompt, pil, return_tensors="pt").to(
            self._device, dtype=self._dtype
        )
        with torch.no_grad():
            action_tokens = self._model.predict_action(**inputs)

        # action_tokens is (7,) for the original OpenVLA
        raw = np.array(action_tokens, dtype=np.float32).flatten()
        action = np.zeros(self._robot_action_dim, dtype=np.float32)
        action[: min(len(raw), self._robot_action_dim)] = raw[: self._robot_action_dim]
        return action

    def ground(self, instruction: str) -> dict:
        # The 7B VLA could ground via generation; keyword rules are a safe,
        # deterministic fallback that doesn't require a decode pass per task.
        return ground_instruction(instruction)


# ─────────────────────────────────────────────────────────────── factory

class OpenVLALanguageModel:
    def __new__(cls, config: dict) -> BaseLanguageModel:  # type: ignore[misc]
        if not config.get("enabled", True):
            log.info("LanguageModel: disabled — using zero embeddings")
            return _DisabledLanguageModel(config.get("embedding_dim", 512))

        if config.get("stub", True):
            log.info("LanguageModel: running in STUB mode")
            rng = np.random.default_rng(seed=13)
            return _StubLanguageModel(config.get("embedding_dim", 512), rng)

        return _RealLanguageModel(
            model_id=config["model_id"],
            device=config.get("device", "cuda"),
            dtype_str=config.get("dtype", "bfloat16"),
            embedding_dim=config.get("embedding_dim", 512),
        )