// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Hyperelastic material models for large deformation analysis.

use super::{Material, MaterialProperties, MaterialResponse, MaterialState};
use crate::error::{FeaError, FeaResult};
use nalgebra::{DMatrix, DVector, Matrix3, Vector6};
use serde::{Deserialize, Serialize};

/// Neo-Hookean hyperelastic material.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NeoHookean {
    properties: MaterialProperties,
    /// First Lamé parameter
    lambda: f64,
    /// Second Lamé parameter (shear modulus)
    mu: f64,
}

impl NeoHookean {
    /// Create a new Neo-Hookean material.
    pub fn new(elastic_modulus: f64, poisson_ratio: f64) -> Self {
        let properties =
            MaterialProperties::isotropic_elastic(elastic_modulus, poisson_ratio, 1000.0);
        let (lambda, mu) = properties.lame_parameters();

        Self {
            properties,
            lambda,
            mu,
        }
    }

    /// Set density.
    pub fn with_density(mut self, density: f64) -> Self {
        self.properties.density = density;
        self
    }

    /// Compute strain energy density.
    pub fn strain_energy_density(&self, deformation_gradient: &Matrix3<f64>) -> f64 {
        let i1 = deformation_gradient.trace();
        let j = deformation_gradient.determinant();

        if j <= 0.0 {
            return f64::INFINITY; // Invalid deformation
        }

        let ln_j = j.ln();
        0.5 * self.mu * (i1 - 3.0) - self.mu * ln_j + 0.5 * self.lambda * ln_j * ln_j
    }
}

impl Material for NeoHookean {
    fn properties(&self) -> &MaterialProperties {
        &self.properties
    }

    fn compute_response(
        &self,
        strain_increment: &Vector6<f64>,
        current_state: &MaterialState,
        _dt: f64,
    ) -> FeaResult<MaterialResponse> {
        // Full implementation: proper finite strain hyperelastic formulation
        let mut new_state = current_state.clone();
        new_state.total_strain += strain_increment;

        // Convert strain vector to Green-Lagrange strain tensor
        let e_gl = self.strain_vector_to_tensor(&new_state.total_strain);

        // Compute right Cauchy-Green deformation tensor: C = 2*E + I
        let identity = DMatrix::identity(3, 3);
        let c = 2.0 * &e_gl + &identity;

        // Compute deformation invariants
        let i1 = c.trace();
        let _i2 = 0.5 * (i1.powi(2) - (&c * &c).trace());
        let i3 = c.determinant();

        // Lame parameters
        let lambda = self.compute_lame_lambda();
        let mu = self.compute_shear_modulus();

        // Neo-Hookean strain energy derivatives
        let j = i3.sqrt(); // Jacobian of deformation

        // Second Piola-Kirchhoff stress using Neo-Hookean model
        let c_inv = c.clone().try_inverse().ok_or_else(|| {
            FeaError::ComputationFailed("Singular deformation tensor".to_string())
        })?;

        let s = mu * (&identity - &c_inv) + (lambda / 2.0) * (j.powi(2) - 1.0) * &c_inv;

        // Convert tensor stress back to vector form
        let stress_vector = self.stress_tensor_to_vector(&s);

        // Compute material tangent moduli (elasticity tensor)
        let tangent = self.compute_hyperelastic_tangent(&c, &c_inv, j, lambda, mu)?;

        new_state.stress = stress_vector;

        Ok(MaterialResponse::new(stress_vector, tangent, new_state))
    }

    fn elastic_tangent(&self) -> FeaResult<DMatrix<f64>> {
        // Hyperelastic tangent at reference configuration (undeformed state)
        // For Neo-Hookean, this reduces to the linear elastic tangent at small strains
        let lambda = self.compute_lame_lambda();
        let mu = self.compute_shear_modulus();

        let mut d = DMatrix::zeros(6, 6);

        // Diagonal terms (normal stresses)
        d[(0, 0)] = lambda + 2.0 * mu;
        d[(1, 1)] = lambda + 2.0 * mu;
        d[(2, 2)] = lambda + 2.0 * mu;

        // Off-diagonal coupling terms
        d[(0, 1)] = lambda;
        d[(1, 0)] = lambda;
        d[(0, 2)] = lambda;
        d[(2, 0)] = lambda;
        d[(1, 2)] = lambda;
        d[(2, 1)] = lambda;

        // Shear terms
        d[(3, 3)] = mu; // γ_xy
        d[(4, 4)] = mu; // γ_xz
        d[(5, 5)] = mu; // γ_yz

        Ok(d)
    }

    fn material_type(&self) -> &'static str {
        "NeoHookean"
    }
}

impl NeoHookean {
    fn strain_vector_to_tensor(&self, strain_vec: &Vector6<f64>) -> DMatrix<f64> {
        let mut tensor = DMatrix::zeros(3, 3);
        tensor[(0, 0)] = strain_vec[0];
        tensor[(1, 1)] = strain_vec[1];
        tensor[(2, 2)] = strain_vec[2];
        tensor[(0, 1)] = strain_vec[3] / 2.0;
        tensor[(1, 0)] = strain_vec[3] / 2.0;
        tensor[(0, 2)] = strain_vec[4] / 2.0;
        tensor[(2, 0)] = strain_vec[4] / 2.0;
        tensor[(1, 2)] = strain_vec[5] / 2.0;
        tensor[(2, 1)] = strain_vec[5] / 2.0;
        tensor
    }

    fn stress_tensor_to_vector(&self, stress_tensor: &DMatrix<f64>) -> Vector6<f64> {
        Vector6::new(
            stress_tensor[(0, 0)],
            stress_tensor[(1, 1)],
            stress_tensor[(2, 2)],
            stress_tensor[(0, 1)],
            stress_tensor[(0, 2)],
            stress_tensor[(1, 2)],
        )
    }

    fn compute_lame_lambda(&self) -> f64 {
        let e = self.properties.elastic_modulus;
        let nu = self.properties.poisson_ratio;
        e * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
    }

    fn compute_shear_modulus(&self) -> f64 {
        let e = self.properties.elastic_modulus;
        let nu = self.properties.poisson_ratio;
        e / (2.0 * (1.0 + nu))
    }

    fn compute_hyperelastic_tangent(
        &self,
        _c: &DMatrix<f64>,
        c_inv: &DMatrix<f64>,
        j: f64,
        lambda: f64,
        mu: f64,
    ) -> FeaResult<DMatrix<f64>> {
        // Compute fourth-order elasticity tensor for Neo-Hookean material
        let mut tangent = DMatrix::zeros(6, 6);

        // Helper function to map tensor indices to Voigt notation
        let voigt_map = [(0, 0), (1, 1), (2, 2), (0, 1), (0, 2), (1, 2)];

        for i in 0..6 {
            for k in 0..6 {
                let (i1, i2) = voigt_map[i];
                let (j1, j2) = voigt_map[k];

                // Neo-Hookean tangent moduli components
                let delta_ij = if i1 == j1 && i2 == j2 { 1.0 } else { 0.0 };
                let _delta_i1j1 = if i1 == j1 { 1.0 } else { 0.0 };
                let _delta_i2j2 = if i2 == j2 { 1.0 } else { 0.0 };

                let c_inv_i1j1 = c_inv[(i1, j1)];
                let c_inv_i2j2 = c_inv[(i2, j2)];
                let c_inv_i1j2 = c_inv[(i1, j2)];
                let c_inv_i2j1 = c_inv[(i2, j1)];

                // Material tangent for Neo-Hookean model
                tangent[(i, k)] = lambda * j.powi(2) * c_inv_i1j1 * c_inv_i2j2
                    + (lambda * (j.powi(2) - 1.0) - 2.0 * mu)
                        * 0.5
                        * (c_inv_i1j2 * c_inv_i2j1 + c_inv_i1j1 * c_inv_i2j2)
                    + 2.0 * mu * delta_ij;

                // Apply Voigt factor for shear components
                if i >= 3 || k >= 3 {
                    tangent[(i, k)] *= if i >= 3 && k >= 3 { 4.0 } else { 2.0 };
                }
            }
        }

        Ok(tangent)
    }
}

/// Mooney-Rivlin hyperelastic material.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MooneyRivlin {
    properties: MaterialProperties,
    /// Material parameter C10
    c10: f64,
    /// Material parameter C01
    c01: f64,
    /// Bulk modulus parameter
    d1: f64,
}

impl MooneyRivlin {
    /// Create a new Mooney-Rivlin material.
    pub fn new(c10: f64, c01: f64, d1: f64) -> Self {
        // Approximate elastic modulus
        let elastic_modulus = 6.0 * (c10 + c01);
        let poisson_ratio = 0.495; // Nearly incompressible

        let properties =
            MaterialProperties::isotropic_elastic(elastic_modulus, poisson_ratio, 1000.0);

        Self {
            properties,
            c10,
            c01,
            d1,
        }
    }

    /// Set density.
    pub fn with_density(mut self, density: f64) -> Self {
        self.properties.density = density;
        self
    }
}

impl Material for MooneyRivlin {
    fn properties(&self) -> &MaterialProperties {
        &self.properties
    }

    fn compute_response(
        &self,
        strain_increment: &Vector6<f64>,
        current_state: &MaterialState,
        _dt: f64,
    ) -> FeaResult<MaterialResponse> {
        // Simplified implementation using small strain approximation
        let mut new_state = current_state.clone();
        new_state.total_strain += strain_increment;

        let elastic_modulus = self.properties.elastic_modulus;
        let poisson_ratio = self.properties.poisson_ratio;

        let factor = elastic_modulus / ((1.0 + poisson_ratio) * (1.0 - 2.0 * poisson_ratio));
        let mut d = DMatrix::zeros(6, 6);

        d[(0, 0)] = factor * (1.0 - poisson_ratio);
        d[(1, 1)] = factor * (1.0 - poisson_ratio);
        d[(2, 2)] = factor * (1.0 - poisson_ratio);
        d[(0, 1)] = factor * poisson_ratio;
        d[(1, 0)] = d[(0, 1)];
        d[(0, 2)] = factor * poisson_ratio;
        d[(2, 0)] = d[(0, 2)];
        d[(1, 2)] = factor * poisson_ratio;
        d[(2, 1)] = d[(1, 2)];
        d[(3, 3)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;
        d[(4, 4)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;
        d[(5, 5)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;

        let strain_vec = DVector::from_iterator(6, new_state.total_strain.iter().copied());
        let stress_vec = &d * &strain_vec;
        let stress = Vector6::from_iterator(stress_vec.iter().copied());
        new_state.stress = stress;

        Ok(MaterialResponse::new(stress, d, new_state))
    }

    fn elastic_tangent(&self) -> FeaResult<DMatrix<f64>> {
        let elastic_modulus = self.properties.elastic_modulus;
        let poisson_ratio = self.properties.poisson_ratio;

        let factor = elastic_modulus / ((1.0 + poisson_ratio) * (1.0 - 2.0 * poisson_ratio));
        let mut d = DMatrix::zeros(6, 6);

        d[(0, 0)] = factor * (1.0 - poisson_ratio);
        d[(1, 1)] = factor * (1.0 - poisson_ratio);
        d[(2, 2)] = factor * (1.0 - poisson_ratio);
        d[(0, 1)] = factor * poisson_ratio;
        d[(1, 0)] = d[(0, 1)];
        d[(0, 2)] = factor * poisson_ratio;
        d[(2, 0)] = d[(0, 2)];
        d[(1, 2)] = factor * poisson_ratio;
        d[(2, 1)] = d[(1, 2)];
        d[(3, 3)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;
        d[(4, 4)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;
        d[(5, 5)] = factor * (1.0 - 2.0 * poisson_ratio) / 2.0;

        Ok(d)
    }

    fn material_type(&self) -> &'static str {
        "MooneyRivlin"
    }
}

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

    #[test]
    fn test_neo_hookean_creation() {
        let material = NeoHookean::new(1e6, 0.49);
        assert_eq!(material.material_type(), "NeoHookean");
        assert!(!material.is_linear());
    }

    #[test]
    fn test_mooney_rivlin_creation() {
        let material = MooneyRivlin::new(80e3, 20e3, 0.0);
        assert_eq!(material.material_type(), "MooneyRivlin");
    }
}