#include "fesa/assembly/sparse_assembler.hpp"

#include "fesa/analysis/analysis_model.hpp"
#include "fesa/assembly/parallel_for.hpp"
#include "fesa/elements/euler_beam_3d.hpp"
#include "fesa/elements/mitc4_shell.hpp"
#include "fesa/fem/dof_manager.hpp"

#include <array>
#include <limits>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace fesa {
namespace {

constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kBeamElementDofCount = 12U;
constexpr std::size_t kBeamContributionCount =
    kBeamElementDofCount * kBeamElementDofCount;
constexpr std::size_t kShellElementDofCount = 24U;
constexpr std::size_t kShellContributionCount =
    kShellElementDofCount * kShellElementDofCount;

using BeamElementBuffer =
    std::array<CooContribution, kBeamContributionCount>;
using ShellElementBuffer =
    std::array<CooContribution, kShellContributionCount>;

Result<SparseMatrix> assemblyFailure(
    const std::string& code,
    const SourceLocation& location,
    const std::string& identity,
    const std::string& message) {
    return Result<SparseMatrix>::failure(Status::failure(
        FailureCategory::model,
        {{Severity::error,
          code,
          location,
          "*ELEMENT",
          identity,
          message}}));
}

} // namespace

Result<SparseMatrix> SparseAssembler::assembleStiffness(
    const AnalysisModel& model,
    const DofManager& dofs,
    const ParallelFor& parallelFor) {
    const Domain& domain = model.domain();
    if (domain.nodes().size() >
        (std::numeric_limits<std::size_t>::max)() / kDofsPerNode ||
        dofs.fullDofCount() != domain.nodes().size() * kDofsPerNode) {
        return assemblyFailure(
            "invalid-assembly-dimensions",
            {domain.sourcePath(), 0U},
            std::to_string(dofs.fullDofCount()),
            "DofManager dimensions do not match the active model nodes.");
    }
    if (!model.activeElements().empty() && !domain.shellElements().empty()) {
        return assemblyFailure(
            "unsupported-mixed-element-model",
            {domain.sourcePath(), 0U},
            "B33:FESA-MITC4",
            "Sparse assembly does not support mixed beam and shell models.");
    }

    if (!domain.shellElements().empty()) {
        if (domain.shellElements().size() >
            (std::numeric_limits<std::size_t>::max)() /
                kShellContributionCount) {
            return assemblyFailure(
                "invalid-assembly-dimensions",
                {domain.sourcePath(), 0U},
                std::to_string(domain.shellElements().size()),
                "Shell contribution storage exceeds the addressable range.");
        }

        std::vector<std::optional<std::array<double, 3>>> directorsByNode(
            domain.nodes().size());
        for (const auto& frame : domain.shellNodeInitialFrames()) {
            if (frame.nodeIndex >= directorsByNode.size() ||
                directorsByNode[frame.nodeIndex]) {
                return assemblyFailure(
                    "invalid-assembly-element",
                    {domain.sourcePath(), 0U},
                    std::to_string(frame.nodeIndex),
                    "Shell initial frames must map uniquely to model nodes.");
            }
            directorsByNode[frame.nodeIndex] = frame.director;
        }

        struct ShellInput {
            std::array<const Node*, 4> nodes;
            std::array<std::array<double, 3>, 4> directors;
            const ShellSection* section;
            const LinearElasticMaterial* material;
            std::array<std::size_t, kShellElementDofCount> scatter;
        };
        std::vector<ShellInput> inputs;
        inputs.reserve(domain.shellElements().size());
        for (std::size_t elementOrder = 0U;
             elementOrder < domain.shellElements().size();
             ++elementOrder) {
            const auto& element = domain.shellElements()[elementOrder];
            if (element.materialIndex >= domain.materials().size() ||
                element.sectionIndex >= domain.shellSections().size()) {
                return assemblyFailure(
                    "invalid-assembly-element",
                    element.location,
                    element.sourceId.sourceLabelText,
                    "Shell element references an entity outside the Domain.");
            }

            ShellInput input{};
            input.section = &domain.shellSections()[element.sectionIndex];
            input.material = &domain.materials()[element.materialIndex];
            try {
                input.scatter = dofs.shellElementScatter(
                    static_cast<EntityIndex>(elementOrder));
            } catch (const std::out_of_range&) {
                return assemblyFailure(
                    "invalid-assembly-scatter",
                    element.location,
                    element.sourceId.sourceLabelText,
                    "DofManager does not contain the active shell scatter.");
            }
            for (std::size_t nodePosition = 0U;
                 nodePosition < element.nodeIndices.size();
                 ++nodePosition) {
                const EntityIndex nodeIndex = element.nodeIndices[nodePosition];
                if (nodeIndex >= domain.nodes().size() ||
                    !directorsByNode[nodeIndex]) {
                    return assemblyFailure(
                        "invalid-assembly-element",
                        element.location,
                        element.sourceId.sourceLabelText,
                        "Shell element requires a valid node and initial director.");
                }
                input.nodes[nodePosition] = &domain.nodes()[nodeIndex];
                input.directors[nodePosition] = *directorsByNode[nodeIndex];
                for (std::size_t component = 0U;
                     component < kDofsPerNode;
                     ++component) {
                    const std::size_t local =
                        nodePosition * kDofsPerNode + component;
                    const std::size_t expected =
                        static_cast<std::size_t>(nodeIndex) * kDofsPerNode +
                        component;
                    if (input.scatter[local] != expected ||
                        input.scatter[local] >= dofs.fullDofCount()) {
                        return assemblyFailure(
                            "invalid-assembly-scatter",
                            element.location,
                            element.sourceId.sourceLabelText,
                            "Shell scatter does not match the active model topology.");
                    }
                }
            }
            inputs.push_back(input);
        }

        std::vector<ShellElementBuffer> localBuffers(inputs.size());
        std::vector<std::optional<Status>> localFailures(inputs.size());
        parallelFor.execute(
            inputs.size(),
            [&](const std::size_t elementOrder) {
                const auto& input = inputs[elementOrder];
                const auto shell = Mitc4Shell::create(
                    input.nodes,
                    input.directors,
                    *input.section,
                    *input.material);
                if (!shell.hasValue()) {
                    localFailures[elementOrder] = shell.status();
                    return;
                }
                const auto stiffness = shell.value().stiffness();
                if (!stiffness.hasValue()) {
                    localFailures[elementOrder] = stiffness.status();
                    return;
                }

                auto& buffer = localBuffers[elementOrder];
                for (std::size_t localRow = 0U;
                     localRow < kShellElementDofCount;
                     ++localRow) {
                    for (std::size_t localColumn = 0U;
                         localColumn < kShellElementDofCount;
                         ++localColumn) {
                        const std::size_t localOrder =
                            localRow * kShellElementDofCount + localColumn;
                        buffer[localOrder] = {
                            input.scatter[localRow],
                            input.scatter[localColumn],
                            stiffness.value().stabilizedGlobal24(
                                localRow, localColumn),
                            elementOrder,
                            localOrder};
                    }
                }
            });

        for (std::size_t elementOrder = 0U;
             elementOrder < localFailures.size();
             ++elementOrder) {
            if (localFailures[elementOrder]) {
                return Result<SparseMatrix>::failure(
                    *localFailures[elementOrder]);
            }
        }

        std::vector<CooContribution> contributions;
        contributions.reserve(
            localBuffers.size() * kShellContributionCount);
        // Flatten in source-element order after workers complete. The canonical
        // COO reduction remains the sole writer of global CSR values.
        for (const auto& buffer : localBuffers) {
            contributions.insert(
                contributions.end(), buffer.begin(), buffer.end());
        }
        return SparseMatrix::fromCoo(
            dofs.fullDofCount(),
            dofs.fullDofCount(),
            std::move(contributions),
            dofs.sparsePattern());
    }

    if (model.activeElements().size() >
        (std::numeric_limits<std::size_t>::max)() /
            kBeamContributionCount) {
        return assemblyFailure(
            "invalid-assembly-dimensions",
            {domain.sourcePath(), 0U},
            std::to_string(model.activeElements().size()),
            "Element contribution storage exceeds the addressable range.");
    }

    std::vector<std::array<std::size_t, kBeamElementDofCount>> scatters;
    scatters.reserve(model.activeElements().size());
    for (const EntityIndex elementIndex : model.activeElements()) {
        if (elementIndex >= domain.elements().size()) {
            return assemblyFailure(
                "invalid-assembly-element",
                {domain.sourcePath(), 0U},
                std::to_string(elementIndex),
                "Active element index is outside the Domain.");
        }
        const auto& element = domain.elements()[elementIndex];
        if (element.nodeIndices[0U] >= domain.nodes().size() ||
            element.nodeIndices[1U] >= domain.nodes().size() ||
            element.materialIndex >= domain.materials().size() ||
            element.sectionIndex >= domain.sections().size()) {
            return assemblyFailure(
                "invalid-assembly-element",
                element.location,
                element.sourceId.sourceLabelText,
                "Element references an entity outside the Domain.");
        }

        std::array<std::size_t, kBeamElementDofCount> scatter{};
        try {
            scatter = dofs.elementScatter(elementIndex);
        } catch (const std::out_of_range&) {
            return assemblyFailure(
                "invalid-assembly-scatter",
                element.location,
                element.sourceId.sourceLabelText,
                "DofManager does not contain the active element scatter.");
        }
        for (std::size_t endpoint = 0U; endpoint < 2U; ++endpoint) {
            for (std::size_t component = 0U;
                 component < kDofsPerNode;
                 ++component) {
                const std::size_t local = endpoint * kDofsPerNode + component;
                const std::size_t expected =
                    static_cast<std::size_t>(element.nodeIndices[endpoint]) *
                        kDofsPerNode +
                    component;
                if (scatter[local] != expected ||
                    scatter[local] >= dofs.fullDofCount()) {
                    return assemblyFailure(
                        "invalid-assembly-scatter",
                        element.location,
                        element.sourceId.sourceLabelText,
                        "Element scatter does not match the active model topology.");
                }
            }
        }
        scatters.push_back(scatter);
    }

    std::vector<BeamElementBuffer> localBuffers(model.activeElements().size());
    std::vector<std::optional<Status>> localFailures(
        model.activeElements().size());
    parallelFor.execute(
        model.activeElements().size(),
        [&](const std::size_t elementOrder) {
            const EntityIndex elementIndex = model.activeElements()[elementOrder];
            const auto& definition = domain.elements()[elementIndex];
            const auto beam = EulerBeam3D::create(
                domain.nodes()[definition.nodeIndices[0U]],
                domain.nodes()[definition.nodeIndices[1U]],
                domain.sections()[definition.sectionIndex],
                domain.materials()[definition.materialIndex]);
            if (!beam.hasValue()) {
                localFailures[elementOrder] = beam.status();
                return;
            }

            const Matrix stiffness = beam.value().globalStiffness();
            auto& buffer = localBuffers[elementOrder];
            const auto& scatter = scatters[elementOrder];
            for (std::size_t localRow = 0U;
                 localRow < kBeamElementDofCount;
                 ++localRow) {
                for (std::size_t localColumn = 0U;
                     localColumn < kBeamElementDofCount;
                     ++localColumn) {
                    const std::size_t localOrder =
                        localRow * kBeamElementDofCount + localColumn;
                    buffer[localOrder] = {
                        scatter[localRow],
                        scatter[localColumn],
                        stiffness(localRow, localColumn),
                        elementOrder,
                        localOrder};
                }
            }
        });

    for (std::size_t elementOrder = 0U;
         elementOrder < localFailures.size();
         ++elementOrder) {
        if (localFailures[elementOrder]) {
            return Result<SparseMatrix>::failure(
                *localFailures[elementOrder]);
        }
    }

    std::vector<CooContribution> contributions;
    contributions.reserve(
        localBuffers.size() * kBeamContributionCount);
    // Flatten only after all workers complete; workers never share CSR state.
    for (const auto& buffer : localBuffers) {
        contributions.insert(
            contributions.end(), buffer.begin(), buffer.end());
    }
    return SparseMatrix::fromCoo(
        dofs.fullDofCount(),
        dofs.fullDofCount(),
        std::move(contributions),
        dofs.sparsePattern());
}

} // namespace fesa
