import os
import numpy as np
import torch
import matplotlib.pyplot as plt
from PIL import Image
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

# Environment setup (assuming dependencies are installed)
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

# Select device
if torch.cuda.is_available():
    device = torch.device("cuda")
    torch.autocast("cuda", dtype=torch.bfloat16).__enter__()
    if torch.cuda.get_device_properties(0).major >= 8:
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True
elif torch.backends.mps.is_available():
    device = torch.device("mps")
else:
    device = torch.device("cpu")
print(f"using device: {device}")

# Helper functions (modified to save instead of show)
def save_mask_on_image(image, mask, save_path, alpha=0.5, borders=True):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(image)

    color = np.array([30/255, 144/255, 255/255, alpha])
    h, w = mask.shape[-2:]
    mask = mask.astype(np.uint8)
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)

    if borders:
        import cv2
        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)

    ax.imshow(mask_image)
    ax.axis('off')
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close(fig)

def save_points_on_image(image, coords, labels, save_path):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(image)
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    marker_size = 375
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)
    ax.axis('off')
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close(fig)

def save_box_on_image(image, box, save_path):
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(image)
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))
    ax.axis('off')
    plt.savefig(save_path, bbox_inches='tight', pad_inches=0)
    plt.close(fig)


# Load model
sam2_checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
sam2 = build_sam2(model_cfg, sam2_checkpoint, device=device)
predictor = SAM2ImagePredictor(sam2)

# Example 1: Truck image with point prompt
image = Image.open('./notebooks/images/truck.jpg').convert("RGB")
image_np = np.array(image)

with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
    predictor.set_image(image_np)
    input_point = np.array([[825, 619]])
    input_label = np.array([1])
    masks, scores, logits = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=True,
    )

# Save results
os.makedirs("results", exist_ok=True)
for i, (mask, score) in enumerate(zip(masks, scores)):
    save_path = f"results/truck_mask_{i}_score_{score:.3f}.png"
    save_mask_on_image(image_np, mask, save_path)
    np.save(f"results/truck_mask_{i}.npy", mask)

save_points_on_image(image_np, input_point, input_label, "results/truck_points.png")

# Example 2: Truck with box prompt
with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
    input_box = np.array([425, 600, 700, 875])
    masks, scores, _ = predictor.predict(
        point_coords=None,
        point_labels=None,
        box=input_box,
        multimask_output=False,
    )

for i, (mask, score) in enumerate(zip(masks, scores)):
    save_path = f"results/truck_box_mask_{i}_score_{score:.3f}.png"
    save_mask_on_image(image_np, mask, save_path)
    np.save(f"results/truck_box_mask_{i}.npy", mask)

save_box_on_image(image_np, input_box, "results/truck_box.png")

# Example 3: Truck with multiple points
with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
    input_points = np.array([[450, 600], [775, 800]])
    input_labels = np.array([1, 1])
    masks, scores, _ = predictor.predict(
        point_coords=input_points,
        point_labels=input_labels,
        multimask_output=False,
    )

for i, (mask, score) in enumerate(zip(masks, scores)):
    save_path = f"results/truck_multi_points_mask_{i}_score_{score:.3f}.png"
    save_mask_on_image(image_np, mask, save_path)
    np.save(f"results/truck_multi_points_mask_{i}.npy", mask)

save_points_on_image(image_np, input_points, input_labels, "results/truck_multi_points.png")

# Example 4: Groceries image with batch prompts
image2 = Image.open('./notebooks/images/groceries.jpg').convert("RGB")
image2_np = np.array(image2)

with torch.inference_mode(), torch.autocast(device.type, dtype=torch.bfloat16):
    predictor.set_image(image2_np)
    input_points2 = np.array([[350, 450], [400, 450]])
    input_labels2 = np.array([1, 1])
    masks2, scores2, _ = predictor.predict(
        point_coords=input_points2,
        point_labels=input_labels2,
        multimask_output=False,
    )

for i, (mask, score) in enumerate(zip(masks2, scores2)):
    save_path = f"results/groceries_mask_{i}_score_{score:.3f}.png"
    save_mask_on_image(image2_np, mask, save_path)
    np.save(f"results/groceries_mask_{i}.npy", mask)

save_points_on_image(image2_np, input_points2, input_labels2, "results/groceries_points.png")

print("All results saved in 'results' folder.")