#!/opt/homebrew/bin/python3.11
"""
Deep Analysis of 112:36 Frame
1. Detailed Captioning
2. Search for "Envelope" and "Hand holding object"
"""

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

UUID = "384b0ff44aaaa1f1"
BASE_DIR = f"output/{UUID}/florence2_results"
IMG_NAME = "scan_6756.jpg"  # 112:36
IMG_PATH = os.path.join(BASE_DIR, IMG_NAME)


# 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
    )


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

image = Image.open(IMG_PATH).convert("RGB")

print("🧠 Loading Florence-2 model...")
try:
    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)

    # 1. Detailed Caption
    print("\n📝 Generating Detailed Caption...")
    prompt = "<DETAILED_CAPTION>"
    inputs = processor(text=prompt, images=image, return_tensors="pt")
    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=True)[0]
    print(f"🗣️ Caption: {generated_text}")

    # 2. Object Detection for specific items
    search_terms = ["envelope", "letter", "hand holding paper", "stamp", "small paper"]
    img_cv = cv2.imread(IMG_PATH)

    for term in search_terms:
        print(f"\n🔍 Detecting '{term}'...")
        prompt_ovd = "<OPEN_VOCABULARY_DETECTION>"
        # Note: OVD usually takes text input differently or relies on generation.
        # For Florence-2, OVD often requires text_input in processor or prompt format.
        # We will try the standard way first.

        inputs = processor(text=prompt_ovd, images=image, return_tensors="pt")
        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]

        try:
            parsed_answer = processor.post_process_generation(
                generated_text, task=prompt_ovd, image_size=(image.width, image.height)
            )
            results = parsed_answer.get("<OPEN_VOCABULARY_DETECTION>", {})
            bboxes = results.get("bboxes", [])
            labels = results.get("bboxes_labels", [])

            if bboxes:
                print(f"   ✅ Found '{term}': {labels}")
                for i, (box, label) in enumerate(zip(bboxes, labels)):
                    if term.lower() in label.lower() or (
                        term == "envelope" and "paper" in label.lower()
                    ):
                        x1, y1, x2, y2 = map(int, box)
                        print(f"      📍 Box: ({x1},{y1}) -> ({x2},{y2})")

                        # Crop
                        crop = img_cv[y1:y2, x1:x2]
                        crop_path = os.path.join(
                            BASE_DIR, f"crop_deep_{term.replace(' ', '_')}_{i}.jpg"
                        )
                        cv2.imwrite(crop_path, crop)

                        # Draw
                        cv2.rectangle(img_cv, (x1, y1), (x2, y2), (0, 255, 0), 3)
                        cv2.putText(
                            img_cv,
                            label,
                            (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            1,
                            (0, 255, 0),
                            2,
                        )
            else:
                print("   ❌ Not found.")
        except Exception as e:
            print(f"   ⚠️ Error: {e}")

    res_path = os.path.join(BASE_DIR, "deep_analysis_result.jpg")
    cv2.imwrite(res_path, img_cv)
    print(f"\n🎨 Result saved to {res_path}")

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