module;

#include <algorithm>
#include <cstddef>
#include <mfem.hpp>
#include <stdexcept>

module mean_field;
import :mapping.prepared_cache;
import :mapping.types;

namespace mean_field::mapping {
    namespace {
        void pack_vector(
            double *&destination,
            const mfem::Vector &vector,
            const int dimension
        ) {
            if (vector.Size() != dimension)
                throw std::invalid_argument("Prepared mapping vector dimension mismatch.");
            std::copy_n(vector.HostRead(), dimension, destination);
            destination += dimension;
        }

        void pack_matrix(
            double *&destination,
            const mfem::DenseMatrix &matrix,
            const int dimension
        ) {
            if (matrix.Height() != dimension || matrix.Width() != dimension)
                throw std::invalid_argument("Prepared mapping matrix dimension mismatch.");
            std::copy_n(matrix.HostRead(), dimension * dimension, destination);
            destination += dimension * dimension;
        }

        void unpack_vector(
            const double *&source,
            mfem::Vector &vector,
            const int dimension
        ) {
            vector.SetSize(dimension);
            std::copy_n(source, dimension, vector.HostWrite());
            source += dimension;
        }

        void unpack_matrix(
            const double *&source,
            mfem::DenseMatrix &matrix,
            const int dimension
        ) {
            matrix.SetSize(dimension);
            std::copy_n(source, dimension * dimension, matrix.HostWrite());
            source += dimension * dimension;
        }
    } // namespace

    void VolumeMappingCache::SetSize(
        const int point_count,
        const int dimension
    ) {
        if (point_count < 0 || dimension < 1 || dimension > 3)
            throw std::invalid_argument("Prepared mapping storage requires nonnegative point count and dimension 1-3.");
        const int stride = 3 * dimension + 4 * dimension * dimension + 4;
        m_data.resize(static_cast<std::size_t>(point_count) * stride);
        m_point_count  = point_count;
        m_dimension    = dimension;
        m_point_stride = stride;
    }

    const double *VolumeMappingCache::GetPointData(const int point) const {
        if (point < 0 || point >= m_point_count)
            throw std::out_of_range("Prepared mapping quadrature point is out of range.");
        return m_data.data() + static_cast<std::size_t>(point) * m_point_stride;
    }

    void VolumeMappingCache::Store(
        const int point,
        const VolumeMappingContext &context
    ) {
        // Validate the index through the same checked accessor used by readers.
        (void)GetPointData(point);
        double *data = m_data.data() + static_cast<std::size_t>(point) * m_point_stride;
        pack_vector(data, context.mapping.reference_position, m_dimension);
        pack_vector(data, context.mapping.displaced_position, m_dimension);
        pack_vector(data, context.mapping.physical_position, m_dimension);
        pack_matrix(data, context.mapping.displacement_jacobian, m_dimension);
        pack_matrix(data, context.mapping.mapping_jacobian, m_dimension);
        pack_matrix(data, context.mapping.inverse_mapping_jacobian, m_dimension);
        pack_matrix(data, context.quadrature.J_inv, m_dimension);
        *data++ = context.mapping.mapping_determinant;
        *data++ = context.mapping.compactified ? 1.0 : 0.0;
        *data++ = context.quadrature.detJ;
        *data   = context.quadrature.weight;
    }

    void VolumeMappingCache::Load(
        const int point,
        VolumeMappingContext &context
    ) const {
        const double *data = GetPointData(point);
        unpack_vector(data, context.mapping.reference_position, m_dimension);
        unpack_vector(data, context.mapping.displaced_position, m_dimension);
        unpack_vector(data, context.mapping.physical_position, m_dimension);
        unpack_matrix(data, context.mapping.displacement_jacobian, m_dimension);
        unpack_matrix(data, context.mapping.mapping_jacobian, m_dimension);
        unpack_matrix(data, context.mapping.inverse_mapping_jacobian, m_dimension);
        unpack_matrix(data, context.quadrature.J_inv, m_dimension);
        context.mapping.mapping_determinant = *data++;
        context.mapping.compactified        = *data++ != 0.0;
        context.quadrature.detJ             = *data++;
        context.quadrature.weight           = *data;
    }

    void VolumeMappingCache::LoadInverseJacobian(
        const int point,
        mfem::DenseMatrix &inverse
    ) const {
        const double *data = GetPointData(point) + 3 * m_dimension + 3 * m_dimension * m_dimension;
        unpack_matrix(data, inverse, m_dimension);
    }

    int VolumeMappingCache::GetPointCount() const {
        return m_point_count;
    }
    int VolumeMappingCache::GetDimension() const {
        return m_dimension;
    }
} // namespace mean_field::mapping