cfsem 11.1.0

use crate::mesh::{QuadMeshView2d, QuadratureRule};
use crate::physics::solenoid_stress::family::QuadElementFamily;
use crate::physics::solenoid_stress::geometry::{FaceSample, validate_structural_2d_mesh};
use crate::physics::solenoid_stress::types::{
    DOF_PER_NODE, Structural2dFormulation, TractionLoad, local_dofs, scatter_local_matrix,
};

use super::{SparseOperator, collect_sparse_operator_chunks, concat_sparse_operators};

/// Build the local dense traction operator for one loaded face.
///
/// The returned block has shape `(DOF_PER_ELEMENT, 2)`. Its two columns correspond to:
/// - column `0`: unit radial traction `[force / area]`
/// - column `1`: unit axial traction `[force / area]`
///
/// Multiplying this block by `[t_r, t_z]^T` gives the element's consistent nodal load vector
/// `[energy / distance]`.
///
/// Each block entry therefore has units of area.
fn traction_face_kernel<const NODES_PER_ELEMENT: usize, const DOF_PER_ELEMENT: usize>(
    samples: &[FaceSample<f64, NODES_PER_ELEMENT>],
    formulation: Structural2dFormulation,
) -> Result<[[f64; 2]; DOF_PER_ELEMENT], String> {
    const {
        assert!(DOF_PER_ELEMENT == DOF_PER_NODE * NODES_PER_ELEMENT);
    }
    let mut local = [[0.0; 2]; DOF_PER_ELEMENT];

    for sample in samples {
        // `|dx/ds|` is the physical line Jacobian for the face quadrature parameter.
        let tangent_norm =
            (sample.tangent[0] * sample.tangent[0] + sample.tangent[1] * sample.tangent[1]).sqrt();
        let scale = formulation.face_scale(sample.point, tangent_norm, sample.weight)?;
        for local_node in 0..NODES_PER_ELEMENT {
            // The two columns encode independent unit tractions in the global radial and axial
            // directions, so the block is diagonal in those two traction components.
            local[2 * local_node][0] += scale * sample.n[local_node];
            local[2 * local_node + 1][1] += scale * sample.n[local_node];
        }
    }

    Ok(local)
}

/// Assemble the global traction-to-RHS operator for one quadrilateral family.
///
/// Output shape: `(2 * mesh.num_nodes(), 2 * traction_faces.len())`.
///
/// Row meaning:
/// - row `2*a` is the radial generalized-force equation for node `a`,
/// - row `2*a + 1` is the axial generalized-force equation for node `a`.
///
/// Column meaning:
/// - column `2*j` is unit radial traction on `traction_faces[j]`,
/// - column `2*j + 1` is unit axial traction on `traction_faces[j]`.
///
/// Entry units: `[area]`.
pub(crate) fn traction_operator_for_family<
    Family,
    const NODES_PER_ELEMENT: usize,
    const DOF_PER_ELEMENT: usize,
>(
    mesh: QuadMeshView2d<'_, f64, NODES_PER_ELEMENT>,
    traction_faces: &[TractionLoad],
    formulation: Structural2dFormulation,
    quadrature: QuadratureRule,
) -> Result<SparseOperator, String>
where
    Family: QuadElementFamily<NODES_PER_ELEMENT>,
{
    const {
        assert!(DOF_PER_ELEMENT == DOF_PER_NODE * NODES_PER_ELEMENT);
    }
    validate_structural_2d_mesh(mesh, formulation)?;
    traction_operator_range_for_family::<Family, NODES_PER_ELEMENT, DOF_PER_ELEMENT>(
        mesh,
        traction_faces,
        formulation,
        quadrature,
        0,
        traction_faces.len(),
    )
}

/// Assemble traction loads using traction-face ranges split across Rayon workers.
pub(crate) fn traction_operator_for_family_par<
    Family,
    const NODES_PER_ELEMENT: usize,
    const DOF_PER_ELEMENT: usize,
>(
    mesh: QuadMeshView2d<'_, f64, NODES_PER_ELEMENT>,
    traction_faces: &[TractionLoad],
    formulation: Structural2dFormulation,
    quadrature: QuadratureRule,
) -> Result<SparseOperator, String>
where
    Family: QuadElementFamily<NODES_PER_ELEMENT>,
{
    const {
        assert!(DOF_PER_ELEMENT == DOF_PER_NODE * NODES_PER_ELEMENT);
    }
    validate_structural_2d_mesh(mesh, formulation)?;
    let chunks = collect_sparse_operator_chunks(traction_faces.len(), |start, end| {
        traction_operator_range_for_family::<Family, NODES_PER_ELEMENT, DOF_PER_ELEMENT>(
            mesh,
            traction_faces,
            formulation,
            quadrature,
            start,
            end,
        )
    })?;

    Ok(concat_sparse_operators(
        chunks,
        mesh.num_nodes() * 2,
        2 * traction_faces.len(),
        traction_faces.len() * DOF_PER_ELEMENT * 2,
    ))
}

/// Apply traction amplitudes directly to a reduced RHS without building a sparse operator.
pub(crate) fn apply_traction_rhs_for_family<
    Family,
    const NODES_PER_ELEMENT: usize,
    const DOF_PER_ELEMENT: usize,
>(
    mesh: QuadMeshView2d<'_, f64, NODES_PER_ELEMENT>,
    traction_faces: &[TractionLoad],
    formulation: Structural2dFormulation,
    quadrature: QuadratureRule,
    values: &[f64],
    global_to_reduced: &[usize],
    rhs: &mut [f64],
) -> Result<(), String>
where
    Family: QuadElementFamily<NODES_PER_ELEMENT>,
{
    const {
        assert!(DOF_PER_ELEMENT == DOF_PER_NODE * NODES_PER_ELEMENT);
    }
    validate_structural_2d_mesh(mesh, formulation)?;
    let expected = 2 * traction_faces.len();
    if values.len() != expected {
        return Err(format!(
            "traction_values has length {}, but operator expects {expected} values",
            values.len()
        ));
    }
    for_traction_face_blocks::<Family, NODES_PER_ELEMENT, DOF_PER_ELEMENT, _>(
        mesh,
        traction_faces,
        formulation,
        quadrature,
        0,
        traction_faces.len(),
        |load_index, global_rows, local| {
            let radial = values[2 * load_index];
            let axial = values[2 * load_index + 1];
            for local_dof in 0..DOF_PER_ELEMENT {
                let reduced_row = global_to_reduced[global_rows[local_dof]];
                if reduced_row != usize::MAX {
                    rhs[reduced_row] += local[local_dof][0] * radial + local[local_dof][1] * axial;
                }
            }
            Ok(())
        },
    )
}

fn for_traction_face_blocks<
    Family,
    const NODES_PER_ELEMENT: usize,
    const DOF_PER_ELEMENT: usize,
    Consume,
>(
    mesh: QuadMeshView2d<'_, f64, NODES_PER_ELEMENT>,
    traction_faces: &[TractionLoad],
    formulation: Structural2dFormulation,
    quadrature: QuadratureRule,
    load_start: usize,
    load_end: usize,
    mut consume: Consume,
) -> Result<(), String>
where
    Family: QuadElementFamily<NODES_PER_ELEMENT>,
    Consume:
        FnMut(usize, [usize; DOF_PER_ELEMENT], [[f64; 2]; DOF_PER_ELEMENT]) -> Result<(), String>,
{
    let nloads = load_end - load_start;
    for (load_index, load) in traction_faces
        .iter()
        .enumerate()
        .skip(load_start)
        .take(nloads)
    {
        if load.element >= mesh.num_elements() {
            return Err(format!(
                "traction load references element {}, but mesh has only {} elements",
                load.element,
                mesh.num_elements()
            ));
        }
        let coords = mesh.element_coords(load.element)?;
        let nodes = mesh.element_nodes(load.element)?;
        let samples = Family::face_samples(&coords, load.local_face, quadrature)?;
        let local =
            traction_face_kernel::<NODES_PER_ELEMENT, DOF_PER_ELEMENT>(&samples, formulation)?;
        let global_rows = local_dofs::<NODES_PER_ELEMENT, DOF_PER_ELEMENT>(&nodes);
        consume(load_index, global_rows, local)?;
    }
    Ok(())
}

fn traction_operator_range_for_family<
    Family,
    const NODES_PER_ELEMENT: usize,
    const DOF_PER_ELEMENT: usize,
>(
    mesh: QuadMeshView2d<'_, f64, NODES_PER_ELEMENT>,
    traction_faces: &[TractionLoad],
    formulation: Structural2dFormulation,
    quadrature: QuadratureRule,
    load_start: usize,
    load_end: usize,
) -> Result<SparseOperator, String>
where
    Family: QuadElementFamily<NODES_PER_ELEMENT>,
{
    let ndof = mesh.num_nodes() * 2;
    let ncol = 2 * traction_faces.len();
    let nloads = load_end - load_start;
    let mut rows = Vec::with_capacity(nloads * DOF_PER_ELEMENT * 2);
    let mut cols = Vec::with_capacity(nloads * DOF_PER_ELEMENT * 2);
    let mut vals = Vec::with_capacity(nloads * DOF_PER_ELEMENT * 2);

    for_traction_face_blocks::<Family, NODES_PER_ELEMENT, DOF_PER_ELEMENT, _>(
        mesh,
        traction_faces,
        formulation,
        quadrature,
        load_start,
        load_end,
        |load_index, global_rows, local| {
            let global_cols = [2 * load_index, 2 * load_index + 1];
            scatter_local_matrix(
                &mut rows,
                &mut cols,
                &mut vals,
                &global_rows,
                &global_cols,
                &local,
            );
            Ok(())
        },
    )?;

    Ok(SparseOperator {
        rows,
        cols,
        vals,
        nrow: ndof,
        ncol,
    })
}