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

#include "mfem/linalg/densemat.hpp"
#include "mfem/linalg/vector.hpp"

namespace serif::mapping {
    using namespace utils::errors;
    namespace {
        // Replace all of these with a fast serialization class at some point. Probably in something like a serif/serialize module
        void pack_vector(double*& destination, const mfem::Vector& vector) {
            const int dim = vector.Size();
            std::copy_n(vector.HostRead(), dim, destination);
            destination += dim;
        }

        void pack_matrix(double*& destination, const mfem::DenseMatrix& matrix) {
            serif_assert<SERiFErrorCode::MATRIX_SHAPE_ERROR>(matrix.Height() == matrix.Width(), "Matrix must be square to pack");
            const int dim = matrix.Height();
            std::copy_n(matrix.HostRead(), dim * dim, destination);
            destination += dim * dim;
        }

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

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

    void VolumeMappingCache::set_size(int point_count, int dimension) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(point_count < 0 || dimension < 1 || dimension > 3, "Invalid point count or dimension for volume mapping cache. Point count: {}, dimension: {}", point_count, dimension);
        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;
    }

    void VolumeMappingCache::validate_point(const int point) const {
        serif_assert<SERiFErrorCode::OUT_OF_RANGE_ERROR>(point < 0 || point >= m_point_count, "Point index {} is out of range for volume mapping cache with {} points", point, m_point_count);
    }

    const double *VolumeMappingCache::get_point_data(const int point) const {
        validate_point(point);
        return m_data.data() + static_cast<std::size_t>(point) * m_point_stride;
    }

    void VolumeMappingCache::store(int point, VolumeMappingContext &context) {
        validate_point(point);
        double *data = m_data.data() + static_cast<std::size_t>(point) * m_point_stride;

        // What we build here is a FIFO data structure. we pack
        // all the raw data from the mapping context into a single contiguous array of doubles.
        // This allows us to avoid repeated allocations as the vector will grow to some appropriate size and then stay roughly there.
        pack_vector(data, context.mapping.reference_position);
        pack_vector(data, context.mapping.displaced_position);
        pack_vector(data, context.mapping.physical_position);

        pack_matrix(data, context.mapping.displacement_jacobian);
        pack_matrix(data, context.mapping.mapping_jacobian);
        pack_matrix(data, context.mapping.inverse_mapping_jacobian);
        pack_matrix(data, context.quadrature.J_element_to_physical_inverse);

        *data++ = context.mapping.mapping_determinate;
        *data++ = context.mapping.compactified ? 1.0 : 0.0; // its possible this could cause issues, have not run into but when unpacked I worry that exact equality to some double value could be a problem. If there are wierd issues this may be a place to look,
        *data++ = context.quadrature.mapping_determinate;
        *data   = context.quadrature.physical_integration_weight;
    }

    void VolumeMappingCache::load(int point, VolumeMappingContext &context) const {
        const double *data = get_point_data(point);

        // FIFO data structure so need to unpack in the same order as we packed.
        // TODO: Redo this with some sort of fast serialization class so we do not need to worry about ordering
        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_element_to_physical_inverse, m_dimension);

        context.mapping.mapping_determinate            = *data++;
        context.mapping.compactified                   = *data++ != 0.0; // See comment in store()
        context.quadrature.mapping_determinate         = *data++;
        context.quadrature.physical_integration_weight = *data;
    }

    void VolumeMappingCache::load_inverse_jacobian(int point, mfem::DenseMatrix &inverse) const {
        // Again this is unsafe, and really should be replaced. This was a quick cache scheme that solved the direct issue I had; however, it should be replaced
        // with something that has at least a modicum of safety built in (we are pretty good about this in other parts of SERiF but its about 8pm and
        // Im quite tired....
        const double *data = get_point_data(point) + 3 * m_dimension + 3 * m_dimension * m_dimension; // skip to the inverse jacobian, this is just pointer math...Think back to CS-101. Do that have that in the future, are we still learning about memory like this is a PDP-11?
        unpack_matrix(data, inverse, m_dimension);
    }

    int VolumeMappingCache::get_point_count() const noexcept {
        return m_point_count;
    }

    int VolumeMappingCache::get_dimension() const noexcept {
        return m_dimension;
    }
}