import os
import torch
import torchvision.transforms as transforms
from torch.utils.data import Dataset
from pycocotools.coco import COCO
from PIL import Image

class COCODataset(Dataset):
    def __init__(self, root, annotation_file, transforms=None):
        self.root = root
        self.coco = COCO(annotation_file)
        self.transforms = transforms
        
        self.image_ids = list(self.coco.imgs.keys())

    def __getitem__(self, idx):
        image_id = self.image_ids[idx]
        image_info = self.coco.imgs[image_id]
        image_path = os.path.join(self.root, image_info["file_name"])
        
        image = Image.open(image_path).convert("RGB")

        annotations = self.coco.loadAnns(self.coco.getAnnIds(imgIds=image_id))
        
        boxes = []
        labels = []
        
        for ann in annotations:
            xmin = ann["bbox"][0]
            ymin = ann["bbox"][1]
            width = ann["bbox"][2]
            height = ann["bbox"][3]
            
            boxes.append([xmin, ymin, xmin + width, ymin + height])
            labels.append(ann["category_id"])  # Ensure category_id starts from 1
        
        if len(boxes) == 0:
            boxes = torch.zeros((0, 4), dtype=torch.float32)
            labels = torch.zeros((0,), dtype=torch.int64)
        else:
            boxes = torch.tensor(boxes, dtype=torch.float32)
            labels = torch.tensor(labels, dtype=torch.int64)

        target = {"boxes": boxes, "labels": labels, "image_id": torch.tensor([image_id])}

        if self.transforms:
            image = self.transforms(image)

        return image, target

    def __len__(self):
        return len(self.image_ids)

def get_transforms():
    return transforms.Compose([
        transforms.ToTensor()
    ])