dem_atom 0.1.3

//! DEM atom properties: named material types, per-pair mixing tables, and per-atom radius/density.

use std::any::{Any, TypeId};
use std::f64::consts::PI;

use mddem_app::prelude::*;
use nalgebra::{Quaternion, UnitQuaternion};
use serde::Deserialize;

use mddem_core::{AtomData, AtomDataRegistry, AtomPlugin, Config};

// ── Config structs ──────────────────────────────────────────────────────────

fn default_friction() -> f64 {
    0.4
}

#[derive(Deserialize, Clone)]
#[serde(deny_unknown_fields)]
/// A single material definition from `[[dem.materials]]`.
pub struct MaterialConfig {
    /// Material name, referenced by particle insert blocks.
    pub name: String,
    /// Young's modulus (Pa).
    pub youngs_mod: f64,
    /// Poisson's ratio (dimensionless, 0–0.5).
    pub poisson_ratio: f64,
    /// Coefficient of restitution (0–1).
    pub restitution: f64,
    /// Coulomb friction coefficient.
    #[serde(default = "default_friction")]
    pub friction: f64,
}

/// TOML `[dem]` — top-level DEM configuration containing material definitions.
#[derive(Deserialize, Clone, Default)]
pub struct DemConfig {
    pub materials: Option<Vec<MaterialConfig>>,
}

// ── MaterialTable — per-material and per-pair precomputed properties ────────

/// Per-material properties and per-pair precomputed mixing tables (geometric-mean).
pub struct MaterialTable {
    pub names: Vec<String>,
    pub youngs_mod: Vec<f64>,
    pub poisson_ratio: Vec<f64>,
    pub friction: Vec<f64>,
    pub restitution: Vec<f64>,
    pub beta_ij: Vec<Vec<f64>>,
    pub friction_ij: Vec<Vec<f64>>,
    /// Precomputed effective Young's modulus for each material pair (Hertz contact).
    pub e_eff_ij: Vec<Vec<f64>>,
    /// Precomputed effective shear modulus for each material pair (Mindlin contact).
    pub g_eff_ij: Vec<Vec<f64>>,
}

impl Default for MaterialTable {
    fn default() -> Self {
        Self::new()
    }
}

impl MaterialTable {
    pub fn new() -> Self {
        MaterialTable {
            names: Vec::new(),
            youngs_mod: Vec::new(),
            poisson_ratio: Vec::new(),
            friction: Vec::new(),
            restitution: Vec::new(),
            beta_ij: Vec::new(),
            friction_ij: Vec::new(),
            e_eff_ij: Vec::new(),
            g_eff_ij: Vec::new(),
        }
    }

    pub fn add_material(
        &mut self,
        name: &str,
        youngs_mod: f64,
        poisson_ratio: f64,
        restitution: f64,
        friction: f64,
    ) -> u32 {
        let idx = self.names.len() as u32;
        self.names.push(name.to_string());
        self.youngs_mod.push(youngs_mod);
        self.poisson_ratio.push(poisson_ratio);
        self.restitution.push(restitution);
        self.friction.push(friction);
        idx
    }

    pub fn find_material(&self, name: &str) -> Option<u32> {
        self.names.iter().position(|n| n == name).map(|i| i as u32)
    }

    pub fn build_pair_tables(&mut self) {
        let n = self.names.len();
        self.beta_ij = vec![vec![0.0; n]; n];
        self.friction_ij = vec![vec![0.0; n]; n];
        self.e_eff_ij = vec![vec![0.0; n]; n];
        self.g_eff_ij = vec![vec![0.0; n]; n];
        for i in 0..n {
            for j in 0..n {
                // Geometric mean mixing for restitution
                let e_ij = (self.restitution[i] * self.restitution[j]).sqrt();
                let log_e = e_ij.ln();
                self.beta_ij[i][j] = -log_e / (PI * PI + log_e * log_e).sqrt();

                // Geometric mean mixing for friction
                self.friction_ij[i][j] = (self.friction[i] * self.friction[j]).sqrt();

                // Effective Young's modulus (Hertz)
                let nu_i = self.poisson_ratio[i];
                let nu_j = self.poisson_ratio[j];
                self.e_eff_ij[i][j] = 1.0
                    / ((1.0 - nu_i * nu_i) / self.youngs_mod[i]
                        + (1.0 - nu_j * nu_j) / self.youngs_mod[j]);

                // Effective shear modulus (Mindlin)
                self.g_eff_ij[i][j] = 1.0
                    / (2.0 * (2.0 - nu_i) * (1.0 + nu_i) / self.youngs_mod[i]
                        + 2.0 * (2.0 - nu_j) * (1.0 + nu_j) / self.youngs_mod[j]);
            }
        }
    }
}

// ── DemAtom per-atom data ────────────────────────────────────────────────────

/// Per-atom DEM extension data: particle radius, density, inverse inertia, and rotational fields.
pub struct DemAtom {
    pub radius: Vec<f64>,
    pub density: Vec<f64>,
    pub inv_inertia: Vec<f64>,
    pub quaternion: Vec<UnitQuaternion<f64>>,
    pub omega: Vec<[f64; 3]>,
    pub ang_mom: Vec<[f64; 3]>,
    pub torque: Vec<[f64; 3]>,
}

impl Default for DemAtom {
    fn default() -> Self {
        Self::new()
    }
}

impl DemAtom {
    pub fn new() -> Self {
        DemAtom {
            radius: Vec::new(),
            density: Vec::new(),
            inv_inertia: Vec::new(),
            quaternion: Vec::new(),
            omega: Vec::new(),
            ang_mom: Vec::new(),
            torque: Vec::new(),
        }
    }
}

impl AtomData for DemAtom {
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn truncate(&mut self, n: usize) {
        self.radius.truncate(n);
        self.density.truncate(n);
        self.inv_inertia.truncate(n);
        self.quaternion.truncate(n);
        self.omega.truncate(n);
        self.ang_mom.truncate(n);
        self.torque.truncate(n);
    }

    fn swap_remove(&mut self, i: usize) {
        self.radius.swap_remove(i);
        self.density.swap_remove(i);
        self.inv_inertia.swap_remove(i);
        self.quaternion.swap_remove(i);
        self.omega.swap_remove(i);
        self.ang_mom.swap_remove(i);
        self.torque.swap_remove(i);
    }

    fn pack(&self, i: usize, buf: &mut Vec<f64>) {
        buf.push(self.radius[i]);
        buf.push(self.density[i]);
        buf.push(self.inv_inertia[i]);
        let q = self.quaternion[i];
        buf.push(q.w);
        buf.push(q.i);
        buf.push(q.j);
        buf.push(q.k);
        buf.push(self.omega[i][0]);
        buf.push(self.omega[i][1]);
        buf.push(self.omega[i][2]);
        buf.push(self.ang_mom[i][0]);
        buf.push(self.ang_mom[i][1]);
        buf.push(self.ang_mom[i][2]);
        buf.push(self.torque[i][0]);
        buf.push(self.torque[i][1]);
        buf.push(self.torque[i][2]);
    }

    fn unpack(&mut self, buf: &[f64]) -> usize {
        self.radius.push(buf[0]);
        self.density.push(buf[1]);
        self.inv_inertia.push(buf[2]);
        self.quaternion.push(UnitQuaternion::from_quaternion(
            Quaternion::new(buf[3], buf[4], buf[5], buf[6]),
        ));
        self.omega.push([buf[7], buf[8], buf[9]]);
        self.ang_mom.push([buf[10], buf[11], buf[12]]);
        self.torque.push([buf[13], buf[14], buf[15]]);
        16
    }

    fn apply_permutation(&mut self, perm: &[usize], n: usize) {
        {
            let scratch: Vec<f64> = perm.iter().map(|&p| self.radius[p]).collect();
            self.radius[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<f64> = perm.iter().map(|&p| self.density[p]).collect();
            self.density[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<f64> = perm.iter().map(|&p| self.inv_inertia[p]).collect();
            self.inv_inertia[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<UnitQuaternion<f64>> = perm.iter().map(|&p| self.quaternion[p]).collect();
            self.quaternion[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<[f64; 3]> = perm.iter().map(|&p| self.omega[p]).collect();
            self.omega[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<[f64; 3]> = perm.iter().map(|&p| self.ang_mom[p]).collect();
            self.ang_mom[..n].copy_from_slice(&scratch);
        }
        {
            let scratch: Vec<[f64; 3]> = perm.iter().map(|&p| self.torque[p]).collect();
            self.torque[..n].copy_from_slice(&scratch);
        }
    }

    fn pack_forward(&self, i: usize, buf: &mut Vec<f64>) {
        buf.push(self.omega[i][0]);
        buf.push(self.omega[i][1]);
        buf.push(self.omega[i][2]);
    }

    fn unpack_forward(&mut self, i: usize, buf: &[f64]) -> usize {
        self.omega[i] = [buf[0], buf[1], buf[2]];
        3
    }

    fn pack_reverse(&self, i: usize, buf: &mut Vec<f64>) {
        buf.push(self.torque[i][0]);
        buf.push(self.torque[i][1]);
        buf.push(self.torque[i][2]);
    }

    fn unpack_reverse(&mut self, i: usize, buf: &[f64]) -> usize {
        self.torque[i][0] += buf[0];
        self.torque[i][1] += buf[1];
        self.torque[i][2] += buf[2];
        3
    }

    fn zero(&mut self, n: usize) {
        self.torque[..n].fill([0.0; 3]);
    }

    fn forward_comm_size(&self) -> usize { 3 } // omega
    fn reverse_comm_size(&self) -> usize { 3 } // torque
}

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

    #[test]
    fn single_material_beta_and_friction() {
        let mut mt = MaterialTable::new();
        mt.add_material("glass", 8.7e9, 0.3, 0.95, 0.4);
        mt.build_pair_tables();

        let e = 0.95_f64;
        let log_e = e.ln();
        let expected_beta = -log_e / (PI * PI + log_e * log_e).sqrt();
        assert!(
            (mt.beta_ij[0][0] - expected_beta).abs() < 1e-12,
            "beta should be {}, got {}",
            expected_beta,
            mt.beta_ij[0][0]
        );
        assert!(
            (mt.friction_ij[0][0] - 0.4).abs() < 1e-12,
            "friction should be 0.4, got {}",
            mt.friction_ij[0][0]
        );
    }

    #[test]
    fn multi_material_mixing_symmetry() {
        let mut mt = MaterialTable::new();
        mt.add_material("glass", 8.7e9, 0.3, 0.95, 0.4);
        mt.add_material("steel", 200e9, 0.28, 0.8, 0.3);
        mt.build_pair_tables();

        // Symmetry
        assert!(
            (mt.beta_ij[0][1] - mt.beta_ij[1][0]).abs() < 1e-15,
            "beta_ij should be symmetric"
        );
        assert!(
            (mt.friction_ij[0][1] - mt.friction_ij[1][0]).abs() < 1e-15,
            "friction_ij should be symmetric"
        );

        // Geometric mean mixing for friction
        let expected_friction = (0.4_f64 * 0.3).sqrt();
        assert!(
            (mt.friction_ij[0][1] - expected_friction).abs() < 1e-12,
            "friction_ij should be geometric mean {}, got {}",
            expected_friction,
            mt.friction_ij[0][1]
        );

        // Geometric mean mixing for restitution -> beta
        let e_mix = (0.95_f64 * 0.8).sqrt();
        let log_e = e_mix.ln();
        let expected_beta = -log_e / (PI * PI + log_e * log_e).sqrt();
        assert!(
            (mt.beta_ij[0][1] - expected_beta).abs() < 1e-12,
            "beta_ij should use geometric mean restitution"
        );

        // e_eff and g_eff symmetry
        assert!(
            (mt.e_eff_ij[0][1] - mt.e_eff_ij[1][0]).abs() < 1e-6,
            "e_eff_ij should be symmetric"
        );
        assert!(
            (mt.g_eff_ij[0][1] - mt.g_eff_ij[1][0]).abs() < 1e-6,
            "g_eff_ij should be symmetric"
        );
        assert!(mt.e_eff_ij[0][0] > 0.0, "e_eff should be positive");
        assert!(mt.g_eff_ij[0][0] > 0.0, "g_eff should be positive");
    }

    #[test]
    fn e_eff_matches_manual_computation() {
        let mut mt = MaterialTable::new();
        mt.add_material("glass", 8.7e9, 0.3, 0.95, 0.4);
        mt.build_pair_tables();

        let nu = 0.3_f64;
        let e = 8.7e9_f64;
        let expected = 1.0 / (2.0 * (1.0 - nu * nu) / e);
        assert!(
            (mt.e_eff_ij[0][0] - expected).abs() < 1.0,
            "e_eff_ij[0][0] should be {}, got {}",
            expected,
            mt.e_eff_ij[0][0]
        );
    }
}

// ── Plugin ───────────────────────────────────────────────────────────────────

/// Registers [`DemAtom`] extension and [`MaterialTable`] from `[[dem.materials]]` config.
pub struct DemAtomPlugin;

impl Plugin for DemAtomPlugin {
    fn default_config(&self) -> Option<&str> {
        Some(
            r#"# Material definitions for DEM particles
[[dem.materials]]
name = "glass"
youngs_mod = 8.7e9
poisson_ratio = 0.3
restitution = 0.95
friction = 0.4

# Additional materials can be added:
# [[dem.materials]]
# name = "steel"
# youngs_mod = 200e9
# poisson_ratio = 0.28
# restitution = 0.8
# friction = 0.3"#,
        )
    }

    fn build(&self, app: &mut App) {
        app.add_plugins(AtomPlugin);

        if let Some(registry_option) = app.get_mut_resource(TypeId::of::<AtomDataRegistry>()) {
            let mut registry_binder = registry_option.borrow_mut();
            let registry = registry_binder.downcast_mut::<AtomDataRegistry>().unwrap();
            registry.register(DemAtom::new());
        } else {
            panic!("AtomDataRegistry not found — AtomPlugin must be added first");
        }

        let dem_config = Config::load::<DemConfig>(app, "dem");

        // Build MaterialTable from config at plugin build time
        let mut material_table = MaterialTable::new();

        if let Some(ref materials) = dem_config.materials {
            for mat in materials {
                material_table.add_material(
                    &mat.name,
                    mat.youngs_mod,
                    mat.poisson_ratio,
                    mat.restitution,
                    mat.friction,
                );
            }
            material_table.build_pair_tables();
        }

        app.add_resource(material_table);
    }
}