#!/usr/bin/env python3
"""Finite element solution of a 1D string vibration eigenvalue problem.

Problem:
    u'' + k^2 u = 0,  0 < x < L
    u(0) = u(L) = 0

Weak form:
    integral_0^L u' v' dx = k^2 integral_0^L u v dx

This script uses linear finite elements on a uniform mesh.  The command-line
argument N is the number of elements, so the number of nodes is N + 1.
The boundary nodes are fixed, and the generalized eigenvalue problem is solved
for the internal degrees of freedom.
"""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np


LENGTH = 1.0
OUTPUT_FILE = "fem1d_results.txt"


def build_global_matrices(num_elements: int, length: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Build global stiffness and mass matrices for linear 1D elements."""
    num_nodes = num_elements + 1
    nodes = np.linspace(0.0, length, num_nodes)

    stiffness = np.zeros((num_nodes, num_nodes))
    mass = np.zeros((num_nodes, num_nodes))

    for elem in range(num_elements):
        h = nodes[elem + 1] - nodes[elem]

        elem_stiffness = (1.0 / h) * np.array(
            [
                [1.0, -1.0],
                [-1.0, 1.0],
            ]
        )
        elem_mass = (h / 6.0) * np.array(
            [
                [2.0, 1.0],
                [1.0, 2.0],
            ]
        )

        global_nodes = [elem, elem + 1]
        for local_i, global_i in enumerate(global_nodes):
            for local_j, global_j in enumerate(global_nodes):
                stiffness[global_i, global_j] += elem_stiffness[local_i, local_j]
                mass[global_i, global_j] += elem_mass[local_i, local_j]

    return nodes, stiffness, mass


def apply_fixed_end_boundary_conditions(
    stiffness: np.ndarray,
    mass: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Remove both fixed boundary degrees of freedom."""
    return stiffness[1:-1, 1:-1], mass[1:-1, 1:-1]


def solve_generalized_symmetric_eigenproblem(
    stiffness: np.ndarray,
    mass: np.ndarray,
    num_modes: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Solve K u = lambda M u using only NumPy.

    Since M is symmetric positive definite, write M = L L^T and transform
    the problem into a standard symmetric eigenvalue problem:

        L^{-1} K L^{-T} y = lambda y,  u = L^{-T} y.
    """

    cholesky = np.linalg.cholesky(mass)
    temp = np.linalg.solve(cholesky, stiffness)
    standard_matrix = np.linalg.solve(cholesky, temp.T).T

    eigenvalues, transformed_vectors = np.linalg.eigh(standard_matrix)
    order = np.argsort(eigenvalues)
    eigenvalues = eigenvalues[order][:num_modes]
    transformed_vectors = transformed_vectors[:, order][:, :num_modes]

    eigenvectors = np.linalg.solve(cholesky.T, transformed_vectors)

    for mode in range(num_modes):
        norm = np.sqrt(eigenvectors[:, mode].T @ mass @ eigenvectors[:, mode])
        eigenvectors[:, mode] /= norm

    return eigenvalues, eigenvectors


def expand_internal_modes(eigenvectors_internal: np.ndarray, num_elements: int) -> np.ndarray:
    """Add zero boundary values to internal eigenvectors."""
    num_nodes = num_elements + 1
    full_vectors = np.zeros((num_nodes, eigenvectors_internal.shape[1]))
    full_vectors[1:-1, :] = eigenvectors_internal
    return full_vectors


def write_results(
    path: Path,
    nodes: np.ndarray,
    eigenvalues: np.ndarray,
    eigenvectors: np.ndarray,
) -> None:
    """Write eigenvalues and eigenvectors to a text file."""
    with path.open("w", encoding="utf-8") as file:
        file.write("# 1D FEM string vibration eigenvalue problem\n")
        file.write("# Equation: u'' + k^2 u = 0, u(0)=u(L)=0\n")
        file.write("# Eigenvalue lambda = k^2\n")
        file.write(f"# number_of_elements = {len(nodes) - 1}\n")
        file.write(f"# number_of_nodes = {len(nodes)}\n")
        file.write(f"# length = {nodes[-1] - nodes[0]:.16e}\n\n")

        file.write("# Eigenvalues\n")
        file.write("# mode  lambda=k^2              k\n")
        for index, eigenvalue in enumerate(eigenvalues, start=1):
            file.write(f"{index:5d}  {eigenvalue:.16e}  {np.sqrt(eigenvalue):.16e}\n")

        file.write("\n# Eigenvectors\n")
        file.write("# Each eigenvector includes the fixed boundary values u(0)=u(L)=0.\n")
        header = "node_index  x" + "".join(f"  mode_{i}" for i in range(1, len(eigenvalues) + 1))
        file.write(f"# {header}\n")
        for node_index, x in enumerate(nodes, start=1):
            values = "".join(f"  {eigenvectors[node_index - 1, mode]: .16e}" for mode in range(len(eigenvalues)))
            file.write(f"{node_index:10d}  {x:.16e}{values}\n")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Solve the 1D FEM string vibration eigenvalue problem."
    )
    parser.add_argument("N", type=int, help="number of elements")
    parser.add_argument("m", type=int, help="number of modes, counted from the smallest eigenvalue")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.N < 2:
        raise SystemExit("N must be at least 2 because both ends are fixed.")

    max_modes = args.N - 1
    if args.m < 1 or args.m > max_modes:
        raise SystemExit(f"m must satisfy 1 <= m <= {max_modes}.")

    nodes, stiffness, mass = build_global_matrices(args.N, LENGTH)
    stiffness_internal, mass_internal = apply_fixed_end_boundary_conditions(stiffness, mass)
    eigenvalues, eigenvectors_internal = solve_generalized_symmetric_eigenproblem(
        stiffness_internal,
        mass_internal,
        args.m,
    )
    eigenvectors = expand_internal_modes(eigenvectors_internal, args.N)

    output_path = Path(OUTPUT_FILE)
    write_results(output_path, nodes, eigenvalues, eigenvectors)
    print(f"Wrote {args.m} modes to {output_path}")


if __name__ == "__main__":
    main()