//! Image interpolation for transformed coordinates.

use nalgebra::Point3;
use rtx_medical_io::Volume;

/// Interpolation method.
#[derive(Debug, Clone, Copy, Default)]
pub enum InterpolationMethod {
    /// Nearest neighbor interpolation.
    NearestNeighbor,
    /// Trilinear interpolation.
    #[default]
    Trilinear,
}

/// Interpolate a volume at a given physical coordinate.
pub fn interpolate(volume: &Volume, point: &Point3<f64>, method: InterpolationMethod) -> f64 {
    match method {
        InterpolationMethod::NearestNeighbor => interpolate_nearest(volume, point),
        InterpolationMethod::Trilinear => interpolate_trilinear(volume, point),
    }
}

/// Nearest neighbor interpolation.
pub fn interpolate_nearest(volume: &Volume, point: &Point3<f64>) -> f64 {
    let shape = volume.shape();
    let spacing = volume.spacing();

    // Convert physical to voxel coordinates
    let x = (point.x / spacing[0]).round() as i64;
    let y = (point.y / spacing[1]).round() as i64;
    let z = (point.z / spacing[2]).round() as i64;

    // Bounds check
    if x < 0
        || x >= shape[0] as i64
        || y < 0
        || y >= shape[1] as i64
        || z < 0
        || z >= shape[2] as i64
    {
        return 0.0;
    }

    volume
        .get(x as usize, y as usize, z as usize)
        .unwrap_or(0.0)
}

/// Trilinear interpolation.
pub fn interpolate_trilinear(volume: &Volume, point: &Point3<f64>) -> f64 {
    let shape = volume.shape();
    let spacing = volume.spacing();

    // Convert physical to voxel coordinates
    let x = point.x / spacing[0];
    let y = point.y / spacing[1];
    let z = point.z / spacing[2];

    // Get integer and fractional parts
    let x0 = x.floor() as i64;
    let y0 = y.floor() as i64;
    let z0 = z.floor() as i64;

    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;

    // Bounds check
    if x0 < 0
        || x1 >= shape[0] as i64
        || y0 < 0
        || y1 >= shape[1] as i64
        || z0 < 0
        || z1 >= shape[2] as i64
    {
        return 0.0;
    }

    let x0 = x0 as usize;
    let y0 = y0 as usize;
    let z0 = z0 as usize;
    let x1 = x1 as usize;
    let y1 = y1 as usize;
    let z1 = z1 as usize;

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

    // 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
    c0 * (1.0 - zd) + c1 * zd
}

/// Apply a transform and resample a volume.
pub fn resample_volume<T: crate::transform::Transform>(
    volume: &Volume,
    transform: &T,
    output_shape: [usize; 3],
    output_spacing: [f64; 3],
    method: InterpolationMethod,
) -> Volume {
    let mut output = Volume::zeros(output_shape);
    output.set_spacing(output_spacing);

    for z in 0..output_shape[2] {
        for y in 0..output_shape[1] {
            for x in 0..output_shape[0] {
                // Physical coordinate in output space
                let point = Point3::new(
                    x as f64 * output_spacing[0],
                    y as f64 * output_spacing[1],
                    z as f64 * output_spacing[2],
                );

                // Transform to input space (use inverse transform)
                if let Some(inv) = transform.try_inverse() {
                    let input_point = inv.transform_point(&point);
                    let value = interpolate(volume, &input_point, method);
                    output.set(x, y, z, value);
                }
            }
        }
    }

    output
}

/// Apply a transform and resample a volume in parallel.
pub fn resample_volume_parallel<T: crate::transform::Transform>(
    volume: &Volume,
    transform: &T,
    output_shape: [usize; 3],
    output_spacing: [f64; 3],
    method: InterpolationMethod,
) -> Volume {
    use rayon::prelude::*;

    let inv = match transform.try_inverse() {
        Some(inv) => inv,
        None => return Volume::zeros(output_shape),
    };

    let total_voxels = output_shape[0] * output_shape[1] * output_shape[2];

    let values: Vec<f64> = (0..total_voxels)
        .into_par_iter()
        .map(|idx| {
            let x = idx % output_shape[0];
            let y = (idx / output_shape[0]) % output_shape[1];
            let z = idx / (output_shape[0] * output_shape[1]);

            let point = Point3::new(
                x as f64 * output_spacing[0],
                y as f64 * output_spacing[1],
                z as f64 * output_spacing[2],
            );

            let input_point = inv.transform_point(&point);
            interpolate(volume, &input_point, method)
        })
        .collect();

    let mut output = Volume::zeros(output_shape);
    output.set_spacing(output_spacing);
    for (idx, value) in values.into_iter().enumerate() {
        let x = idx % output_shape[0];
        let y = (idx / output_shape[0]) % output_shape[1];
        let z = idx / (output_shape[0] * output_shape[1]);
        output.set(x, y, z, value);
    }

    output
}

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

    #[test]
    fn test_nearest_neighbor() {
        let mut vol = Volume::zeros([10, 10, 10]);
        vol.set(5, 5, 5, 100.0);

        // Exact voxel center
        let p = Point3::new(5.0, 5.0, 5.0);
        let v = interpolate_nearest(&vol, &p);
        assert_eq!(v, 100.0);

        // Slightly off should still round to (5,5,5)
        let p = Point3::new(5.2, 4.8, 5.3);
        let v = interpolate_nearest(&vol, &p);
        assert_eq!(v, 100.0);

        // Far off should be 0
        let p = Point3::new(0.0, 0.0, 0.0);
        let v = interpolate_nearest(&vol, &p);
        assert_eq!(v, 0.0);
    }

    #[test]
    fn test_trilinear_center() {
        let mut vol = Volume::zeros([10, 10, 10]);
        vol.set(5, 5, 5, 100.0);

        // Exact voxel position
        let p = Point3::new(5.0, 5.0, 5.0);
        let v = interpolate_trilinear(&vol, &p);
        assert_eq!(v, 100.0);
    }

    #[test]
    fn test_trilinear_midpoint() {
        let mut vol = Volume::zeros([10, 10, 10]);
        // Set two adjacent voxels
        vol.set(5, 5, 5, 0.0);
        vol.set(6, 5, 5, 100.0);

        // Midpoint along x
        let p = Point3::new(5.5, 5.0, 5.0);
        let v = interpolate_trilinear(&vol, &p);
        assert!((v - 50.0).abs() < 1e-10);
    }

    #[test]
    fn test_trilinear_corner() {
        let mut vol = Volume::zeros([10, 10, 10]);
        // Set all 8 corners of a cube
        for dz in 0..2 {
            for dy in 0..2 {
                for dx in 0..2 {
                    vol.set(2 + dx, 2 + dy, 2 + dz, (dx + dy + dz) as f64 * 100.0);
                }
            }
        }

        // Center of cube
        let p = Point3::new(2.5, 2.5, 2.5);
        let v = interpolate_trilinear(&vol, &p);
        // Average of 0, 100, 100, 200, 100, 200, 200, 300 = 150
        assert!((v - 150.0).abs() < 1e-10);
    }

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

        let p = Point3::new(-1.0, 5.0, 5.0);
        assert_eq!(interpolate_trilinear(&vol, &p), 0.0);

        let p = Point3::new(10.0, 5.0, 5.0);
        assert_eq!(interpolate_trilinear(&vol, &p), 0.0);
    }
}