#include "fesa/assembly/load_assembler.hpp"

#include "fesa/constraints/essential_constraints.hpp"

#include <algorithm>
#include <charconv>
#include <cmath>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

namespace fesa {
namespace {

constexpr std::size_t dofsPerNode = 6U;
constexpr double shellMomentProjectionTolerance = 1.0e-12;

Status loadFailure(
    const std::string& code,
    const SourceLocation& location,
    const std::string& keyword,
    const std::string& identity,
    const std::string& message) {
    return Status::Failure(
        FailureCategory::kModel,
        {{Severity::kError, code, location, keyword, identity, message}});
}

char asciiLower(const char value) {
    if (value >= 'A' && value <= 'Z') {
        return static_cast<char>(value + ('a' - 'A'));
    }
    return value;
}

bool equalName(const std::string& left, const std::string& right) {
    return left.size() == right.size() &&
           std::equal(
               left.begin(),
               left.end(),
               right.begin(),
               [](const char leftValue, const char rightValue) {
                   return asciiLower(leftValue) == asciiLower(rightValue);
               });
}

bool tryPositiveInteger(const std::string& text, std::int64_t& value) {
    const char* const first = text.data();
    const char* const last = first + text.size();
    const auto parsed = std::from_chars(first, last, value);
    return parsed.ec == std::errc{} && parsed.ptr == last && value > 0;
}

bool isStrictlyIncreasing(const std::vector<std::size_t>& values) {
    return std::adjacent_find(
               values.begin(),
               values.end(),
               [](const std::size_t left, const std::size_t right) {
                   return left >= right;
               }) == values.end();
}

Status validateDofOrder(
    const DofManager& dofs,
    const std::size_t expectedFullCount,
    const SourceLocation& location) {
    const std::size_t fullCount = dofs.fullDofCount();
    const auto& freeDofs = dofs.freeDofs();
    const auto& constrainedDofs = dofs.constrainedDofs();
    if (fullCount != expectedFullCount ||
        freeDofs.size() != dofs.freeDofCount() ||
        constrainedDofs.size() != dofs.constrainedDofCount() ||
        dofs.prescribedValues().Size() != constrainedDofs.size() ||
        constrainedDofs.size() > fullCount ||
        freeDofs.size() != fullCount - constrainedDofs.size()) {
        return loadFailure(
            "invalid-load-dimensions",
            location,
            "LOAD_ASSEMBLER",
            std::to_string(fullCount),
            "Full, free, constrained, prescribed, and model dimensions must agree.");
    }
    if (!isStrictlyIncreasing(freeDofs) ||
        !isStrictlyIncreasing(constrainedDofs)) {
        return loadFailure(
            "invalid-load-order",
            location,
            "LOAD_ASSEMBLER",
            std::to_string(fullCount),
            "Free and constrained DOFs must use stable increasing full-DOF order.");
    }

    std::vector<unsigned char> ownership(fullCount, 0U);
    try {
        for (std::size_t equation = 0U;
             equation < freeDofs.size();
             ++equation) {
            const std::size_t fullDof = freeDofs[equation];
            if (fullDof >= fullCount || ownership[fullDof] != 0U ||
                dofs.freeEquation(fullDof) != equation) {
                return loadFailure(
                    "invalid-load-order",
                    location,
                    "LOAD_ASSEMBLER",
                    std::to_string(fullDof),
                    "Free equation numbering must match stable full-DOF order.");
            }
            ownership[fullDof] = 1U;
        }
        for (const std::size_t fullDof : constrainedDofs) {
            if (fullDof >= fullCount || ownership[fullDof] != 0U ||
                dofs.freeEquation(fullDof).has_value()) {
                return loadFailure(
                    "invalid-load-order",
                    location,
                    "LOAD_ASSEMBLER",
                    std::to_string(fullDof),
                    "Constrained DOFs must be unique and absent from free equations.");
            }
            ownership[fullDof] = 2U;
        }
    } catch (const std::out_of_range&) {
        return loadFailure(
            "invalid-load-dimensions",
            location,
            "LOAD_ASSEMBLER",
            std::to_string(fullCount),
            "DofManager equation storage must cover every full DOF.");
    }
    if (std::find(ownership.begin(), ownership.end(), 0U) != ownership.end()) {
        return loadFailure(
            "invalid-load-order",
            location,
            "LOAD_ASSEMBLER",
            std::to_string(fullCount),
            "Free and constrained DOFs must partition the full range.");
    }
    return Status::Ok();
}

Result<std::vector<EntityIndex>> resolveTarget(
    const Domain& domain,
    const NodalLoad& load) {
    std::vector<const NodeSet*> matchingSets;
    for (const auto& set : domain.nodeSets()) {
        if (equalName(set.name, load.target)) {
            matchingSets.push_back(&set);
        }
    }

    std::vector<EntityIndex> matchingNodes;
    std::int64_t label = 0;
    if (tryPositiveInteger(load.target, label)) {
        for (std::size_t index = 0U; index < domain.nodes().size(); ++index) {
            if (domain.nodes()[index].sourceId.source_label == label) {
                matchingNodes.push_back(static_cast<EntityIndex>(index));
            }
        }
    }

    if (matchingSets.size() > 1U || matchingNodes.size() > 1U ||
        (!matchingSets.empty() && !matchingNodes.empty())) {
        return Result<std::vector<EntityIndex>>::Failure(loadFailure(
            "invalid-load-target",
            load.location,
            "CLOAD",
            load.target,
            "The load target must resolve unambiguously to one node or one expanded node set."));
    }
    if (!matchingSets.empty()) {
        const auto& nodes = matchingSets.front()->nodeIndices;
        std::vector<unsigned char> seen(domain.nodes().size(), 0U);
        for (const EntityIndex node : nodes) {
            if (node >= domain.nodes().size() || seen[node] != 0U) {
                return Result<std::vector<EntityIndex>>::Failure(loadFailure(
                    "invalid-load-target",
                    load.location,
                    "CLOAD",
                    load.target,
                    "The expanded node set must contain unique in-range stable node identities."));
            }
            seen[node] = 1U;
        }
        return Result<std::vector<EntityIndex>>::Success(nodes);
    }
    if (!matchingNodes.empty()) {
        return Result<std::vector<EntityIndex>>::Success(
            std::move(matchingNodes));
    }
    return Result<std::vector<EntityIndex>>::Failure(loadFailure(
        "invalid-load-target",
        load.location,
        "CLOAD",
        load.target,
        "The load target must resolve to one semantic node or node set."));
}

Status validateFiniteVector(
    const Vector& values,
    const SourceLocation& location,
    const std::string& identity) {
    for (std::size_t index = 0U; index < values.Size(); ++index) {
        if (!std::isfinite(values[index])) {
            return loadFailure(
                "nonfinite-load-value",
                location,
                "LOAD_ASSEMBLER",
                identity + ":" + std::to_string(index),
                "Load and prescribed displacement vectors must contain finite values.");
        }
    }
    return Status::Ok();
}

Status validateShellMoments(
    const Domain& domain,
    const Vector& fullLoad) {
    if (domain.shellElements().empty()) {
        return Status::Ok();
    }

    std::vector<const ShellNodeInitialFrame*> frameByNode(
        domain.nodes().size(), nullptr);
    for (const auto& frame : domain.shellNodeInitialFrames()) {
        if (frame.nodeIndex >= frameByNode.size() ||
            frameByNode[frame.nodeIndex] != nullptr) {
            return loadFailure(
                "invalid-shell-director",
                {domain.sourcePath(), 0U},
                "NODE",
                std::to_string(frame.nodeIndex),
                "Shell nodal directors must have unique in-range node identities.");
        }
        frameByNode[frame.nodeIndex] = &frame;
    }

    for (std::size_t node = 0U; node < domain.nodes().size(); ++node) {
        const double momentX = fullLoad[node * dofsPerNode + 3U];
        const double momentY = fullLoad[node * dofsPerNode + 4U];
        const double momentZ = fullLoad[node * dofsPerNode + 5U];
        if (momentX == 0.0 && momentY == 0.0 && momentZ == 0.0) {
            continue;
        }

        const auto* const frame = frameByNode[node];
        if (frame == nullptr) {
            return loadFailure(
                "invalid-shell-director",
                domain.nodes()[node].location,
                "NODE",
                domain.nodes()[node].sourceId.source_label_text,
                "A loaded shell node must have an approved initial director.");
        }

        const double momentScale = std::max(
            std::abs(momentX),
            std::max(std::abs(momentY), std::abs(momentZ)));
        const double scaledX = momentX / momentScale;
        const double scaledY = momentY / momentScale;
        const double scaledZ = momentZ / momentScale;
        const double scaledNorm = std::hypot(scaledX, scaledY, scaledZ);
        const double scaledDot =
            frame->director[0U] * scaledX +
            frame->director[1U] * scaledY +
            frame->director[2U] * scaledZ;
        const double projectionRatio = std::abs(scaledDot) / scaledNorm;
        if (!(projectionRatio <= shellMomentProjectionTolerance)) {
            return loadFailure(
                "unsupported-drilling-load",
                domain.nodes()[node].location,
                "CLOAD",
                domain.nodes()[node].sourceId.source_label_text,
                "The aggregate nodal moment has an unsupported director-parallel component.");
        }
    }
    return Status::Ok();
}

} // namespace

Result<Vector> LoadAssembler::assembleFullNodalLoad(
    const AnalysisModel& model,
    const DofManager& dofs) {
    const Domain& domain = model.domain();
    if (domain.nodes().size() >
        (std::numeric_limits<std::size_t>::max)() / dofsPerNode) {
        return Result<Vector>::Failure(loadFailure(
            "invalid-load-dimensions",
            {domain.sourcePath(), 0U},
            "LOAD_ASSEMBLER",
            domain.sourceContentIdentity(),
            "The semantic node count cannot be represented in full-DOF order."));
    }
    const std::size_t expectedFullCount =
        domain.nodes().size() * dofsPerNode;
    const Status dofStatus = validateDofOrder(
        dofs, expectedFullCount, {domain.sourcePath(), 0U});
    if (!dofStatus.IsOk()) {
        return Result<Vector>::Failure(dofStatus);
    }
    for (std::size_t node = 0U; node < domain.nodes().size(); ++node) {
        for (std::size_t component = 0U;
             component < dofsPerNode;
             ++component) {
            try {
                if (dofs.fullDof(
                        static_cast<EntityIndex>(node),
                        static_cast<DofComponent>(component)) !=
                    node * dofsPerNode + component) {
                    return Result<Vector>::Failure(loadFailure(
                        "invalid-load-order",
                        domain.nodes()[node].location,
                        "LOAD_ASSEMBLER",
                        domain.nodes()[node].sourceId.source_label_text,
                        "DofManager node/component identity must match full-DOF order."));
                }
            } catch (const std::out_of_range&) {
                return Result<Vector>::Failure(loadFailure(
                    "invalid-load-dimensions",
                    domain.nodes()[node].location,
                    "LOAD_ASSEMBLER",
                    domain.nodes()[node].sourceId.source_label_text,
                    "DofManager must provide all six DOFs for every semantic node."));
            }
        }
    }

    const auto& activeLoads = model.activeLoads();
    const auto& loads = model.step().loads;
    if (activeLoads.size() != loads.size()) {
        return Result<Vector>::Failure(loadFailure(
            "invalid-load-order",
            model.step().location,
            "CLOAD",
            model.step().name,
            "The active load view must include every sole-step load once."));
    }

    Vector fullLoad{expectedFullCount};
    // Active load indices are required to be the original source order; this
    // loop is therefore also the fixed floating-point accumulation order.
    for (std::size_t sourceOrder = 0U;
         sourceOrder < activeLoads.size();
         ++sourceOrder) {
        const EntityIndex loadIndex = activeLoads[sourceOrder];
        if (static_cast<std::size_t>(loadIndex) != sourceOrder ||
            loadIndex >= loads.size()) {
            return Result<Vector>::Failure(loadFailure(
                "invalid-load-order",
                model.step().location,
                "CLOAD",
                std::to_string(sourceOrder),
                "Active loads must retain complete stable source order."));
        }
        const auto& load = loads[loadIndex];
        if (load.dof < 1 || load.dof > static_cast<int>(dofsPerNode)) {
            return Result<Vector>::Failure(loadFailure(
                "invalid-load-dof",
                load.location,
                "CLOAD",
                load.target,
                "A nodal load component must be in the range 1 through 6."));
        }
        if (!std::isfinite(load.magnitude)) {
            return Result<Vector>::Failure(loadFailure(
                "nonfinite-load-value",
                load.location,
                "CLOAD",
                load.target,
                "A nodal load magnitude must be finite."));
        }

        auto target = resolveTarget(domain, load);
        if (!target.HasValue()) {
            return Result<Vector>::Failure(target.GetStatus());
        }
        const auto component = static_cast<DofComponent>(load.dof - 1);
        for (const EntityIndex node : target.Value()) {
            const std::size_t fullDof = dofs.fullDof(node, component);
            const double accumulated = fullLoad[fullDof] + load.magnitude;
            if (!std::isfinite(accumulated)) {
                return Result<Vector>::Failure(loadFailure(
                    "nonfinite-load-accumulation",
                    load.location,
                    "CLOAD",
                    load.target,
                    "Source-order load accumulation produced a nonfinite value."));
            }
            fullLoad[fullDof] = accumulated;
        }
    }
    const Status shellMomentStatus = validateShellMoments(domain, fullLoad);
    if (!shellMomentStatus.IsOk()) {
        return Result<Vector>::Failure(shellMomentStatus);
    }
    return Result<Vector>::Success(std::move(fullLoad));
}

Result<Vector> LoadAssembler::effectiveFreeRhs(
    const Vector& fullLoad,
    const SparseMatrix& kfc,
    const Vector& prescribedValues,
    const DofManager& dofs) {
    const SourceLocation location{{}, 0U};
    const Status dofStatus =
        validateDofOrder(dofs, fullLoad.Size(), location);
    if (!dofStatus.IsOk()) {
        return Result<Vector>::Failure(dofStatus);
    }
    if (kfc.Rows() != dofs.freeDofCount() ||
        kfc.Columns() != dofs.constrainedDofCount() ||
        prescribedValues.Size() != dofs.constrainedDofCount()) {
        return Result<Vector>::Failure(loadFailure(
            "invalid-load-dimensions",
            location,
            "LOAD_ASSEMBLER",
            std::to_string(kfc.Rows()) + "x" +
                std::to_string(kfc.Columns()),
            "Kfc rows/columns and prescribed values must match free/constrained order."));
    }
    const Status matrixStatus = kfc.Validate();
    if (!matrixStatus.IsOk()) {
        return Result<Vector>::Failure(matrixStatus);
    }
    const Status loadStatus =
        validateFiniteVector(fullLoad, location, "full-load");
    if (!loadStatus.IsOk()) {
        return Result<Vector>::Failure(loadStatus);
    }
    const Status prescribedStatus = validateFiniteVector(
        prescribedValues, location, "prescribed-values");
    if (!prescribedStatus.IsOk()) {
        return Result<Vector>::Failure(prescribedStatus);
    }

    Vector correction{kfc.Rows()};
    for (std::size_t row = 0U; row < kfc.Rows(); ++row) {
        double sum = 0.0;
        for (std::size_t position = kfc.RowOffsets()[row];
             position < kfc.RowOffsets()[row + 1U];
             ++position) {
            const double product = kfc.Values()[position] *
                prescribedValues[kfc.ColumnIndices()[position]];
            if (!std::isfinite(product)) {
                return Result<Vector>::Failure(loadFailure(
                    "nonfinite-load-accumulation",
                    location,
                    "LOAD_ASSEMBLER",
                    std::to_string(row),
                    "Kfc times prescribed displacement produced a nonfinite product."));
            }
            sum += product;
            if (!std::isfinite(sum)) {
                return Result<Vector>::Failure(loadFailure(
                    "nonfinite-load-accumulation",
                    location,
                    "LOAD_ASSEMBLER",
                    std::to_string(row),
                    "Kfc times prescribed displacement produced a nonfinite row sum."));
            }
        }
        correction[row] = sum;
    }

    Vector rhs = EssentialConstraints::gatherFree(fullLoad, dofs);
    // The constrained vector is already in DofManager order, so this is the
    // approved elimination equation rhs = Ff - Kfc*dc without reordering dc.
    for (std::size_t row = 0U; row < rhs.Size(); ++row) {
        const double value = rhs[row] - correction[row];
        if (!std::isfinite(value)) {
            return Result<Vector>::Failure(loadFailure(
                "nonfinite-load-accumulation",
                location,
                "LOAD_ASSEMBLER",
                std::to_string(row),
                "Effective RHS subtraction produced a nonfinite value."));
        }
        rhs[row] = value;
    }
    return Result<Vector>::Success(std::move(rhs));
}

} // namespace fesa