medrs 0.1.2

#![allow(clippy::doc_markdown)]
//! Python bindings for medrs using PyO3.

use ndarray::{ArrayD, IxDyn};
use numpy::{IntoPyArray, PyArrayDyn, PyArrayMethods, PyReadonlyArrayDyn};
use pyo3::exceptions::{
    PyFileNotFoundError, PyIOError, PyMemoryError, PyStopIteration, PyValueError,
};
use pyo3::prelude::*;
use pyo3::types::IntoPyDict;

use crate::error::Error as MedrsError;
use crate::nifti::image::ArrayData;

/// Convert a medrs Error to the appropriate Python exception.
/// This provides proper exception types for different error conditions.
fn to_py_err(e: MedrsError, context: &str) -> PyErr {
    match &e {
        MedrsError::Io(io_err) => {
            // Map I/O errors to PyIOError
            PyIOError::new_err(format!("{}: {}", context, io_err))
        }
        MedrsError::MemoryAllocation(msg) => {
            PyMemoryError::new_err(format!("{}: {}", context, msg))
        }
        MedrsError::InvalidDimensions(msg)
        | MedrsError::InvalidAffine(msg)
        | MedrsError::InvalidCropRegion(msg)
        | MedrsError::ShapeMismatch(msg)
        | MedrsError::InvalidFileFormat(msg)
        | MedrsError::InvalidOrientation(msg)
        | MedrsError::NonContiguousArray(msg)
        | MedrsError::Configuration(msg)
        | MedrsError::Decompression(msg)
        | MedrsError::Exhausted(msg) => PyValueError::new_err(format!("{}: {}", context, msg)),
        MedrsError::InvalidMagic(magic) => PyValueError::new_err(format!(
            "{}: invalid NIfTI magic bytes {:?}",
            context, magic
        )),
        MedrsError::UnsupportedDataType(code) => {
            PyValueError::new_err(format!("{}: unsupported data type code {}", context, code))
        }
        MedrsError::DataTypeMismatch { expected, got } => PyValueError::new_err(format!(
            "{}: data type mismatch (expected {}, got {})",
            context, expected, got
        )),
        MedrsError::TransformError { operation, reason } => {
            PyValueError::new_err(format!("{}: {} failed: {}", context, operation, reason))
        }
    }
}
use crate::nifti::DataType;
use crate::nifti::{self, NiftiImage as RustNiftiImage};
use crate::pipeline::TransformPipeline as RustTransformPipeline;
use crate::transforms::crop::{
    compute_center_crop_regions, compute_label_aware_crop_regions,
    compute_random_spatial_crop_regions,
};
use crate::transforms::{self, Interpolation, Orientation};

// ============================================================================
// Validation helpers for Python boundary
// ============================================================================

/// Validate shape array has positive dimensions.
fn validate_shape(shape: &[usize; 3], name: &str) -> PyResult<()> {
    for (i, &dim) in shape.iter().enumerate() {
        if dim == 0 {
            return Err(PyValueError::new_err(format!(
                "{} dimension {} must be positive (got 0)",
                name, i
            )));
        }
    }
    Ok(())
}

/// Validate spacing array has positive values.
fn validate_spacing(spacing: &[f32; 3], name: &str) -> PyResult<()> {
    for (i, &s) in spacing.iter().enumerate() {
        if s <= 0.0 {
            return Err(PyValueError::new_err(format!(
                "{} dimension {} must be positive (got {})",
                name, i, s
            )));
        }
        if !s.is_finite() {
            return Err(PyValueError::new_err(format!(
                "{} dimension {} must be finite (got {})",
                name, i, s
            )));
        }
    }
    Ok(())
}

/// Validate a 3-element shape vector and return it as an array.
fn parse_shape3(values: &[usize], name: &str) -> PyResult<[usize; 3]> {
    if values.len() != 3 {
        return Err(PyValueError::new_err(format!(
            "{} must be a 3-element sequence (got {})",
            name,
            values.len()
        )));
    }
    let shape = [values[0], values[1], values[2]];
    validate_shape(&shape, name)?;
    Ok(shape)
}

/// Validate file path is safe and exists.
fn validate_file_path(path: &str, operation: &str) -> PyResult<std::path::PathBuf> {
    if path.is_empty() {
        return Err(PyValueError::new_err(format!(
            "{}: file path cannot be empty",
            operation
        )));
    }

    // Check for null bytes (prevents injection in C APIs)
    if path.contains('\0') {
        return Err(PyValueError::new_err(format!(
            "{}: file path cannot contain null bytes",
            operation
        )));
    }

    let path_buf = std::path::PathBuf::from(path);

    // For loading operations, check if file exists
    if operation.contains("load") || operation.contains("read") {
        if !path_buf.exists() {
            return Err(PyFileNotFoundError::new_err(format!(
                "{}: file not found: {}",
                operation, path
            )));
        }

        if !path_buf.is_file() {
            return Err(PyValueError::new_err(format!(
                "{}: path is not a file: {}",
                operation, path
            )));
        }
    }

    // For saving operations, check if parent directory exists
    if operation.contains("save") || operation.contains("write") {
        if let Some(parent) = path_buf.parent() {
            if !parent.exists() {
                return Err(PyFileNotFoundError::new_err(format!(
                    "{}: parent directory does not exist: {}",
                    operation,
                    parent.display()
                )));
            }
        }
    }

    Ok(path_buf)
}

/// Validate intensity range parameters.
fn validate_intensity_range(min: f64, max: f64, param_name: &str) -> PyResult<()> {
    if !min.is_finite() {
        return Err(PyValueError::new_err(format!(
            "{}: min value must be finite (got {})",
            param_name, min
        )));
    }
    if !max.is_finite() {
        return Err(PyValueError::new_err(format!(
            "{}: max value must be finite (got {})",
            param_name, max
        )));
    }
    if min > max {
        return Err(PyValueError::new_err(format!(
            "{}: min ({}) cannot be greater than max ({})",
            param_name, min, max
        )));
    }
    Ok(())
}

/// Validate probability value (0.0 to 1.0).
fn validate_probability(p: f64, param_name: &str) -> PyResult<()> {
    if !p.is_finite() {
        return Err(PyValueError::new_err(format!(
            "{}: probability must be finite (got {})",
            param_name, p
        )));
    }
    if !(0.0..=1.0).contains(&p) {
        return Err(PyValueError::new_err(format!(
            "{}: probability must be between 0.0 and 1.0 (got {})",
            param_name, p
        )));
    }
    Ok(())
}

/// Shared helper for creating NiftiImage from numpy array.
/// Handles F-order conversion and validation.
fn create_nifti_from_numpy_array(
    arr: ndarray::ArrayViewD<'_, f32>,
    affine: Option<[[f32; 4]; 4]>,
) -> PyResult<RustNiftiImage> {
    use ndarray::ShapeBuilder;

    let shape = arr.shape();

    if shape.len() < 3 {
        return Err(PyValueError::new_err(
            "Array must have at least 3 dimensions (D,H,W)",
        ));
    }

    // Validate that no dimension is zero
    for (i, &dim) in shape.iter().enumerate() {
        if dim == 0 {
            return Err(PyValueError::new_err(format!(
                "Array dimension {} cannot be 0",
                i
            )));
        }
    }

    // Check for integer overflow when casting to u16 (NIfTI header limitation)
    for (i, &dim) in shape.iter().enumerate() {
        if dim > u16::MAX as usize {
            return Err(PyValueError::new_err(format!(
                "Array dimension {} ({}) exceeds maximum NIfTI dimension size ({})",
                i,
                dim,
                u16::MAX
            )));
        }
    }

    // Create F-order array to match NIfTI convention
    // Use as_slice_memory_order to get data in physical layout
    #[allow(clippy::option_if_let_else)]
    let data_vec: Vec<f32> = if let Some(slice) = arr.as_slice_memory_order() {
        slice.to_vec()
    } else {
        // Fallback: iterate in logical order
        arr.iter().copied().collect()
    };

    // Determine if input is F-order
    let is_f_order = !arr.is_standard_layout() && arr.as_slice_memory_order().is_some();

    let array = if is_f_order {
        // Input was F-order, data_vec is in F-order
        ArrayD::from_shape_vec(IxDyn(shape).f(), data_vec)
            .map_err(|e| PyValueError::new_err(format!("Invalid array shape: {}", e)))?
    } else {
        // Input was C-order, convert to F-order
        let c_order = ArrayD::from_shape_vec(shape.to_vec(), data_vec)
            .map_err(|e| PyValueError::new_err(format!("Invalid array shape: {}", e)))?;
        let mut f_order = ArrayD::zeros(IxDyn(shape).f());
        f_order.assign(&c_order);
        f_order
    };

    let affine = affine.unwrap_or([
        [1.0, 0.0, 0.0, 0.0],
        [0.0, 1.0, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0],
    ]);

    Ok(RustNiftiImage::from_array(array, affine))
}

/// A NIfTI image with header metadata and voxel data.
///
/// Supports method chaining for transform operations.
///
/// Example:
///     >>> img = medrs.load("brain.nii.gz")
///     >>> processed = img.resample([1.0, 1.0, 1.0]).z_normalize().clamp(0, 1)
///     >>> processed.save("output.nii.gz")
#[pyclass(name = "NiftiImage")]
struct PyNiftiImage {
    inner: RustNiftiImage,
}

#[pymethods]
impl PyNiftiImage {
    /// Create a new NIfTI image from a numpy array (>=3D).
    ///
    /// Args:
    ///     data: numpy array of voxel values (last 3 dims are spatial)
    ///     affine: 4x4 affine transformation matrix (optional)
    #[new]
    #[pyo3(signature = (data, affine=None))]
    fn new(data: PyReadonlyArrayDyn<'_, f32>, affine: Option<[[f32; 4]; 4]>) -> PyResult<Self> {
        let inner = create_nifti_from_numpy_array(data.as_array(), affine)?;
        Ok(Self { inner })
    }

    /// Image shape as (depth, height, width).
    #[getter]
    fn shape(&self) -> Vec<usize> {
        self.inner.shape().to_vec()
    }

    /// Number of dimensions.
    #[getter]
    fn ndim(&self) -> usize {
        self.inner.ndim()
    }

    /// Data type as string.
    #[getter]
    fn dtype(&self) -> &'static str {
        self.inner.dtype().type_name()
    }

    /// Voxel spacing in mm.
    #[getter]
    fn spacing(&self) -> Vec<f32> {
        self.inner.spacing().to_vec()
    }

    /// 4x4 affine transformation matrix.
    #[getter]
    fn affine(&self) -> [[f32; 4]; 4] {
        self.inner.affine()
    }

    /// Set the affine transformation matrix.
    #[setter]
    fn set_affine(&mut self, affine: [[f32; 4]; 4]) {
        self.inner.set_affine(affine);
    }

    /// Image orientation code (e.g., "RAS", "LPS").
    #[getter]
    fn orientation(&self) -> String {
        let affine = self.inner.affine();
        transforms::orientation_from_affine(&affine).to_string()
    }

    /// Raw voxel data as a numpy array (float32).
    ///
    /// Uses zero-copy views when possible; may fall back to a copy for non-contiguous data.
    #[getter]
    fn data<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
        self.to_numpy(py)
    }

    /// Get image data as numpy array (float32).
    ///
    /// Similar to nibabel's get_fdata(). Applies scaling factors if present.
    /// Supports arbitrary ndim.
    fn to_numpy<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
        if let Some(arr) = to_numpy_view(py, &self.inner) {
            return Ok(arr);
        }
        to_numpy_array(py, &self.inner)
    }

    /// Get image data as numpy view when possible (may fall back to copy).
    #[allow(clippy::option_if_let_else)]
    fn to_numpy_view<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
        if let Some(arr) = to_numpy_view(py, &self.inner) {
            Ok(arr)
        } else {
            to_numpy_array(py, &self.inner)
        }
    }

    /// Get image data as a torch tensor (shares memory when possible).
    fn to_torch(&self, py: Python<'_>) -> PyResult<PyObject> {
        let torch = py.import("torch")?;
        let dtype = torch_dtype(py, self.inner.dtype());
        let has_dtype = dtype.is_some();

        // Prefer view when dtype is torch-supported and contiguous
        let np_obj = if has_dtype {
            if let Some(np_view) = to_numpy_view_native(py, &self.inner) {
                np_view
            } else {
                arraydata_to_numpy(
                    py,
                    &self
                        .inner
                        .owned_data()
                        .map_err(|e| to_py_err(e, "to_torch owned_data"))?,
                    self.inner.shape(),
                )?
            }
        } else {
            to_numpy_array(py, &self.inner)?
                .into_pyobject(py)?
                .into_any()
                .unbind()
        };

        let tensor = torch.getattr("from_numpy")?.call1((np_obj,))?;

        if let Some(dt) = dtype {
            let tensor = tensor.call_method1("to", (dt,))?;
            Ok(tensor.unbind())
        } else {
            Ok(tensor.unbind())
        }
    }

    /// Get image data as a JAX array (shares memory via numpy when possible).
    fn to_jax(&self, py: Python<'_>) -> PyResult<PyObject> {
        let jnp = py.import("jax.numpy")?;
        let np_obj = to_numpy_array(py, &self.inner)?;
        let arr = jnp.getattr("array")?.call1((np_obj,))?;
        Ok(arr.unbind())
    }

    /// Convert to PyTorch tensor with custom dtype and device.
    ///
    /// This is the most efficient way to load medical imaging data directly
    /// into PyTorch with the target precision and device placement.
    #[pyo3(signature = (dtype=None, device=None))]
    fn to_torch_with_dtype_and_device(
        &self,
        py: Python<'_>,
        dtype: Option<PyObject>,
        device: Option<&str>,
    ) -> PyResult<PyObject> {
        let device_str = device.unwrap_or("cpu");
        let torch = py.import("torch")?;

        // Get numpy array with maximum efficiency
        let np_obj = if let Some(np_view) = to_numpy_view_native(py, &self.inner) {
            np_view
        } else {
            to_numpy_array(py, &self.inner)?
                .into_pyobject(py)?
                .into_any()
                .unbind()
        };

        // Convert to tensor
        let tensor = torch.getattr("from_numpy")?.call1((np_obj,))?;

        // Apply device placement
        let tensor = tensor.call_method1("to", (device_str,))?;

        // Apply dtype if specified
        if let Some(dt) = dtype {
            let tensor = tensor.call_method1("to", (dt,))?;
            Ok(tensor.unbind())
        } else {
            Ok(tensor.unbind())
        }
    }

    /// Convert to JAX array with custom dtype and device.
    ///
    /// This is the most efficient way to load medical imaging data directly
    /// into JAX with the target precision and device placement.
    #[allow(clippy::useless_let_if_seq)]
    #[pyo3(signature = (dtype=None, device=None))]
    fn to_jax_with_dtype_and_device(
        &self,
        py: Python<'_>,
        dtype: Option<PyObject>,
        device: Option<&str>,
    ) -> PyResult<PyObject> {
        let device_str = device.unwrap_or("cpu");
        let jax = py.import("jax")?;
        let _jnp = py.import("jax.numpy")?;

        // Get device object using correct JAX API
        let device_obj = if device_str == "cpu" {
            let cpu_devices = jax.getattr("devices")?.call1(("cpu",))?;
            cpu_devices.get_item(0)?
        } else if device_str.starts_with("cuda") || device_str.starts_with("gpu") {
            // JAX uses "gpu" not "cuda" for CUDA devices
            let gpu_devices = jax.getattr("devices")?.call1(("gpu",))?;
            let device_id: usize = device_str
                .strip_prefix("cuda:")
                .or_else(|| device_str.strip_prefix("gpu:"))
                .unwrap_or("0")
                .parse()
                .unwrap_or(0);
            gpu_devices.get_item(device_id)?
        } else {
            return Err(PyValueError::new_err(format!(
                "Unsupported device: {}. Use 'cpu', 'cuda', 'cuda:N', 'gpu', or 'gpu:N'",
                device_str
            )));
        };

        // Use optimized on-device array creation
        let jax = py.import("jax")?;
        let jnp = py.import("jax.numpy")?;

        // Check if we can use zero-copy numpy path for maximum efficiency
        if let Some(np_view) = to_numpy_view_native(py, &self.inner) {
            // Use numpy view for maximum efficiency
            let mut arr = jnp.getattr("array")?.call1((np_view,))?;

            // Apply dtype if specified
            if let Some(dt) = dtype {
                arr = jnp.getattr("astype")?.call1((dt,))?;
            }

            // Use jax.device_put for efficient async device placement
            let device_put = jax.getattr("device_put")?;
            let arr = device_put.call1((arr, &device_obj))?;

            return Ok(arr.into());
        }

        // Fallback to regular numpy array with efficient transfer
        let np_obj = to_numpy_array(py, &self.inner)?;
        let mut arr = jnp.getattr("array")?.call1((np_obj,))?;

        // Apply dtype if specified
        if let Some(dt) = dtype {
            arr = jnp.getattr("astype")?.call1((dt,))?;
        }

        // Use jax.device_put for efficient async device placement
        let device_put = jax.getattr("device_put")?;
        let arr = device_put.call1((arr, &device_obj))?;

        Ok(arr.into())
    }

    /// Get image data as a numpy array with native dtype.
    ///
    /// Half/bfloat16 are returned as float32 for compatibility.
    fn to_numpy_native(&self, py: Python<'_>) -> PyResult<PyObject> {
        arraydata_to_numpy(
            py,
            &self
                .inner
                .owned_data()
                .map_err(|e| to_py_err(e, "to_numpy_native"))?,
            self.inner.shape(),
        )
    }

    /// Create a new NiftiImage from a numpy array.
    ///
    /// This is a convenience method for creating NIfTI images from numpy arrays.
    /// For more control, use the NiftiImage constructor directly.
    ///
    /// Args:
    ///     data: numpy array of voxel values (last 3 dims are spatial)
    ///     affine: 4x4 affine transformation matrix (optional)
    ///
    /// Returns:
    ///     NiftiImage instance
    ///
    /// Example:
    ///     >>> import numpy as np
    ///     >>> data = np.random.rand(64, 64, 32).astype(np.float32)
    ///     >>> img = NiftiImage.from_numpy(data)
    #[staticmethod]
    #[pyo3(signature = (data, affine=None))]
    fn from_numpy<'py>(
        _py: Python<'py>,
        data: PyReadonlyArrayDyn<'py, f32>,
        affine: Option<[[f32; 4]; 4]>,
    ) -> PyResult<Self> {
        let inner = create_nifti_from_numpy_array(data.as_array(), affine)?;
        Ok(Self { inner })
    }

    /// Save image to file.
    ///
    /// Format is determined by extension (.nii or .nii.gz).
    fn save(&self, path: &str) -> PyResult<()> {
        let validated_path = validate_file_path(path, "save")?;
        let path_str = validated_path
            .to_str()
            .ok_or_else(|| PyValueError::new_err("path contains invalid UTF-8"))?;
        nifti::save(&self.inner, path_str)
            .map_err(|e| to_py_err(e, &format!("Failed to save {}", path)))
    }

    /// Convert image to a different data type.
    ///
    /// This is useful for reducing file size when saving. For example,
    /// converting from float32 to bfloat16 reduces storage by 50%.
    ///
    /// Args:
    ///     dtype: Target dtype as string. Supported values:
    ///         - "float32", "f32" - 32-bit float (default)
    ///         - "float64", "f64" - 64-bit float
    ///         - "float16", "f16" - IEEE 754 half precision
    ///         - "bfloat16", "bf16" - Brain floating point 16-bit
    ///         - "int8", "i8" - Signed 8-bit integer
    ///         - "uint8", "u8" - Unsigned 8-bit integer
    ///         - "int16", "i16" - Signed 16-bit integer
    ///         - "uint16", "u16" - Unsigned 16-bit integer
    ///         - "int32", "i32" - Signed 32-bit integer
    ///         - "uint32", "u32" - Unsigned 32-bit integer
    ///         - "int64", "i64" - Signed 64-bit integer
    ///         - "uint64", "u64" - Unsigned 64-bit integer
    ///
    /// Returns:
    ///     New MedicalImage with converted dtype
    ///
    /// Example:
    ///     >>> img = medrs.load("volume.nii.gz")
    ///     >>> img_bf16 = img.with_dtype("bfloat16")
    ///     >>> img_bf16.save("volume_bf16.nii.gz")  # 50% smaller file
    fn with_dtype(&self, dtype: &str) -> PyResult<Self> {
        let target_dtype = match dtype.to_lowercase().as_str() {
            "float32" | "f32" => nifti::DataType::Float32,
            "float64" | "f64" => nifti::DataType::Float64,
            "float16" | "f16" => nifti::DataType::Float16,
            "bfloat16" | "bf16" => nifti::DataType::BFloat16,
            "int8" | "i8" => nifti::DataType::Int8,
            "uint8" | "u8" => nifti::DataType::UInt8,
            "int16" | "i16" => nifti::DataType::Int16,
            "uint16" | "u16" => nifti::DataType::UInt16,
            "int32" | "i32" => nifti::DataType::Int32,
            "uint32" | "u32" => nifti::DataType::UInt32,
            "int64" | "i64" => nifti::DataType::Int64,
            "uint64" | "u64" => nifti::DataType::UInt64,
            _ => return Err(PyValueError::new_err(format!(
                "Unsupported dtype '{}'. Use: float32, float64, float16, bfloat16, int8, uint8, int16, uint16, int32, uint32, int64, uint64",
                dtype
            ))),
        };

        Ok(Self {
            inner: self
                .inner
                .with_dtype(target_dtype)
                .map_err(|e| to_py_err(e, "with_dtype"))?,
        })
    }

    /// Resample to target voxel spacing.
    ///
    /// Args:
    ///     spacing: Target spacing as [x, y, z] in mm
    ///     method: Interpolation method ("trilinear" or "nearest", default: "trilinear")
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    #[pyo3(signature = (spacing, method=None))]
    fn resample(&self, spacing: [f32; 3], method: Option<&str>) -> PyResult<Self> {
        let method_str = method.unwrap_or("trilinear");
        let interp = match method_str {
            "trilinear" | "linear" => Interpolation::Trilinear,
            "nearest" => Interpolation::Nearest,
            _ => {
                return Err(PyValueError::new_err(
                    "method must be 'trilinear' or 'nearest'",
                ))
            }
        };

        let resampled = transforms::resample_to_spacing(&self.inner, spacing, interp)
            .map_err(|e| PyValueError::new_err(format!("Resampling failed: {}", e)))?;
        Ok(Self { inner: resampled })
    }

    /// Resample to target shape.
    ///
    /// Args:
    ///     shape: Target shape as [depth, height, width]
    ///     method: Interpolation method ("trilinear" or "nearest", default: "trilinear")
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    #[pyo3(signature = (shape, method=None))]
    fn resample_to_shape(&self, shape: [usize; 3], method: Option<&str>) -> PyResult<Self> {
        let method_str = method.unwrap_or("trilinear");
        let interp = match method_str {
            "trilinear" | "linear" => Interpolation::Trilinear,
            "nearest" => Interpolation::Nearest,
            _ => {
                return Err(PyValueError::new_err(
                    "method must be 'trilinear' or 'nearest'",
                ))
            }
        };

        let resampled = transforms::resample_to_shape(&self.inner, shape, interp)
            .map_err(|e| PyValueError::new_err(format!("Resampling failed: {}", e)))?;
        Ok(Self { inner: resampled })
    }

    /// Reorient to target orientation.
    ///
    /// Args:
    ///     orientation: Target orientation code (e.g., "RAS", "LPS")
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn reorient(&self, orientation: &str) -> PyResult<Self> {
        let target: Orientation = orientation
            .parse()
            .map_err(|e| PyValueError::new_err(format!("{}", e)))?;

        Ok(Self {
            inner: transforms::reorient(&self.inner, target)
                .map_err(|e| PyValueError::new_err(format!("Reorientation failed: {}", e)))?,
        })
    }

    /// Z-score normalization (zero mean, unit variance).
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn z_normalize(&self) -> PyResult<Self> {
        Ok(Self {
            inner: transforms::z_normalization(&self.inner)
                .map_err(|e| to_py_err(e, "z_normalize"))?,
        })
    }

    /// Rescale intensity to range [min, max].
    ///
    /// Args:
    ///     out_min: Minimum output value
    ///     out_max: Maximum output value
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn rescale(&self, out_min: f64, out_max: f64) -> PyResult<Self> {
        Ok(Self {
            inner: transforms::rescale_intensity(&self.inner, out_min, out_max)
                .map_err(|e| to_py_err(e, "rescale"))?,
        })
    }

    /// Clamp intensity values to range [min, max].
    ///
    /// Args:
    ///     min: Minimum value
    ///     max: Maximum value
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn clamp(&self, min: f64, max: f64) -> PyResult<Self> {
        Ok(Self {
            inner: transforms::clamp(&self.inner, min, max).map_err(|e| to_py_err(e, "clamp"))?,
        })
    }

    /// Crop or pad to target shape.
    ///
    /// Args:
    ///     target_shape: Target shape as [depth, height, width]
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn crop_or_pad(&self, target_shape: Vec<usize>) -> PyResult<Self> {
        Ok(Self {
            inner: transforms::crop_or_pad(&self.inner, &target_shape)
                .map_err(|e| PyValueError::new_err(e.to_string()))?,
        })
    }

    /// Flip along specified axes.
    ///
    /// Args:
    ///     axes: List of axes to flip (0=depth, 1=height, 2=width)
    ///
    /// Returns:
    ///     New NiftiImage (supports method chaining)
    fn flip(&self, axes: Vec<usize>) -> PyResult<Self> {
        Ok(Self {
            inner: transforms::flip(&self.inner, &axes)
                .map_err(|e| PyValueError::new_err(e.to_string()))?,
        })
    }

    /// Check if the image data is already materialized in memory.
    ///
    /// When False, the data is mmap'd from disk and will be materialized
    /// on each transform call.
    fn is_materialized(&self) -> bool {
        self.inner.is_materialized()
    }

    /// Convert mmap'd data to owned memory.
    ///
    /// Call this once before running multiple transforms to avoid
    /// re-materializing the data on each transform call.
    ///
    /// Returns:
    ///     New NiftiImage with data in memory (supports method chaining)
    ///
    /// Example:
    ///     >>> img = medrs.load("brain.nii.gz").materialize()
    ///     >>> # Now transforms are fast as data is in memory
    ///     >>> processed = img.z_normalize().rescale(0, 1).flip([0])
    fn materialize(&self) -> PyResult<Self> {
        Ok(Self {
            inner: self
                .inner
                .materialize()
                .map_err(|e| to_py_err(e, "materialize"))?,
        })
    }

    fn __repr__(&self) -> String {
        format!(
            "NiftiImage(shape={:?}, dtype={}, spacing={:?}, orientation={})",
            self.shape(),
            self.dtype(),
            self.spacing(),
            self.orientation()
        )
    }
}

/// Z-score normalize an image (zero mean, unit variance).
#[pyfunction]
fn z_normalization(image: &PyNiftiImage) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::z_normalization(&image.inner)
            .map_err(|e| to_py_err(e, "z_normalization"))?,
    })
}

/// Rescale intensity to the provided range.
#[pyfunction]
#[pyo3(signature = (image, output_range=(0.0, 1.0)))]
fn rescale_intensity(image: &PyNiftiImage, output_range: (f64, f64)) -> PyResult<PyNiftiImage> {
    let (out_min, out_max) = output_range;
    Ok(PyNiftiImage {
        inner: transforms::rescale_intensity(&image.inner, out_min, out_max)
            .map_err(|e| to_py_err(e, "rescale_intensity"))?,
    })
}

/// Clamp intensity values into a fixed range.
#[pyfunction]
fn clamp(image: &PyNiftiImage, min_value: f64, max_value: f64) -> PyResult<PyNiftiImage> {
    validate_intensity_range(min_value, max_value, "clamp")?;
    Ok(PyNiftiImage {
        inner: transforms::clamp(&image.inner, min_value, max_value)
            .map_err(|e| to_py_err(e, "clamp"))?,
    })
}

/// Crop or pad an image to the target shape.
#[pyfunction]
fn crop_or_pad(image: &PyNiftiImage, target_shape: Vec<usize>) -> PyResult<PyNiftiImage> {
    if target_shape.len() != 3 {
        return Err(PyValueError::new_err(
            "target_shape must be a 3-element sequence",
        ));
    }

    Ok(PyNiftiImage {
        inner: transforms::crop_or_pad(&image.inner, &target_shape)
            .map_err(|e| PyValueError::new_err(e.to_string()))?,
    })
}

/// Resample to target voxel spacing.
#[pyfunction]
#[pyo3(signature = (image, target_spacing, method=None))]
fn resample(
    image: &PyNiftiImage,
    target_spacing: (f32, f32, f32),
    method: Option<&str>,
) -> PyResult<PyNiftiImage> {
    let interp = match method.unwrap_or("trilinear") {
        "trilinear" | "linear" => Interpolation::Trilinear,
        "nearest" => Interpolation::Nearest,
        other => {
            return Err(PyValueError::new_err(format!(
                "method must be 'trilinear' or 'nearest', got {}",
                other
            )))
        }
    };

    let spacing = [target_spacing.0, target_spacing.1, target_spacing.2];

    let resampled = transforms::resample_to_spacing(&image.inner, spacing, interp)
        .map_err(|e| PyValueError::new_err(format!("Resampling failed: {}", e)))?;
    Ok(PyNiftiImage { inner: resampled })
}

/// Reorient an image to the target orientation code (e.g., RAS or LPS).
#[pyfunction]
fn reorient(image: &PyNiftiImage, orientation: &str) -> PyResult<PyNiftiImage> {
    let target: Orientation = orientation
        .parse()
        .map_err(|e| PyValueError::new_err(format!("{}", e)))?;

    Ok(PyNiftiImage {
        inner: transforms::reorient(&image.inner, target)
            .map_err(|e| PyValueError::new_err(format!("Reorientation failed: {}", e)))?,
    })
}

/// Load a NIfTI image from file.
///
/// Supports both .nii and .nii.gz formats.
///
/// Args:
///     path: Path to the NIfTI file
///
/// Returns:
///     NiftiImage instance
///
/// Example:
///     >>> img = medrs.load("brain.nii.gz")
#[pyfunction]
fn load(path: &str) -> PyResult<PyNiftiImage> {
    let validated_path = validate_file_path(path, "load")?;
    let path_str = validated_path
        .to_str()
        .ok_or_else(|| PyValueError::new_err("path contains invalid UTF-8"))?;
    nifti::load(path_str)
        .map(|inner| PyNiftiImage { inner })
        .map_err(|e| to_py_err(e, &format!("Failed to load {}", path)))
}

/// Load a NIfTI image directly to a PyTorch tensor.
///
/// This is the most efficient way to load medical imaging data into PyTorch.
/// Eliminates memory copies and supports half-precision tensors directly.
///
/// Args:
///     path: Path to the NIfTI file
///     dtype: PyTorch dtype (default: torch.float32)
///     device: PyTorch device (default: "cpu")
///
/// Returns:
///     PyTorch tensor with shape matching the image
///
/// Example:
///     >>> import torch
///     >>> tensor = medrs.load_to_torch("volume.nii", dtype=torch.float16, device="cuda")
#[pyfunction]
#[pyo3(signature = (path, dtype=None, device="cpu"))]
fn load_to_torch(
    py: Python<'_>,
    path: &str,
    dtype: Option<PyObject>,
    device: &str,
) -> PyResult<PyObject> {
    // Load image using medrs
    let img = nifti::load(path).map_err(|e| to_py_err(e, &format!("Failed to load {}", path)))?;

    // Convert to PyTorch tensor directly
    let py_img = PyNiftiImage { inner: img };
    py_img.to_torch_with_dtype_and_device(py, dtype, Some(device))
}

/// Load only a cropped region from a NIfTI file without loading the entire volume.
///
/// This is extremely efficient for training pipelines that load large volumes
/// just to crop small patches (e.g., loading 64^3 patch from 256^3 volume).
///
/// Args:
///     path: Path to NIfTI file (must be uncompressed .nii)
///     crop_offset: Starting coordinates of crop region [d, h, w]
///     crop_shape: Size of crop region [d, h, w]
///
/// Returns:
///     NiftiImage instance with cropped data
///
/// Example:
///     >>> # Load 64^3 patch starting at (32, 32, 32)
///     >>> img = medrs.load_cropped("volume.nii", [32, 32, 32], [64, 64, 64])
#[pyfunction]
fn load_cropped(
    path: &str,
    crop_offset: [usize; 3],
    crop_shape: [usize; 3],
) -> PyResult<PyNiftiImage> {
    nifti::load_cropped(path, crop_offset, crop_shape)
        .map(|inner| PyNiftiImage { inner })
        .map_err(|e| to_py_err(e, &format!("Failed to load cropped {}", path)))
}

/// Load a cropped region with optional reorientation and resampling.
///
/// Advanced version that computes the minimal region needed from the raw file
/// to achieve the desired output after transforms.
///
/// Args:
///     path: Path to NIfTI file (must be uncompressed .nii)
///     output_shape: Desired output shape after all transforms [d, h, w]
///     target_spacing: Optional target voxel spacing [mm] (None = keep original)
///     target_orientation: Optional target orientation code (None = keep original)
///     output_offset: Optional offset in output space for non-centered crops [d, h, w]
///
/// Returns:
///     NiftiImage instance with processed cropped data
///
/// Example:
///     >>> # Load 64^3 RAS 1mm isotropic patch from any orientation/spacing
///     >>> img = medrs.load_resampled(
///     ...     "volume.nii",
///     ...     output_shape=[64, 64, 64],
///     ...     target_spacing=[1.0, 1.0, 1.0],
///     ...     target_orientation="RAS"
///     ... )
#[pyfunction]
#[pyo3(signature = (path, output_shape, target_spacing=None, target_orientation=None, output_offset=None))]
fn load_resampled(
    path: &str,
    output_shape: [usize; 3],
    target_spacing: Option<[f32; 3]>,
    target_orientation: Option<String>,
    output_offset: Option<[usize; 3]>,
) -> PyResult<PyNiftiImage> {
    use crate::transforms::Orientation;

    let orientation = match target_orientation {
        Some(s) => Some(
            s.parse::<Orientation>()
                .map_err(|e| PyValueError::new_err(format!("{}", e)))?,
        ),
        None => None,
    };

    let config = nifti::LoadCroppedConfig {
        output_shape,
        target_spacing,
        target_orientation: orientation,
        output_offset,
    };

    nifti::load_cropped_config(path, config)
        .map(|inner| PyNiftiImage { inner })
        .map_err(|e| to_py_err(e, &format!("Failed to load cropped {}", path)))
}

/// Load a cropped region directly into a PyTorch tensor without numpy intermediate.
///
/// This is the most efficient way to load medical imaging data into PyTorch.
/// Eliminates memory copies and supports half-precision tensors directly.
///
/// Args:
///     path: Path to NIfTI file (must be uncompressed .nii)
///     output_shape: Desired output shape after all transforms [d, h, w]
///     target_spacing: Optional target voxel spacing [mm] (None = keep original)
///     target_orientation: Optional target orientation code (None = keep original)
///     output_offset: Optional offset in output space for non-centered crops [d, h, w]
///     dtype: PyTorch dtype (default: torch.float32)
///     device: PyTorch device (default: "cpu")
///
/// Returns:
///     PyTorch tensor with shape [d, h, w]
///
/// Example:
///     >>> # Load 64^3 RAS 1mm isotropic patch directly as f16 tensor
///     >>> import torch
///     >>> tensor = medrs.load_cropped_to_torch(
///     ...     "volume.nii",
///     ...     output_shape=[64, 64, 64],
///     ...     target_spacing=[1.0, 1.0, 1.0],
///     ...     target_orientation="RAS",
///     ...     dtype=torch.float16,  # Direct f16 loading!
///     ...     device="cuda"
///     ... )
#[pyfunction]
#[pyo3(signature = (path, output_shape, target_spacing=None, target_orientation=None, output_offset=None, dtype=None, device="cpu"))]
#[allow(clippy::too_many_arguments)]
fn load_cropped_to_torch(
    py: Python<'_>,
    path: &str,
    output_shape: [usize; 3],
    target_spacing: Option<[f32; 3]>,
    target_orientation: Option<String>,
    output_offset: Option<[usize; 3]>,
    dtype: Option<PyObject>,
    device: &str,
) -> PyResult<PyObject> {
    use crate::transforms::Orientation;

    // Validate inputs at Python boundary
    validate_shape(&output_shape, "output_shape")?;
    if let Some(ref spacing) = target_spacing {
        validate_spacing(spacing, "target_spacing")?;
    }

    // Load image using our I/O-optimized function
    let orientation = match target_orientation {
        Some(s) => Some(
            s.parse::<Orientation>()
                .map_err(|e| PyValueError::new_err(format!("{}", e)))?,
        ),
        None => None,
    };

    let config = nifti::LoadCroppedConfig {
        output_shape,
        target_spacing,
        target_orientation: orientation,
        output_offset,
    };

    let img = nifti::load_cropped_config(path, config)
        .map_err(|e| to_py_err(e, &format!("Failed to load cropped {}", path)))?;

    // Convert to PyTorch tensor directly using our optimized I/O + tensor conversion
    let py_img = PyNiftiImage { inner: img };
    py_img.to_torch_with_dtype_and_device(py, dtype, Some(device))
}

/// Load a cropped region directly into a JAX array without numpy intermediate.
///
/// This is the most efficient way to load medical imaging data into JAX.
/// Eliminates memory copies and supports bfloat16/f16 directly.
///
/// Args:
///     path: Path to NIfTI file (must be uncompressed .nii)
///     output_shape: Desired output shape after all transforms [d, h, w]
///     target_spacing: Optional target voxel spacing [mm] (None = keep original)
///     target_orientation: Optional target orientation code (None = keep original)
///     output_offset: Optional offset in output space for non-centered crops [d, h, w]
///     dtype: JAX dtype (default: jax.numpy.float32)
///     device: JAX device (default: "cpu")
///
/// Returns:
///     JAX array with shape [d, h, w]
///
/// Example:
///     >>> # Load 64^3 RAS 1mm isotropic patch directly as bfloat16
///     >>> import jax
///     >>> array = medrs.load_cropped_to_jax(
///     ...     "volume.nii",
///     ...     output_shape=[64, 64, 64],
///     ...     target_spacing=[1.0, 1.0, 1.0],
///     ...     target_orientation="RAS",
///     ...     dtype=jax.numpy.bfloat16,  # Direct bfloat16 loading!
///     ...     device="cuda:0"
///     ... )
#[pyfunction]
#[pyo3(signature = (path, output_shape, target_spacing=None, target_orientation=None, output_offset=None, dtype=None, device="cpu"))]
#[allow(clippy::too_many_arguments)]
fn load_cropped_to_jax(
    py: Python<'_>,
    path: &str,
    output_shape: [usize; 3],
    target_spacing: Option<[f32; 3]>,
    target_orientation: Option<String>,
    output_offset: Option<[usize; 3]>,
    dtype: Option<PyObject>,
    device: &str,
) -> PyResult<PyObject> {
    use crate::transforms::Orientation;

    // Validate inputs at Python boundary
    validate_shape(&output_shape, "output_shape")?;
    if let Some(ref spacing) = target_spacing {
        validate_spacing(spacing, "target_spacing")?;
    }

    // Load image using our I/O-optimized function
    let orientation = match target_orientation {
        Some(s) => Some(
            s.parse::<Orientation>()
                .map_err(|e| PyValueError::new_err(format!("{}", e)))?,
        ),
        None => None,
    };

    let config = nifti::LoadCroppedConfig {
        output_shape,
        target_spacing,
        target_orientation: orientation,
        output_offset,
    };

    let img = nifti::load_cropped_config(path, config)
        .map_err(|e| to_py_err(e, &format!("Failed to load cropped {}", path)))?;

    // Convert to JAX array directly using our optimized I/O + array conversion
    let py_img = PyNiftiImage { inner: img };
    py_img.to_jax_with_dtype_and_device(py, dtype, Some(device))
}

/// High-performance training data loader with prefetching and caching.
///
/// This is the most efficient way to load training patches from multiple volumes.
/// Maintains an LRU cache and prefetches upcoming data to maximize throughput.
#[pyclass(name = "TrainingDataLoader")]
pub struct PyTrainingDataLoader {
    loader: nifti::TrainingDataLoader,
}

#[pymethods]
impl PyTrainingDataLoader {
    /// Create a new training data loader.
    ///
    /// Args:
    ///     volumes: List of NIfTI file paths
    ///     patch_size: Patch size to extract [d, h, w]
    ///     patches_per_volume: Number of patches per volume
    ///     patch_overlap: Overlap between patches [d, h, w] in voxels
    ///     randomize: Whether to randomize patch positions
    ///     cache_size: Maximum number of patches to cache
    ///
    /// Example:
    ///     ```python
    ///     loader = medrs.TrainingDataLoader(
    ///         volumes=["vol1.nii", "vol2.nii"],
    ///         patch_size=[64, 64, 64],
    ///         patches_per_volume=4,
    ///         patch_overlap=[0, 0, 0],
    ///         randomize=True,
    ///         cache_size=1000
    ///     )
    ///     patch = loader.next_patch()
    ///     ```
    #[new]
    #[pyo3(signature = (volumes, patch_size, patches_per_volume, patch_overlap, randomize, cache_size=None))]
    fn new(
        volumes: Vec<String>,
        patch_size: [usize; 3],
        patches_per_volume: usize,
        patch_overlap: [usize; 3],
        randomize: bool,
        cache_size: Option<usize>,
    ) -> PyResult<Self> {
        for i in 0..3 {
            if patch_overlap[i] >= patch_size[i] {
                return Err(PyValueError::new_err(
                    "patch_overlap must be smaller than patch_size in all dimensions",
                ));
            }
        }

        let config = nifti::CropLoaderConfig {
            patch_size,
            patches_per_volume,
            patch_overlap,
            randomize,
        };

        let cache_size = cache_size.unwrap_or(1000);
        let loader = nifti::TrainingDataLoader::new(volumes, config, cache_size)
            .map_err(|e| PyValueError::new_err(format!("Failed to create loader: {}", e)))?;

        Ok(Self { loader })
    }

    /// Get next training patch.
    ///
    /// Returns the next training patch with automatic prefetching.
    /// Raises StopIteration when all patches are processed.
    fn next_patch(&mut self) -> PyResult<PyNiftiImage> {
        match self.loader.next_patch() {
            Ok(inner) => Ok(PyNiftiImage { inner }),
            Err(MedrsError::Exhausted(msg)) => Err(PyStopIteration::new_err(msg)),
            Err(e) => Err(to_py_err(e, "next_patch")),
        }
    }

    fn __len__(&self) -> usize {
        self.loader.volumes_len() * self.loader.patches_per_volume()
    }

    fn __iter__(mut slf: PyRefMut<'_, Self>) -> PyResult<PyRefMut<'_, Self>> {
        slf.loader
            .reset()
            .map_err(|e| PyValueError::new_err(format!("Failed to reset loader: {}", e)))?;
        Ok(slf)
    }

    fn __next__(mut slf: PyRefMut<'_, Self>) -> PyResult<Option<PyNiftiImage>> {
        match slf.loader.next_patch() {
            Ok(img) => Ok(Some(PyNiftiImage { inner: img })),
            Err(MedrsError::Exhausted(_)) => Ok(None),
            Err(e) => Err(to_py_err(e, "iterator")),
        }
    }

    /// Get performance statistics.
    fn stats(&self) -> String {
        format!("{}", self.loader.stats())
    }

    /// Reset the loader to start from the beginning.
    fn reset(&mut self) -> PyResult<()> {
        self.loader
            .reset()
            .map_err(|e| PyValueError::new_err(format!("Failed to reset loader: {}", e)))
    }
}

/// Composable transform pipeline with lazy evaluation.
///
/// Build transformation chains that are optimized and applied efficiently.
/// Supports method chaining for a fluent API similar to MONAI's Compose.
///
/// Example:
///     >>> pipeline = medrs.TransformPipeline()
///     >>>     .z_normalize()
///     >>>     .clamp(-1.0, 1.0)
///     >>>     .resample_to_shape([64, 64, 64])
///     >>> processed = pipeline.apply(img)
#[pyclass(name = "TransformPipeline")]
pub struct PyTransformPipeline {
    inner: RustTransformPipeline,
}

#[pymethods]
impl PyTransformPipeline {
    /// Create a new transform pipeline.
    ///
    /// Args:
    ///     lazy: Enable lazy evaluation (default: True). When True, transforms
    ///           are composed and optimized before execution.
    #[new]
    #[pyo3(signature = (lazy=true))]
    fn new(lazy: bool) -> Self {
        let inner = if lazy {
            RustTransformPipeline::new()
        } else {
            RustTransformPipeline::new().lazy(false)
        };
        Self { inner }
    }

    /// Add z-score normalization (zero mean, unit variance).
    ///
    /// Returns:
    ///     Self for method chaining
    fn z_normalize(mut self_: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).z_normalize();
        self_
    }

    /// Add intensity rescaling to range [out_min, out_max].
    ///
    /// Args:
    ///     out_min: Minimum output value
    ///     out_max: Maximum output value
    ///
    /// Returns:
    ///     Self for method chaining
    fn rescale(mut self_: PyRefMut<'_, Self>, out_min: f32, out_max: f32) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).rescale(out_min, out_max);
        self_
    }

    /// Add intensity clamping to range [min, max].
    ///
    /// Args:
    ///     min: Minimum value
    ///     max: Maximum value
    ///
    /// Returns:
    ///     Self for method chaining
    fn clamp(mut self_: PyRefMut<'_, Self>, min: f32, max: f32) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).clamp(min, max);
        self_
    }

    /// Add resampling to target voxel spacing.
    ///
    /// Args:
    ///     spacing: Target spacing as [x, y, z] in mm
    ///
    /// Returns:
    ///     Self for method chaining
    fn resample_to_spacing(mut self_: PyRefMut<'_, Self>, spacing: [f32; 3]) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).resample_to_spacing(spacing);
        self_
    }

    /// Add resampling to target shape.
    ///
    /// Args:
    ///     shape: Target shape as [depth, height, width]
    ///
    /// Returns:
    ///     Self for method chaining
    fn resample_to_shape(mut self_: PyRefMut<'_, Self>, shape: [usize; 3]) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).resample_to_shape(shape);
        self_
    }

    /// Add flip along specified axes.
    ///
    /// Args:
    ///     axes: List of axes to flip (0=depth, 1=height, 2=width)
    ///
    /// Returns:
    ///     Self for method chaining
    fn flip(mut self_: PyRefMut<'_, Self>, axes: Vec<usize>) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).flip(&axes);
        self_
    }

    /// Enable or disable lazy evaluation.
    ///
    /// Args:
    ///     lazy: Whether to use lazy evaluation
    ///
    /// Returns:
    ///     Self for method chaining
    fn set_lazy(mut self_: PyRefMut<'_, Self>, lazy: bool) -> PyRefMut<'_, Self> {
        self_.inner = std::mem::take(&mut self_.inner).lazy(lazy);
        self_
    }

    /// Apply the pipeline to an image.
    ///
    /// Args:
    ///     image: Input NiftiImage
    ///
    /// Returns:
    ///     Transformed NiftiImage
    ///
    /// Raises:
    ///     ValueError: If the pipeline fails to apply
    fn apply(&self, image: &PyNiftiImage) -> PyResult<PyNiftiImage> {
        let result = self
            .inner
            .apply(&image.inner)
            .map_err(|e| PyValueError::new_err(e.to_string()))?;
        Ok(PyNiftiImage { inner: result })
    }

    #[allow(clippy::unused_self)]
    fn __repr__(&self) -> String {
        "TransformPipeline(...)".to_string()
    }
}

// Random augmentation functions

/// Apply random flip along specified axes with given probability.
///
/// Args:
///     image: Input image
///     axes: Axes that may be flipped (0=depth, 1=height, 2=width)
///     prob: Probability of flipping each axis (default: 0.5)
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
///
/// Example:
///     >>> augmented = medrs.random_flip(img, [0, 1, 2], prob=0.5)
#[pyfunction]
#[pyo3(signature = (image, axes, prob=None, seed=None))]
fn random_flip(
    image: &PyNiftiImage,
    axes: Vec<usize>,
    prob: Option<f32>,
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    // Validate probability if provided
    if let Some(p) = prob {
        validate_probability(p as f64, "random_flip")?;
    }

    // Validate axes
    for &axis in &axes {
        if axis >= 3 {
            return Err(PyValueError::new_err(format!(
                "random_flip: axis {} is out of range (must be 0, 1, or 2)",
                axis
            )));
        }
    }

    Ok(PyNiftiImage {
        inner: transforms::random_flip(&image.inner, &axes, prob, seed)
            .map_err(|e| PyValueError::new_err(e.to_string()))?,
    })
}

/// Add random Gaussian noise to the image.
///
/// Args:
///     image: Input image
///     std: Standard deviation of the noise (default: 0.1)
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, std=None, seed=None))]
fn random_gaussian_noise(
    image: &PyNiftiImage,
    std: Option<f32>,
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_gaussian_noise(&image.inner, std, seed)
            .map_err(|e| to_py_err(e, "random_gaussian_noise"))?,
    })
}

/// Randomly scale image intensity.
///
/// Multiplies intensity by a random factor sampled from [1-scale_range, 1+scale_range].
///
/// Args:
///     image: Input image
///     scale_range: Range for random scaling factor (default: 0.1, meaning 0.9-1.1)
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, scale_range=None, seed=None))]
fn random_intensity_scale(
    image: &PyNiftiImage,
    scale_range: Option<f32>,
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_intensity_scale(&image.inner, scale_range, seed)
            .map_err(|e| to_py_err(e, "random_intensity_scale"))?,
    })
}

/// Randomly shift image intensity.
///
/// Adds a random offset sampled from [-shift_range, shift_range].
///
/// Args:
///     image: Input image
///     shift_range: Range for random shift (default: 0.1)
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, shift_range=None, seed=None))]
fn random_intensity_shift(
    image: &PyNiftiImage,
    shift_range: Option<f32>,
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_intensity_shift(&image.inner, shift_range, seed)
            .map_err(|e| to_py_err(e, "random_intensity_shift"))?,
    })
}

/// Randomly rotate the image by 90-degree increments.
///
/// Performs random rotation in the specified plane by 0, 90, 180, or 270 degrees.
///
/// Args:
///     image: Input image
///     axes: Tuple of two axes defining the rotation plane (e.g., (0, 1) for depth-height plane)
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, axes, seed=None))]
fn random_rotate_90(
    image: &PyNiftiImage,
    axes: (usize, usize),
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_rotate_90(&image.inner, axes, seed)
            .map_err(|e| PyValueError::new_err(e.to_string()))?,
    })
}

/// Apply random gamma correction to image intensity.
///
/// Applies the transform: output = input^gamma where gamma is randomly sampled.
///
/// Args:
///     image: Input image (should be normalized to [0, 1] for best results)
///     gamma_range: Range for gamma sampling as (min, max) (default: (0.7, 1.5))
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, gamma_range=None, seed=None))]
fn random_gamma(
    image: &PyNiftiImage,
    gamma_range: Option<(f32, f32)>,
    seed: Option<u64>,
) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_gamma(&image.inner, gamma_range, seed)
            .map_err(|e| to_py_err(e, "random_gamma"))?,
    })
}

/// Apply a random combination of common augmentations.
///
/// This is a convenience function that applies multiple augmentations:
/// - Random flip (prob=0.5 per axis)
/// - Random intensity scale
/// - Random intensity shift
/// - Random Gaussian noise
///
/// Args:
///     image: Input image
///     seed: Optional random seed for reproducibility
///
/// Returns:
///     Augmented image
#[pyfunction]
#[pyo3(signature = (image, seed=None))]
fn random_augment(image: &PyNiftiImage, seed: Option<u64>) -> PyResult<PyNiftiImage> {
    Ok(PyNiftiImage {
        inner: transforms::random_augment(&image.inner, seed)
            .map_err(|e| PyValueError::new_err(e.to_string()))?,
    })
}

/// High-performance medical image I/O and processing.
///
/// Fast NIfTI reading/writing and transforms optimized for deep learning.
///
/// Example:
///     >>> import medrs
///     >>> img = medrs.load("brain.nii.gz")
///     >>> data = img.to_numpy()
///     >>> processed = img.resample([1.0, 1.0, 1.0]).z_normalize().clamp(0, 1)
///     >>> processed.save("output.nii.gz")
#[pymodule]
fn _medrs(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Classes
    m.add_class::<PyNiftiImage>()?;
    m.add_class::<PyTrainingDataLoader>()?;
    m.add_class::<PyTransformPipeline>()?;

    // Basic transforms
    m.add_function(wrap_pyfunction!(z_normalization, m)?)?;
    m.add_function(wrap_pyfunction!(rescale_intensity, m)?)?;
    m.add_function(wrap_pyfunction!(clamp, m)?)?;
    m.add_function(wrap_pyfunction!(crop_or_pad, m)?)?;
    m.add_function(wrap_pyfunction!(resample, m)?)?;
    m.add_function(wrap_pyfunction!(reorient, m)?)?;

    // I/O functions
    m.add_function(wrap_pyfunction!(load, m)?)?;
    m.add_function(wrap_pyfunction!(load_to_torch, m)?)?;
    m.add_function(wrap_pyfunction!(load_cropped, m)?)?;
    m.add_function(wrap_pyfunction!(load_resampled, m)?)?;
    m.add_function(wrap_pyfunction!(load_cropped_to_torch, m)?)?;
    m.add_function(wrap_pyfunction!(load_cropped_to_jax, m)?)?;

    // Crop-first transform functions
    m.add_function(wrap_pyfunction!(load_label_aware_cropped, m)?)?;
    m.add_function(wrap_pyfunction!(compute_crop_regions, m)?)?;
    m.add_function(wrap_pyfunction!(compute_random_spatial_crops, m)?)?;
    m.add_function(wrap_pyfunction!(compute_center_crop, m)?)?;

    // Random augmentation functions
    m.add_function(wrap_pyfunction!(random_flip, m)?)?;
    m.add_function(wrap_pyfunction!(random_gaussian_noise, m)?)?;
    m.add_function(wrap_pyfunction!(random_intensity_scale, m)?)?;
    m.add_function(wrap_pyfunction!(random_intensity_shift, m)?)?;
    m.add_function(wrap_pyfunction!(random_rotate_90, m)?)?;
    m.add_function(wrap_pyfunction!(random_gamma, m)?)?;
    m.add_function(wrap_pyfunction!(random_augment, m)?)?;

    Ok(())
}

fn arraydata_to_numpy(py: Python<'_>, data: &ArrayData, shape: &[usize]) -> PyResult<PyObject> {
    let dyn_shape = IxDyn(shape);
    Ok(match data {
        ArrayData::U8(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::I8(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::I16(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::U16(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::I32(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::U32(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::I64(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::U64(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::F16(a) => a
            .mapv(|v| v.to_f32())
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::BF16(a) => a
            .mapv(|v| v.to_f32())
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::F32(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
        ArrayData::F64(a) => a
            .to_owned()
            .into_pyarray(py)
            .reshape(dyn_shape)?
            .unbind()
            .into_any(),
    })
}

fn to_numpy_view<'py>(
    py: Python<'py>,
    image: &RustNiftiImage,
) -> Option<Bound<'py, PyArrayDyn<f32>>> {
    if let Some(view) = image.as_view_f32() {
        // numpy will copy; this still avoids extra materialization in Rust
        let arr = PyArrayDyn::from_array(py, &view);
        return Some(arr);
    }
    None
}

fn to_numpy_array<'py>(
    py: Python<'py>,
    image: &RustNiftiImage,
) -> PyResult<Bound<'py, PyArrayDyn<f32>>> {
    let data = image.to_f32().map_err(|e| to_py_err(e, "to_f32"))?;
    let shape: Vec<usize> = data.shape().to_vec();

    // Data is in F-order (column-major) layout in memory
    // Use as_slice_memory_order to get bytes in physical order
    let slice = data
        .as_slice_memory_order()
        .ok_or_else(|| PyValueError::new_err("Array is not contiguous in memory"))?;
    let np = py.import("numpy")?;

    // Create 1D array from F-order bytes
    let flat = slice.to_vec().into_pyarray(py);

    // Reshape with order='F' to interpret the flat data as F-order
    // This creates an F-contiguous array matching NIfTI convention
    let kwargs = [("order", "F")].into_py_dict(py)?;
    let reshaped = np
        .call_method("reshape", (flat, &shape), Some(&kwargs))
        .map_err(|e| PyValueError::new_err(format!("Failed to reshape array: {}", e)))?;

    reshaped
        .extract::<Bound<'py, PyArrayDyn<f32>>>()
        .map_err(|e| PyValueError::new_err(format!("Failed to extract array: {}", e)))
}

fn to_numpy_view_native(py: Python<'_>, image: &RustNiftiImage) -> Option<PyObject> {
    let arr_obj = match image.dtype() {
        DataType::UInt8 => image
            .as_view_t::<u8>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::Int8 => image
            .as_view_t::<i8>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::Int16 => image
            .as_view_t::<i16>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::UInt16 => image
            .as_view_t::<u16>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::Int32 => image
            .as_view_t::<i32>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::UInt32 => image
            .as_view_t::<u32>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::Int64 => image
            .as_view_t::<i64>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::UInt64 => image
            .as_view_t::<u64>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        // numpy crate lacks f16/bf16 Element impl, fall back to f32
        DataType::Float16 | DataType::BFloat16 => None,
        DataType::Float32 => image
            .as_view_t::<f32>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
        DataType::Float64 => image
            .as_view_t::<f64>()
            .map(|v| PyArrayDyn::from_array(py, &v).unbind().into_any()),
    };
    arr_obj
}

fn torch_dtype(py: Python<'_>, dtype: DataType) -> Option<PyObject> {
    let torch = py.import("torch").ok()?;
    let dt = match dtype {
        DataType::UInt8 => "uint8",
        DataType::Int8 => "int8",
        DataType::Int16 => "int16",
        DataType::Int32 => "int32",
        DataType::Int64 => "int64",
        DataType::Float16 => "float16",
        DataType::BFloat16 => "bfloat16",
        DataType::Float32 => "float32",
        DataType::Float64 => "float64",
        // PyTorch doesn't support unsigned 16/32/64-bit integers
        DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => return None,
    };
    torch.getattr(dt).ok().map(|o| o.unbind())
}

/// Load a NIfTI file with byte-exact cropping for label-aware training.
///
/// This function combines MONAI's `RandCropByPosNegLabeld` with medrs's
/// byte-exact loading for maximum performance. It computes optimal crop
/// regions containing both positive and negative labels, then loads only
/// the required bytes.
///
/// Args:
///     image_path: Path to the image file
///     label_path: Path to the label file
///     patch_size: Target patch size [x, y, z]
///     pos_neg_ratio: Ratio of positive to negative samples (default: 1.0)
///     min_pos_samples: Minimum positive samples per volume (default: 4)
///     seed: Random seed for reproducibility (optional)
///
/// Returns:
///     Tuple of (cropped_image, cropped_label) as PyNiftiImages
#[pyfunction]
#[pyo3(signature = (image_path, label_path, patch_size, pos_neg_ratio=None, min_pos_samples=None, seed=None))]
fn load_label_aware_cropped(
    _py: Python<'_>,
    image_path: &str,
    label_path: &str,
    patch_size: Vec<usize>,
    pos_neg_ratio: Option<f64>,
    min_pos_samples: Option<usize>,
    seed: Option<u64>,
) -> PyResult<(PyNiftiImage, PyNiftiImage)> {
    let patch_size = parse_shape3(&patch_size, "patch_size")?;

    // Load full images first (this could be optimized further)
    let image = nifti::load(image_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load image: {}", e))
    })?;
    let label = nifti::load(label_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load label: {}", e))
    })?;

    // Configure cropping
    let config = transforms::RandCropByPosNegLabelConfig {
        patch_size,
        pos_neg_ratio: pos_neg_ratio.unwrap_or(1.0) as f32,
        min_pos_samples: min_pos_samples.unwrap_or(4),
        seed,
        background_label: 0.0,
    };

    // Compute crop regions (this is fast, operates on labels only)
    let crop_regions = compute_label_aware_crop_regions(&config, &image, &label, 1)
        .map_err(|e| to_py_err(e, "compute_label_aware_crop_regions"))?;

    if crop_regions.is_empty() {
        return Err(PyValueError::new_err("No valid crop regions found"));
    }

    let region = &crop_regions[0];

    // Load cropped regions (this is the fast byte-exact loading)
    let cropped_image =
        nifti::load_cropped(image_path, region.start, region.size).map_err(|e| {
            PyErr::new::<pyo3::exceptions::PyIOError, _>(format!(
                "Failed to load cropped image: {}",
                e
            ))
        })?;
    let cropped_label =
        nifti::load_cropped(label_path, region.start, region.size).map_err(|e| {
            PyErr::new::<pyo3::exceptions::PyIOError, _>(format!(
                "Failed to load cropped label: {}",
                e
            ))
        })?;

    Ok((
        PyNiftiImage {
            inner: cropped_image,
        },
        PyNiftiImage {
            inner: cropped_label,
        },
    ))
}

/// Compute crop regions for smart loading without actually loading the data.
///
/// This function allows users to compute optimal crop regions first,
/// then use them with `load_cropped()` for maximum control.
/// This is useful for batch processing and advanced training pipelines.
///
/// Args:
///     image_path: Path to the image file
///     label_path: Path to the label file
///     patch_size: Target patch size [x, y, z]
///     num_samples: Number of crop regions to compute
///     pos_neg_ratio: Ratio of positive to negative samples (default: 1.0)
///     min_pos_samples: Minimum positive samples per volume (default: 4)
///     seed: Random seed for reproducibility (optional)
///
/// Returns:
///     List of crop regions as dictionaries with 'start', 'end', and 'size' keys
#[pyfunction]
#[pyo3(signature = (image_path, label_path, patch_size, num_samples, pos_neg_ratio=None, min_pos_samples=None, seed=None))]
#[allow(clippy::too_many_arguments)]
fn compute_crop_regions(
    py: Python<'_>,
    image_path: &str,
    label_path: &str,
    patch_size: Vec<usize>,
    num_samples: usize,
    pos_neg_ratio: Option<f64>,
    min_pos_samples: Option<usize>,
    seed: Option<u64>,
) -> PyResult<Vec<PyObject>> {
    let patch_size = parse_shape3(&patch_size, "patch_size")?;

    // Load images to get shape information
    let image = nifti::load(image_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load image: {}", e))
    })?;
    let label = nifti::load(label_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load label: {}", e))
    })?;

    // Configure cropping
    let config = transforms::RandCropByPosNegLabelConfig {
        patch_size,
        pos_neg_ratio: pos_neg_ratio.unwrap_or(1.0) as f32,
        min_pos_samples: min_pos_samples.unwrap_or(4),
        seed,
        background_label: 0.0,
    };

    // Compute crop regions
    let crop_regions = compute_label_aware_crop_regions(&config, &image, &label, num_samples)
        .map_err(|e| to_py_err(e, "compute_label_aware_crop_regions"))?;

    // Convert to Python dictionaries
    let mut regions_py = Vec::new();
    for region in crop_regions {
        let region_dict = pyo3::types::PyDict::new(py);
        region_dict.set_item("start", region.start.to_vec())?;
        region_dict.set_item("end", region.end.to_vec())?;
        region_dict.set_item("size", region.size.to_vec())?;
        regions_py.push(region_dict.into());
    }

    Ok(regions_py)
}

/// Compute random spatial crop regions.
///
/// This function implements MONAI's `RandSpatialCropd` functionality
/// optimized for medrs's crop-first approach.
///
/// Args:
///     image_path: Path to the image file
///     patch_size: Target patch size [x, y, z]
///     num_samples: Number of crop regions to compute
///     seed: Random seed for reproducibility (optional)
///     allow_smaller: Whether to allow smaller crops at boundaries (default: false)
///
/// Returns:
///     List of crop regions as dictionaries
#[pyfunction]
#[pyo3(signature = (image_path, patch_size, num_samples, seed=None, allow_smaller=None))]
fn compute_random_spatial_crops(
    py: Python<'_>,
    image_path: &str,
    patch_size: Vec<usize>,
    num_samples: usize,
    seed: Option<u64>,
    allow_smaller: Option<bool>,
) -> PyResult<Vec<PyObject>> {
    let patch_size = parse_shape3(&patch_size, "patch_size")?;

    // Load image to get shape information
    let image = nifti::load(image_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load image: {}", e))
    })?;

    // Configure cropping
    let config = transforms::SpatialCropConfig {
        patch_size,
        seed,
        allow_smaller: allow_smaller.unwrap_or(false),
    };

    // Compute crop regions
    let crop_regions = compute_random_spatial_crop_regions(&config, &image, num_samples);

    // Convert to Python dictionaries
    let mut regions_py = Vec::new();
    for region in crop_regions {
        let region_dict = pyo3::types::PyDict::new(py);
        region_dict.set_item("start", region.start.to_vec())?;
        region_dict.set_item("end", region.end.to_vec())?;
        region_dict.set_item("size", region.size.to_vec())?;
        regions_py.push(region_dict.into());
    }

    Ok(regions_py)
}

/// Compute center crop region.
///
/// This function implements MONAI's `CenterSpatialCropd` functionality
/// optimized for medrs's crop-first approach.
///
/// Args:
///     image_path: Path to the image file
///     patch_size: Target patch size [x, y, z]
///
/// Returns:
///     Crop region as dictionary
#[pyfunction]
fn compute_center_crop(
    py: Python<'_>,
    image_path: &str,
    patch_size: Vec<usize>,
) -> PyResult<PyObject> {
    let patch_size = parse_shape3(&patch_size, "patch_size")?;

    // Load image to get shape information
    let image = nifti::load(image_path).map_err(|e| {
        PyErr::new::<pyo3::exceptions::PyIOError, _>(format!("Failed to load image: {}", e))
    })?;

    // Compute center crop
    let region = compute_center_crop_regions(patch_size, &image);

    // Convert to Python dictionary
    let region_dict = pyo3::types::PyDict::new(py);
    region_dict.set_item("start", region.start.to_vec())?;
    region_dict.set_item("end", region.end.to_vec())?;
    region_dict.set_item("size", region.size.to_vec())?;

    Ok(region_dict.into())
}