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

//! Plasticity models for finite element analysis.

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

/// Von Mises plasticity with isotropic hardening.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct IsotropicPlasticity {
    properties: MaterialProperties,
    elastic_tangent: DMatrix<f64>,
    yield_stress: f64,
    hardening_modulus: f64,
}

impl IsotropicPlasticity {
    /// Create a new isotropic plasticity material.
    pub fn new(
        elastic_modulus: f64,
        poisson_ratio: f64,
        yield_stress: f64,
        hardening_modulus: f64,
    ) -> Self {
        let properties =
            MaterialProperties::isotropic_elastic(elastic_modulus, poisson_ratio, 7850.0);

        // Elastic tangent matrix
        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;

        Self {
            properties,
            elastic_tangent: d,
            yield_stress,
            hardening_modulus,
        }
    }

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

    /// Compute von Mises stress.
    fn von_mises_stress(&self, stress: &Vector6<f64>) -> f64 {
        let s11 = stress[0];
        let s22 = stress[1];
        let s33 = stress[2];
        let s12 = stress[3];
        let s13 = stress[4];
        let s23 = stress[5];

        let diff1 = s11 - s22;
        let diff2 = s22 - s33;
        let diff3 = s33 - s11;

        let vm_squared = 0.5 * (diff1 * diff1 + diff2 * diff2 + diff3 * diff3)
            + 3.0 * (s12 * s12 + s13 * s13 + s23 * s23);

        vm_squared.sqrt()
    }

    /// Check yield condition.
    fn yield_function(&self, stress: &Vector6<f64>, equivalent_plastic_strain: f64) -> f64 {
        let vm_stress = self.von_mises_stress(stress);
        let current_yield_stress =
            self.yield_stress + self.hardening_modulus * equivalent_plastic_strain;
        vm_stress - current_yield_stress
    }
}

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

    fn compute_response(
        &self,
        strain_increment: &Vector6<f64>,
        current_state: &MaterialState,
        _dt: f64,
    ) -> FeaResult<MaterialResponse> {
        let mut new_state = current_state.clone();

        // Trial elastic step
        let elastic_strain_increment = strain_increment;
        let trial_stress = current_state.stress + &self.elastic_tangent * elastic_strain_increment;

        // Check yield condition
        let yield_value =
            self.yield_function(&trial_stress, current_state.equivalent_plastic_strain);

        if yield_value <= 0.0 {
            // Elastic response
            new_state.total_strain += strain_increment;
            new_state.stress = trial_stress;

            Ok(MaterialResponse::new(
                trial_stress,
                self.elastic_tangent.clone(),
                new_state,
            ))
        } else {
            // Plastic response - return mapping
            let vm_stress = self.von_mises_stress(&trial_stress);
            let _current_yield_stress = self.yield_stress
                + self.hardening_modulus * current_state.equivalent_plastic_strain;

            // Plastic multiplier
            let shear_modulus = self.properties.shear_modulus();
            let plastic_multiplier = yield_value / (3.0 * shear_modulus + self.hardening_modulus);

            // Update plastic strain
            let plastic_strain_increment = plastic_multiplier * 1.5 * trial_stress / vm_stress;
            new_state.plastic_strain += plastic_strain_increment;
            new_state.equivalent_plastic_strain += plastic_multiplier;

            // Update stress
            let stress_correction = &self.elastic_tangent * plastic_strain_increment;
            new_state.stress = trial_stress - stress_correction;
            new_state.total_strain += strain_increment;

            // Consistent tangent (simplified)
            let mut consistent_tangent = self.elastic_tangent.clone();
            let beta = 3.0 * shear_modulus / (3.0 * shear_modulus + self.hardening_modulus);
            consistent_tangent *= beta;

            Ok(MaterialResponse::new(
                new_state.stress,
                consistent_tangent,
                new_state,
            ))
        }
    }

    fn elastic_tangent(&self) -> FeaResult<DMatrix<f64>> {
        Ok(self.elastic_tangent.clone())
    }

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

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

    #[test]
    fn test_isotropic_plasticity_creation() {
        let material = IsotropicPlasticity::new(200e9, 0.3, 250e6, 1e9);
        assert_eq!(material.material_type(), "IsotropicPlasticity");
        assert!(!material.is_linear());
    }

    #[test]
    fn test_elastic_response() {
        let material = IsotropicPlasticity::new(200e9, 0.3, 250e6, 1e9);
        let small_strain = Vector6::new(0.0001, 0.0, 0.0, 0.0, 0.0, 0.0);
        let state = MaterialState::default();

        let response = material
            .compute_response(&small_strain, &state, 1.0)
            .unwrap();
        assert!(response.is_valid);
        assert!(!response.updated_state.has_yielded());
    }
}