#!/usr/bin/env python3
# coding=utf-8
"""ConCor-1 inference example: an image + its text in, correspondences out.

Given an image and a paired text, ConCor-1 predicts the full set of
correspondences between visually referential text spans and instance-level
image segments — the text spans are *not* given as queries; the model decides
which parts of the text are grounded.

Usage (the bundled `example.png` with its COCONut-PanCap caption):
    python example_inference.py --image example.png \
        --text "This image depicts a close-up of a brown bear in a natural outdoor setting. \
The background consists of lush green grass. In the foreground, a large brown bear is \
positioned centrally."

    # a category list works just as well as a caption
    python example_inference.py --image example.png --text "bear . grass . tree . person"

    # write a mask overlay next to the printed correspondences
    python example_inference.py --image example.png --text "..." --output overlay.png
"""

from __future__ import annotations

import argparse
from pathlib import Path
from typing import List

import numpy as np
import torch
from PIL import Image, ImageDraw
from transformers import AutoModel, AutoProcessor

# Visually distinct overlay colours.
COLORS = [
    (239, 64, 64), (64, 204, 64), (64, 115, 255), (255, 204, 0), (255, 115, 0),
    (217, 51, 217), (0, 204, 217), (153, 89, 13), (102, 255, 102), (140, 26, 255),
    (255, 153, 204), (0, 140, 0), (179, 179, 0), (0, 26, 153), (204, 140, 51),
    (128, 128, 128), (255, 0, 128), (0, 255, 140), (140, 0, 0), (0, 140, 140),
]


def overlay_masks(image: Image.Image, correspondences: List[dict], alpha: float = 0.5) -> Image.Image:
    """Blend each correspondence's mask over the image and label it with its phrases."""
    canvas = np.array(image.convert("RGB"), dtype=np.float32)
    for index, correspondence in enumerate(correspondences):
        mask = correspondence.get("mask")
        if mask is None or not mask.any():
            continue
        color = np.array(COLORS[index % len(COLORS)], dtype=np.float32)
        canvas[mask] = (1.0 - alpha) * canvas[mask] + alpha * color

    overlaid = Image.fromarray(canvas.astype(np.uint8))
    draw = ImageDraw.Draw(overlaid)
    for index, correspondence in enumerate(correspondences):
        mask = correspondence.get("mask")
        if mask is None or not mask.any():
            continue
        rows, columns = np.nonzero(mask)
        label = " / ".join(correspondence["text_phrases"]) or "(no text span)"
        label = f"{label}  {correspondence['presence_score']:.2f}"
        anchor = (int(columns.min()), max(int(rows.min()) - 12, 0))
        draw.text(anchor, label, fill=COLORS[index % len(COLORS)])
    return overlaid


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--model", default=str(Path(__file__).parent), help="model repo id or local path")
    parser.add_argument("--image", required=True, type=Path)
    parser.add_argument("--text", required=True, help="the text to ground (caption, category list, referring expression)")
    parser.add_argument("--output", type=Path, default=None, help="write a mask overlay here (PNG)")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument(
        "--attn_implementation",
        default="flash_attention_2",
        choices=["flash_attention_2", "sdpa", "eager"],
        help="flash_attention_2 reproduces the paper's numbers exactly; sdpa needs no extra dependency",
    )
    parser.add_argument("--presence_threshold", type=float, default=None, help="default 0.1")
    parser.add_argument("--text_threshold", type=float, default=None, help="default 0.45")
    parser.add_argument("--image_threshold", type=float, default=None, help="default 0.45")
    parser.add_argument("--nms_iou_threshold", type=float, default=None, help="default 0.5")
    args = parser.parse_args()

    processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
    model = AutoModel.from_pretrained(
        args.model,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        attn_implementation=args.attn_implementation,
    ).to(args.device).eval()

    image = Image.open(args.image).convert("RGB")
    inputs = processor(images=image, text=args.text, return_tensors="pt").to(args.device)

    with torch.inference_mode(), torch.autocast(args.device, dtype=torch.bfloat16):
        outputs = model(**inputs)

    correspondences = processor.post_process_correspondences(
        outputs,
        text=args.text,
        target_sizes=[(image.height, image.width)],
        presence_threshold=args.presence_threshold,
        text_threshold=args.text_threshold,
        image_threshold=args.image_threshold,
        nms_iou_threshold=args.nms_iou_threshold,
    )[0]

    print(f"\nimage: {args.image}  ({image.width}x{image.height})")
    print(f"text:  {args.text}")
    print(f"\n{len(correspondences)} correspondence(s):")
    for correspondence in correspondences:
        mask = correspondence["mask"]
        phrases = " / ".join(correspondence["text_phrases"]) or "(no text span)"
        area = 100.0 * mask.mean() if mask is not None else 0.0
        print(
            f"  presence={correspondence['presence_score']:.3f}  "
            f"bridge={correspondence['bridge_index']:3d}  "
            f"mask={area:5.2f}% of image  spans={correspondence['text_spans']}  |  {phrases}"
        )

    if args.output is not None:
        overlay_masks(image, correspondences).save(args.output)
        print(f"\nwrote {args.output}")


if __name__ == "__main__":
    main()