//! Organ geometry representation for digital twins.
//!
//! This module provides structures for representing patient-specific organ
//! geometry from segmented medical imaging data.

use serde::{Deserialize, Serialize};

use crate::error::{DigitalTwinError, Result};
use crate::tissue::{TissueDatabase, TissueProperties, TissueType};

/// Tissue label at a voxel position.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub struct TissueLabel(pub u8);

impl TissueLabel {
    /// Create a new tissue label.
    pub fn new(label: u8) -> Self {
        Self(label)
    }

    /// Get the tissue type for this label.
    pub fn tissue_type(&self) -> TissueType {
        TissueType::from_label(self.0)
    }

    /// Get the raw label value.
    pub fn value(&self) -> u8 {
        self.0
    }
}

impl From<TissueType> for TissueLabel {
    fn from(tt: TissueType) -> Self {
        Self(tt.label())
    }
}

/// Per-voxel data including tissue label and physical state.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct VoxelData {
    /// Tissue label (from segmentation)
    pub label: TissueLabel,
    /// Temperature [°C] - for bioheat simulations
    pub temperature: f32,
    /// Damage parameter (Arrhenius) [-]
    pub damage: f32,
    /// Custom scalar field (user-defined)
    pub scalar_field: f32,
}

impl Default for VoxelData {
    fn default() -> Self {
        Self {
            label: TissueLabel(0), // Air
            temperature: 37.0,     // Body temperature
            damage: 0.0,
            scalar_field: 0.0,
        }
    }
}

impl VoxelData {
    /// Create voxel data with specific tissue label.
    pub fn with_label(label: TissueLabel) -> Self {
        Self {
            label,
            ..Default::default()
        }
    }

    /// Check if voxel is inside the body (not air).
    pub fn is_tissue(&self) -> bool {
        self.label.0 != 0
    }
}

/// 3D organ geometry representation.
///
/// Stores the patient-specific geometry as a voxelized representation
/// with tissue labels and physical state at each voxel.
#[derive(Debug, Clone)]
pub struct OrganGeometry {
    /// Voxel data array [x * y * z]
    data: Vec<VoxelData>,
    /// Shape [x, y, z]
    shape: [usize; 3],
    /// Voxel spacing [dx, dy, dz] in mm
    spacing: [f32; 3],
    /// Origin position in world coordinates [x, y, z] in mm
    origin: [f32; 3],
    /// Tissue property database
    tissue_db: TissueDatabase,
}

impl OrganGeometry {
    /// Create new organ geometry filled with air.
    pub fn new(shape: [usize; 3], spacing: [f32; 3]) -> Self {
        let size = shape[0] * shape[1] * shape[2];
        Self {
            data: vec![VoxelData::default(); size],
            shape,
            spacing,
            origin: [0.0, 0.0, 0.0],
            tissue_db: TissueDatabase::standard(),
        }
    }

    /// Create geometry from segmentation labels.
    ///
    /// # Arguments
    /// * `labels` - Flattened array of tissue labels [x * y * z]
    /// * `shape` - Volume shape [x, y, z]
    /// * `spacing` - Voxel spacing in mm
    pub fn from_labels(labels: &[u8], shape: [usize; 3], spacing: [f32; 3]) -> Result<Self> {
        let expected_size = shape[0] * shape[1] * shape[2];
        if labels.len() != expected_size {
            return Err(DigitalTwinError::InvalidGeometry(format!(
                "Label array size {} doesn't match shape {:?} (expected {})",
                labels.len(),
                shape,
                expected_size
            )));
        }

        let data: Vec<VoxelData> = labels
            .iter()
            .map(|&label| VoxelData::with_label(TissueLabel(label)))
            .collect();

        Ok(Self {
            data,
            shape,
            spacing,
            origin: [0.0, 0.0, 0.0],
            tissue_db: TissueDatabase::standard(),
        })
    }

    /// Set the origin position.
    pub fn set_origin(&mut self, origin: [f32; 3]) {
        self.origin = origin;
    }

    /// Set custom tissue database.
    pub fn set_tissue_database(&mut self, db: TissueDatabase) {
        self.tissue_db = db;
    }

    /// Get the shape of the geometry.
    pub fn shape(&self) -> [usize; 3] {
        self.shape
    }

    /// Get voxel spacing in mm.
    pub fn spacing(&self) -> [f32; 3] {
        self.spacing
    }

    /// Get origin position.
    pub fn origin(&self) -> [f32; 3] {
        self.origin
    }

    /// Get total number of voxels.
    pub fn num_voxels(&self) -> usize {
        self.data.len()
    }

    /// Get physical dimensions in mm.
    pub fn dimensions(&self) -> [f32; 3] {
        [
            self.shape[0] as f32 * self.spacing[0],
            self.shape[1] as f32 * self.spacing[1],
            self.shape[2] as f32 * self.spacing[2],
        ]
    }

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

    /// Convert 3D coordinates to 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
    }

    /// Check if coordinates are within bounds.
    pub fn in_bounds(&self, x: usize, y: usize, z: usize) -> bool {
        x < self.shape[0] && y < self.shape[1] && z < self.shape[2]
    }

    /// Get voxel data at coordinates.
    pub fn get(&self, x: usize, y: usize, z: usize) -> Option<&VoxelData> {
        if self.in_bounds(x, y, z) {
            Some(&self.data[self.coords_to_index(x, y, z)])
        } else {
            None
        }
    }

    /// Get mutable voxel data at coordinates.
    pub fn get_mut(&mut self, x: usize, y: usize, z: usize) -> Option<&mut VoxelData> {
        if self.in_bounds(x, y, z) {
            let idx = self.coords_to_index(x, y, z);
            Some(&mut self.data[idx])
        } else {
            None
        }
    }

    /// Get voxel data by flat index.
    pub fn get_by_index(&self, index: usize) -> Option<&VoxelData> {
        self.data.get(index)
    }

    /// Get mutable voxel data by flat index.
    pub fn get_by_index_mut(&mut self, index: usize) -> Option<&mut VoxelData> {
        self.data.get_mut(index)
    }

    /// Get tissue properties at coordinates.
    pub fn tissue_properties(&self, x: usize, y: usize, z: usize) -> Option<TissueProperties> {
        self.get(x, y, z)
            .map(|voxel| self.tissue_db.get_or_default(voxel.label.tissue_type()))
    }

    /// Get all voxel data as a slice.
    pub fn data(&self) -> &[VoxelData] {
        &self.data
    }

    /// Get mutable access to all voxel data.
    pub fn data_mut(&mut self) -> &mut [VoxelData] {
        &mut self.data
    }

    /// Get reference to tissue database.
    pub fn tissue_db(&self) -> &TissueDatabase {
        &self.tissue_db
    }

    /// Set tissue label at coordinates.
    pub fn set_label(&mut self, x: usize, y: usize, z: usize, label: TissueLabel) -> bool {
        if let Some(voxel) = self.get_mut(x, y, z) {
            voxel.label = label;
            true
        } else {
            false
        }
    }

    /// Set temperature at coordinates.
    pub fn set_temperature(&mut self, x: usize, y: usize, z: usize, temp: f32) -> bool {
        if let Some(voxel) = self.get_mut(x, y, z) {
            voxel.temperature = temp;
            true
        } else {
            false
        }
    }

    /// Get temperature field as flat array.
    pub fn temperature_field(&self) -> Vec<f32> {
        self.data.iter().map(|v| v.temperature).collect()
    }

    /// Set temperature field from flat array.
    pub fn set_temperature_field(&mut self, field: &[f32]) -> Result<()> {
        if field.len() != self.data.len() {
            return Err(DigitalTwinError::ShapeMismatch {
                expected: self.shape,
                got: [field.len(), 1, 1],
            });
        }

        for (voxel, &temp) in self.data.iter_mut().zip(field.iter()) {
            voxel.temperature = temp;
        }
        Ok(())
    }

    /// Get thermal conductivity field (k) from tissue properties.
    pub fn thermal_conductivity_field(&self) -> Vec<f32> {
        self.data
            .iter()
            .map(|v| {
                self.tissue_db
                    .get_or_default(v.label.tissue_type())
                    .thermal_conductivity
            })
            .collect()
    }

    /// Get density field (ρ) from tissue properties.
    pub fn density_field(&self) -> Vec<f32> {
        self.data
            .iter()
            .map(|v| self.tissue_db.get_or_default(v.label.tissue_type()).density)
            .collect()
    }

    /// Get specific heat field (c) from tissue properties.
    pub fn specific_heat_field(&self) -> Vec<f32> {
        self.data
            .iter()
            .map(|v| {
                self.tissue_db
                    .get_or_default(v.label.tissue_type())
                    .specific_heat
            })
            .collect()
    }

    /// Get perfusion rate field (ω) from tissue properties.
    pub fn perfusion_field(&self) -> Vec<f32> {
        self.data
            .iter()
            .map(|v| {
                self.tissue_db
                    .get_or_default(v.label.tissue_type())
                    .perfusion_rate
            })
            .collect()
    }

    /// Count voxels of each tissue type.
    pub fn tissue_histogram(&self) -> std::collections::HashMap<TissueType, usize> {
        let mut counts = std::collections::HashMap::new();
        for voxel in &self.data {
            *counts.entry(voxel.label.tissue_type()).or_insert(0) += 1;
        }
        counts
    }

    /// Create a spherical region centered at world coordinates.
    ///
    /// # Arguments
    /// * `center` - Center position in mm [x, y, z]
    /// * `radius` - Radius in mm
    /// * `label` - Tissue label to assign
    pub fn create_sphere(&mut self, center: [f32; 3], radius: f32, label: TissueLabel) {
        let radius_sq = radius * radius;

        for z in 0..self.shape[2] {
            for y in 0..self.shape[1] {
                for x in 0..self.shape[0] {
                    // World position of this voxel
                    let pos = [
                        self.origin[0] + x as f32 * self.spacing[0],
                        self.origin[1] + y as f32 * self.spacing[1],
                        self.origin[2] + z as f32 * self.spacing[2],
                    ];

                    // Distance squared from center
                    let dist_sq = (pos[0] - center[0]).powi(2)
                        + (pos[1] - center[1]).powi(2)
                        + (pos[2] - center[2]).powi(2);

                    if dist_sq <= radius_sq {
                        self.set_label(x, y, z, label);
                    }
                }
            }
        }
    }
}

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

    #[test]
    fn test_geometry_creation() {
        let geom = OrganGeometry::new([10, 10, 10], [1.0, 1.0, 1.0]);

        assert_eq!(geom.shape(), [10, 10, 10]);
        assert_eq!(geom.num_voxels(), 1000);
        assert_eq!(geom.dimensions(), [10.0, 10.0, 10.0]);
    }

    #[test]
    fn test_from_labels() {
        let labels: Vec<u8> = (0..27).map(|i| if i == 13 { 6 } else { 0 }).collect();
        let geom = OrganGeometry::from_labels(&labels, [3, 3, 3], [1.0, 1.0, 1.0]).unwrap();

        // Center voxel should be liver (label 6)
        let center = geom.get(1, 1, 1).unwrap();
        assert_eq!(center.label.value(), 6);
        assert_eq!(center.label.tissue_type(), TissueType::Liver);
    }

    #[test]
    fn test_coordinate_conversion() {
        let geom = OrganGeometry::new([5, 6, 7], [1.0, 1.0, 1.0]);

        let index = geom.coords_to_index(2, 3, 4);
        let coords = geom.index_to_coords(index);

        assert_eq!(coords, [2, 3, 4]);
    }

    #[test]
    fn test_tissue_properties_lookup() {
        let labels = vec![6u8; 8]; // All liver
        let geom = OrganGeometry::from_labels(&labels, [2, 2, 2], [1.0, 1.0, 1.0]).unwrap();

        let props = geom.tissue_properties(0, 0, 0).unwrap();
        assert_eq!(props.tissue_type, TissueType::Liver);
        assert!(props.thermal_conductivity > 0.5);
    }

    #[test]
    fn test_temperature_field() {
        let mut geom = OrganGeometry::new([3, 3, 3], [1.0, 1.0, 1.0]);

        // Set center to 42°C
        geom.set_temperature(1, 1, 1, 42.0);

        let temp_field = geom.temperature_field();
        let center_idx = geom.coords_to_index(1, 1, 1);
        assert_eq!(temp_field[center_idx], 42.0);
    }

    #[test]
    fn test_create_sphere() {
        let mut geom = OrganGeometry::new([20, 20, 20], [1.0, 1.0, 1.0]);

        // Create tumor sphere at center
        let center = [10.0, 10.0, 10.0];
        geom.create_sphere(center, 5.0, TissueLabel::from(TissueType::Tumor));

        // Center should be tumor
        let center_voxel = geom.get(10, 10, 10).unwrap();
        assert_eq!(center_voxel.label.tissue_type(), TissueType::Tumor);

        // Corner should be air
        let corner_voxel = geom.get(0, 0, 0).unwrap();
        assert_eq!(corner_voxel.label.tissue_type(), TissueType::Air);
    }

    #[test]
    fn test_tissue_histogram() {
        let mut labels = vec![0u8; 100]; // Mostly air
        labels[50] = 6; // One liver voxel
        labels[51] = 6;
        labels[52] = 3; // One muscle voxel

        let geom = OrganGeometry::from_labels(&labels, [10, 10, 1], [1.0, 1.0, 1.0]).unwrap();
        let hist = geom.tissue_histogram();

        assert_eq!(hist.get(&TissueType::Air), Some(&97));
        assert_eq!(hist.get(&TissueType::Liver), Some(&2));
        assert_eq!(hist.get(&TissueType::Muscle), Some(&1));
    }
}