import os
import pandas as pd
import torch
from torch.utils.data import Dataset
from PIL import Image
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np

class NiftiSegmentationDataset(Dataset):
    def __init__(self, csv_file, channel_keys = ['first_path'], augment=True):
        """
        Args:
            csv_file (str): CSV file with columns: 'image_path', 'label_path'
            use_mask (bool): Whether to concatenate mask as extra channel (not typical for segmentation targets)
            augment (bool): Apply data augmentation
        """
        self.df = pd.read_csv(csv_file)
        self.augment = augment
        source_directory = '/workspace/data/mips3'
        breast_mask_source_directory = '/workspace/Segmentation/breastmaskmamamia'
        self.samples = []
        for _, row in self.df.iterrows():
            image_paths = []
            for key in channel_keys:
                image_paths.append(os.path.join(source_directory, row['patient_id'], key + '.png'))
                
            
            mask_path = os.path.join(source_directory, row['patient_id'], 'label.png') 
            breast_mask_path = os.path.join(breast_mask_source_directory, row['patient_id']+'.png') 
            if os.path.exists(mask_path):
                self.samples.append({
                    "image_path": image_paths,
                    "breast_mask_path": breast_mask_path,
                    "mask_path": mask_path,
            })
            else:
                self.samples.append({
                    "image_path": image_paths,
                    "mask_path": 'does not exist',
                })


        self.transform = self.build_transforms()


    def build_transforms(self):
        additional_targets = {'breast_mask': 'mask'}
        if self.augment:
            return A.Compose([
                A.HorizontalFlip(p=0.5),
                A.VerticalFlip(p=0.5),
                A.RandomRotate90(p=0.5),
                A.Rotate(limit=15, p=0.5),
                A.Affine(
                    scale=(0.95, 1.05),
                    shear=5,
                    translate_percent=(0.05, 0.05),
                    p=0.5
                ),
                # A.ColorJitter(brightness=0.1, contrast=0.1, p=0.5),
                A.GaussianBlur(blur_limit=3, p=0.3),
                A.GaussNoise(var_limit=(10.0, 50.0), p=0.3),
                A.Normalize(mean=(0.5,), std=(0.5,)),
                ToTensorV2()
            ], additional_targets=additional_targets)
        else:
            return A.Compose([
                A.Normalize(mean=(0.5,), std=(0.5,)),
                ToTensorV2()
            ], additional_targets=additional_targets)

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

    def __getitem__(self, idx):
        sample = self.samples[idx]

        # Load images as numpy arrays

        image_channels = []
        for image_path in sample["image_path"]:
            img = np.array(Image.open(image_path).convert('L'))  # Convert to grayscale
            image_channels.append(img)

        image = np.stack(image_channels, axis=-1)  # Shape: [H, W, C]
        if os.path.exists(sample["mask_path"]):
            mask = np.array(Image.open(sample["mask_path"]).convert('L'))

        breast_mask =np.array(Image.open(sample["breast_mask_path"]).convert('L'))

        augmented = self.transform(image=image, mask=mask, breast_mask = breast_mask)
        image = augmented['image']           # Tensor [C, H, W]
        # return image
        mask = augmented['mask'].long()      # Tensor [H, W] as long tensor
        mask = mask.unsqueeze(0)             # Tensor [1, H, W] if needed
        breast_mask = augmented['breast_mask'].unsqueeze(0)  # [1, H, W]

        # Combine breast mask as additional input channel
        image = image * breast_mask
        # image = torch.cat([image, breast_mask], dim=0) 
        # else:
        #     augmented = self.transform(image=image)
        #     image = augmented['image']
        #     mask = torch.zeros((1, image.shape[1], image.shape[2]), dtype=torch.long)   
        return image, mask