#include <fesa/analysis/linear_static_analysis.hpp>

#include <algorithm>
#include <array>
#include <cstddef>
#include <exception>
#include <stdexcept>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include <fesa/assembly/equation_system.hpp>
#include <fesa/assembly/serial_assembler.hpp>
#include <fesa/constraints/essential_bc.hpp>
#include <fesa/elements/beam/beam3d2.hpp>
#include <fesa/fem/dof_manager.hpp>
#include <fesa/solvers/linear/pardiso_linear_solver.hpp>

namespace fesa {
namespace {

AnalysisRunResult failure(std::vector<Diagnostic> diagnostics) {
    return {false, std::nullopt, std::move(diagnostics)};
}

AnalysisRunResult equation_failure(
    std::string code,
    std::string message) {
    return failure({{
        DiagnosticStage::equation,
        Severity::error,
        std::move(code),
        std::move(message),
        std::nullopt,
    }});
}

AnalysisRunResult results_failure(
    std::string code,
    std::string message) {
    return failure({{
        DiagnosticStage::results,
        Severity::error,
        std::move(code),
        std::move(message),
        std::nullopt,
    }});
}

NodalFrame build_nodal_frame(
    const Domain& domain,
    const DofManager& dofs,
    const std::vector<double>& displacement,
    const std::vector<double>& reaction) {
    std::vector<NodeId> node_ids;
    node_ids.reserve(domain.nodes().size());
    for (const Node& node : domain.nodes()) {
        node_ids.push_back(node.id);
    }
    std::ranges::sort(node_ids);

    NodalFrame nodal;
    nodal.node_ids = std::move(node_ids);
    nodal.origins.reserve(nodal.node_ids.size());
    nodal.displacement.reserve(nodal.node_ids.size());
    nodal.reaction.reserve(nodal.node_ids.size());
    for (const NodeId node_id : nodal.node_ids) {
        nodal.origins.push_back(domain.node(node_id).origin);
        std::array<double, 6> node_displacement{};
        std::array<double, 6> node_reaction{};
        for (std::size_t component = 0; component < 6; ++component) {
            const std::size_t full_dof = dofs.full_dof({
                node_id,
                static_cast<NodeDof>(component),
            });
            node_displacement[component] = displacement[full_dof];
            node_reaction[component] = reaction[full_dof];
        }
        nodal.displacement.push_back(node_displacement);
        nodal.reaction.push_back(node_reaction);
    }
    return nodal;
}

ElementFrame build_element_frame(
    const Domain& domain,
    const DofManager& dofs,
    const std::vector<double>& displacement) {
    std::vector<const BeamElement*> elements;
    elements.reserve(domain.beam_elements().size());
    for (const BeamElement& element : domain.beam_elements()) {
        elements.push_back(&element);
    }
    std::ranges::sort(
        elements,
        {},
        [](const BeamElement* element) {
            return element->id.value();
        });

    ElementFrame frame;
    frame.beams.reserve(elements.size());
    for (const BeamElement* element : elements) {
        const Beam3D2Input input{
            {
                domain.node(element->nodes[0]).position,
                domain.node(element->nodes[1]).position,
            },
            element->nodes,
            domain.material(element->material),
            domain.section(element->section),
        };
        const BeamKernelResult kernel = compute_beam3d2(input);
        if (!kernel.contribution.has_value()) {
            const std::string message = kernel.diagnostics.empty()
                ? "Beam recovery requires a valid element input."
                : kernel.diagnostics.front().message;
            throw std::runtime_error{message};
        }

        std::array<double, 12> element_displacement{};
        const std::array<std::size_t, 12> full_dofs =
            dofs.element_full_dofs(*element);
        for (std::size_t local = 0; local < full_dofs.size(); ++local) {
            element_displacement[local] = displacement[full_dofs[local]];
        }
        std::vector<BeamSectionResult> recovered = recover_beam3d2(
            input,
            element_displacement,
            input.section.recovery_points);
        if (recovered.size() != 2U) {
            throw std::logic_error{
                "Beam recovery must return exactly two end results."};
        }

        frame.beams.push_back({
            element->id,
            element->origin,
            kernel.contribution->frame,
            {
                std::move(recovered[0]),
                std::move(recovered[1]),
            },
        });
    }
    return frame;
}

}  // namespace

AnalysisRunResult LinearStaticAnalysis::run(const Domain& domain) const {
    const DofManager dofs = DofManager::build(domain);

    std::optional<EquationSystem> original;
    try {
        original = assemble_serial(domain, dofs);
    } catch (const std::exception& error) {
        return equation_failure(
            "equation.assembly_failed", error.what());
    }

    ConstraintResult constrained =
        eliminate_essential_bcs(*original, dofs);
    if (!constrained.reduced_system.has_value()) {
        return failure(std::move(constrained.diagnostics));
    }

    std::vector<double> reduced_solution;
    if (dofs.free_equation_count() != 0) {
        PardisoLinearSolver solver;
        LinearSolveResult solved = solver.solve(
            constrained.reduced_system->stiffness,
            constrained.reduced_system->force);
        if (!solved.diagnostics.empty()) {
            return failure(std::move(solved.diagnostics));
        }
        reduced_solution = std::move(solved.solution);
    }

    std::vector<double> displacement;
    try {
        displacement = dofs.reconstruct_full(reduced_solution);
    } catch (const std::exception& error) {
        return equation_failure(
            "equation.solution_reconstruction_failed", error.what());
    }

    std::vector<double> reaction;
    try {
        reaction = recover_reaction(*original, displacement);
    } catch (const std::exception& error) {
        return equation_failure(
            "equation.reaction_recovery_failed", error.what());
    }

    ElementFrame element;
    try {
        element = build_element_frame(domain, dofs, displacement);
    } catch (const std::exception& error) {
        return results_failure(
            "results.element_recovery_failed", error.what());
    }

    ResultDatabase database{
        "2.0.0",
        {{
            domain.step().name,
            {{
                1.0,
                build_nodal_frame(
                    domain, dofs, displacement, reaction),
                std::move(element),
                {},
            }},
        }},
    };
    Status status = validate_result_database(database);
    if (!status.succeeded) {
        return failure(std::move(status.diagnostics));
    }

    return {true, std::move(database), {}};
}

}  // namespace fesa