qtruss 0.9.2

//! Describes a 2D truss.

use raqote::*;

use maria_linalg::{
    Matrix,
    Vector,
};

use super::{
    Element,
    Material,
    PlottableElement,
};

/// Width of canvas (in px) for plotting.
const CANVAS_WIDTH: i32 = 1024;

/// Height of canvas (in px) for plotting.
const CANVAS_HEIGHT: i32 = 768;

#[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.
    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.
    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 => 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 {
                    return false;
                }
            } else if let Some (s) = material.maxstress {
                if stress > s {
                    return false;
                }
            }
        }

        true
    }

    /// Computes fabrication complexity.
    pub fn complexity(&self, areas: [f64; K], maxarea: f64, beta: f64) -> f64 {
        let mut complexity = 0.0;

        for k in 0..K {
            let area = areas[k];

            if area < maxarea {
                complexity += 1.0 - (-beta * area).exp() + area * (-beta).exp();
            } else {
                complexity += 1.0;
            }
        }

        complexity
    }

    /// 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
    }

    /// Computes the *magnitude* of the maximum internal stress and the maximum area in this truss.
    /// 
    /// *Note*: the maximum stress does not necessarily correspond to the maximum area.  This function
    /// enables truss plotting.
    fn maximum(&self, areas: [f64; K], materials: [Material; K]) -> (f64, f64) {
        let mut maxstress: f64 = 0.0;
        let mut maxarea: f64 = 0.0;

        for k in 0..K {
            if let Some (s) = self.stress(areas, materials, k) {
                if s.abs() > maxstress.abs() {
                    maxstress = s;
                }
            }

            if areas[k] > maxarea {
                maxarea = areas[k];
            }
        }

        (maxstress, maxarea)
    }

    /// Returns the top-left corner and dimensions of the rectangular bounding box around the truss.
    fn bounding_box(&self) -> (f64, f64, f64, f64) {
        let mut topleft = Vector::zero();
        let mut bottomright = Vector::zero();

        let score = |vec: Vector<2>| {
            -vec[0] + vec[1]
        };

        for k in 0..K {
            let one = self.nodes[self.elements[k].one].location;
            let two = self.nodes[self.elements[k].two].location;

            if score(one) < score(topleft) {
                topleft = one;
            } else if score(two) < score(topleft) {
                topleft = two;
            }

            if score(one) > score(bottomright) {
                bottomright = one;
            } else if score(two) > score(bottomright) {
                bottomright = two;
            }
        }

        println!("Bottom right: {}, {}", bottomright[0], bottomright[1],);
        println!("Top left: {}, {}", topleft[0], topleft[1],);

        (
            topleft[0],
            topleft[1],
            bottomright[0] - topleft[0],
            bottomright[1] - topleft[1],
        )
    }

    /// Returns a list of plottable elements.
    fn make_plottable_elements(&self, areas: [f64; K], materials: [Material; K]) -> [PlottableElement; K] {
        let mut plottable = [PlottableElement::zero(); K];

        let (top, left, width, height) = self.bounding_box();
        let (maxstress, maxarea) = self.maximum(areas, materials);

        for k in 0..K {
            let stress = self.stress(areas, materials, k).unwrap_or(0.0);
            let area = areas[k];

            plottable[k] = self.elements[k].to_plottable(
                top,
                left,
                width,
                height,
                stress,
                maxstress,
                area,
                maxarea,
            );
        }

        plottable
    }

    /// Plots this truss.
    pub fn plot(&self, areas: [f64; K], materials: [Material; K], filename: &str) -> Option<()> {
        let elements = self.make_plottable_elements(areas, materials);
    
        // Construct canvas
        let mut dt = DrawTarget::new(CANVAS_WIDTH, CANVAS_HEIGHT);

        // Construct white background
        dt.fill_rect(
            0.0,
            0.0,
            CANVAS_WIDTH as f32,
            CANVAS_HEIGHT as f32,
            &Source::Solid(SolidSource {
                r: 0xff,
                g: 0xff,
                b: 0xff,
                a: 0xff,
            }),
            &DrawOptions::new(),
        );

        for element in elements {
            let mut pb = PathBuilder::new();
            pb.move_to(CANVAS_WIDTH as f32 * element.x1, CANVAS_HEIGHT as f32 * element.y1);
            pb.line_to(CANVAS_WIDTH as f32 * element.x2, CANVAS_HEIGHT as f32 * element.y2);
            let path = pb.finish();

            let (r, g, b, a) = element.rgba();

            println!("{}, {}", CANVAS_WIDTH as f32 * element.x1, CANVAS_HEIGHT as f32 * element.y1);
            println!("{}. {}", CANVAS_WIDTH as f32 * element.x2, CANVAS_HEIGHT as f32 * element.y2);

            println!("Color: {}, {}, {}, {}", r, g, b, a);
    
            dt.stroke(
                &path,
                &Source::Solid(SolidSource {
                    r,
                    g,
                    b,
                    a,
                }),
                &StrokeStyle {
                    cap: LineCap::Square,
                    join: LineJoin::Bevel,
                    width: 10.,
                    miter_limit: 2.,
                    dash_array: vec![10.,],
                    dash_offset: 0.,
                },
                &DrawOptions::new(),
            );
        }
    
        dt.write_png(filename).ok()
    }
}