import os
import cv2
import json
import torch
import numpy as np
import supervision as sv
import pycocotools.mask as mask_util
from torchvision.ops import box_convert

# Assuming you have these directories and files in your project
# If not, you'll need to adjust the paths accordingly
try:
    from sam2.build_sam import build_sam2
    from sam2.sam2_image_predictor import SAM2ImagePredictor
    from GroundingDINO.groundingdino.util.inference import load_model, load_image, predict
except ImportError as e:
    print(f"Error: {e}")
    print("Please ensure that the 'sam2' and 'GroundingDINO' modules are correctly installed and in your PYTHONPATH.")
    exit()

"""
Hyper parameters
"""
TEXT_PROMPT = "eyes. lips. ears."
IMG_PATH = os.path.abspath("b.jpg")  # Example image path, make it absolute
SAM2_CHECKPOINT = os.path.abspath("models/sam2.1_hiera_large.pt")
SAM2_MODEL_CONFIG = os.path.abspath("models/sam2.1_hiera_l.yaml")
GROUNDING_DINO_CONFIG = os.path.abspath("models/GroundingDINO_SwinT_OGC.py")
GROUNDING_DINO_CHECKPOINT = os.path.abspath("models/groundingdino_swint_ogc.pth")
BOX_THRESHOLD = 0.45
TEXT_THRESHOLD = 0.25
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
OUTPUT_DIR = os.path.abspath("outputs/grounded_sam2_local_demo")
DUMP_JSON_RESULTS = True

# Create output directory
os.makedirs(OUTPUT_DIR, exist_ok=True)

# Check if files exist
if not os.path.exists(IMG_PATH):
    print(f"Error: Image file not found at {IMG_PATH}")
    exit()
if not os.path.exists(SAM2_CHECKPOINT):
    print(f"Error: SAM2 checkpoint file not found at {SAM2_CHECKPOINT}")
    exit()
if not os.path.exists(SAM2_MODEL_CONFIG):
    print(f"Error: SAM2 model config file not found at {SAM2_MODEL_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CONFIG):
    print(f"Error: Grounding DINO config file not found at {GROUNDING_DINO_CONFIG}")
    exit()
if not os.path.exists(GROUNDING_DINO_CHECKPOINT):
    print(f"Error: Grounding DINO checkpoint file not found at {GROUNDING_DINO_CHECKPOINT}")
    exit()

# Build SAM2 image predictor
sam2_model = build_sam2(SAM2_MODEL_CONFIG, SAM2_CHECKPOINT, device=DEVICE)

sam2_predictor = SAM2ImagePredictor(sam2_model)

# Build grounding dino model
grounding_model = load_model(
    model_config_path=GROUNDING_DINO_CONFIG,
    model_checkpoint_path=GROUNDING_DINO_CHECKPOINT,
    device=DEVICE
)

# Setup the input image and text prompt for SAM 2 and Grounding DINO
text = TEXT_PROMPT
image_source, image = load_image(IMG_PATH)

sam2_predictor.set_image(image_source)

boxes, confidences, labels = predict(
    model=grounding_model,
    image=image,
    caption=text,
    box_threshold=BOX_THRESHOLD,
    text_threshold=TEXT_THRESHOLD,
)

# Process the box prompt for SAM 2
h, w, _ = image_source.shape
boxes = boxes * torch.Tensor([w, h, w, h])
input_boxes = box_convert(boxes=boxes, in_fmt="cxcywh", out_fmt="xyxy").numpy()

# FIXME: figure how does this influence the G-DINO model (This part seems unnecessary for the main issue, keeping it as is)
torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()

if torch.cuda.is_available() and torch.cuda.get_device_properties(0).major >= 8:
    # turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

masks, scores, logits = sam2_predictor.predict(
    point_coords=None,
    point_labels=None,
    box=input_boxes,
    multimask_output=False,
)

"""
Post-process the output
"""
if masks.ndim == 4:
    masks = masks.squeeze(1)

confidences = confidences.numpy().tolist()
class_names = labels
class_ids = np.array(list(range(len(class_names))))
labels = [
    f"{class_name} {confidence:.2f}"
    for class_name, confidence
    in zip(class_names, confidences)
]

"""
Visualize image with supervision
"""
img = cv2.imread(IMG_PATH)
detections = sv.Detections(
    xyxy=input_boxes,
    mask=masks.astype(bool),
    class_id=class_ids
)

box_annotator = sv.BoxAnnotator()
annotated_frame = box_annotator.annotate(scene=img.copy(), detections=detections)

label_annotator = sv.LabelAnnotator()
annotated_frame = label_annotator.annotate(scene=annotated_frame, detections=detections, labels=labels)
cv2.imwrite(os.path.join(OUTPUT_DIR, "groundingdino_annotated_image.jpg"), annotated_frame)

mask_annotator = sv.MaskAnnotator()
annotated_frame = mask_annotator.annotate(scene=annotated_frame, detections=detections)
cv2.imwrite(os.path.join(OUTPUT_DIR, "grounded_sam2_annotated_image_with_mask.jpg"), annotated_frame)

"""
Dump the results in standard format and save as json files
"""

def single_mask_to_rle(mask):
    rle = mask_util.encode(np.array(mask[:, :, None], order="F", dtype="uint8"))[0]
    rle["counts"] = rle["counts"].decode("utf-8")
    return rle

if DUMP_JSON_RESULTS:
    mask_rles = [single_mask_to_rle(mask) for mask in masks]
    input_boxes = input_boxes.tolist()
    scores = scores.tolist()
    results = {
        "image_path": IMG_PATH,
        "annotations": [
            {
                "class_name": class_name,
                "bbox": box,
                "segmentation": mask_rle,
                "score": score,
            }
            for class_name, box, mask_rle, score in zip(class_names, input_boxes, mask_rles, scores)
        ],
        "box_format": "xyxy",
        "img_width": w,
        "img_height": h,
    }

    with open(os.path.join(OUTPUT_DIR, "grounded_sam2_local_image_demo_results.json"), "w") as f:
        json.dump(results, f, indent=4)