//! MRE-to-mesh mapping utilities.
//!
//! This module provides functionality to map MRE (Magnetic Resonance Elastography)
//! property fields onto finite element mesh nodes.

use crate::error::{Mri2FeError, Result};
use nalgebra::{DVector, Point3};
use rayon::prelude::*;
use rtx_medical_io::Volume;
use rtx_mesh_gen::TetrahedralMesh;
use rtx_segmentation::prelude::*;
use serde::{Deserialize, Serialize};

/// Configuration for MRE mapping.
#[derive(Debug, Clone)]
pub struct MreMappingConfig {
    /// Number of K-means segments for spatial grouping.
    pub n_segments: usize,
    /// Interpolation method.
    pub interpolation: InterpolationMethod,
    /// Search radius for local averaging (mm).
    pub search_radius: Option<f64>,
    /// Mask threshold for valid MRE values.
    pub mask_threshold: f64,
}

impl Default for MreMappingConfig {
    fn default() -> Self {
        Self {
            n_segments: 5,
            interpolation: InterpolationMethod::Trilinear,
            search_radius: None,
            mask_threshold: 0.0,
        }
    }
}

/// Interpolation method for MRE values.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum InterpolationMethod {
    /// Nearest neighbor interpolation.
    NearestNeighbor,
    /// Trilinear interpolation.
    Trilinear,
    /// Inverse distance weighting.
    InverseDistance,
}

/// MRE property type.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum MrePropertyType {
    /// Complex shear modulus (G' and G'').
    ComplexShear,
    /// Wave amplitude.
    Amplitude,
    /// Phase velocity.
    PhaseVelocity,
    /// Stiffness (μ) and damping (ξ).
    StiffnessDamping,
}

/// Result of MRE mapping operation.
#[derive(Debug, Clone)]
pub struct MreMappingResult {
    /// Mapped storage modulus values at mesh nodes.
    pub storage_modulus: Vec<f64>,
    /// Mapped loss modulus values at mesh nodes.
    pub loss_modulus: Vec<f64>,
    /// Mask indicating valid values (1.0) vs extrapolated (0.0).
    pub valid_mask: Vec<f64>,
    /// Segment assignments for each node.
    pub segments: Option<Vec<usize>>,
    /// Statistics about the mapping.
    pub stats: MappingStats,
}

/// Statistics about the MRE mapping.
#[derive(Debug, Clone, Default)]
pub struct MappingStats {
    /// Number of nodes mapped.
    pub nodes_mapped: usize,
    /// Number of nodes with valid MRE data.
    pub nodes_valid: usize,
    /// Number of nodes extrapolated.
    pub nodes_extrapolated: usize,
    /// Mean storage modulus.
    pub mean_storage: f64,
    /// Mean loss modulus.
    pub mean_loss: f64,
    /// Standard deviation of storage modulus.
    pub std_storage: f64,
    /// Standard deviation of loss modulus.
    pub std_loss: f64,
}

/// Map MRE properties onto mesh nodes.
pub fn map_mre_to_mesh(
    mesh: &TetrahedralMesh,
    storage_volume: &Volume,
    loss_volume: &Volume,
    mask: Option<&Volume>,
    config: &MreMappingConfig,
) -> Result<MreMappingResult> {
    let n_nodes = mesh.vertices.len();
    if n_nodes == 0 {
        return Err(Mri2FeError::MappingError("Empty mesh".to_string()));
    }

    // Get volume properties
    let shape = storage_volume.shape();
    let spacing = storage_volume.spacing();
    let origin = storage_volume.origin();

    // Map each node to volume coordinates and sample
    let node_values: Vec<(f64, f64, bool)> = mesh
        .vertices
        .par_iter()
        .map(|vertex| {
            let node = vertex.position;
            // Convert world coordinates to voxel coordinates
            let vx = ((node.x - origin[0]) / spacing[0]).floor() as i64;
            let vy = ((node.y - origin[1]) / spacing[1]).floor() as i64;
            let vz = ((node.z - origin[2]) / spacing[2]).floor() as i64;

            // Check if within bounds
            if vx < 0
                || vy < 0
                || vz < 0
                || vx >= shape[0] as i64
                || vy >= shape[1] as i64
                || vz >= shape[2] as i64
            {
                return (0.0, 0.0, false);
            }

            let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);

            // Check mask if provided
            if let Some(m) = mask {
                if m.get(vx, vy, vz).unwrap_or(0.0) <= config.mask_threshold {
                    return (0.0, 0.0, false);
                }
            }

            match config.interpolation {
                InterpolationMethod::NearestNeighbor => {
                    let g_prime = storage_volume.get(vx, vy, vz).unwrap_or(0.0);
                    let g_double_prime = loss_volume.get(vx, vy, vz).unwrap_or(0.0);
                    (g_prime, g_double_prime, true)
                }
                InterpolationMethod::Trilinear => {
                    let g_prime = trilinear_interp(storage_volume, &node, &origin, &spacing);
                    let g_double_prime = trilinear_interp(loss_volume, &node, &origin, &spacing);
                    (g_prime, g_double_prime, true)
                }
                InterpolationMethod::InverseDistance => {
                    let radius = config.search_radius.unwrap_or(spacing[0] * 2.0);
                    let g_prime = idw_interp(storage_volume, &node, &origin, &spacing, radius);
                    let g_double_prime = idw_interp(loss_volume, &node, &origin, &spacing, radius);
                    (g_prime, g_double_prime, true)
                }
            }
        })
        .collect();

    // Separate results
    let mut storage_modulus = Vec::with_capacity(n_nodes);
    let mut loss_modulus = Vec::with_capacity(n_nodes);
    let mut valid_mask = Vec::with_capacity(n_nodes);
    let mut nodes_valid = 0;

    for (g_prime, g_double_prime, valid) in &node_values {
        storage_modulus.push(*g_prime);
        loss_modulus.push(*g_double_prime);
        valid_mask.push(if *valid { 1.0 } else { 0.0 });
        if *valid {
            nodes_valid += 1;
        }
    }

    // Optionally perform K-means segmentation
    let segments = if config.n_segments > 1 && nodes_valid > config.n_segments {
        Some(segment_nodes(
            &storage_modulus,
            &loss_modulus,
            &valid_mask,
            config.n_segments,
        )?)
    } else {
        None
    };

    // Calculate statistics
    let stats = calculate_stats(&storage_modulus, &loss_modulus, &valid_mask, nodes_valid);

    Ok(MreMappingResult {
        storage_modulus,
        loss_modulus,
        valid_mask,
        segments,
        stats,
    })
}

/// Trilinear interpolation in a volume.
fn trilinear_interp(
    volume: &Volume,
    point: &Point3<f64>,
    origin: &[f64; 3],
    spacing: &[f64; 3],
) -> f64 {
    let shape = volume.shape();

    // Get fractional voxel coordinates
    let fx = (point.x - origin[0]) / spacing[0];
    let fy = (point.y - origin[1]) / spacing[1];
    let fz = (point.z - origin[2]) / spacing[2];

    let x0 = fx.floor() as i64;
    let y0 = fy.floor() as i64;
    let z0 = fz.floor() as i64;

    let x1 = x0 + 1;
    let y1 = y0 + 1;
    let z1 = z0 + 1;

    // Check bounds
    if x0 < 0
        || y0 < 0
        || z0 < 0
        || x1 >= shape[0] as i64
        || y1 >= shape[1] as i64
        || z1 >= shape[2] as i64
    {
        // Fall back to nearest neighbor
        let vx = fx.round().clamp(0.0, (shape[0] - 1) as f64) as usize;
        let vy = fy.round().clamp(0.0, (shape[1] - 1) as f64) as usize;
        let vz = fz.round().clamp(0.0, (shape[2] - 1) as f64) as usize;
        return volume.get(vx, vy, vz).unwrap_or(0.0);
    }

    let (x0, y0, z0) = (x0 as usize, y0 as usize, z0 as usize);
    let (x1, y1, z1) = (x1 as usize, y1 as usize, z1 as usize);

    // Interpolation weights
    let xd = fx - fx.floor();
    let yd = fy - fy.floor();
    let zd = fz - fz.floor();

    // Sample 8 corners
    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);

    // Trilinear interpolation
    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;

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

    c0 * (1.0 - zd) + c1 * zd
}

/// Inverse distance weighted interpolation.
fn idw_interp(
    volume: &Volume,
    point: &Point3<f64>,
    origin: &[f64; 3],
    spacing: &[f64; 3],
    radius: f64,
) -> f64 {
    let shape = volume.shape();

    // Center voxel
    let cx = ((point.x - origin[0]) / spacing[0]).round() as i64;
    let cy = ((point.y - origin[1]) / spacing[1]).round() as i64;
    let cz = ((point.z - origin[2]) / spacing[2]).round() as i64;

    // Search radius in voxels
    let rx = (radius / spacing[0]).ceil() as i64;
    let ry = (radius / spacing[1]).ceil() as i64;
    let rz = (radius / spacing[2]).ceil() as i64;

    let mut sum = 0.0;
    let mut weight_sum = 0.0;

    for dz in -rz..=rz {
        for dy in -ry..=ry {
            for dx in -rx..=rx {
                let vx = cx + dx;
                let vy = cy + dy;
                let vz = cz + dz;

                if vx < 0
                    || vy < 0
                    || vz < 0
                    || vx >= shape[0] as i64
                    || vy >= shape[1] as i64
                    || vz >= shape[2] as i64
                {
                    continue;
                }

                let (vx, vy, vz) = (vx as usize, vy as usize, vz as usize);

                // World position of this voxel
                let vox_x = origin[0] + vx as f64 * spacing[0];
                let vox_y = origin[1] + vy as f64 * spacing[1];
                let vox_z = origin[2] + vz as f64 * spacing[2];

                let dist = ((point.x - vox_x).powi(2)
                    + (point.y - vox_y).powi(2)
                    + (point.z - vox_z).powi(2))
                .sqrt();

                if dist <= radius {
                    let weight = if dist < 1e-10 {
                        1e10
                    } else {
                        1.0 / (dist * dist)
                    };
                    sum += weight * volume.get(vx, vy, vz).unwrap_or(0.0);
                    weight_sum += weight;
                }
            }
        }
    }

    if weight_sum > 0.0 {
        sum / weight_sum
    } else {
        0.0
    }
}

/// Segment nodes using K-means on MRE values.
fn segment_nodes(
    storage: &[f64],
    loss: &[f64],
    valid_mask: &[f64],
    n_segments: usize,
) -> Result<Vec<usize>> {
    // Collect valid data points
    let mut data: Vec<DVector<f64>> = Vec::new();
    let mut valid_indices: Vec<usize> = Vec::new();

    for i in 0..storage.len() {
        if valid_mask[i] > 0.5 {
            data.push(DVector::from_vec(vec![storage[i], loss[i]]));
            valid_indices.push(i);
        }
    }

    if data.len() < n_segments {
        return Err(Mri2FeError::MappingError(format!(
            "Not enough valid points ({}) for {} segments",
            data.len(),
            n_segments
        )));
    }

    // Run K-means
    let kmeans = KMeans::with_k(n_segments);
    let result = kmeans.fit(&data)?;

    // Map back to all nodes
    let mut segments = vec![0; storage.len()];
    for (i, &orig_idx) in valid_indices.iter().enumerate() {
        segments[orig_idx] = result.labels[i];
    }

    Ok(segments)
}

/// Calculate mapping statistics.
fn calculate_stats(
    storage: &[f64],
    loss: &[f64],
    valid_mask: &[f64],
    nodes_valid: usize,
) -> MappingStats {
    let n_nodes = storage.len();

    // Calculate means
    let (sum_storage, sum_loss): (f64, f64) = storage
        .iter()
        .zip(loss.iter())
        .zip(valid_mask.iter())
        .filter(|(_, m)| **m > 0.5)
        .map(|((s, l), _)| (*s, *l))
        .fold((0.0, 0.0), |(ss, sl), (s, l)| (ss + s, sl + l));

    let mean_storage = if nodes_valid > 0 {
        sum_storage / nodes_valid as f64
    } else {
        0.0
    };
    let mean_loss = if nodes_valid > 0 {
        sum_loss / nodes_valid as f64
    } else {
        0.0
    };

    // Calculate standard deviations
    let (var_storage, var_loss): (f64, f64) = storage
        .iter()
        .zip(loss.iter())
        .zip(valid_mask.iter())
        .filter(|(_, m)| **m > 0.5)
        .map(|((s, l), _)| ((*s - mean_storage).powi(2), (*l - mean_loss).powi(2)))
        .fold((0.0, 0.0), |(vs, vl), (s, l)| (vs + s, vl + l));

    let std_storage = if nodes_valid > 1 {
        (var_storage / (nodes_valid - 1) as f64).sqrt()
    } else {
        0.0
    };
    let std_loss = if nodes_valid > 1 {
        (var_loss / (nodes_valid - 1) as f64).sqrt()
    } else {
        0.0
    };

    MappingStats {
        nodes_mapped: n_nodes,
        nodes_valid,
        nodes_extrapolated: n_nodes - nodes_valid,
        mean_storage,
        mean_loss,
        std_storage,
        std_loss,
    }
}

/// Convert stiffness (μ) and damping (ξ) to complex shear modulus.
pub fn stiffness_damping_to_complex_shear(mu: f64, xi: f64) -> (f64, f64) {
    // G' = μ(1 - ξ²) / (1 + ξ²)
    // G'' = 2μξ / (1 + ξ²)
    let xi_sq = xi * xi;
    let denom = 1.0 + xi_sq;
    let g_prime = mu * (1.0 - xi_sq) / denom;
    let g_double_prime = 2.0 * mu * xi / denom;
    (g_prime, g_double_prime)
}

/// Convert complex shear modulus to stiffness and damping.
pub fn complex_shear_to_stiffness_damping(g_prime: f64, g_double_prime: f64) -> (f64, f64) {
    // μ = sqrt(G'² + G''²)
    // ξ = G'' / (G' + μ)
    let mu = (g_prime * g_prime + g_double_prime * g_double_prime).sqrt();
    let xi = if (g_prime + mu).abs() > 1e-10 {
        g_double_prime / (g_prime + mu)
    } else {
        0.0
    };
    (mu, xi)
}

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

    #[test]
    fn test_stiffness_damping_conversion() {
        let mu = 1000.0;
        let xi = 0.1;

        let (g_prime, g_double_prime) = stiffness_damping_to_complex_shear(mu, xi);
        let (mu_back, xi_back) = complex_shear_to_stiffness_damping(g_prime, g_double_prime);

        assert!((mu - mu_back).abs() < 1e-6);
        assert!((xi - xi_back).abs() < 1e-6);
    }

    #[test]
    fn test_trilinear_interp() {
        // Create a simple volume with identity affine
        let identity_affine = [
            [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],
        ];
        let volume = Volume::new(
            vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
            [2, 2, 2],
            [1.0, 1.0, 1.0],
            [0.0, 0.0, 0.0],
            identity_affine,
        );

        let origin = [0.0, 0.0, 0.0];
        let spacing = [1.0, 1.0, 1.0];

        // Center of volume
        let point = Point3::new(0.5, 0.5, 0.5);
        let val = trilinear_interp(&volume, &point, &origin, &spacing);

        // Should be average of all 8 corners = 4.5
        assert!((val - 4.5).abs() < 1e-6);
    }

    #[test]
    fn test_mapping_config_default() {
        let config = MreMappingConfig::default();
        assert_eq!(config.n_segments, 5);
        assert_eq!(config.interpolation, InterpolationMethod::Trilinear);
        assert!(config.search_radius.is_none());
    }

    #[test]
    fn test_mapping_stats() {
        let storage = vec![100.0, 200.0, 300.0, 0.0];
        let loss = vec![10.0, 20.0, 30.0, 0.0];
        let valid_mask = vec![1.0, 1.0, 1.0, 0.0];

        let stats = calculate_stats(&storage, &loss, &valid_mask, 3);

        assert_eq!(stats.nodes_mapped, 4);
        assert_eq!(stats.nodes_valid, 3);
        assert_eq!(stats.nodes_extrapolated, 1);
        assert!((stats.mean_storage - 200.0).abs() < 1e-6);
        assert!((stats.mean_loss - 20.0).abs() < 1e-6);
    }
}