//! Mean Squared Error similarity metric.

use super::SimilarityMetric;
use rayon::prelude::*;
use rtx_medical_io::Volume;

/// Mean Squared Error (MSE) metric.
///
/// Computes: MSE = (1/N) * Σ(fixed[i] - moving[i])²
///
/// Lower values indicate better alignment.
#[derive(Debug, Clone, Default)]
pub struct MeanSquaredError;

impl MeanSquaredError {
    /// Create a new MSE metric.
    pub fn new() -> Self {
        Self
    }
}

impl SimilarityMetric for MeanSquaredError {
    fn compute(&self, fixed: &Volume, moving: &Volume, mask: Option<&Volume>) -> f64 {
        let shape = fixed.shape();

        if shape != moving.shape() {
            return f64::MAX;
        }

        // Collect all voxel pairs
        let total_voxels = shape[0] * shape[1] * shape[2];
        let results: Vec<(f64, usize)> = (0..total_voxels)
            .into_par_iter()
            .map(|idx| {
                let x = idx % shape[0];
                let y = (idx / shape[0]) % shape[1];
                let z = idx / (shape[0] * shape[1]);

                // Check mask
                if let Some(m) = mask {
                    if m.get(x, y, z).unwrap_or(0.0) <= 0.0 {
                        return (0.0, 0);
                    }
                }

                let f = fixed.get(x, y, z).unwrap_or(0.0);
                let m = moving.get(x, y, z).unwrap_or(0.0);
                let diff = f - m;
                (diff * diff, 1)
            })
            .collect();

        let (sum_sq, count): (f64, usize) = results
            .iter()
            .fold((0.0, 0), |acc, &(sq, c)| (acc.0 + sq, acc.1 + c));

        if count == 0 {
            return f64::MAX;
        }

        sum_sq / count as f64
    }

    fn is_similarity(&self) -> bool {
        false // Lower is better
    }
}

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

    #[test]
    fn test_identical_images() {
        let mut vol = Volume::zeros([10, 10, 10]);
        for z in 0..10 {
            for y in 0..10 {
                for x in 0..10 {
                    vol.set(x, y, z, (x + y + z) as f64);
                }
            }
        }

        let mse = MeanSquaredError::new();
        let value = mse.compute(&vol, &vol, None);
        assert!(value < 1e-10);
    }

    #[test]
    fn test_constant_difference() {
        let mut vol1 = Volume::zeros([10, 10, 10]);
        let mut vol2 = Volume::zeros([10, 10, 10]);

        for z in 0..10 {
            for y in 0..10 {
                for x in 0..10 {
                    vol1.set(x, y, z, 5.0);
                    vol2.set(x, y, z, 7.0); // Constant difference of 2
                }
            }
        }

        let mse = MeanSquaredError::new();
        let value = mse.compute(&vol1, &vol2, None);
        assert!((value - 4.0).abs() < 1e-10); // MSE = (2)² = 4
    }

    #[test]
    fn test_with_mask() {
        let mut vol1 = Volume::zeros([10, 10, 10]);
        let mut vol2 = Volume::zeros([10, 10, 10]);
        let mut mask = Volume::zeros([10, 10, 10]);

        // Only set mask for a small region
        for z in 0..5 {
            for y in 0..5 {
                for x in 0..5 {
                    vol1.set(x, y, z, 5.0);
                    vol2.set(x, y, z, 7.0);
                    mask.set(x, y, z, 1.0);
                }
            }
        }

        // Rest of volumes have different values but should be ignored
        for z in 5..10 {
            for y in 5..10 {
                for x in 5..10 {
                    vol1.set(x, y, z, 100.0);
                    vol2.set(x, y, z, 0.0);
                }
            }
        }

        let mse = MeanSquaredError::new();
        let value = mse.compute(&vol1, &vol2, Some(&mask));
        assert!((value - 4.0).abs() < 1e-10); // MSE should only consider masked region
    }

    #[test]
    fn test_to_cost() {
        let mse = MeanSquaredError::new();
        assert!(!mse.is_similarity());
        assert_eq!(mse.to_cost(5.0), 5.0); // Already a cost
    }
}