"""
Medical Imaging Dataset - Image Preprocessing
==============================================

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 image preprocessing functions for medical images
including normalization, resizing, and enhancement.
"""

import numpy as np
import cv2
from PIL import Image
from typing import Tuple, Optional


def preprocess_image(image_path: str, 
                    image_type: str = 'xray',
                    target_size: Optional[Tuple[int, int]] = None,
                    normalize: bool = True) -> np.ndarray:
    """
    Preprocess a medical image.
    
    Args:
        image_path (str): Path to the image file
        image_type (str): Type of medical image ('xray', 'ct_scan', 'mri')
        target_size (tuple): Target size (width, height) for resizing
        normalize (bool): Whether to normalize pixel values to [0, 1]
        
    Returns:
        numpy.ndarray: Preprocessed image array
    """
    # Load image
    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    
    if image is None:
        raise ValueError(f"Could not load image from {image_path}")
    
    # Resize if target size is specified
    if target_size:
        image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA)
    
    # Apply image type specific preprocessing
    if image_type == 'xray':
        image = preprocess_xray(image)
    elif image_type == 'ct_scan':
        image = preprocess_ct_scan(image)
    elif image_type == 'mri':
        image = preprocess_mri(image)
    
    # Normalize if requested
    if normalize:
        image = normalize_image(image)
    
    return image


def preprocess_xray(image: np.ndarray) -> np.ndarray:
    """
    Preprocess X-ray images.
    
    Args:
        image (numpy.ndarray): Input X-ray image
        
    Returns:
        numpy.ndarray: Preprocessed X-ray image
    """
    # Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    image = clahe.apply(image)
    
    # Apply Gaussian blur to reduce noise
    image = cv2.GaussianBlur(image, (3, 3), 0)
    
    return image


def preprocess_ct_scan(image: np.ndarray) -> np.ndarray:
    """
    Preprocess CT scan images.
    
    Args:
        image (numpy.ndarray): Input CT scan image
        
    Returns:
        numpy.ndarray: Preprocessed CT scan image
    """
    # Apply histogram equalization
    image = cv2.equalizeHist(image)
    
    # Apply median filter to reduce noise
    image = cv2.medianBlur(image, 5)
    
    return image


def preprocess_mri(image: np.ndarray) -> np.ndarray:
    """
    Preprocess MRI images.
    
    Args:
        image (numpy.ndarray): Input MRI image
        
    Returns:
        numpy.ndarray: Preprocessed MRI image
    """
    # Apply CLAHE for better contrast
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    image = clahe.apply(image)
    
    # Apply bilateral filter to preserve edges while reducing noise
    image = cv2.bilateralFilter(image, 9, 75, 75)
    
    return image


def normalize_image(image: np.ndarray) -> np.ndarray:
    """
    Normalize image pixel values to [0, 1] range.
    
    Args:
        image (numpy.ndarray): Input image
        
    Returns:
        numpy.ndarray: Normalized image
    """
    # Convert to float
    image = image.astype(np.float32)
    
    # Normalize to [0, 1]
    image = (image - image.min()) / (image.max() - image.min() + 1e-8)
    
    return image


def enhance_contrast(image: np.ndarray, alpha: float = 1.5, beta: int = 0) -> np.ndarray:
    """
    Enhance image contrast.
    
    Args:
        image (numpy.ndarray): Input image
        alpha (float): Contrast control (1.0-3.0)
        beta (int): Brightness control (0-100)
        
    Returns:
        numpy.ndarray: Enhanced image
    """
    enhanced = cv2.convertScaleAbs(image, alpha=alpha, beta=beta)
    return enhanced


def apply_window_level(image: np.ndarray, window: int = 400, level: int = 50) -> np.ndarray:
    """
    Apply window/level transformation (common in medical imaging).
    
    Args:
        image (numpy.ndarray): Input image
        window (int): Window width
        level (int): Window center/level
        
    Returns:
        numpy.ndarray: Windowed image
    """
    min_val = level - window // 2
    max_val = level + window // 2
    
    # Clip values
    image = np.clip(image, min_val, max_val)
    
    # Normalize to [0, 255]
    image = ((image - min_val) / (max_val - min_val) * 255).astype(np.uint8)
    
    return image


def batch_preprocess(images: list, 
                    image_type: str = 'xray',
                    target_size: Optional[Tuple[int, int]] = None,
                    normalize: bool = True) -> list:
    """
    Preprocess a batch of images.
    
    Args:
        images (list): List of image paths or arrays
        image_type (str): Type of medical image
        target_size (tuple): Target size for resizing
        normalize (bool): Whether to normalize
        
    Returns:
        list: List of preprocessed images
    """
    preprocessed = []
    
    for image in images:
        if isinstance(image, str):
            # If it's a path, load and preprocess
            processed = preprocess_image(image, image_type, target_size, normalize)
        else:
            # If it's already an array, preprocess directly
            if target_size:
                image = cv2.resize(image, target_size, interpolation=cv2.INTER_AREA)
            
            if image_type == 'xray':
                processed = preprocess_xray(image)
            elif image_type == 'ct_scan':
                processed = preprocess_ct_scan(image)
            elif image_type == 'mri':
                processed = preprocess_mri(image)
            else:
                processed = image
            
            if normalize:
                processed = normalize_image(processed)
        
        preprocessed.append(processed)
    
    return preprocessed


if __name__ == '__main__':
    # Example usage
    print("Medical Imaging Dataset - Preprocessing Module")
    print("This module provides preprocessing functions for medical images.")
    print("\nAvailable functions:")
    print("- preprocess_image(): Preprocess a single image")
    print("- preprocess_xray(): Preprocess X-ray images")
    print("- preprocess_ct_scan(): Preprocess CT scan images")
    print("- preprocess_mri(): Preprocess MRI images")
    print("- normalize_image(): Normalize image pixel values")
    print("- enhance_contrast(): Enhance image contrast")
    print("- apply_window_level(): Apply window/level transformation")
    print("- batch_preprocess(): Preprocess a batch of images")