import os
import cv2
import json
import torch
import numpy as np

from segment_anything import sam_model_registry, SamPredictor

# ---------------- CONFIG ----------------
SAM_CHECKPOINT = "/media/bs5/New Volume/Brinda/Codes/Florida/SAM/sam_vit_h_4b8939.pth"
MODEL_TYPE = "vit_h"

IMAGE_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/images_seg"
LABEL_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/New_Dataset/set2/labels_seg/"
OUTPUT_DIR = "/media/bs5/New Volume/Brinda/Dataset/Florida/classification_data/"

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(os.path.join(OUTPUT_DIR, "masks"), exist_ok=True)
os.makedirs(os.path.join(OUTPUT_DIR, "overlay"), exist_ok=True)
os.makedirs(os.path.join(OUTPUT_DIR, "polygons"), exist_ok=True)

# ---------------- LOAD SAM ----------------
sam = sam_model_registry[MODEL_TYPE](checkpoint=SAM_CHECKPOINT)
sam.to(DEVICE)
predictor = SamPredictor(sam)

# ---------------- HELPERS ----------------
def yolo_to_xyxy(label, img_w, img_h):
    """
    YOLO format: class cx cy w h (normalized)
    Convert → pixel [x1, y1, x2, y2]
    """
    cls, cx, cy, w, h = label

    x1 = int((cx - w / 2) * img_w)
    y1 = int((cy - h / 2) * img_h)
    x2 = int((cx + w / 2) * img_w)
    y2 = int((cy + h / 2) * img_h)

    return [x1, y1, x2, y2]


def expand_box(x1, y1, x2, y2, pad, w, h):
    return [
        max(0, x1 - pad),
        max(0, y1 - pad),
        min(w, x2 + pad),
        min(h, y2 + pad),
    ]


def mask_to_polygon(mask):
    contours, _ = cv2.findContours(
        mask.astype(np.uint8),
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )

    polygons = []
    for cnt in contours:
        cnt = cnt.squeeze()
        if len(cnt.shape) == 2 and len(cnt) > 2:
            polygons.append(cnt.tolist())

    return polygons


# ---------------- PROCESS ----------------
for img_file in os.listdir(IMAGE_DIR):

    if not img_file.lower().endswith((".jpg", ".png", ".jpeg")):
        continue

    image_path = os.path.join(IMAGE_DIR, img_file)
    label_path = os.path.join(LABEL_DIR, img_file.replace(".jpg", ".txt").replace(".png", ".txt"))

    if not os.path.exists(label_path):
        continue

    image = cv2.imread(image_path)
    if image is None:
        continue

    h, w = image.shape[:2]
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    predictor.set_image(image_rgb)

    overlay = image.copy()
    output_data = []

    # -------- Read YOLO labels --------
    with open(label_path, "r") as f:
        lines = f.readlines()

    for idx, line in enumerate(lines):
        parts = list(map(float, line.strip().split()))
        cls_id = int(parts[0])
        cx, cy, bw, bh = parts[1:]

        # YOLO → XYXY
        x1, y1, x2, y2 = yolo_to_xyxy([cls_id, cx, cy, bw, bh], w, h)

        # Expand bbox (important for SAM)
        x1, y1, x2, y2 = expand_box(x1, y1, x2, y2, pad=10, w=w, h=h)

        box = np.array([x1, y1, x2, y2])

        # -------- SAM Prediction --------
        masks, scores, _ = predictor.predict(
            box=box,
            multimask_output=False
        )

        mask = masks[0]

        # -------- Save mask --------
        mask_name = f"{img_file.split('.')[0]}_{idx}.png"
        mask_path = os.path.join(OUTPUT_DIR, "masks", mask_name)
        cv2.imwrite(mask_path, mask.astype(np.uint8) * 255)

        # -------- Overlay --------
        color = np.random.randint(0, 255, (3,))
        overlay[mask] = (overlay[mask] * 0.5 + color * 0.5).astype(np.uint8)

        # -------- Polygon --------
        polygons = mask_to_polygon(mask)

        output_data.append({
            "id": idx,
            "class": cls_id,
            "bbox": [x1, y1, x2, y2],
            "polygon": polygons
        })

    # -------- Save overlay --------
    cv2.imwrite(
        os.path.join(OUTPUT_DIR, "overlay", img_file),
        overlay
    )

    # -------- Save JSON --------
    json_path = os.path.join(
        OUTPUT_DIR, "polygons", img_file.replace(".jpg", ".json").replace(".png", ".json")
    )

    with open(json_path, "w") as f:
        json.dump(output_data, f, indent=4)

print(" Processing Complete")