#include "serif/solver/normalization/detail/packing.hpp"

#include <cstddef>
#include <utility>
#include <vector>

#include "serif/utils/error/errors.hpp"

namespace serif::solver::normalization::detail
{
    namespace
    {
        /** @brief Preserve the declared block order while checking the total before allocating packed storage. */
        [[nodiscard]] std::expected<std::vector<double>, std::string> pack_factors(std::span<const std::span<const double>> coordinates)
        {
            using namespace utils::errors;

            std::size_t count = 0;
            for (const std::span<const double> factors : coordinates)
            {
                if (factors.size() > std::vector<double>{}.max_size() - count)
                {
                    return std::unexpected(format_error_message<SERiFErrorCode::INVALID_DOF_COUNT>("The packed normalization layout exceeds the supported vector size. Reduce its coordinate counts."));
                }
                count += factors.size();
            }

            std::vector<double> packed;
            packed.reserve(count);
            for (const std::span<const double> factors : coordinates)
            {
                packed.insert(packed.end(), factors.begin(), factors.end());
            }
            return packed;
        }
    }

    std::expected<DiagonalNormalization, std::string> pack_normalization(std::span<const std::span<const double>> states, std::span<const std::span<const double>> residuals)
    {
        auto state_factors = pack_factors(states);
        if (!state_factors)
        {
            return std::unexpected(state_factors.error());
        }
        auto residual_factors = pack_factors(residuals);
        if (!residual_factors)
        {
            return std::unexpected(residual_factors.error());
        }
        return make_diagonal_normalization(std::move(*state_factors), std::move(*residual_factors));
    }
}