import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import json
import base64
from io import BytesIO
from PIL import Image
import re

import torch
from transformers import Qwen2VLForConditionalGeneration, GenerationConfig, AutoProcessor
import spaces


def extract_answer_content(text: str) -> str:
    """
    Extracts the content between <answer> and </answer> tags.
    If no closing tag is found, extracts everything after the first <answer>.

    Returns:
        str: The extracted content.
    """
    text = text.replace("```", " ").replace("json", " ").strip()

    # Try to find full <answer>...</answer>
    match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
    if match:
        return match.group(1).strip()
    # Fallback: everything after the first <answer>
    match = re.search(r"<answer>(.*)", text, re.DOTALL)
    return match.group(1).strip() if match else text


def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
    """
    Encode a PIL Image to a base64 string.

    Args:
        image (PIL.Image): The image to encode.
        format (str): Image format to use (e.g., "PNG", "JPEG"). Default is "PNG".

    Returns:
        str: Base64-encoded string of the image.
    """
    buffer = BytesIO()
    image.save(buffer, format=format)
    buffer.seek(0)
    encoded_string = base64.b64encode(buffer.read()).decode("utf-8")
    return encoded_string



SYSTEM_PROMPT = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
    "<think> reasoning process here </think><answer> answer here </answer>"
)



# Load model and processor
processor = AutoProcessor.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg", max_pixels=1024*28*28)

device='cuda' if torch.cuda.is_available() else "cpu"
model = Qwen2VLForConditionalGeneration.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg",
             torch_dtype=torch.bfloat16,
             device_map=device)

generation_config=GenerationConfig(
        do_sample=True,
        temperature=0.01,
        top_k=1,
        top_p=0.001,
        repetition_penalty=1.0,
        max_new_tokens=2048,
        use_cache=True 
)

def build_prompt(image, user_text):
    #base64_image = encode_image_to_base64(image)
    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT
        },
        {
            "role": "user",
            "content": [
                #{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
                 {"type": "image"},
                {"type": "text", "text": user_text},
            ],
        },
    ]    
    return messages

PROMPT_CLOSE="""Generate a structured scene graph for an image using the specified object and relationship categories.

### **Output Format:**
```json
{
  "objects": [
    {"id": "object_name.number", "bbox": [x1, y1, x2, y2]},
    ...
  ],
  "relationships": [
    {"subject": "object_name.number", "predicate": "relationship_type", "object": "object_name.number"},
    ...
  ]
}
```

### **Guidelines:**
- **Objects:**
  - Assign unique IDs in the format `"object_name.number"` (e.g., `"person.1"`). The **object_name** must belong to the predefined object set: `["airplane", "animal", "arm", "bag", "banana", "basket", "beach", "bear", "bed", "bench", "bike", "bird", "board", "boat", "book", "boot", "bottle", "bowl", "box", "boy", "branch", "building", "bus", "cabinet", "cap", "car", "cat", "chair", "child", "clock", "coat", "counter", "cow", "cup", "curtain", "desk", "dog", "door", "drawer", "ear", "elephant", "engine", "eye", "face", "fence", "finger", "flag", "flower", "food", "fork", "fruit", "giraffe", "girl", "glass", "glove", "guy", "hair", "hand", "handle", "hat", "head", "helmet", "hill", "horse", "house", "jacket", "jean", "kid", "kite", "lady", "lamp", "laptop", "leaf", "leg", "letter", "light", "logo", "man", "men", "motorcycle", "mountain", "mouth", "neck", "nose", "number", "orange", "pant", "paper", "paw", "people", "person", "phone", "pillow", "pizza", "plane", "plant", "plate", "player", "pole", "post", "pot", "racket", "railing", "rock", "roof", "room", "screen", "seat", "sheep", "shelf", "shirt", "shoe", "short", "sidewalk", "sign", "sink", "skateboard", "ski", "skier", "sneaker", "snow", "sock", "stand", "street", "surfboard", "table", "tail", "tie", "tile", "tire", "toilet", "towel", "tower", "track", "train", "tree", "truck", "trunk", "umbrella", "vase", "vegetable", "vehicle", "wave", "wheel", "window", "windshield", "wing", "wire", "woman", "zebra"]`.
  - Provide a bounding box `[x1, y1, x2, y2]` in integer pixel format.
  - Include all visible objects, even if they have no relationships.

- **Relationships:**
  - Define relationships using `"subject"`, `"predicate"`, and `"object"`.
  - The **predicate** must belong to the predefined relationship set: `["above", "across", "against", "along", "and", "at", "attached to", "behind", "belonging to", "between", "carrying", "covered in", "covering", "eating", "flying in", "for", "from", "growing on", "hanging from", "has", "holding", "in", "in front of", "laying on", "looking at", "lying on", "made of", "mounted on", "near", "of", "on", "on back of", "over", "painted on", "parked on", "part of", "playing", "riding", "says", "sitting on", "standing on", "to", "under", "using", "walking in", "walking on", "watching", "wearing", "wears", "with"]`.
  - Omit relationships for orphan objects.

### **Example Output:**
```json
{
  "objects": [
    {"id": "person.1", "bbox": [120, 200, 350, 700]},
    {"id": "bike.2", "bbox": [100, 600, 400, 800]},
    {"id": "helmet.3", "bbox": [150, 150, 280, 240]},
    {"id": "tree.4", "bbox": [500, 100, 750, 700]}
  ],
  "relationships": [
    {"subject": "person.1", "predicate": "riding", "object": "bike.2"},
    {"subject": "person.1", "predicate": "wearing", "object": "helmet.3"}
  ]
}
```

Now, generate the complete scene graph for the provided image:
"""

def is_box(item):
    return (
        isinstance(item, (list, tuple)) and
        len(item) == 4 and
        all(isinstance(e, (int, float)) for e in item)
    )

def scale_box(box, scale):
    sw, sh = scale
    return [int(box[0]*sw), int(box[1]*sh), int(box[2]*sw), int(box[3]*sh)]

@spaces.GPU
def generate_sgg(image):
    global model
    
    device='cuda' if torch.cuda.is_available() else "cpu"
    if next(model.parameters()).device != torch.device(device):
        model = model.to(device)
    
    iw, ih = image.size
    scale_factors = (iw / 1000.0, ih / 1000.0)

    conversation = build_prompt(image, PROMPT_CLOSE)
    text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)

    inputs = processor(
        text=[text_prompt], images=[image], padding=True, return_tensors="pt"
    )
    inputs = inputs.to(model.device)
    with torch.no_grad():
        output_ids = model.generate(**inputs, generation_config=generation_config)
    generated_ids = [
        output_ids[len(input_ids) :]
        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
    ]
    output_text = processor.batch_decode(
        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
    )[0]    

    resp = extract_answer_content(output_text)
    
    try:
        resp = json.loads(resp)
        pred_objs = resp['objects']
        pred_rels = resp['relationships']
        new_objs = []
        for obj in pred_objs:
            assert len(obj['bbox']) == 4, "len(obj['bbox']) != 4"
            assert is_box(obj['bbox']), "invalid box :{}".format(obj['bbox'])
            assert 'id' in obj, "invalid obj:{}".format(obj)        
            obj['bbox'] = scale_box(obj['bbox'], scale_factors)
            new_objs.append(obj)
        pred_objs = new_objs
        resp = {"objects": pred_objs, "relationships": pred_rels}

        # visualize pred_objs
        draw = ImageDraw.Draw(image)
        try:
            font = ImageFont.truetype("arial.ttf", 16)
        except:
            font = ImageFont.load_default()

        for obj in pred_objs:
            bbox = obj["bbox"]
            label = obj["id"]
            draw.rectangle(bbox, outline="red", width=3)
            draw.text((bbox[0], bbox[1] - 15), label, fill="red", font=font)

        return image, f"objects:{pred_objs},\n relationships:{pred_rels}\n"
    except:
        return image, resp

gr.Interface(
    fn=generate_sgg,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Textbox(label="Scene Graph")],
    title="R1-SGG: Compile Scene Graphs with Reinforcement Learning",
    description="Upload an image and generate a structured scene graph in JSON format."
).launch(share=True)