#include <fesa/constraints/essential_bc.hpp>

#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace fesa {
namespace {

std::optional<std::string> validate_system(
    const EquationSystem& system) {
    const SymmetricCsr& matrix = system.stiffness;
    if (system.force.size() != matrix.order) {
        return "Force size must equal the matrix order.";
    }
    if (matrix.order >
        static_cast<std::size_t>(
            std::numeric_limits<std::int32_t>::max())) {
        return "Matrix order exceeds the 32-bit CSR index range.";
    }
    if (matrix.row_offsets.size() != matrix.order + 1 ||
        matrix.row_offsets.empty() ||
        matrix.row_offsets.front() != 0) {
        return "CSR row offsets must contain order + 1 entries "
               "starting at zero.";
    }
    if (matrix.column_indices.size() != matrix.values.size()) {
        return "CSR column and value counts must match.";
    }

    std::int32_t previous_offset = 0;
    for (const std::int32_t offset : matrix.row_offsets) {
        if (offset < previous_offset || offset < 0 ||
            static_cast<std::size_t>(offset) >
                matrix.column_indices.size()) {
            return "CSR row offsets must be nondecreasing and in range.";
        }
        previous_offset = offset;
    }
    if (static_cast<std::size_t>(matrix.row_offsets.back()) !=
        matrix.column_indices.size()) {
        return "The final CSR row offset must equal the entry count.";
    }

    for (std::size_t row = 0; row < matrix.order; ++row) {
        std::int32_t previous_column = -1;
        const std::size_t begin =
            static_cast<std::size_t>(matrix.row_offsets[row]);
        const std::size_t end =
            static_cast<std::size_t>(matrix.row_offsets[row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const std::int32_t column = matrix.column_indices[entry];
            if (column < static_cast<std::int32_t>(row) ||
                column >= static_cast<std::int32_t>(matrix.order) ||
                column <= previous_column) {
                return "CSR rows must contain sorted unique "
                       "upper-triangle columns.";
            }
            previous_column = column;
        }
    }
    return std::nullopt;
}

Diagnostic equation_error(std::string code, std::string message) {
    return {
        DiagnosticStage::equation,
        Severity::error,
        std::move(code),
        std::move(message),
        std::nullopt,
    };
}

std::int32_t csr_index(const std::size_t value) {
    if (value >
        static_cast<std::size_t>(
            std::numeric_limits<std::int32_t>::max())) {
        throw std::overflow_error{
            "Reduced system exceeds the 32-bit CSR index range."};
    }
    return static_cast<std::int32_t>(value);
}

}  // namespace

ConstraintResult eliminate_essential_bcs(
    const EquationSystem& original,
    const DofManager& dofs,
    const std::span<const PrescribedDof> prescribed) {
    ConstraintResult result;
    if (const auto error = validate_system(original);
        error.has_value()) {
        result.diagnostics.push_back(equation_error(
            "equation.invalid_system", *error));
        return result;
    }
    if (original.stiffness.order != dofs.full_dof_count()) {
        result.diagnostics.push_back(equation_error(
            "equation.dof_count_mismatch",
            "Matrix order must equal the DofManager full DOF count."));
        return result;
    }

    const std::size_t full_count = original.stiffness.order;
    std::vector<bool> constrained(full_count, false);
    std::vector<double> prescribed_full(full_count, 0.0);
    for (const PrescribedDof& value : prescribed) {
        if (value.dof < 1 || value.dof > 6 ||
            !std::isfinite(value.value)) {
            result.diagnostics.push_back(equation_error(
                "equation.invalid_prescribed_dof",
                "Prescribed DOF data must contain a DOF in [1, 6] "
                "and a finite value."));
            continue;
        }

        const DofAddress address{
            value.node,
            static_cast<NodeDof>(value.dof - std::uint8_t{1}),
        };
        std::size_t full = 0;
        std::optional<double> owned_value;
        try {
            full = dofs.full_dof(address);
            owned_value = dofs.prescribed_value(address);
        } catch (const std::out_of_range&) {
            result.diagnostics.push_back(equation_error(
                "equation.invalid_prescribed_dof",
                "Prescribed DOF references a node absent from the "
                "DofManager."));
            continue;
        }

        if (constrained[full]) {
            result.diagnostics.push_back(equation_error(
                "equation.conflicting_prescribed_dof",
                "A full DOF has more than one prescribed value."));
            continue;
        }
        constrained[full] = true;

        if (!owned_value.has_value() ||
            *owned_value != value.value) {
            result.diagnostics.push_back(equation_error(
                "equation.prescribed_dof_mismatch",
                "Prescribed DOFs must match the state owned by the "
                "DofManager."));
            continue;
        }
        prescribed_full[full] = *owned_value;
    }

    const std::size_t constrained_count =
        full_count - dofs.free_equation_count();
    if (prescribed.size() != constrained_count) {
        result.diagnostics.push_back(equation_error(
            "equation.prescribed_dof_mismatch",
            "Prescribed DOF count must match the DofManager "
            "constraint state."));
    }
    if (!result.diagnostics.empty()) {
        return result;
    }

    std::vector<std::size_t> full_to_free(full_count, full_count);
    std::vector<std::size_t> free_to_full;
    free_to_full.reserve(dofs.free_equation_count());
    for (std::size_t full = 0; full < full_count; ++full) {
        if (!constrained[full]) {
            full_to_free[full] = free_to_full.size();
            free_to_full.push_back(full);
        }
    }
    if (free_to_full.size() != dofs.free_equation_count()) {
        result.diagnostics.push_back(equation_error(
            "equation.prescribed_dof_mismatch",
            "Prescribed DOFs do not identify the DofManager constraints."));
        return result;
    }

    std::vector<double> reduced_force(
        free_to_full.size(), 0.0);
    for (std::size_t free = 0; free < free_to_full.size(); ++free) {
        reduced_force[free] = original.force[free_to_full[free]];
    }

    SymmetricCsr reduced{
        free_to_full.size(),
        std::vector<std::int32_t>(free_to_full.size() + 1, 0),
        {},
        {},
    };
    reduced.column_indices.reserve(
        original.stiffness.column_indices.size());
    reduced.values.reserve(original.stiffness.values.size());

    for (std::size_t full_row = 0; full_row < full_count; ++full_row) {
        const std::size_t begin = static_cast<std::size_t>(
            original.stiffness.row_offsets[full_row]);
        const std::size_t end = static_cast<std::size_t>(
            original.stiffness.row_offsets[full_row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const std::size_t full_column =
                static_cast<std::size_t>(
                    original.stiffness.column_indices[entry]);
            const double stiffness = original.stiffness.values[entry];

            if (constrained[full_row]) {
                if (!constrained[full_column]) {
                    reduced_force[full_to_free[full_column]] -=
                        stiffness * prescribed_full[full_row];
                }
            } else if (constrained[full_column]) {
                reduced_force[full_to_free[full_row]] -=
                    stiffness * prescribed_full[full_column];
            } else {
                reduced.column_indices.push_back(
                    csr_index(full_to_free[full_column]));
                reduced.values.push_back(stiffness);
            }
        }
        if (!constrained[full_row]) {
            reduced.row_offsets[full_to_free[full_row] + 1] =
                csr_index(reduced.column_indices.size());
        }
    }

    result.reduced_system = ReducedSystem{
        std::move(reduced),
        std::move(reduced_force),
        std::move(free_to_full),
        std::move(prescribed_full),
    };
    return result;
}

std::vector<double> recover_reaction(
    const EquationSystem& original,
    const std::span<const double> full_displacement) {
    if (const auto error = validate_system(original);
        error.has_value()) {
        throw std::invalid_argument{*error};
    }
    if (full_displacement.size() != original.stiffness.order) {
        throw std::invalid_argument{
            "Full displacement size must equal the matrix order."};
    }

    std::vector<double> reaction = original.force;
    for (double& value : reaction) {
        value = -value;
    }
    for (std::size_t row = 0;
         row < original.stiffness.order;
         ++row) {
        const std::size_t begin = static_cast<std::size_t>(
            original.stiffness.row_offsets[row]);
        const std::size_t end = static_cast<std::size_t>(
            original.stiffness.row_offsets[row + 1]);
        for (std::size_t entry = begin; entry < end; ++entry) {
            const std::size_t column =
                static_cast<std::size_t>(
                    original.stiffness.column_indices[entry]);
            const double stiffness = original.stiffness.values[entry];
            reaction[row] +=
                stiffness * full_displacement[column];
            if (column != row) {
                reaction[column] +=
                    stiffness * full_displacement[row];
            }
        }
    }
    return reaction;
}

}  // namespace fesa
