//! Magnetic Resonance Elastography (MRE) material property utilities.
//!
//! This module provides utilities for working with MRE-derived material
//! properties, including conversions between different representations
//! and multi-frequency data handling.

use crate::error::{MaterialError, Result};
use crate::viscoelastic::PronyCoefficients;
use serde::{Deserialize, Serialize};

/// MRE material property type.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MrePropertyType {
    /// Complex shear modulus (G', G'')
    ComplexShear,
    /// Stiffness and damping ratio (μ, ξ)
    StiffnessDamping,
}

/// Multi-frequency MRE measurement data.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MreData {
    /// Property type
    pub property_type: MrePropertyType,
    /// Frequencies in Hz
    pub frequencies: Vec<f64>,
    /// Property values (storage modulus G' or stiffness μ)
    pub primary_values: Vec<f64>,
    /// Secondary property values (loss modulus G'' or damping ratio ξ)
    pub secondary_values: Vec<f64>,
}

impl MreData {
    /// Create new MRE data from complex shear modulus measurements.
    ///
    /// # Arguments
    /// * `frequencies` - MRE frequencies in Hz
    /// * `storage_modulus` - Storage modulus G' values in Pa
    /// * `loss_modulus` - Loss modulus G'' values in Pa
    pub fn from_complex_shear(
        frequencies: Vec<f64>,
        storage_modulus: Vec<f64>,
        loss_modulus: Vec<f64>,
    ) -> Result<Self> {
        if frequencies.len() != storage_modulus.len() || frequencies.len() != loss_modulus.len() {
            return Err(MaterialError::InvalidInput(
                "All arrays must have the same length".to_string(),
            ));
        }

        Ok(Self {
            property_type: MrePropertyType::ComplexShear,
            frequencies,
            primary_values: storage_modulus,
            secondary_values: loss_modulus,
        })
    }

    /// Create new MRE data from stiffness and damping ratio measurements.
    ///
    /// # Arguments
    /// * `frequencies` - MRE frequencies in Hz
    /// * `stiffness` - Shear stiffness μ values in Pa
    /// * `damping_ratio` - Damping ratio ξ values (dimensionless)
    pub fn from_stiffness_damping(
        frequencies: Vec<f64>,
        stiffness: Vec<f64>,
        damping_ratio: Vec<f64>,
    ) -> Result<Self> {
        if frequencies.len() != stiffness.len() || frequencies.len() != damping_ratio.len() {
            return Err(MaterialError::InvalidInput(
                "All arrays must have the same length".to_string(),
            ));
        }

        Ok(Self {
            property_type: MrePropertyType::StiffnessDamping,
            frequencies,
            primary_values: stiffness,
            secondary_values: damping_ratio,
        })
    }

    /// Get angular frequencies (ω = 2πf).
    pub fn angular_frequencies(&self) -> Vec<f64> {
        self.frequencies
            .iter()
            .map(|&f| 2.0 * std::f64::consts::PI * f)
            .collect()
    }

    /// Convert to complex shear modulus representation.
    ///
    /// If already in complex shear format, returns (gp, gpp) directly.
    /// If in stiffness/damping format, performs conversion.
    pub fn to_complex_shear(&self) -> (Vec<f64>, Vec<f64>) {
        match self.property_type {
            MrePropertyType::ComplexShear => {
                (self.primary_values.clone(), self.secondary_values.clone())
            }
            MrePropertyType::StiffnessDamping => {
                let mut gp = Vec::with_capacity(self.frequencies.len());
                let mut gpp = Vec::with_capacity(self.frequencies.len());

                for i in 0..self.frequencies.len() {
                    let (g, gpp_val) = crate::viscoelastic::mu_xi_to_complex(
                        self.primary_values[i],
                        self.secondary_values[i],
                    );
                    gp.push(g);
                    gpp.push(gpp_val);
                }

                (gp, gpp)
            }
        }
    }

    /// Fit Prony series coefficients to this MRE data.
    pub fn fit_prony(&self) -> Result<PronyCoefficients> {
        let (gp, gpp) = self.to_complex_shear();
        let omega = self.angular_frequencies();

        crate::viscoelastic::calculate_prony(&gp, &gpp, &omega)
    }

    /// Calculate mean storage modulus.
    pub fn mean_storage_modulus(&self) -> f64 {
        let (gp, _) = self.to_complex_shear();
        gp.iter().sum::<f64>() / gp.len() as f64
    }

    /// Calculate mean loss modulus.
    pub fn mean_loss_modulus(&self) -> f64 {
        let (_, gpp) = self.to_complex_shear();
        gpp.iter().sum::<f64>() / gpp.len() as f64
    }

    /// Calculate mean complex magnitude |G*|.
    pub fn mean_complex_magnitude(&self) -> f64 {
        let (gp, gpp) = self.to_complex_shear();
        let mags: Vec<f64> = gp
            .iter()
            .zip(gpp.iter())
            .map(|(&g, &gpp)| (g * g + gpp * gpp).sqrt())
            .collect();
        mags.iter().sum::<f64>() / mags.len() as f64
    }
}

/// MRE region properties for segmented data.
///
/// Represents average MRE properties within a segmented region.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MreRegion {
    /// Region label/ID
    pub label: i64,
    /// Number of voxels in region
    pub voxel_count: usize,
    /// Mean storage modulus (G') in Pa
    pub mean_storage_modulus: f64,
    /// Mean loss modulus (G'') in Pa
    pub mean_loss_modulus: f64,
    /// Standard deviation of storage modulus
    pub std_storage_modulus: f64,
    /// Standard deviation of loss modulus
    pub std_loss_modulus: f64,
    /// Fitted Prony coefficients (if available)
    pub prony: Option<PronyCoefficients>,
}

impl MreRegion {
    /// Create a new MRE region with basic properties.
    pub fn new(label: i64, mean_gp: f64, mean_gpp: f64) -> Self {
        Self {
            label,
            voxel_count: 0,
            mean_storage_modulus: mean_gp,
            mean_loss_modulus: mean_gpp,
            std_storage_modulus: 0.0,
            std_loss_modulus: 0.0,
            prony: None,
        }
    }

    /// Calculate mean stiffness from complex modulus.
    pub fn mean_stiffness(&self) -> f64 {
        let (mu, _) = crate::viscoelastic::complex_to_mu_xi(
            self.mean_storage_modulus,
            self.mean_loss_modulus,
        );
        mu
    }

    /// Calculate mean damping ratio from complex modulus.
    pub fn mean_damping_ratio(&self) -> f64 {
        let (_, xi) = crate::viscoelastic::complex_to_mu_xi(
            self.mean_storage_modulus,
            self.mean_loss_modulus,
        );
        xi
    }

    /// Calculate complex magnitude |G*|.
    pub fn complex_magnitude(&self) -> f64 {
        (self.mean_storage_modulus.powi(2) + self.mean_loss_modulus.powi(2)).sqrt()
    }

    /// Calculate loss tangent tan(δ) = G''/G'.
    pub fn loss_tangent(&self) -> f64 {
        self.mean_loss_modulus / self.mean_storage_modulus
    }
}

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

    #[test]
    fn test_mre_data_complex_shear() {
        let data =
            MreData::from_complex_shear(vec![10.0, 20.0], vec![1500.0, 1800.0], vec![500.0, 400.0])
                .unwrap();

        assert_eq!(data.property_type, MrePropertyType::ComplexShear);
        assert_eq!(data.frequencies.len(), 2);
    }

    #[test]
    fn test_mre_data_stiffness_damping() {
        let data = MreData::from_stiffness_damping(
            vec![10.0, 20.0],
            vec![2000.0, 2500.0],
            vec![0.2, 0.15],
        )
        .unwrap();

        assert_eq!(data.property_type, MrePropertyType::StiffnessDamping);

        // Should be able to convert to complex shear
        let (gp, gpp) = data.to_complex_shear();
        assert_eq!(gp.len(), 2);
        assert_eq!(gpp.len(), 2);
    }

    #[test]
    fn test_angular_frequencies() {
        let data =
            MreData::from_complex_shear(vec![10.0, 20.0], vec![1.0, 2.0], vec![0.5, 1.0]).unwrap();

        let omega = data.angular_frequencies();
        assert!((omega[0] - 62.83185307179586).abs() < 1e-10);
        assert!((omega[1] - 125.66370614359172).abs() < 1e-10);
    }

    #[test]
    fn test_mre_region() {
        let region = MreRegion::new(1, 1500.0, 500.0);

        assert_eq!(region.label, 1);
        assert_eq!(region.mean_storage_modulus, 1500.0);
        assert_eq!(region.mean_loss_modulus, 500.0);

        // Check derived quantities
        let magnitude = region.complex_magnitude();
        assert!((magnitude - (1500.0_f64.powi(2) + 500.0_f64.powi(2)).sqrt()).abs() < 1e-10);

        let loss_tan = region.loss_tangent();
        assert!((loss_tan - 500.0 / 1500.0).abs() < 1e-10);
    }
}