//! 3D volumetric data structure for medical imaging.
//!
//! The `Volume` struct provides a container for 3D image data with spatial
//! metadata (spacing, origin, affine transform).

use crate::nifti::transform::{self, Affine4};
use rayon::prelude::*;
use serde::{Deserialize, Serialize};

/// A 3D volume of medical imaging data.
///
/// Data is stored in row-major order (C-order): z varies slowest, x varies fastest.
/// The layout is `data[z * (y_size * x_size) + y * x_size + x]`.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Volume {
    /// Voxel data in row-major order
    data: Vec<f64>,

    /// Dimensions [x, y, z]
    shape: [usize; 3],

    /// Voxel spacing in mm [dx, dy, dz]
    spacing: [f64; 3],

    /// Origin in world coordinates [ox, oy, oz]
    origin: [f64; 3],

    /// 4x4 affine transformation matrix (voxel to world)
    affine: Affine4,
}

impl Volume {
    /// Create a new volume from raw data and metadata.
    ///
    /// # Arguments
    /// * `data` - Voxel data in row-major order
    /// * `shape` - Dimensions [x, y, z]
    /// * `spacing` - Voxel spacing in mm [dx, dy, dz]
    /// * `origin` - Origin in world coordinates [ox, oy, oz]
    /// * `affine` - 4x4 affine transformation matrix
    ///
    /// # Panics
    /// Panics if data length doesn't match shape.
    pub fn new(
        data: Vec<f64>,
        shape: [usize; 3],
        spacing: [f64; 3],
        origin: [f64; 3],
        affine: Affine4,
    ) -> Self {
        let expected_len = shape[0] * shape[1] * shape[2];
        assert_eq!(
            data.len(),
            expected_len,
            "Data length {} doesn't match shape {:?} (expected {})",
            data.len(),
            shape,
            expected_len
        );

        Self {
            data,
            shape,
            spacing,
            origin,
            affine,
        }
    }

    /// Create a volume filled with zeros.
    pub fn zeros(shape: [usize; 3]) -> Self {
        let size = shape[0] * shape[1] * shape[2];
        let spacing = [1.0, 1.0, 1.0];
        let origin = [0.0, 0.0, 0.0];
        let affine = transform::from_spacing_origin(spacing, origin);

        Self {
            data: vec![0.0; size],
            shape,
            spacing,
            origin,
            affine,
        }
    }

    /// Create a volume filled with a constant value.
    pub fn filled(shape: [usize; 3], value: f64) -> Self {
        let size = shape[0] * shape[1] * shape[2];
        let spacing = [1.0, 1.0, 1.0];
        let origin = [0.0, 0.0, 0.0];
        let affine = transform::from_spacing_origin(spacing, origin);

        Self {
            data: vec![value; size],
            shape,
            spacing,
            origin,
            affine,
        }
    }

    /// Get the volume dimensions [x, y, z].
    pub fn shape(&self) -> [usize; 3] {
        self.shape
    }

    /// Get the voxel spacing in mm [dx, dy, dz].
    pub fn spacing(&self) -> [f64; 3] {
        self.spacing
    }

    /// Get the origin in world coordinates [ox, oy, oz].
    pub fn origin(&self) -> [f64; 3] {
        self.origin
    }

    /// Get the affine transformation matrix.
    pub fn affine(&self) -> Affine4 {
        self.affine
    }

    /// Get a reference to the raw voxel data.
    pub fn data(&self) -> &[f64] {
        &self.data
    }

    /// Get a mutable reference to the raw voxel data.
    pub fn data_mut(&mut self) -> &mut [f64] {
        &mut self.data
    }

    /// Get the total number of voxels.
    pub fn num_voxels(&self) -> usize {
        self.shape[0] * self.shape[1] * self.shape[2]
    }

    /// Convert flat index to 3D coordinates (x, y, z).
    pub fn index_to_coords(&self, index: usize) -> [usize; 3] {
        let x = index % self.shape[0];
        let y = (index / self.shape[0]) % self.shape[1];
        let z = index / (self.shape[0] * self.shape[1]);
        [x, y, z]
    }

    /// Convert 3D coordinates to flat index.
    pub fn coords_to_index(&self, x: usize, y: usize, z: usize) -> usize {
        z * (self.shape[0] * self.shape[1]) + y * self.shape[0] + x
    }

    /// Get voxel value at coordinates (x, y, z).
    ///
    /// Returns None if coordinates are out of bounds.
    pub fn get(&self, x: usize, y: usize, z: usize) -> Option<f64> {
        if x < self.shape[0] && y < self.shape[1] && z < self.shape[2] {
            Some(self.data[self.coords_to_index(x, y, z)])
        } else {
            None
        }
    }

    /// Set voxel value at coordinates (x, y, z).
    ///
    /// Returns false if coordinates are out of bounds.
    pub fn set(&mut self, x: usize, y: usize, z: usize, value: f64) -> bool {
        if x < self.shape[0] && y < self.shape[1] && z < self.shape[2] {
            let idx = self.coords_to_index(x, y, z);
            self.data[idx] = value;
            true
        } else {
            false
        }
    }

    /// Get voxel value at coordinates without bounds checking.
    ///
    /// # Safety
    /// Caller must ensure coordinates are valid.
    #[allow(unsafe_op_in_unsafe_fn)]
    pub unsafe fn get_unchecked(&self, x: usize, y: usize, z: usize) -> f64 {
        let idx = self.coords_to_index(x, y, z);
        // SAFETY: caller guarantees coordinates are valid
        unsafe { *self.data.get_unchecked(idx) }
    }

    /// Set voxel value at coordinates without bounds checking.
    ///
    /// # Safety
    /// Caller must ensure coordinates are valid.
    #[allow(unsafe_op_in_unsafe_fn)]
    pub unsafe fn set_unchecked(&mut self, x: usize, y: usize, z: usize, value: f64) {
        let idx = self.coords_to_index(x, y, z);
        // SAFETY: caller guarantees coordinates are valid
        unsafe { *self.data.get_unchecked_mut(idx) = value }
    }

    /// Convert voxel coordinates to world coordinates.
    pub fn voxel_to_world(&self, voxel: [f64; 3]) -> [f64; 3] {
        transform::voxel_to_world(&self.affine, voxel)
    }

    /// Convert world coordinates to voxel coordinates.
    pub fn world_to_voxel(&self, world: [f64; 3]) -> Option<[f64; 3]> {
        transform::world_to_voxel(&self.affine, world)
    }

    /// Get the bounding box in world coordinates.
    ///
    /// Returns (min_corner, max_corner).
    pub fn world_bounds(&self) -> ([f64; 3], [f64; 3]) {
        let corners = [
            [0.0, 0.0, 0.0],
            [self.shape[0] as f64, 0.0, 0.0],
            [0.0, self.shape[1] as f64, 0.0],
            [0.0, 0.0, self.shape[2] as f64],
            [self.shape[0] as f64, self.shape[1] as f64, 0.0],
            [self.shape[0] as f64, 0.0, self.shape[2] as f64],
            [0.0, self.shape[1] as f64, self.shape[2] as f64],
            [
                self.shape[0] as f64,
                self.shape[1] as f64,
                self.shape[2] as f64,
            ],
        ];

        let world_corners: Vec<_> = corners.iter().map(|&c| self.voxel_to_world(c)).collect();

        let min = [
            world_corners
                .iter()
                .map(|c| c[0])
                .fold(f64::INFINITY, f64::min),
            world_corners
                .iter()
                .map(|c| c[1])
                .fold(f64::INFINITY, f64::min),
            world_corners
                .iter()
                .map(|c| c[2])
                .fold(f64::INFINITY, f64::min),
        ];

        let max = [
            world_corners
                .iter()
                .map(|c| c[0])
                .fold(f64::NEG_INFINITY, f64::max),
            world_corners
                .iter()
                .map(|c| c[1])
                .fold(f64::NEG_INFINITY, f64::max),
            world_corners
                .iter()
                .map(|c| c[2])
                .fold(f64::NEG_INFINITY, f64::max),
        ];

        (min, max)
    }

    /// Get minimum voxel value.
    pub fn min(&self) -> f64 {
        self.data
            .par_iter()
            .copied()
            .reduce(|| f64::INFINITY, f64::min)
    }

    /// Get maximum voxel value.
    pub fn max(&self) -> f64 {
        self.data
            .par_iter()
            .copied()
            .reduce(|| f64::NEG_INFINITY, f64::max)
    }

    /// Get mean voxel value.
    pub fn mean(&self) -> f64 {
        let sum: f64 = self.data.par_iter().sum();
        sum / self.data.len() as f64
    }

    /// Get standard deviation of voxel values.
    pub fn std(&self) -> f64 {
        let mean = self.mean();
        let variance: f64 = self
            .data
            .par_iter()
            .map(|&v| (v - mean).powi(2))
            .sum::<f64>()
            / self.data.len() as f64;
        variance.sqrt()
    }

    /// Get unique label values (for segmentation volumes).
    pub fn unique_labels(&self) -> Vec<i64> {
        use std::collections::BTreeSet;
        let labels: BTreeSet<i64> = self.data.par_iter().map(|&v| v.round() as i64).collect();
        labels.into_iter().collect()
    }

    /// Apply a function to each voxel value (in-place).
    pub fn apply<F>(&mut self, f: F)
    where
        F: Fn(f64) -> f64 + Sync + Send,
    {
        self.data.par_iter_mut().for_each(|v| *v = f(*v));
    }

    /// Apply a function to each voxel, creating a new volume.
    pub fn map<F>(&self, f: F) -> Self
    where
        F: Fn(f64) -> f64 + Sync + Send,
    {
        let new_data: Vec<f64> = self.data.par_iter().map(|&v| f(v)).collect();
        Self {
            data: new_data,
            shape: self.shape,
            spacing: self.spacing,
            origin: self.origin,
            affine: self.affine,
        }
    }

    /// Threshold the volume: values >= threshold become 1.0, others become 0.0.
    pub fn threshold(&self, threshold: f64) -> Self {
        self.map(|v| if v >= threshold { 1.0 } else { 0.0 })
    }

    /// Create a binary mask where values equal to the label are 1.0.
    pub fn mask_label(&self, label: i64) -> Self {
        self.map(|v| if v.round() as i64 == label { 1.0 } else { 0.0 })
    }

    /// Extract a sub-volume (crop).
    ///
    /// # Arguments
    /// * `start` - Start coordinates [x, y, z]
    /// * `size` - Size of sub-volume [sx, sy, sz]
    pub fn crop(&self, start: [usize; 3], size: [usize; 3]) -> Self {
        let [sx, sy, sz] = size;
        let mut new_data = vec![0.0; sx * sy * sz];

        for z in 0..sz {
            for y in 0..sy {
                for x in 0..sx {
                    let src_idx = self.coords_to_index(start[0] + x, start[1] + y, start[2] + z);
                    let dst_idx = z * (sx * sy) + y * sx + x;
                    new_data[dst_idx] = self.data[src_idx];
                }
            }
        }

        // Adjust origin for the crop
        let new_origin = self.voxel_to_world([start[0] as f64, start[1] as f64, start[2] as f64]);

        Self {
            data: new_data,
            shape: size,
            spacing: self.spacing,
            origin: new_origin,
            affine: self.affine, // Affine stays the same, origin adjusted
        }
    }

    /// Set spacing and update the affine accordingly.
    pub fn set_spacing(&mut self, spacing: [f64; 3]) {
        self.spacing = spacing;
        self.affine = transform::from_spacing_origin(spacing, self.origin);
    }

    /// Set origin and update the affine accordingly.
    pub fn set_origin(&mut self, origin: [f64; 3]) {
        self.origin = origin;
        self.affine = transform::from_spacing_origin(self.spacing, origin);
    }

    /// Set the affine transformation and extract spacing/origin.
    pub fn set_affine(&mut self, affine: Affine4) {
        self.affine = affine;
        self.spacing = transform::get_spacing(&affine);
        self.origin = transform::get_origin(&affine);
    }

    /// Trilinear interpolation at a floating-point voxel coordinate.
    ///
    /// Returns None if the coordinate is outside the volume.
    pub fn interpolate(&self, x: f64, y: f64, z: f64) -> Option<f64> {
        if x < 0.0
            || y < 0.0
            || z < 0.0
            || x >= (self.shape[0] - 1) as f64
            || y >= (self.shape[1] - 1) as f64
            || z >= (self.shape[2] - 1) as f64
        {
            return None;
        }

        let x0 = x.floor() as usize;
        let y0 = y.floor() as usize;
        let z0 = z.floor() as usize;
        let x1 = x0 + 1;
        let y1 = y0 + 1;
        let z1 = z0 + 1;

        let xd = x - x0 as f64;
        let yd = y - y0 as f64;
        let zd = z - z0 as f64;

        // Get 8 corner values
        let c000 = self.get(x0, y0, z0)?;
        let c100 = self.get(x1, y0, z0)?;
        let c010 = self.get(x0, y1, z0)?;
        let c110 = self.get(x1, y1, z0)?;
        let c001 = self.get(x0, y0, z1)?;
        let c101 = self.get(x1, y0, z1)?;
        let c011 = self.get(x0, y1, z1)?;
        let c111 = self.get(x1, y1, z1)?;

        // Interpolate along x
        let c00 = c000 * (1.0 - xd) + c100 * xd;
        let c01 = c001 * (1.0 - xd) + c101 * xd;
        let c10 = c010 * (1.0 - xd) + c110 * xd;
        let c11 = c011 * (1.0 - xd) + c111 * xd;

        // Interpolate along y
        let c0 = c00 * (1.0 - yd) + c10 * yd;
        let c1 = c01 * (1.0 - yd) + c11 * yd;

        // Interpolate along z
        Some(c0 * (1.0 - zd) + c1 * zd)
    }

    /// Sample the volume at a world coordinate using trilinear interpolation.
    pub fn sample_world(&self, world: [f64; 3]) -> Option<f64> {
        let voxel = self.world_to_voxel(world)?;
        self.interpolate(voxel[0], voxel[1], voxel[2])
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_zeros() {
        let vol = Volume::zeros([10, 20, 30]);
        assert_eq!(vol.shape(), [10, 20, 30]);
        assert_eq!(vol.num_voxels(), 6000);
        assert_eq!(vol.min(), 0.0);
        assert_eq!(vol.max(), 0.0);
    }

    #[test]
    fn test_get_set() {
        let mut vol = Volume::zeros([10, 10, 10]);
        vol.set(5, 5, 5, 42.0);
        assert_eq!(vol.get(5, 5, 5), Some(42.0));
        assert_eq!(vol.get(0, 0, 0), Some(0.0));
        assert_eq!(vol.get(100, 100, 100), None);
    }

    #[test]
    fn test_coords_conversion() {
        let vol = Volume::zeros([10, 20, 30]);

        // Test round-trip
        for z in 0..5 {
            for y in 0..5 {
                for x in 0..5 {
                    let idx = vol.coords_to_index(x, y, z);
                    let [rx, ry, rz] = vol.index_to_coords(idx);
                    assert_eq!([x, y, z], [rx, ry, rz]);
                }
            }
        }
    }

    #[test]
    fn test_statistics() {
        let vol = Volume::filled([10, 10, 10], 5.0);
        assert_eq!(vol.mean(), 5.0);
        assert_eq!(vol.std(), 0.0);
    }

    #[test]
    fn test_interpolate() {
        let mut vol = Volume::zeros([10, 10, 10]);
        vol.set(0, 0, 0, 0.0);
        vol.set(1, 0, 0, 1.0);

        // Interpolate at midpoint
        let val = vol.interpolate(0.5, 0.0, 0.0).unwrap();
        assert!((val - 0.5).abs() < 1e-10);
    }
}