#!/opt/homebrew/bin/python3.11
"""
Crop the detected stamp from the 112:36 frame (with Patch).
"""

from PIL import Image
import os
import cv2
import types
from transformers import AutoProcessor, AutoModelForCausalLM

UUID = "384b0ff44aaaa1f1"
BASE_DIR = f"output/{UUID}/florence2_results"
IMG_NAME = "frame_6756.jpg"
img_path = os.path.join(BASE_DIR, IMG_NAME)

print(f"📷 Loading image: {img_path}")
if not os.path.exists(img_path):
    print("❌ Image not found.")
    exit()


# Patch for compatibility
def patch_model(model):
    inner_model = model.language_model
    original_prepare = inner_model.prepare_inputs_for_generation

    def patched_prepare(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        is_valid_cache = False
        if past_key_values is not None:
            if isinstance(past_key_values, (list, tuple)) and len(past_key_values) > 0:
                first_layer = past_key_values[0]
                if first_layer is not None and (
                    not isinstance(first_layer, (list, tuple)) or len(first_layer) > 0
                ):
                    is_valid_cache = True

        if not is_valid_cache:
            return {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "past_key_values": None,
                "use_cache": True,
            }
        else:
            return original_prepare(
                input_ids,
                past_key_values=past_key_values,
                attention_mask=attention_mask,
                inputs_embeds=inputs_embeds,
                **kwargs,
            )

    inner_model.prepare_inputs_for_generation = types.MethodType(
        patched_prepare, inner_model
    )


try:
    img = Image.open(img_path).convert("RGB")
    print(f"📐 Image Size: {img.width}x{img.height}")

    print("🧠 Running detection to get coordinates...")
    processor = AutoProcessor.from_pretrained(
        "microsoft/Florence-2-base", trust_remote_code=True
    )
    model = AutoModelForCausalLM.from_pretrained(
        "microsoft/Florence-2-base", trust_remote_code=True, attn_implementation="eager"
    )
    patch_model(model)

    prompt = "<OPEN_VOCABULARY_DETECTION>"
    inputs = processor(text=prompt, images=img, return_tensors="pt")

    # Generate
    generated_ids = model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"],
        max_new_tokens=1024,
        num_beams=3,
    )
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]

    # Parse
    parsed_answer = processor.post_process_generation(
        generated_text, task=prompt, image_size=(img.width, img.height)
    )
    results = parsed_answer.get("<OPEN_VOCABULARY_DETECTION>", {})
    bboxes = results.get("bboxes", [])

    if bboxes:
        box = bboxes[0]  # Take the first detected stamp
        print(f"📦 Detected Box: {box}")

        # Crop
        box_int = [int(x) for x in box]
        cropped = img.crop(box_int)

        out_path = os.path.join(BASE_DIR, "stamp_from_112_36.jpg")
        cropped.save(out_path)
        print(f"✅ Successfully saved cropped stamp to {out_path}")

        # Also save a visualization
        img_cv = cv2.imread(img_path)
        x1, y1, x2, y2 = map(int, box)
        cv2.rectangle(img_cv, (x1, y1), (x2, y2), (0, 255, 0), 3)
        cv2.putText(
            img_cv, "STAMP", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2
        )
        vis_path = os.path.join(BASE_DIR, "stamp_detection_112_36.jpg")
        cv2.imwrite(vis_path, img_cv)
        print(f"🎨 Visualization saved to {vis_path}")

    else:
        print("❌ No stamp found in this frame.")

except Exception as e:
    print(f"❌ Error: {e}")
    import traceback

    traceback.print_exc()