#include <fesa/constraints/essential_bc.hpp>

#include <algorithm>
#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.";
    }
    if (!std::ranges::all_of(
            matrix.values, [](const double value) {
                return std::isfinite(value);
            }) ||
        !std::ranges::all_of(
            system.force, [](const double value) {
                return std::isfinite(value);
            })) {
        return "Matrix and force values must be finite.";
    }

    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) {
    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<std::size_t> full_to_free(full_count, full_count);
    std::vector<double> prescribed_full(full_count, 0.0);
    for (std::size_t full = 0; full < full_count; ++full) {
        const std::optional<std::size_t> equation =
            dofs.equation(full);
        if (equation.has_value()) {
            full_to_free[full] = *equation;
        } else {
            prescribed_full[full] =
                dofs.prescribed_value(full).value();
        }
    }

    std::vector<double> reduced_force(
        dofs.free_equation_count(), 0.0);
    for (std::size_t full = 0; full < full_count; ++full) {
        if (full_to_free[full] != full_count) {
            reduced_force[full_to_free[full]] = original.force[full];
        }
    }

    SymmetricCsr reduced{
        dofs.free_equation_count(),
        std::vector<std::int32_t>(
            dofs.free_equation_count() + 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 (full_to_free[full_row] == full_count) {
                if (full_to_free[full_column] != full_count) {
                    reduced_force[full_to_free[full_column]] -=
                        stiffness * prescribed_full[full_row];
                }
            } else if (full_to_free[full_column] == full_count) {
                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 (full_to_free[full_row] != full_count) {
            reduced.row_offsets[full_to_free[full_row] + 1] =
                csr_index(reduced.column_indices.size());
        }
    }
    if (!std::ranges::all_of(
            reduced_force, [](const double value) {
                return std::isfinite(value);
            })) {
        result.diagnostics.push_back(equation_error(
            "equation.nonfinite_result",
            "Essential BC elimination produced a nonfinite force."));
        return result;
    }

    result.reduced_system = ReducedSystem{
        std::move(reduced),
        std::move(reduced_force),
    };
    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."};
    }
    if (!std::ranges::all_of(
            full_displacement, [](const double value) {
                return std::isfinite(value);
            })) {
        throw std::invalid_argument{
            "Full displacement values must be finite."};
    }

    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];
            }
        }
    }
    if (!std::ranges::all_of(
            reaction, [](const double value) {
                return std::isfinite(value);
            })) {
        throw std::invalid_argument{
            "Reaction recovery produced a nonfinite value."};
    }
    return reaction;
}

}  // namespace fesa
