import numpy as np
from scipy.sparse.linalg import spsolve

from femsurrogate.fea.assembly import (
    assemble_global_stiffness,
    assemble_load_vector,
    constrained_dofs,
    node_ids,
)
from femsurrogate.fea.model import BeamModel, Displacement


def solve_linear_static(model: BeamModel) -> dict[int, Displacement]:
    stiffness = assemble_global_stiffness(model)
    loads = assemble_load_vector(model)
    fixed = np.array(constrained_dofs(model), dtype=int)
    all_dofs = np.arange(loads.size)
    free = np.setdiff1d(all_dofs, fixed, assume_unique=True)

    displacements = np.zeros_like(loads)
    displacements[free] = spsolve(stiffness[free][:, free], loads[free])

    result: dict[int, Displacement] = {}
    for position, node_id in enumerate(node_ids(model)):
        base = position * 3
        result[node_id] = Displacement(
            node_id=node_id,
            ux=float(displacements[base]),
            uy=float(displacements[base + 1]),
            rz=float(displacements[base + 2]),
        )

    return result