qtruss 0.7.2

//! Describes a 2D truss.

use maria_linalg::{
    Matrix,
    Vector,
};

use super::{
    Element,
    Material,
};

#[derive(Clone, Copy, PartialEq, Debug)]
/// Enumerates the types of constraints available on
/// truss nodes.
pub enum Constraint {
    /// A free node
    Free (Vector<2>),
    /// A pinned node
    Pin,
    /// A node free to slide in the X direction, but not in the Y direction.
    HorizontalSlide (Vector<2>),
    /// A node free to slide in the Y direction, but not in the X direction.
    VerticalSlide (Vector<2>),
}

#[derive(Clone, Copy, PartialEq, Debug)]
/// A 2D truss node.
pub struct Node {
    pub location: Vector<2>,
    pub constraint: Constraint,
}

impl Node {
    /// Constructs an empty "zero" node.
    pub fn zero() -> Self {
        Self {
            location: Vector::zero(),
            constraint: Constraint::Free (Vector::zero()),
        }
    }

    /// Constructs a new node.
    pub fn new(
        location: Vector<2>,
        constraint: Constraint,
    ) -> Self {
        Self {
            location,
            constraint,
        }
    }
}

#[derive(Clone, Copy, Debug)]
/// A 2D truss with `N` nodes, `K` elements, and `F` degrees of freedom.
pub struct Truss<const N: usize, const K: usize, const F: usize> {
    nodes: [Node; N],
    pub elements: [Element; K],
    i: usize,
    pub forces: Option<Vector<F>>,
    pub displacements: Option<Vector<F>>,
}

impl<const N: usize, const K: usize, const F: usize> Truss<N, K, F> {
    pub fn new(nodes: [Node; N]) -> Self {
        Self {
            nodes,
            elements: [Element::zero(); K],
            i: 0,
            forces: None,
            displacements: None,
        }
    }

    /// Add an element to this truss and returns the index of this element.
    /// 
    /// Element indices begin at `0` and end at `K - 1`.
    pub fn add(&mut self, one: usize, two: usize) -> Option<usize> {
        // Quit if the node numbers are invalid
        if one >= N {
            println!("Invalid node number: {}", one);
            return None;
        }
        if two >= N {
            println!("Invalid node number: {}", two);
            return None;
        }

        let node_one = self.nodes[one];
        let node_two = self.nodes[two];
        self.elements[self.i] = Element::new(
            one,
            node_one.location,
            two,
            node_two.location,
        );

        let output = self.i;
        self.i += 1;

        Some (output)
    }

    /// Assemble the global stiffness matrix for this truss.
    /// 
    /// The global stiffness matrix has dimensions `(F, F)`.
    pub fn global_stiffness_matrix(&self, areas: [f64; K], materials: [Material; K]) -> Option<Matrix<F>> {
        // If we don't have enough elements, quit
        if self.i < K {
            return None;
        }

        // Create a new `F` by `F` matrix
        let mut matrix = Matrix::zero();

        // Assemble a list of free degrees (up to two per node)
        // `degrees[i]` represents the global indices (if they exist) at node `i`
        let mut degrees: [(Option<usize>, Option<usize>); N] = [(None, None); N];
        let mut f = 0;
        for (n, node) in self.nodes.iter().enumerate() {
            // Create an X degree of freedom, if necessary
            if matches!(node.constraint, Constraint::Free (_))
                || matches!(node.constraint, Constraint::HorizontalSlide (_))
            {
                degrees[n].0 = Some (f);
                f += 1;
            }

            // Create a Y degree of freedom, if necessary
            if matches!(node.constraint, Constraint::Free (_))
                || matches!(node.constraint, Constraint::VerticalSlide (_))
            {
                degrees[n].1 = Some (f);
                f += 1;
            }
        }

        for k in 0..K {
            let element = self.elements[k];
            let area = areas[k];
            let e = materials[k].e;

            let attached: [Option<usize>; 4] = [
                degrees[element.one].0,
                degrees[element.one].1,
                degrees[element.two].0,
                degrees[element.two].1,
            ];

            for (i, a0) in attached.iter().enumerate() {
                for (j, a1) in attached.iter().enumerate() {
                    if let Some (arow) = a0 {
                        if let Some (acol) = a1 {
                            matrix[(*arow, *acol)] += area * e * element.stiffness[(i, j)];
                        }
                    }
                }
            }
        }

        Some (matrix)
    }
    
    /// Solves this truss.
    pub fn solve(&self, areas: [f64; K], materials: [Material; K]) -> Option<(Vector<F>, Vector<F>)> {
        // Get forces applied
        let mut f = Vector::zero();
        let mut i = 0;

        for node in self.nodes {
            match node.constraint {
                Constraint::Pin => (),
                Constraint::Free (applied) => {
                    f[i] = applied[0];
                    f[i + 1] = applied[1];
                    i += 2;
                }
                Constraint::HorizontalSlide (applied) => {
                    f[i] = applied[0];
                    i += 1;
                }
                Constraint::VerticalSlide (applied) => {
                    f[i] = applied[0];
                    i += 1;
                }
            }
        }

        // Get global stiffness matrix
        let k = match self.global_stiffness_matrix(areas, materials) {
            Some (mx) => mx,
            None => {
                println!("Global stiffness matrix is singular.");
                return None;
            },
        };

        // Compute displacements
        let displacements = k.inverse().mult(f);

        Some ((f, displacements))
    }

    /// Gets the displacements of a provided node, given its index.
    pub fn displacement(&self, areas: [f64; K], materials: [Material; K], node: usize) -> Option<Vector<2>> {
        // Get global displacements
        let displacements = match self.solve(areas, materials) {
            Some ((_, d)) => d,
            None => {
                println!("Could not compute displacement of node `{}`.", node);
                return None;
            },
        };

        // Start a counter
        let mut i = 0;

        // Iterate through the nodes, up to the current node index
        for k in 0..node {
            // Get the current node
            let n = self.nodes[k];

            i += match n.constraint {
                Constraint::Pin => 0,
                Constraint::Free (_) => 2,
                Constraint::HorizontalSlide (_) => 1,
                Constraint::VerticalSlide (_) => 1,
            };
        }

        let mut output: Vector<2> = Vector::zero();

        match self.nodes[node].constraint {
            // A pin doesn't move
            Constraint::Pin => return Some (output),
            
            // A free node can move in either dimension
            Constraint::Free (_) => {
                output[0] = displacements[i];
                output[1] = displacements[i + 1];
            },

            // A slider can move in only one dimension
            Constraint::HorizontalSlide (_) => output[0] = displacements[i],
            Constraint::VerticalSlide (_) => output[1] = displacements[i],
        }

        Some (output)
    }

    /// Computes the internal force in a provided element, given its index.
    pub fn internal_force(&self, areas: [f64; K], materials: [Material; K], elt: usize) -> Option<f64> {
        // Get this element
        let element = self.elements[elt];

        // Get this element's displacements
        let mut u: Vector<4> = Vector::zero();

        let left_constraint = self.nodes[element.one].constraint;
        let right_constraint = self.nodes[element.two].constraint;

        let left_displacement = match self.displacement(areas, materials, element.one) {
            Some (d) => d,
            None => {
                println!("Could not compute internal force of element `{}`.", elt);
                return None;
            },
        };
        match left_constraint {
            Constraint::Pin => {
                u[0] = 0.0;
                u[1] = 0.0;
            },
            Constraint::Free (_) => {
                u[0] = left_displacement[0];
                u[1] = left_displacement[1];
            },
            Constraint::HorizontalSlide (_) => {
                u[0] = left_displacement[0];
                u[1] = 0.0;
            },
            Constraint::VerticalSlide (_) => {
                u[0] = 0.0;
                u[1] = left_displacement[1];
            },
        }

        let right_displacement = match self.displacement(areas, materials, element.two) {
            Some (d) => d,
            None => {
                println!("Could not compute internal force of element `{}`.", elt);
                return None;
            },
        };
        match right_constraint {
            Constraint::Pin => {
                u[2] = 0.0;
                u[3] = 0.0;
            },
            Constraint::Free (_) => {
                u[2] = right_displacement[0];
                u[3] = right_displacement[1];
            },
            Constraint::HorizontalSlide (_) => {
                u[2] = right_displacement[0];
                u[3] = 0.0;
            },
            Constraint::VerticalSlide (_) => {
                u[2] = 0.0;
                u[3] = right_displacement[1];
            },
        }
        
        // Get this element's stiffness matrix
        let k = element.stiffness;

        // Get the reaction forces on this element
        let force = k.mult(u);

        // Get the internal force in this element
        // 
        // Note: for reasons of convention, we always take the last two
        // values of the reaction forces vector
        let internal = Vector::new([
            force[2],
            force[3],
        ]);

        // Get the direction of this element
        let direction = element.direction;

        // Return the internal force of this element
        Some (internal.dot(direction))
    }

    /// Computes the internal stress in a provided element, in Pa.
    pub fn stress(&self, areas: [f64; K], materials: [Material; K], elt: usize) -> Option<f64> {
        let force = match self.internal_force(areas, materials, elt) {
            Some (d) => d,
            None => {
                println!("Could not compute internal stress of element `{}`.", elt);
                return None;
            },
        };
        let area = areas[elt];
        Some (force / area)
    }

    /// Computes the compliance of this truss.
    /// 
    /// *Note*: you must solve this truss by calling `Truss::solve()` before calling this function.
    pub fn compliance(&self, areas: [f64; K], materials: [Material; K]) -> Option<f64> {
        let (forces, displacements) = match self.solve(areas, materials) {
            Some ((f, d)) => (f, d),
            None => {
                println!("Could not compute truss compliance.");
                return None;
            },
        };

        Some (forces.dot(displacements))
    }

    /// Checks if truss materials can hold stress applied.
    pub fn check(&self, areas: [f64; K], materials: [Material; K]) -> bool {
        // Check internal stress in each element
        for k in 0..K {
            let material = materials[k];

            let stress = match self.stress(areas, materials, k) {
                Some (s) => s,
                None => return false,
            };

            // Does the element fail?
            if let Some (s) = material.minstress {
                if stress < s {
                    println!("Element {} fails in compression", k);
                    return false;
                }
            } else if let Some (s) = material.maxstress {
                if stress > s {
                    println!("Element {} fails in tension", k);
                    return false;
                }
            }
        }

        true
    }

    /// Computes the total volume of material contained in this truss.
    pub fn volume(&self, areas: [f64; K]) -> f64 {
        let mut volume = 0.0;

        for k in 0..K {
            let element = self.elements[k];
            let area = areas[k];
            volume += area * element.length;
        }

        volume
    }
}