#!/usr/bin/env python3
"""Run one image-text matching example from this Hugging Face model bundle."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from lab.common import resolve_device, resolve_dtype  # noqa: E402
from lab.multimodal import MatchingHead, VisionTokenAdapter  # noqa: E402
from lab.vision import fuse_prefix, run_vision_model  # noqa: E402


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, default=ROOT)
    parser.add_argument("--task-adapter", type=Path)
    parser.add_argument("--vision-model", default="google/siglip2-base-patch16-256")
    parser.add_argument("--image", type=Path, required=True)
    parser.add_argument("--text", required=True)
    parser.add_argument("--dtype", default="float32")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--output", type=Path)
    return parser.parse_args()


def main():
    args = parse_args()
    import torch
    from PIL import Image
    from transformers import AutoImageProcessor, AutoModel, AutoModelForMaskedLM, AutoTokenizer

    device = resolve_device(args.device)
    dtype = resolve_dtype(args.dtype) or torch.float32
    task_path = args.task_adapter or args.model / "multimodal_task_adapter.pt"
    task = torch.load(task_path, map_location="cpu", weights_only=True)

    processor = AutoImageProcessor.from_pretrained(args.vision_model)
    vision = AutoModel.from_pretrained(args.vision_model, dtype=dtype).to(device).eval()
    model = AutoModelForMaskedLM.from_pretrained(
        args.model, trust_remote_code=True, local_files_only=True, dtype=dtype
    ).to(device).eval()
    tokenizer = AutoTokenizer.from_pretrained(
        args.model, trust_remote_code=True, local_files_only=True
    )

    if (args.model / "recipe.yaml").exists():
        warmup = tokenizer("compression warmup", return_tensors="pt")
        warmup = {key: value.to(device) for key, value in warmup.items()}
        with torch.inference_mode():
            model(**warmup)

    body = model.lfm2
    adapter = VisionTokenAdapter(
        task["vision_dim"], task["lm_dim"], task["num_image_tokens"]
    ).to(device=device, dtype=dtype)
    adapter.load_state_dict(task["projector"])
    head = MatchingHead(body.config.hidden_size).to(device=device, dtype=dtype)
    head.load_state_dict(task["matching_head"])

    image = Image.open(args.image).convert("RGB")
    pixels = processor(images=image, return_tensors="pt")["pixel_values"].to(
        device=device, dtype=dtype
    )
    text = tokenizer(args.text, return_tensors="pt", truncation=True, max_length=128)
    text = {key: value.to(device) for key, value in text.items()}

    with torch.inference_mode():
        visual = run_vision_model(vision, pixels)
        image_tokens = adapter(visual)
        fused, mask, marker = fuse_prefix(
            body,
            tokenizer,
            text["input_ids"],
            text["attention_mask"],
            image_tokens,
            insert_position=1,
        )
        output = body(inputs_embeds=fused, attention_mask=mask)
        matching_score = head(output.last_hidden_state[:, marker]).item()

    result = {
        "model": str(args.model),
        "task_adapter": str(task_path),
        "vision_model": args.vision_model,
        "device": device,
        "dtype": str(dtype),
        "visual_features": list(visual.shape),
        "image_tokens": list(image_tokens.shape),
        "fused_embeddings": list(fused.shape),
        "hidden_state": list(output.last_hidden_state.shape),
        "matching_score": matching_score,
        "finite_hidden": bool(torch.isfinite(output.last_hidden_state).all()),
        "generation": "not supported: LFM2.5 is a bidirectional encoder",
    }
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(json.dumps(result, indent=2) + "\n")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()