#include "fesa/constraints/essential_constraints.hpp"

#include "fesa/fem/dof_manager.hpp"

#include <algorithm>
#include <limits>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace fesa {
namespace {

Status constraintFailure(
    const std::string& code,
    const std::string& identity,
    const std::string& message) {
    return Status::failure(
        FailureCategory::model,
        {{Severity::error,
          code,
          {{}, 0U},
          "ESSENTIAL_CONSTRAINTS",
          identity,
          message}});
}

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 fullCount = dofs.fullDofCount();
    const auto& freeDofs = dofs.freeDofs();
    const auto& constrainedDofs = dofs.constrainedDofs();
    if (freeDofs.size() != dofs.freeDofCount() ||
        constrainedDofs.size() != dofs.constrainedDofCount() ||
        dofs.prescribedValues().size() != constrainedDofs.size() ||
        constrainedDofs.size() > fullCount ||
        freeDofs.size() != fullCount - constrainedDofs.size()) {
        return constraintFailure(
            "invalid-constraint-dimensions",
            std::to_string(fullCount),
            "DofManager full, free, constrained, and prescribed dimensions must agree.");
    }
    if (!isStrictlyIncreasing(freeDofs) ||
        !isStrictlyIncreasing(constrainedDofs)) {
        return constraintFailure(
            "invalid-constraint-order",
            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 constraintFailure(
                    "invalid-constraint-order",
                    std::to_string(fullDof),
                    "Free equation numbering must match the stable free-DOF order.");
            }
            ownership[fullDof] = 1U;
        }
        for (const std::size_t fullDof : constrainedDofs) {
            if (fullDof >= fullCount || ownership[fullDof] != 0U ||
                dofs.freeEquation(fullDof).has_value()) {
                return constraintFailure(
                    "invalid-constraint-order",
                    std::to_string(fullDof),
                    "Constrained DOFs must be unique and absent from free equations.");
            }
            ownership[fullDof] = 2U;
        }
    } catch (const std::out_of_range&) {
        return constraintFailure(
            "invalid-constraint-dimensions",
            std::to_string(fullCount),
            "DofManager equation storage does not cover every full DOF.");
    }
    if (std::find(ownership.begin(), ownership.end(), 0U) != ownership.end()) {
        return constraintFailure(
            "invalid-constraint-order",
            std::to_string(fullCount),
            "Free and constrained DOFs must partition the complete full-DOF range.");
    }
    return Status::ok();
}

Result<SparseMatrix> extractBlock(
    const SparseMatrix& full,
    const std::vector<std::size_t>& rowDofs,
    const std::vector<std::size_t>& columnDofs) {
    const std::size_t absent = (std::numeric_limits<std::size_t>::max)();
    std::vector<std::size_t> localColumn(full.columns(), absent);
    for (std::size_t column = 0U; column < columnDofs.size(); ++column) {
        localColumn[columnDofs[column]] = column;
    }

    SparsePattern pattern;
    pattern.rowOffsets.reserve(rowDofs.size() + 1U);
    pattern.rowOffsets.push_back(0U);
    std::vector<CooContribution> contributions;
    contributions.reserve(full.values().size());
    for (std::size_t localRow = 0U;
         localRow < rowDofs.size();
         ++localRow) {
        const std::size_t fullRow = rowDofs[localRow];
        for (std::size_t position = full.rowOffsets()[fullRow];
             position < full.rowOffsets()[fullRow + 1U];
             ++position) {
            const std::size_t column =
                localColumn[full.columnIndices()[position]];
            if (column == absent) {
                continue;
            }
            pattern.columnIndices.push_back(column);
            // One source CSR entry maps to one block slot, so exact numeric
            // values and structural zeros survive without a new reduction.
            contributions.push_back({
                localRow,
                column,
                full.values()[position],
                localRow,
                position});
        }
        pattern.rowOffsets.push_back(pattern.columnIndices.size());
    }
    return SparseMatrix::fromCoo(
        rowDofs.size(),
        columnDofs.size(),
        std::move(contributions),
        pattern);
}

void requireDofOrder(const DofManager& dofs) {
    if (!validateDofOrder(dofs).isOk()) {
        throw std::invalid_argument{
            "DofManager constraint dimensions or order are invalid."};
    }
}

} // namespace

Result<PartitionedStiffness> EssentialConstraints::partition(
    const SparseMatrix& full,
    const DofManager& dofs) {
    const Status matrixStatus = full.validate();
    if (!matrixStatus.isOk()) {
        return Result<PartitionedStiffness>::failure(matrixStatus);
    }
    if (full.rows() != full.columns() ||
        full.rows() != dofs.fullDofCount()) {
        return Result<PartitionedStiffness>::failure(constraintFailure(
            "invalid-constraint-dimensions",
            std::to_string(full.rows()) + "x" +
                std::to_string(full.columns()),
            "Full stiffness must be square and match the DofManager full dimension."));
    }
    const Status dofStatus = validateDofOrder(dofs);
    if (!dofStatus.isOk()) {
        return Result<PartitionedStiffness>::failure(dofStatus);
    }

    auto kff = extractBlock(full, dofs.freeDofs(), dofs.freeDofs());
    if (!kff.hasValue()) {
        return Result<PartitionedStiffness>::failure(kff.status());
    }
    auto kfc = extractBlock(full, dofs.freeDofs(), dofs.constrainedDofs());
    if (!kfc.hasValue()) {
        return Result<PartitionedStiffness>::failure(kfc.status());
    }
    auto kcf = extractBlock(full, dofs.constrainedDofs(), dofs.freeDofs());
    if (!kcf.hasValue()) {
        return Result<PartitionedStiffness>::failure(kcf.status());
    }
    auto kcc = extractBlock(
        full, dofs.constrainedDofs(), dofs.constrainedDofs());
    if (!kcc.hasValue()) {
        return Result<PartitionedStiffness>::failure(kcc.status());
    }

    return Result<PartitionedStiffness>::success({
        std::move(kff.value()),
        std::move(kfc.value()),
        std::move(kcf.value()),
        std::move(kcc.value())});
}

Vector EssentialConstraints::gatherFree(
    const Vector& full,
    const DofManager& dofs) {
    requireDofOrder(dofs);
    if (full.size() != dofs.fullDofCount()) {
        throw std::invalid_argument{
            "Full vector size must match the DofManager full dimension."};
    }
    Vector reduced{dofs.freeDofCount()};
    for (std::size_t equation = 0U;
         equation < dofs.freeDofs().size();
         ++equation) {
        reduced[equation] = full[dofs.freeDofs()[equation]];
    }
    return reduced;
}

Vector EssentialConstraints::gatherConstrained(
    const Vector& full,
    const DofManager& dofs) {
    requireDofOrder(dofs);
    if (full.size() != dofs.fullDofCount()) {
        throw std::invalid_argument{
            "Full vector size must match the DofManager full dimension."};
    }
    Vector reduced{dofs.constrainedDofCount()};
    for (std::size_t index = 0U;
         index < dofs.constrainedDofs().size();
         ++index) {
        reduced[index] = full[dofs.constrainedDofs()[index]];
    }
    return reduced;
}

Vector EssentialConstraints::reconstructFull(
    const Vector& freeValues,
    const Vector& constrainedValues,
    const DofManager& dofs) {
    requireDofOrder(dofs);
    if (freeValues.size() != dofs.freeDofCount() ||
        constrainedValues.size() != dofs.constrainedDofCount()) {
        throw std::invalid_argument{
            "Reduced vector sizes must match the DofManager order."};
    }

    Vector full{dofs.fullDofCount()};
    for (std::size_t equation = 0U;
         equation < dofs.freeDofs().size();
         ++equation) {
        full[dofs.freeDofs()[equation]] = freeValues[equation];
    }
    // Preserve caller-supplied dc exactly; nonzero prescribed displacement is
    // never replaced with an implicit homogeneous constraint.
    for (std::size_t index = 0U;
         index < dofs.constrainedDofs().size();
         ++index) {
        full[dofs.constrainedDofs()[index]] = constrainedValues[index];
    }
    return full;
}

} // namespace fesa
