"""
Image and mask loading module for MedTDA.

This module provides functions to load medical images and masks from various
file formats, supporting 2D (PNG, JPG, JPEG, TIFF) and 3D/4D (NIfTI, NRRD, MHA, MHD) formats.
"""

import numpy as np
from pathlib import Path
from typing import Union, Tuple, Optional

try:
    import SimpleITK as sitk
    SITK_AVAILABLE = True
except ImportError:
    SITK_AVAILABLE = False

try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False


def load_image(path: Union[str, Path]) -> Tuple[np.ndarray, dict]:
    """
    Load image from file (auto-detects format).
    
    Supports:
    - 2D formats: PNG, JPG, JPEG, TIFF
    - 3D/4D formats: NIfTI (.nii, .nii.gz), NRRD (.nrrd), MHA (.mha), MHD (.mhd)
    
    Parameters
    ----------
    path : str or Path
        Path to the image file.
        
    Returns
    -------
    image_array : numpy.ndarray
        Loaded image as numpy array.
    metadata : dict
        Dictionary containing image metadata:
        - 'spacing': Pixel/voxel spacing (tuple)
        - 'origin': Image origin coordinates (tuple)
        - 'direction': Direction cosines (tuple or None for 2D)
        - 'dtype': Original data type
        - 'format': File format
        
    Raises
    ------
    FileNotFoundError
        If the file doesn't exist.
    ValueError
        If the file format is not supported or required library is missing.
    """
    path = Path(path)
    
    if not path.exists():
        raise FileNotFoundError(f"Image file not found: {path}")
    
    # Determine file format by extension
    suffix = path.suffix.lower()
    if path.name.endswith('.nii.gz'):
        suffix = '.nii.gz'
    
    # 2D formats
    if suffix in ['.png', '.jpg', '.jpeg', '.tiff', '.tif', '.bmp']:
        return _load_2d_image(path)
    
    # 3D/4D formats
    elif suffix in ['.nii', '.nii.gz', '.nrrd', '.mha', '.mhd']:
        return _load_3d_image(path)
    
    else:
        raise ValueError(
            f"Unsupported file format: {suffix}. "
            f"Supported formats: PNG, JPG, JPEG, TIFF (2D); "
            f"NIfTI, NRRD, MHA, MHD (3D/4D)"
        )


def _load_2d_image(path: Path) -> Tuple[np.ndarray, dict]:
    """Load 2D image using PIL/Pillow."""
    if not PIL_AVAILABLE:
        raise ImportError(
            "Pillow is required for loading 2D images. "
            "Install with: pip install Pillow"
        )
    
    try:
        img = Image.open(path)
        image_array = np.array(img)
        
        # Convert to grayscale if RGB
        if image_array.ndim == 3:
            if image_array.shape[2] == 3:  # RGB
                image_array = np.mean(image_array, axis=2)
            elif image_array.shape[2] == 4:  # RGBA
                image_array = np.mean(image_array[:, :, :3], axis=2)
        
        # Ensure float type
        image_array = image_array.astype(np.float64)
        
        metadata = {
            'spacing': (1.0, 1.0),  # Default pixel spacing
            'origin': (0.0, 0.0),
            'direction': None,
            'dtype': str(image_array.dtype),
            'format': path.suffix.lower()
        }
        
        return image_array, metadata
        
    except Exception as e:
        raise ValueError(f"Error loading 2D image from {path}: {str(e)}")


def _load_3d_image(path: Path) -> Tuple[np.ndarray, dict]:
    """Load 3D/4D image using SimpleITK."""
    if not SITK_AVAILABLE:
        raise ImportError(
            "SimpleITK is required for loading 3D/4D medical images. "
            "Install with: pip install SimpleITK"
        )
    
    try:
        sitk_image = sitk.ReadImage(str(path))
        image_array = sitk.GetArrayFromImage(sitk_image)
        
        # Convert to float64
        image_array = image_array.astype(np.float64)
        
        metadata = {
            'spacing': sitk_image.GetSpacing(),
            'origin': sitk_image.GetOrigin(),
            'direction': sitk_image.GetDirection(),
            'dtype': str(image_array.dtype),
            'format': path.suffix.lower()
        }
        
        return image_array, metadata
        
    except Exception as e:
        raise ValueError(f"Error loading 3D/4D image from {path}: {str(e)}")


def load_mask(path: Union[str, Path]) -> Tuple[np.ndarray, dict]:
    """
    Load mask from file (auto-detects format).
    
    Masks are loaded preserving all label values for multi-label support.
    
    Parameters
    ----------
    path : str or Path
        Path to the mask file.
        
    Returns
    -------
    mask_array : numpy.ndarray
        Loaded mask as numpy array (dtype: uint8) with original label values preserved.
    metadata : dict
        Dictionary containing mask metadata (same structure as load_image).
        
    Raises
    ------
    FileNotFoundError
        If the file doesn't exist.
    ValueError
        If the file format is not supported.
        
    Notes
    -----
    Multi-label masks are supported. Use the `label` parameter in preprocessing
    and extraction classes to specify which label to work with.
    """
    try:
        mask_array, metadata = load_image(path)
    except FileNotFoundError:
        raise FileNotFoundError(f"Mask file not found: {Path(path)}")
    
    # Convert to uint8 and preserve all label values (supports multi-label masks)
    mask_array = mask_array.astype(np.uint8)
    metadata['dtype'] = 'uint8'
    
    return mask_array, metadata


def validate_image(image: np.ndarray) -> None:
    """
    Validate image array.
    
    Parameters
    ----------
    image : numpy.ndarray
        Image array to validate.
        
    Raises
    ------
    TypeError
        If image is not a numpy array.
    ValueError
        If image has invalid dimensions or contains invalid values.
    """
    if not isinstance(image, np.ndarray):
        raise TypeError(f"Image must be a numpy array. Got {type(image)}")
    
    if image.size == 0:
        raise ValueError("Image array is empty.")
    
    if image.ndim not in [2, 3, 4]:
        raise ValueError(
            f"Image must be 2D, 3D, or 4D. Got {image.ndim}D array."
        )
    
    if not np.isfinite(image).all():
        raise ValueError("Image contains NaN or infinite values.")


def validate_mask(mask: np.ndarray) -> None:
    """
    Validate mask array.
    
    Parameters
    ----------
    mask : numpy.ndarray
        Mask array to validate.
        
    Raises
    ------
    TypeError
        If mask is not a numpy array.
    ValueError
        If mask has invalid dimensions or values.
    """
    if not isinstance(mask, np.ndarray):
        raise TypeError(f"Mask must be a numpy array. Got {type(mask)}")
    
    if mask.ndim not in [2, 3, 4]:
        raise ValueError(
            f"Mask must be 2D, 3D, or 4D. Got {mask.ndim}D array."
        )
    
    if mask.size == 0:
        raise ValueError("Mask array is empty.")
    
    # Check if mask is binary
    unique_values = np.unique(mask)
    if not np.all(np.isin(unique_values, [0, 1])):
        raise ValueError(
            f"Mask must be binary (0 and 1 only). "
            f"Found values: {unique_values}"
        )


def validate_compatibility(image: np.ndarray, mask: np.ndarray) -> None:
    """
    Check image-mask compatibility.
    
    Ensures image and mask have the same shape.
    
    Parameters
    ----------
    image : numpy.ndarray
        Image array.
    mask : numpy.ndarray
        Mask array.
        
    Raises
    ------
    ValueError
        If image and mask shapes don't match.
    """
    if image.shape != mask.shape:
        raise ValueError(
            f"Image shape {image.shape} does not match "
            f"mask shape {mask.shape}."
        )


def save_image(
    image_array: np.ndarray,
    path: Union[str, Path],
    metadata: Optional[dict] = None
) -> None:
    """
    Save image array to file.
    
    Parameters
    ----------
    image_array : numpy.ndarray
        Image array to save.
    path : str or Path
        Output file path.
    metadata : dict, optional
        Image metadata (spacing, origin, direction).
        Only used for 3D/4D formats.
        
    Raises
    ------
    ValueError
        If the file format is not supported.
    """
    path = Path(path)
    suffix = path.suffix.lower()
    if path.name.endswith('.nii.gz'):
        suffix = '.nii.gz'
    
    # 2D formats
    if suffix in ['.png', '.jpg', '.jpeg', '.tiff', '.tif']:
        _save_2d_image(image_array, path)
    
    # 3D/4D formats
    elif suffix in ['.nii', '.nii.gz', '.nrrd', '.mha', '.mhd']:
        _save_3d_image(image_array, path, metadata)
    
    else:
        raise ValueError(
            f"Unsupported file format for saving: {suffix}. "
            f"Supported formats: PNG, JPG, JPEG, TIFF (2D); "
            f"NIfTI, NRRD, MHA, MHD (3D/4D)"
        )


def _save_2d_image(image_array: np.ndarray, path: Path) -> None:
    """Save 2D image using PIL/Pillow."""
    if not PIL_AVAILABLE:
        raise ImportError(
            "Pillow is required for saving 2D images. "
            "Install with: pip install Pillow"
        )
    
    # Normalize to 0-255 range for saving
    image_min = image_array.min()
    image_max = image_array.max()
    if image_max > image_min:
        normalized = ((image_array - image_min) / (image_max - image_min) * 255).astype(np.uint8)
    else:
        normalized = np.zeros_like(image_array, dtype=np.uint8)
    
    img = Image.fromarray(normalized)
    img.save(path)


def _save_3d_image(
    image_array: np.ndarray,
    path: Path,
    metadata: Optional[dict] = None
) -> None:
    """Save 3D/4D image using SimpleITK."""
    if not SITK_AVAILABLE:
        raise ImportError(
            "SimpleITK is required for saving 3D/4D medical images. "
            "Install with: pip install SimpleITK"
        )
    
    sitk_image = sitk.GetImageFromArray(image_array)
    
    # Set metadata if provided
    if metadata:
        if 'spacing' in metadata:
            sitk_image.SetSpacing(metadata['spacing'])
        if 'origin' in metadata:
            sitk_image.SetOrigin(metadata['origin'])
        if 'direction' in metadata and metadata['direction'] is not None:
            sitk_image.SetDirection(metadata['direction'])
    
    sitk.WriteImage(sitk_image, str(path))
