#!/usr/bin/env python3
"""Minimal OpenVLA-style overfitting experiment on a 0.5B Qwen VLM.

This intentionally small experiment uses the Hugging Face checkpoint
``lmms-lab/llava-onevision-qwen2-0.5b-ov`` and the local LLaVA-OneVision
implementation under ``qwen2-openvla``.  It demonstrates the core VLA loop:

    image + language instruction -> autoregressive action tokens -> 7-D action

The script generates a few deterministic synthetic samples, discretizes actions
into 256 bins, fine-tunes a small subset of the VLM, and reports teacher-forced
action-token accuracy plus constrained autoregressive predictions.

Authentication is read from HF_TOKEN/HUGGING_FACE_HUB_TOKEN.  Never put a token
in this file.
"""

from __future__ import annotations

import argparse
import itertools
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable

import numpy as np
import torch
from PIL import Image, ImageDraw


ROOT = Path(__file__).resolve().parent
LLAVA_REPO = ROOT / "qwen2-openvla"
if not LLAVA_REPO.exists():
    raise FileNotFoundError(
        f"Expected the LLaVA implementation at {LLAVA_REPO}. "
        "Clone https://github.com/IP127000/openVLA-Qwen2-0.5B there first."
    )
sys.path.insert(0, str(LLAVA_REPO))

from llava.constants import DEFAULT_IMAGE_TOKEN, IGNORE_INDEX, IMAGE_TOKEN_INDEX  # noqa: E402
from llava.mm_utils import get_model_name_from_path, process_images, tokenizer_image_token  # noqa: E402
from llava.model.builder import load_pretrained_model  # noqa: E402


@dataclass(frozen=True)
class ToySample:
    image: Image.Image
    instruction: str
    action: np.ndarray


class ActionTokenizer:
    """OpenVLA-style uniform discretization over normalized actions in [-1, 1]."""

    def __init__(self, tokenizer, bins: int = 256) -> None:
        self.tokenizer = tokenizer
        self.n_bins = bins
        self.edges = np.linspace(-1.0, 1.0, bins + 1, dtype=np.float32)
        self.centers = (self.edges[:-1] + self.edges[1:]) / 2.0
        self.tokens = [f"<ACTION_{i}>" for i in range(bins)]

        # The local model loader currently creates <ACTION_0>...<ACTION_270>.
        # add_tokens is idempotent, so this also works if that behavior changes.
        tokenizer.add_tokens(self.tokens)
        self.token_ids = tokenizer.convert_tokens_to_ids(self.tokens)
        if len(set(self.token_ids)) != bins:
            raise RuntimeError("Action tokens were not assigned distinct tokenizer IDs")

    def encode(self, action: np.ndarray) -> list[int]:
        action = np.asarray(action, dtype=np.float32)
        indices = np.digitize(np.clip(action, -1.0, 1.0), self.edges) - 1
        indices = np.clip(indices, 0, self.n_bins - 1)
        return [self.token_ids[int(i)] for i in indices]

    def decode(self, token_ids: Iterable[int]) -> np.ndarray:
        id_to_bin = {token_id: i for i, token_id in enumerate(self.token_ids)}
        try:
            indices = [id_to_bin[int(token_id)] for token_id in token_ids]
        except KeyError as exc:
            raise ValueError(f"Generated non-action token ID: {exc.args[0]}") from exc
        return self.centers[np.asarray(indices)]


def make_toy_samples() -> list[ToySample]:
    specs = [
        ((220, 50, 50), (45, 65), "move to the red square", [-0.75, -0.50, 0.20, 0.0, 0.0, 0.0, 1.0]),
        ((50, 190, 70), (130, 70), "move to the green square", [0.10, -0.45, 0.35, 0.0, 0.0, 0.0, -1.0]),
        ((55, 90, 220), (185, 175), "move to the blue square", [0.65, 0.60, -0.15, 0.0, 0.0, 0.0, 1.0]),
    ]
    samples: list[ToySample] = []
    for color, center, instruction, action in specs:
        image = Image.new("RGB", (224, 224), "white")
        draw = ImageDraw.Draw(image)
        x, y = center
        draw.rectangle((x - 24, y - 24, x + 24, y + 24), fill=color)
        samples.append(ToySample(image, instruction, np.asarray(action, dtype=np.float32)))
    return samples


def build_prompt(tokenizer, instruction: str) -> str:
    user_text = f"{DEFAULT_IMAGE_TOKEN}\nWhat action should the robot take to {instruction}?"
    if getattr(tokenizer, "chat_template", None):
        return tokenizer.apply_chat_template(
            [{"role": "user", "content": user_text}],
            tokenize=False,
            add_generation_prompt=True,
        )
    return user_text + "\nASSISTANT:"


def prepare_example(tokenizer, action_tokenizer: ActionTokenizer, sample: ToySample, device: torch.device):
    prompt = build_prompt(tokenizer, sample.instruction)
    prompt_ids = tokenizer_image_token(
        prompt, tokenizer, image_token_index=IMAGE_TOKEN_INDEX, return_tensors="pt"
    )
    action_ids = torch.tensor(action_tokenizer.encode(sample.action), dtype=torch.long)
    input_ids = torch.cat([prompt_ids, action_ids]).unsqueeze(0).to(device)
    labels = torch.full_like(input_ids, IGNORE_INDEX)
    labels[:, -len(action_ids) :] = action_ids.to(device)
    return input_ids, labels


def prepare_image(image, image_processor, model, device: torch.device):
    pixels = process_images([image], image_processor, model.config)
    dtype = next(model.parameters()).dtype
    if isinstance(pixels, torch.Tensor):
        pixels = [item.to(device=device, dtype=dtype) for item in pixels]
    else:
        pixels = [item.to(device=device, dtype=dtype) for item in pixels]
    return pixels, [image.size]


def set_trainable_subset(model) -> list[torch.nn.Parameter]:
    """Train projector, token embeddings/head, and the final Qwen block."""
    model.requires_grad_(False)
    modules = [model.get_model().mm_projector, model.get_model().embed_tokens, model.lm_head]
    layers = getattr(model.get_model(), "layers", None)
    if layers is not None and len(layers):
        modules.append(layers[-1])
    for module in modules:
        module.requires_grad_(True)
    return [parameter for parameter in model.parameters() if parameter.requires_grad]


@torch.no_grad()
def teacher_forced_metrics(model, tokenizer, action_tokenizer, samples, image_processor, device):
    model.eval()
    losses, correct, total = [], 0, 0
    for sample in samples:
        input_ids, labels = prepare_example(tokenizer, action_tokenizer, sample, device)
        images, image_sizes = prepare_image(sample.image, image_processor, model, device)
        output = model(
            input_ids=input_ids,
            labels=labels,
            images=images,
            image_sizes=image_sizes,
            modalities=["image"],
        )
        losses.append(float(output.loss))
        # Image patch embeddings expand the multimodal sequence, so output
        # logits are longer than the original text labels. Action tokens remain
        # the final tokens; the preceding logits predict those actions.
        action_count = len(sample.action)
        action_logits = output.logits[:, -(action_count + 1) : -1]
        action_labels = labels[:, -action_count:]
        predicted = action_logits.argmax(dim=-1)
        correct += int((predicted == action_labels).sum())
        total += action_count
    return float(np.mean(losses)), correct / max(total, 1)


@torch.no_grad()
def generate_action(model, tokenizer, action_tokenizer, sample, image_processor, device):
    input_ids = tokenizer_image_token(
        build_prompt(tokenizer, sample.instruction),
        tokenizer,
        image_token_index=IMAGE_TOKEN_INDEX,
        return_tensors="pt",
    ).unsqueeze(0).to(device)
    allowed = torch.tensor(action_tokenizer.token_ids, device=device)
    generated_ids: list[int] = []
    # Manual greedy decoding avoids a LLaVA-OneVision generation incompatibility
    # where prefix_allowed_tokens_fn receives an empty synthetic input_ids tensor
    # after the wrapper converts multimodal inputs to inputs_embeds.
    for _ in range(len(sample.action)):
        images, image_sizes = prepare_image(sample.image, image_processor, model, device)
        output = model(
            input_ids=input_ids,
            images=images,
            image_sizes=image_sizes,
            modalities=["image"],
            use_cache=False,
        )
        allowed_logits = output.logits[0, -1, allowed]
        next_id = int(allowed[allowed_logits.argmax()].item())
        generated_ids.append(next_id)
        input_ids = torch.cat([input_ids, torch.tensor([[next_id]], device=device)], dim=1)
    return action_tokenizer.decode(generated_ids)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-id", default="lmms-lab/llava-onevision-qwen2-0.5b-ov")
    parser.add_argument("--steps", type=int, default=100)
    parser.add_argument("--learning-rate", type=float, default=5e-5)
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--save-dir", type=Path)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    torch.manual_seed(0)
    np.random.seed(0)
    if torch.cuda.is_available():
        device = torch.device("cuda")
        dtype_name = "bfloat16" if torch.cuda.is_bf16_supported() else "float16"
        load_device_map = {"": device.index or 0}
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        dtype_name = "float16"
        # The bundled LLaVA loader hard-codes a CUDA move whenever device_map
        # is not "auto". Let it load automatically, then move to MPS below.
        load_device_map = "auto"
    else:
        raise RuntimeError("A CUDA GPU or Apple-Silicon MPS device is required.")

    # huggingface_hub/transformers automatically consume HF_TOKEN. We only
    # check for it so private/gated downloads fail with a useful message.
    if not (os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")):
        print("Note: HF_TOKEN is unset; continuing because this checkpoint may be public.")

    model_name = get_model_name_from_path(args.model_id)
    tokenizer, model, image_processor, _ = load_pretrained_model(
        args.model_id,
        model_base=None,
        model_name=model_name,
        device_map=load_device_map,
        torch_dtype=dtype_name,
        attn_implementation="sdpa",
    )
    if device.type == "mps":
        # FP16 AdamW on MPS overflowed after one real optimizer step. The model
        # fits in unified memory at FP32 and is stable enough for this tiny test.
        model.to(device=device, dtype=torch.float32)
    action_tokenizer = ActionTokenizer(tokenizer, bins=256)
    model.resize_token_embeddings(len(tokenizer))
    model.config.use_cache = False

    samples = make_toy_samples()
    trainable = set_trainable_subset(model)
    optimizer = torch.optim.AdamW(trainable, lr=args.learning_rate, weight_decay=0.0)

    initial_loss, initial_accuracy = teacher_forced_metrics(
        model, tokenizer, action_tokenizer, samples, image_processor, device
    )
    print(f"before: loss={initial_loss:.4f}, action_token_accuracy={initial_accuracy:.3f}")

    model.train()
    for step, sample in zip(range(1, args.steps + 1), itertools.cycle(samples)):
        input_ids, labels = prepare_example(tokenizer, action_tokenizer, sample, device)
        images, image_sizes = prepare_image(sample.image, image_processor, model, device)
        output = model(
            input_ids=input_ids,
            labels=labels,
            images=images,
            image_sizes=image_sizes,
            modalities=["image"],
            use_cache=False,
        )
        output.loss.backward()
        torch.nn.utils.clip_grad_norm_(trainable, 1.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        if step == 1 or step % args.log_every == 0:
            print(f"step={step:04d} train_loss={float(output.loss.detach()):.4f}")

    final_loss, final_accuracy = teacher_forced_metrics(
        model, tokenizer, action_tokenizer, samples, image_processor, device
    )
    print(f"after:  loss={final_loss:.4f}, action_token_accuracy={final_accuracy:.3f}")

    for sample in samples:
        predicted = generate_action(model, tokenizer, action_tokenizer, sample, image_processor, device)
        print(f"{sample.instruction}\n  target={sample.action.round(3)}\n  pred  ={predicted.round(3)}")

    if args.save_dir:
        args.save_dir.mkdir(parents=True, exist_ok=True)
        trainable_names = {name for name, parameter in model.named_parameters() if parameter.requires_grad}
        torch.save(
            {name: value.detach().cpu() for name, value in model.state_dict().items() if name in trainable_names},
            args.save_dir / "minimal_openvla_trainable.pt",
        )
        tokenizer.save_pretrained(args.save_dir)
        print(f"Saved trainable state and tokenizer to {args.save_dir}")


if __name__ == "__main__":
    main()