import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset

class FruitVegetableDataset(Dataset):
    def __init__(self, images_dir, labels_dir, transform=None):
        """
        Args:
            images_dir (str): Path to images directory.
            labels_dir (str): Path to labels directory.
            transform (albumentations.Compose, optional): Albumentations transform pipeline.
        """
        self.images_dir = images_dir
        self.labels_dir = labels_dir
        self.transform = transform
        
        # List all image files
        self.image_files = sorted([
            f for f in os.listdir(images_dir) 
            if f.lower().endswith(('.jpg', '.jpeg', '.png'))
        ])

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

    def __getitem__(self, idx):
        img_name = self.image_files[idx]
        img_path = os.path.join(self.images_dir, img_name)
        
        # Load image via OpenCV (BGR) and convert to RGB
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        height, width, _ = image.shape
        
        # Load corresponding label file
        label_name = os.path.splitext(img_name)[0] + ".txt"
        label_path = os.path.join(self.labels_dir, label_name)
        
        boxes = []
        class_ids = []
        
        if os.path.exists(label_path):
            with open(label_path, "r") as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) == 5:
                        class_id = int(parts[0])
                        # YOLO format: x_center, y_center, width, height (normalized)
                        x_c, y_c, w, h = map(float, parts[1:])
                        
                        # Convert to Pascal VOC absolute coordinates [x_min, y_min, x_max, y_max]
                        xmin = (x_c - w / 2.0) * width
                        ymin = (y_c - h / 2.0) * height
                        xmax = (x_c + w / 2.0) * width
                        ymax = (y_c + h / 2.0) * height
                        
                        # Ensure coordinates are within image boundaries
                        xmin = max(0.0, min(xmin, float(width - 1)))
                        ymin = max(0.0, min(ymin, float(height - 1)))
                        xmax = max(xmin + 1.0, min(xmax, float(width)))
                        ymax = max(ymin + 1.0, min(ymax, float(height)))
                        
                        boxes.append([xmin, ymin, xmax, ymax])
                        class_ids.append(class_id)
                        
        boxes = np.array(boxes, dtype=np.float32).reshape(-1, 4)
        class_ids = np.array(class_ids, dtype=np.int64)
        
        # Apply Albumentations transformations
        if self.transform:
            augmented = self.transform(
                image=image,
                bboxes=boxes,
                category_ids=class_ids
            )
            image = augmented["image"]
            boxes = np.array(augmented["bboxes"], dtype=np.float32).reshape(-1, 4)
            class_ids = np.array(augmented["category_ids"], dtype=np.int64)
            
        # Convert image to CHW tensor and scale to [0, 1]
        image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
        
        # Convert targets to PyTorch Tensors
        # NOTE: torchvision models reserve class 0 for background.
        # We must shift our class IDs by +1 (0 -> 1, 1 -> 2, etc.)
        torch_boxes = torch.as_tensor(boxes, dtype=torch.float32)
        torch_labels = torch.as_tensor(class_ids + 1, dtype=torch.int64)
        
        target = {
            "boxes": torch_boxes,
            "labels": torch_labels,
            "image_id": torch.tensor([idx])
        }
        
        return image_tensor, target

def collate_fn(batch):
    """
    Collate function for DataLoader. Returns a tuple of (images, targets).
    """
    return tuple(zip(*batch))