"""
Medical Imaging Dataset - Data Loader
======================================

Project: Medical Imaging Dataset
Website: https://rskworld.in
Contact: help@rskworld.in, support@rskworld.in
Phone: +91 93305 39277
Founder: Molla Samser
Designer & Tester: Rima Khatun

This module provides functionality to load medical imaging data including
X-rays, CT scans, and MRI images with their diagnostic labels.
"""

import os
import numpy as np
from PIL import Image
import cv2
import json
from pathlib import Path
try:
    import pydicom
    DICOM_AVAILABLE = True
except ImportError:
    DICOM_AVAILABLE = False
    print("Warning: pydicom not available. DICOM files cannot be loaded.")


class MedicalImagingDataset:
    """
    Class to load and manage medical imaging dataset.
    
    Attributes:
        data_path (str): Path to the dataset directory
        xray_path (str): Path to X-ray images
        ct_path (str): Path to CT scan images
        mri_path (str): Path to MRI images
    """
    
    def __init__(self, data_path='./data'):
        """
        Initialize the MedicalImagingDataset.
        
        Args:
            data_path (str): Root path to the dataset directory
        """
        self.data_path = Path(data_path)
        self.xray_path = self.data_path / 'xray'
        self.ct_path = self.data_path / 'ct_scan'
        self.mri_path = self.data_path / 'mri'
        
    def load_xray_images(self, image_format='PNG'):
        """
        Load X-ray images from the dataset.
        
        Args:
            image_format (str): Image format to load (PNG, JPG, DICOM)
            
        Returns:
            list: List of image arrays and their corresponding labels
        """
        images = []
        labels = []
        
        images_dir = self.xray_path / 'images'
        labels_dir = self.xray_path / 'labels'
        
        if not images_dir.exists():
            print(f"Warning: X-ray images directory not found at {images_dir}")
            return images, labels
        
        for image_file in images_dir.iterdir():
            if not image_file.is_file():
                continue
            if image_file.suffix.lower() in ['.png', '.jpg', '.jpeg']:
                image = self._load_image(image_file)
                if image is not None:
                    images.append(image)
                    # Load corresponding label
                    label_file = labels_dir / f"{image_file.stem}.json"
                    if label_file.exists():
                        try:
                            with open(label_file, 'r') as f:
                                labels.append(json.load(f))
                        except (json.JSONDecodeError, IOError) as e:
                            print(f"Warning: Could not load label file {label_file}: {e}")
                            labels.append(None)
                    else:
                        labels.append(None)
        
        return images, labels
    
    def load_ct_images(self, image_format='PNG'):
        """
        Load CT scan images from the dataset.
        
        Args:
            image_format (str): Image format to load (PNG, JPG, DICOM)
            
        Returns:
            list: List of image arrays and their corresponding labels
        """
        images = []
        labels = []
        
        images_dir = self.ct_path / 'images'
        labels_dir = self.ct_path / 'labels'
        
        if not images_dir.exists():
            print(f"Warning: CT scan images directory not found at {images_dir}")
            return images, labels
        
        for image_file in images_dir.iterdir():
            if not image_file.is_file():
                continue
            if image_file.suffix.lower() in ['.png', '.jpg', '.jpeg']:
                image = self._load_image(image_file)
                if image is not None:
                    images.append(image)
                    # Load corresponding label
                    label_file = labels_dir / f"{image_file.stem}.json"
                    if label_file.exists():
                        try:
                            with open(label_file, 'r') as f:
                                labels.append(json.load(f))
                        except (json.JSONDecodeError, IOError) as e:
                            print(f"Warning: Could not load label file {label_file}: {e}")
                            labels.append(None)
                    else:
                        labels.append(None)
            elif image_file.suffix.lower() == '.dcm' and DICOM_AVAILABLE:
                image = self._load_dicom(image_file)
                if image is not None:
                    images.append(image)
                    label_file = labels_dir / f"{image_file.stem}.json"
                    if label_file.exists():
                        with open(label_file, 'r') as f:
                            labels.append(json.load(f))
                    else:
                        labels.append(None)
        
        return images, labels
    
    def load_mri_images(self, image_format='PNG'):
        """
        Load MRI images from the dataset.
        
        Args:
            image_format (str): Image format to load (PNG, JPG, DICOM)
            
        Returns:
            list: List of image arrays and their corresponding labels
        """
        images = []
        labels = []
        
        images_dir = self.mri_path / 'images'
        labels_dir = self.mri_path / 'labels'
        
        if not images_dir.exists():
            print(f"Warning: MRI images directory not found at {images_dir}")
            return images, labels
        
        for image_file in images_dir.iterdir():
            if not image_file.is_file():
                continue
            if image_file.suffix.lower() in ['.png', '.jpg', '.jpeg']:
                image = self._load_image(image_file)
                if image is not None:
                    images.append(image)
                    # Load corresponding label
                    label_file = labels_dir / f"{image_file.stem}.json"
                    if label_file.exists():
                        try:
                            with open(label_file, 'r') as f:
                                labels.append(json.load(f))
                        except (json.JSONDecodeError, IOError) as e:
                            print(f"Warning: Could not load label file {label_file}: {e}")
                            labels.append(None)
                    else:
                        labels.append(None)
            elif image_file.suffix.lower() == '.dcm' and DICOM_AVAILABLE:
                image = self._load_dicom(image_file)
                if image is not None:
                    images.append(image)
                    label_file = labels_dir / f"{image_file.stem}.json"
                    if label_file.exists():
                        with open(label_file, 'r') as f:
                            labels.append(json.load(f))
                    else:
                        labels.append(None)
        
        return images, labels
    
    def _load_image(self, image_path):
        """
        Load a standard image file (PNG, JPG).
        
        Args:
            image_path (Path): Path to the image file
            
        Returns:
            numpy.ndarray: Image array or None if loading fails
        """
        try:
            image = Image.open(image_path)
            return np.array(image)
        except Exception as e:
            print(f"Error loading image {image_path}: {e}")
            return None
    
    def _load_dicom(self, dicom_path):
        """
        Load a DICOM file.
        
        Args:
            dicom_path (Path): Path to the DICOM file
            
        Returns:
            numpy.ndarray: Image array or None if loading fails
        """
        if not DICOM_AVAILABLE:
            print("pydicom is not available. Cannot load DICOM files.")
            return None
        
        try:
            dicom_data = pydicom.dcmread(str(dicom_path))
            image = dicom_data.pixel_array
            return image
        except Exception as e:
            print(f"Error loading DICOM file {dicom_path}: {e}")
            return None
    
    def get_dataset_info(self):
        """
        Get information about the dataset.
        
        Returns:
            dict: Dictionary containing dataset statistics
        """
        info = {
            'xray_count': 0,
            'ct_count': 0,
            'mri_count': 0,
            'total_images': 0
        }
        
        # Count X-ray images
        if self.xray_path.exists():
            images_dir = self.xray_path / 'images'
            if images_dir.exists():
                info['xray_count'] = len([f for f in images_dir.iterdir() 
                                         if f.is_file() and f.suffix.lower() in ['.png', '.jpg', '.jpeg']])
        
        # Count CT scan images
        if self.ct_path.exists():
            images_dir = self.ct_path / 'images'
            if images_dir.exists():
                info['ct_count'] = len([f for f in images_dir.iterdir() 
                                       if f.is_file() and f.suffix.lower() in ['.png', '.jpg', '.jpeg', '.dcm']])
        
        # Count MRI images
        if self.mri_path.exists():
            images_dir = self.mri_path / 'images'
            if images_dir.exists():
                info['mri_count'] = len([f for f in images_dir.iterdir() 
                                        if f.is_file() and f.suffix.lower() in ['.png', '.jpg', '.jpeg', '.dcm']])
        
        info['total_images'] = info['xray_count'] + info['ct_count'] + info['mri_count']
        
        return info


if __name__ == '__main__':
    # Example usage
    dataset = MedicalImagingDataset(data_path='./data')
    
    # Get dataset information
    info = dataset.get_dataset_info()
    print("Dataset Information:")
    print(f"X-ray images: {info['xray_count']}")
    print(f"CT scan images: {info['ct_count']}")
    print(f"MRI images: {info['mri_count']}")
    print(f"Total images: {info['total_images']}")
    
    # Load images
    xray_images, xray_labels = dataset.load_xray_images()
    ct_images, ct_labels = dataset.load_ct_images()
    mri_images, mri_labels = dataset.load_mri_images()
    
    print(f"\nLoaded {len(xray_images)} X-ray images")
    print(f"Loaded {len(ct_images)} CT scan images")
    print(f"Loaded {len(mri_images)} MRI images")