// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Modal analysis: natural frequencies and mode shapes.
//!
//! Solves the generalized eigenvalue problem `K φ = λ M φ` on the free
//! degrees of freedom, where `λ = ω²`.
//!
//! This previously returned `DVector::zeros(num_modes)` without assembling
//! anything. Both halves it needs already existed —
//! [`crate::assembly::GlobalAssembler`] for `K` and `M`, and
//! [`crate::solvers::eigenvalue::EigenvalueSolver`] for the eigenproblem —
//! they were simply never connected.

use super::{Analysis, AnalysisConfig, AnalysisData, AnalysisResults};
use crate::assembly::{AdvancedDofNumbering, DofMappingStrategy, GlobalAssembler};
use crate::boundary::{BoundaryCondition, BoundaryConditionSet};
use crate::error::{AnalysisError, FeaResult};
use crate::materials::MaterialDatabase;
use crate::mesh::Mesh;
use crate::solvers::eigenvalue::{EigenvalueSolver, ModalResults};
use nalgebra::{DMatrix, DVector};

/// Modal analysis for computing natural frequencies and mode shapes.
#[derive(Debug)]
pub struct ModalAnalysis {
    mesh: Mesh,
    materials: MaterialDatabase,
    num_modes: usize,
    config: AnalysisConfig,
    boundary_conditions: BoundaryConditionSet,
    shift: Option<f64>,
    progress: f64,
    complete: bool,
}

impl ModalAnalysis {
    pub fn new(
        mesh: Mesh,
        materials: MaterialDatabase,
        num_modes: usize,
        config: AnalysisConfig,
    ) -> Self {
        Self {
            mesh,
            materials,
            num_modes,
            config,
            boundary_conditions: BoundaryConditionSet::new(),
            shift: None,
            progress: 0.0,
            complete: false,
        }
    }

    /// Constrain the structure.
    ///
    /// Without this the stiffness matrix is singular: an unconstrained body
    /// has rigid-body modes at zero frequency. That is a physical fact rather
    /// than a numerical accident, so [`Self::run`] reports it as an error
    /// naming the remedy instead of returning six meaningless near-zero
    /// modes.
    pub fn with_boundary_conditions(mut self, boundary_conditions: BoundaryConditionSet) -> Self {
        self.boundary_conditions = boundary_conditions;
        self
    }

    /// Solve for the modes nearest this frequency-squared value instead of
    /// the lowest ones.
    ///
    /// The one case that genuinely needs it is a free-free structure, where a
    /// small negative shift moves the rigid-body modes off the singularity.
    pub fn with_shift(mut self, shift: f64) -> Self {
        self.shift = Some(shift);
        self
    }

    /// Number the degrees of freedom and apply displacement constraints.
    fn setup_dof_numbering(&self) -> FeaResult<AdvancedDofNumbering> {
        let mut dof_numbering = AdvancedDofNumbering::displacement_only(
            &self.mesh,
            DofMappingStrategy::BandwidthOptimized,
        )?;

        for bc in self.boundary_conditions.conditions() {
            // Only Dirichlet conditions remove degrees of freedom. A modal
            // analysis has no applied load, so Neumann conditions have
            // nothing to contribute and are ignored rather than silently
            // mis-applied.
            if let BoundaryCondition::Dirichlet(dirichlet) = bc {
                for &node_id in &dirichlet.nodes {
                    for &component in &dirichlet.components {
                        if let Some(dof) = dof_numbering.get_dof(node_id, component) {
                            dof_numbering.constrain_dof(dof)?;
                        }
                    }
                }
            }
        }

        Ok(dof_numbering)
    }
}

impl Analysis for ModalAnalysis {
    fn run(&mut self) -> FeaResult<AnalysisResults> {
        if self.num_modes == 0 {
            return Err(AnalysisError::InvalidConfiguration(
                "modal analysis requires at least one mode".to_string(),
            )
            .into());
        }
        if self.mesh.num_elements() == 0 {
            return Err(AnalysisError::InvalidConfiguration(
                "mesh contains no elements".to_string(),
            )
            .into());
        }

        let dof_numbering = self.setup_dof_numbering()?;
        self.progress = 0.2;

        let assembler =
            GlobalAssembler::new(self.config.assembly_options.clone(), self.materials.clone());
        let system =
            assembler.assemble_system(&self.mesh, &dof_numbering.to_dof_numbering(), 0.0)?;
        self.progress = 0.5;

        let mass = system.mass_matrix.as_ref().ok_or_else(|| {
            AnalysisError::InvalidConfiguration(
                "assembly produced no mass matrix; modal analysis needs one".to_string(),
            )
        })?;

        let free_dofs = &system.dof_numbering.free_dofs;
        if free_dofs.is_empty() {
            return Err(AnalysisError::InvalidConfiguration(
                "every degree of freedom is constrained; there is nothing to vibrate".to_string(),
            )
            .into());
        }
        if self.num_modes > free_dofs.len() {
            return Err(AnalysisError::InvalidConfiguration(format!(
                "requested {} modes from a system with {} free degrees of freedom",
                self.num_modes,
                free_dofs.len()
            ))
            .into());
        }

        let free_stiffness = system
            .stiffness_matrix
            .extract_submatrix(free_dofs, free_dofs)?;
        let free_mass = mass.extract_submatrix(free_dofs, free_dofs)?;
        self.progress = 0.7;

        let mut solver = EigenvalueSolver::new(self.num_modes);
        if let Some(shift) = self.shift {
            solver = solver.with_shift(shift);
        }
        let (eigenvalues, free_shapes) = solver.solve(&free_stiffness, &free_mass)?;
        self.progress = 0.9;

        // Expand each mode shape from the free-DOF space back to the full
        // DOF vector, leaving constrained entries at zero.
        let total_dofs = system.dof_numbering.total_dofs;
        let mut mode_shapes = DMatrix::zeros(total_dofs, self.num_modes);
        for mode in 0..self.num_modes {
            for (i, &global_dof) in free_dofs.iter().enumerate() {
                mode_shapes[(global_dof, mode)] = free_shapes[(i, mode)];
            }
        }

        let modal = ModalResults::from_eigenvalues(eigenvalues, free_shapes);

        // The reported solution is the fundamental mode shape, so that a
        // caller treating this like any other analysis gets something
        // physical rather than a zero vector.
        let solution =
            DVector::from_iterator(total_dofs, (0..total_dofs).map(|d| mode_shapes[(d, 0)]));

        let mut results = AnalysisResults::new("Modal".to_string(), solution);
        results.add_data(
            "eigenvalues".to_string(),
            AnalysisData::Vector(modal.eigenvalues.clone()),
        );
        results.add_data(
            "frequencies".to_string(),
            AnalysisData::Vector(modal.frequencies.clone()),
        );
        results.add_data(
            "periods".to_string(),
            AnalysisData::Vector(modal.periods.clone()),
        );
        results.add_data("mode_shapes".to_string(), AnalysisData::Matrix(mode_shapes));

        self.progress = 1.0;
        self.complete = true;

        Ok(results)
    }

    fn analysis_type(&self) -> &'static str {
        "Modal"
    }

    fn set_solver_options(&mut self, options: crate::solvers::SolverOptions) {
        self.config.solver_options = options;
    }

    fn set_assembly_options(&mut self, options: crate::assembly::AssemblyOptions) {
        self.config.assembly_options = options;
    }

    fn progress(&self) -> f64 {
        self.progress
    }

    fn is_complete(&self) -> bool {
        self.complete
    }
}