conspire 0.7.7

use crate::math::{Erase, Quantity};
use crate::units::{Dimensionless, Stress, UnitDiv};
use std::ops::{Div, Mul};

use crate::{
    constitutive::{ConstitutiveError, solid::elastic::internal_variables::ElasticIV},
    fem::block::element::{
        Element, ElementNodalCoordinates, FiniteElement, FiniteElementError, GradientVectors,
        IntegrationWeights,
        solid::{ElementNodalForcesSolid, ElementNodalStiffnessesSolid, SolidFiniteElement},
    },
    math::{
        ContractSecondFourthWithFirst, HessianBlock, Jacobian, Matrix, Scalar, Solution,
        SquareMatrix, Tensor, TensorList, Vector,
        optimize::{EqualityConstraint, FirstOrderRootFinding, NewtonRaphson},
    },
    mechanics::{
        DeformationGradient, FirstPiolaKirchhoffStress, FirstPiolaKirchhoffStressList,
        FirstPiolaKirchhoffTangentStiffness, FirstPiolaKirchhoffTangentStiffnessList,
    },
    units::Volume,
};

/// The indices of the internal variables that are free to move.
fn free_indices<C, V>(constitutive_model: &C, size: usize) -> Vec<usize>
where
    C: ElasticIV<V>,
{
    let mut free = vec![true; size];
    constitutive_model
        .internal_variables_fixed()
        .iter()
        .for_each(|&index| free[index] = false);
    (0..size).filter(|&index| free[index]).collect()
}

/// The internal variables held at each integration point of an element.
pub type InternalVariables<const G: usize, V> = TensorList<V, G>;

pub trait ElasticIVFiniteElement<
    C,
    const G: usize,
    const M: usize,
    const N: usize,
    const P: usize,
    V,
    E,
> where
    C: ElasticIV<V>,
    C::Residual: Erase<Erased = E>,
    Self: SolidFiniteElement<G, M, N, P>,
    V: Erase<Erased = E> + Jacobian + Solution,
    <V as Tensor>::Unit: UnitDiv<<V as Tensor>::Unit, Output = Dimensionless>,
    E: Tensor,
    for<'a> &'a C::Residual: Div<C::TangentVv, Output = V>,
    for<'a> &'a V: Mul<Quantity<Dimensionless>, Output = V> + Mul<Scalar, Output = V>,
    for<'a> &'a Matrix: Mul<&'a V, Output = Vector>,
{
    /// The internal variables an element starts from at every integration point.
    fn internal_variables_initial(&self, constitutive_model: &C) -> InternalVariables<G, V>;
    /// Solves the internal variables at every integration point, holding the
    /// deformation fixed.
    ///
    /// The integration points are independent of one another, the internal
    /// variables of one never entering the residual of another.
    fn internal_variables_root(
        &self,
        local_solver: &NewtonRaphson,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<InternalVariables<G, V>, FiniteElementError>;
    /// Steps the internal variables alongside a decrement of the nodal
    /// coordinates, rather than solving them where the coordinates now are.
    fn internal_variables_increment(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
        nodal_decrement: &ElementNodalCoordinates<N>,
        step: Scalar,
    ) -> Result<InternalVariables<G, V>, FiniteElementError>;
    fn nodal_forces(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalForcesSolid<N>, FiniteElementError>;
    /// The nodal forces with the residual of the internal variables eliminated
    /// into them, for when they are carried rather than solved.
    fn nodal_forces_eliminated(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalForcesSolid<N>, FiniteElementError>;
    /// The tangent stiffnesses with the internal variables condensed out.
    ///
    /// ```math
    /// \mathcal{C} = \mathcal{K}_{uu} - \mathcal{K}_{uv}\mathcal{K}_{vv}^{-1}\mathcal{K}_{vu}
    /// ```
    fn nodal_stiffnesses(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalStiffnessesSolid<N>, FiniteElementError>;
}

/// Solves the internal variables at one integration point.
///
/// The gauge freedom is fixed at the initial values, which are zero over those
/// indices, so the constraint is imposed by leaving them out of the system
/// rather than by a multiplier.
fn root_at_point<C, V, E>(
    local_solver: &NewtonRaphson,
    constitutive_model: &C,
    deformation_gradient: &DeformationGradient,
    internal_variables: &V,
) -> Result<V, ConstitutiveError>
where
    C: ElasticIV<V>,
    C::Residual: Erase<Erased = E>,
    V: Erase<Erased = E> + Jacobian + Solution,
    <V as Tensor>::Unit: UnitDiv<<V as Tensor>::Unit, Output = Dimensionless>,
    E: Tensor,
    for<'a> &'a C::Residual: Div<C::TangentVv, Output = V>,
    for<'a> &'a V: Mul<Quantity<Dimensionless>, Output = V> + Mul<Scalar, Output = V>,
    for<'a> &'a Matrix: Mul<&'a V, Output = Vector>,
{
    local_solver
        .root(
            |root: &V| {
                Ok(constitutive_model.internal_variables_residual(deformation_gradient, root)?)
            },
            |root: &V| Ok(constitutive_model.tangents(deformation_gradient, root)?.3),
            internal_variables.clone(),
            EqualityConstraint::Fixed(constitutive_model.internal_variables_fixed().to_vec()),
            None,
        )
        .map_err(|error| ConstitutiveError::custom(error, deformation_gradient))
}

/// The nodal forces a list of stresses integrates to.
fn assemble_forces<const G: usize, const N: usize>(
    stresses: FirstPiolaKirchhoffStressList<G>,
    gradient_vectors: &GradientVectors<3, G, N>,
    integration_weights: &IntegrationWeights<G, Volume>,
) -> ElementNodalForcesSolid<N> {
    stresses
        .iter()
        .zip(gradient_vectors.iter().zip(integration_weights))
        .map(|(stress, (gradient_vectors_point, integration_weight))| {
            gradient_vectors_point
                .iter()
                .map(|gradient_vector| (stress * gradient_vector) * integration_weight)
                .collect()
        })
        .sum()
}

/// The residual of the internal variables at one integration point, laid flat.
fn local_residual<C, V>(
    constitutive_model: &C,
    deformation_gradient: &DeformationGradient,
    internal_variables: &V,
) -> Result<Vector, ConstitutiveError>
where
    C: ElasticIV<V>,
    V: Jacobian + Solution,
{
    let mut residual = Vector::zero(internal_variables.size());
    constitutive_model
        .internal_variables_residual(deformation_gradient, internal_variables)?
        .fill_into(&mut residual);
    Ok(residual)
}

/// The local block over the free internal variables, the fixed ones leaving it
/// through their rows and columns rather than through the inverse.
fn local_block<K>(tangent_vv: &K, size: usize, unmap: &[usize]) -> SquareMatrix
where
    K: HessianBlock,
{
    let mut block = SquareMatrix::zero(size);
    tangent_vv.fill_into_block(&mut block, 0, 0);
    let mut local = SquareMatrix::zero(unmap.len());
    unmap.iter().enumerate().for_each(|(a, &i)| {
        unmap
            .iter()
            .enumerate()
            .for_each(|(b, &j)| local[a][b] = block[i][j])
    });
    local
}

/// The stress at one integration point with the residual of the internal
/// variables eliminated into it.
///
/// ```math
/// \mathbf{P} - \mathcal{K}_{uv}\mathcal{K}_{vv}^{-1}\mathbf{r}_v
/// ```
///
/// The internal variables are carried rather than solved, so this residual
/// only agrees with the stress alone once they have converged.
fn eliminated_at_point<C, V>(
    constitutive_model: &C,
    deformation_gradient: &DeformationGradient,
    internal_variables: &V,
) -> Result<FirstPiolaKirchhoffStress, ConstitutiveError>
where
    C: ElasticIV<V>,
    V: Jacobian + Solution,
{
    let mut stress = constitutive_model
        .first_piola_kirchhoff_stress(deformation_gradient, internal_variables)?;
    let size = internal_variables.size();
    let unmap = free_indices(constitutive_model, size);
    let residual = local_residual(constitutive_model, deformation_gradient, internal_variables)?;
    let mut reduced = Vector::zero(unmap.len());
    unmap
        .iter()
        .enumerate()
        .for_each(|(a, &i)| reduced[a] = residual[i]);
    let (_, _, tangent_uv, tangent_vv) =
        constitutive_model.tangents(deformation_gradient, internal_variables)?;
    let eliminated = local_block(&tangent_vv, size, &unmap)
        .solve_lu(&reduced)
        .map_err(|error| ConstitutiveError::custom(error, deformation_gradient))?;
    let mut cross = SquareMatrix::zero(size);
    tangent_uv.fill_into_block(&mut cross, 0, 0);
    (0..3).for_each(|i| {
        (0..3).for_each(|j| {
            stress[i][j] -= unmap
                .iter()
                .enumerate()
                .map(|(a, &v)| Quantity::new(cross[3 * i + j][v] * eliminated[a]))
                .sum::<Quantity<Stress>>()
        })
    });
    Ok(stress)
}

/// The internal variables at one integration point stepped alongside a
/// decrement of the deformation.
///
/// ```math
/// \Delta\mathbf{v} = -\mathcal{K}_{vv}^{-1}\left(\mathbf{r}_v + \mathcal{K}_{vu}\Delta\mathbf{u}\right)
/// ```
fn increment_at_point<C, V>(
    constitutive_model: &C,
    deformation_gradient: &DeformationGradient,
    deformation_gradient_decrement: &DeformationGradient,
    internal_variables: &V,
    step: Scalar,
) -> Result<V, ConstitutiveError>
where
    C: ElasticIV<V>,
    V: Jacobian + Solution,
{
    let size = internal_variables.size();
    let unmap = free_indices(constitutive_model, size);
    let residual = local_residual(constitutive_model, deformation_gradient, internal_variables)?;
    let (_, tangent_vu, _, tangent_vv) =
        constitutive_model.tangents(deformation_gradient, internal_variables)?;
    let mut coupling = SquareMatrix::zero(size);
    tangent_vu.fill_into_block(&mut coupling, 0, 0);
    //
    // The deformation is handed over as a decrement, so its contribution to
    // the local residual enters with the opposite sign.
    //
    let mut reduced = Vector::zero(unmap.len());
    unmap.iter().enumerate().for_each(|(a, &i)| {
        reduced[a] = residual[i]
            - (0..3)
                .map(|k| {
                    (0..3)
                        .map(|l| coupling[i][3 * k + l] * deformation_gradient_decrement[k][l])
                        .sum::<Quantity>()
                })
                .sum::<Quantity>()
                .value()
    });
    let solution = local_block(&tangent_vv, size, &unmap)
        .solve_lu(&reduced)
        .map_err(|error| ConstitutiveError::custom(error, deformation_gradient))?;
    let mut decrement = Vector::zero(size);
    unmap
        .iter()
        .enumerate()
        .for_each(|(a, &i)| decrement[i] = solution[a] * step);
    //
    // Eliminating solves one direction for these and the nodal coordinates
    // together, so a shortened step is the same fraction of both. Solving
    // against an already shortened decrement would be a different direction
    // instead of less of this one, and would leave these where the whole of
    // their own residual had been taken out even where nothing else moved.
    //
    let mut incremented = internal_variables.clone();
    incremented.decrement_from(&decrement);
    Ok(incremented)
}

/// The tangent stiffness at one integration point with the internal variables
/// condensed out.
fn condensed_at_point<C, V>(
    constitutive_model: &C,
    deformation_gradient: &DeformationGradient,
    internal_variables: &V,
) -> Result<FirstPiolaKirchhoffTangentStiffness, ConstitutiveError>
where
    C: ElasticIV<V>,
    V: Tensor,
{
    let (tangent_uu, tangent_vu, tangent_uv, tangent_vv) =
        constitutive_model.tangents(deformation_gradient, internal_variables)?;
    let size = internal_variables.size();
    let unmap = free_indices(constitutive_model, size);
    let factorization = local_block(&tangent_vv, size, &unmap)
        .factorize_lu()
        .map_err(|error| ConstitutiveError::custom(error, deformation_gradient))?;
    let mut coupling = SquareMatrix::zero(size);
    tangent_vu.fill_into_block(&mut coupling, 0, 0);
    let mut cross = SquareMatrix::zero(size);
    tangent_uv.fill_into_block(&mut cross, 0, 0);
    let mut column = Vector::zero(unmap.len());
    let mut eliminated = vec![Vector::zero(unmap.len()); size];
    (0..size).for_each(|c| {
        unmap
            .iter()
            .enumerate()
            .for_each(|(a, &i)| column[a] = coupling[i][c]);
        factorization.solve_into(&column, &mut eliminated[c])
    });
    let mut condensed = tangent_uu;
    (0..3).for_each(|i| {
        (0..3).for_each(|j| {
            (0..3).for_each(|k| {
                (0..3).for_each(|l| {
                    condensed[i][j][k][l] -= unmap
                        .iter()
                        .enumerate()
                        .map(|(a, &v)| {
                            Quantity::new(cross[3 * i + j][v] * eliminated[3 * k + l][a])
                        })
                        .sum::<Quantity<Stress>>()
                })
            })
        })
    });
    Ok(condensed)
}

impl<C, const G: usize, const N: usize, const O: usize, const P: usize, V, E>
    ElasticIVFiniteElement<C, G, 3, N, P, V, E> for Element<3, G, N, O>
where
    C: ElasticIV<V>,
    C::Residual: Erase<Erased = E>,
    Self: SolidFiniteElement<G, 3, N, P>,
    V: Erase<Erased = E> + Jacobian + Solution,
    <V as Tensor>::Unit: UnitDiv<<V as Tensor>::Unit, Output = Dimensionless>,
    E: Tensor,
    for<'a> &'a C::Residual: Div<C::TangentVv, Output = V>,
    for<'a> &'a V: Mul<Quantity<Dimensionless>, Output = V> + Mul<Scalar, Output = V>,
    for<'a> &'a Matrix: Mul<&'a V, Output = Vector>,
{
    fn internal_variables_initial(&self, constitutive_model: &C) -> InternalVariables<G, V> {
        std::array::from_fn(|_| constitutive_model.internal_variables_initial()).into()
    }
    fn internal_variables_root(
        &self,
        local_solver: &NewtonRaphson,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<InternalVariables<G, V>, FiniteElementError> {
        self.deformation_gradients(nodal_coordinates)
            .iter()
            .zip(internal_variables)
            .map(|(deformation_gradient, internal_variables_point)| {
                root_at_point(
                    local_solver,
                    constitutive_model,
                    deformation_gradient,
                    internal_variables_point,
                )
            })
            .collect::<Result<InternalVariables<G, V>, _>>()
            .map_err(|error| FiniteElementError::upstream(error, self))
    }
    fn internal_variables_increment(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
        nodal_decrement: &ElementNodalCoordinates<N>,
        step: Scalar,
    ) -> Result<InternalVariables<G, V>, FiniteElementError> {
        self.deformation_gradients(nodal_coordinates)
            .iter()
            .zip(self.deformation_gradients(nodal_decrement).iter())
            .zip(internal_variables)
            .map(
                |((deformation_gradient, decrement), internal_variables_point)| {
                    increment_at_point(
                        constitutive_model,
                        deformation_gradient,
                        decrement,
                        internal_variables_point,
                        step,
                    )
                },
            )
            .collect::<Result<InternalVariables<G, V>, _>>()
            .map_err(|error| FiniteElementError::upstream(error, self))
    }
    fn nodal_forces(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalForcesSolid<N>, FiniteElementError> {
        let stresses = self
            .deformation_gradients(nodal_coordinates)
            .iter()
            .zip(internal_variables)
            .map(|(deformation_gradient, internal_variables_point)| {
                constitutive_model
                    .first_piola_kirchhoff_stress(deformation_gradient, internal_variables_point)
            })
            .collect::<Result<FirstPiolaKirchhoffStressList<G>, _>>()
            .map_err(|error| FiniteElementError::upstream(error, self))?;
        Ok(assemble_forces(
            stresses,
            self.gradient_vectors(),
            self.integration_weights(),
        ))
    }
    fn nodal_forces_eliminated(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalForcesSolid<N>, FiniteElementError> {
        let stresses = self
            .deformation_gradients(nodal_coordinates)
            .iter()
            .zip(internal_variables)
            .map(|(deformation_gradient, internal_variables_point)| {
                eliminated_at_point(
                    constitutive_model,
                    deformation_gradient,
                    internal_variables_point,
                )
            })
            .collect::<Result<FirstPiolaKirchhoffStressList<G>, _>>()
            .map_err(|error| FiniteElementError::upstream(error, self))?;
        Ok(assemble_forces(
            stresses,
            self.gradient_vectors(),
            self.integration_weights(),
        ))
    }
    fn nodal_stiffnesses(
        &self,
        constitutive_model: &C,
        nodal_coordinates: &ElementNodalCoordinates<N>,
        internal_variables: &InternalVariables<G, V>,
    ) -> Result<ElementNodalStiffnessesSolid<N>, FiniteElementError> {
        let condensed = self
            .deformation_gradients(nodal_coordinates)
            .iter()
            .zip(internal_variables)
            .map(|(deformation_gradient, internal_variables_point)| {
                condensed_at_point(
                    constitutive_model,
                    deformation_gradient,
                    internal_variables_point,
                )
            })
            .collect::<Result<FirstPiolaKirchhoffTangentStiffnessList<G>, _>>()
            .map_err(|error| FiniteElementError::upstream(error, self))?;
        Ok(condensed
            .iter()
            .zip(
                self.gradient_vectors()
                    .iter()
                    .zip(self.integration_weights()),
            )
            .map(|(tangent, (gradient_vectors, integration_weight))| {
                gradient_vectors
                    .iter()
                    .map(|gradient_vector_a| {
                        gradient_vectors
                            .iter()
                            .map(|gradient_vector_b| {
                                tangent.contract_second_fourth_with_first(
                                    gradient_vector_a,
                                    gradient_vector_b,
                                ) * integration_weight
                            })
                            .collect()
                    })
                    .collect()
            })
            .sum())
    }
}