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

#include "mfem/linalg/vector.hpp"
#include "mfem/linalg/densemat.hpp"
#include "serif/utils/checks/finite.hpp"

namespace serif::mapping {
    using namespace serif::utils::errors;
    void map_hdiv_flux_to_physical(const MappingPointContext &context, const mfem::Vector &reference_flux, mfem::Vector &physical_flux) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(reference_flux.Size() == context.mapping_jacobian.Width(), "Reference flux vector has incompatible dimension with the mapping jacobian ({} <=> {}).", reference_flux.Size(), context.mapping_jacobian.Width());

        physical_flux.SetSize(reference_flux.Size());
        context.mapping_jacobian.Mult(reference_flux, physical_flux);
        physical_flux /= context.mapping_determinate;
    }

    void map_physical_flux_to_hdiv_reference(const MappingPointContext &context, const mfem::Vector &physical_flux, mfem::Vector &reference_flux) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(physical_flux.Size() == context.inverse_mapping_jacobian.Height(), "Physical flux vector has incompatible dimension with the inverse mapping jacobian ({} <=> {}).", physical_flux.Size(), context.inverse_mapping_jacobian.Height());

        reference_flux.SetSize(physical_flux.Size());
        context.inverse_mapping_jacobian.Mult(physical_flux, reference_flux);
        reference_flux /= context.mapping_determinate;
    }

    void map_reference_gradient_to_physical(const MappingPointContext &context, const mfem::Vector &reference_gradient, mfem::Vector &physical_gradient) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(reference_gradient.Size() == context.inverse_mapping_jacobian.Height(), "Reference gradient vector has incompatible dimension with the inverse mapping jacobian ({} <=> {}).", reference_gradient.Size(), context.inverse_mapping_jacobian.Height());

        physical_gradient.SetSize(reference_gradient.Size());
        context.inverse_mapping_jacobian.Mult(reference_gradient, physical_gradient);
    }

    void map_physical_gradient_to_reference(const MappingPointContext &context, const mfem::Vector &physical_gradient, mfem::Vector &reference_gradient) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(physical_gradient.Size() == context.mapping_jacobian.Height(), "Physical gradient vector has incompatible dimension with the mapping jacobian ({} <=> {}).", physical_gradient.Size(), context.mapping_jacobian.Height());

        reference_gradient.SetSize(physical_gradient.Size());
        context.mapping_jacobian.MultTranspose(physical_gradient, reference_gradient);
    }

    void map_reference_vector_gradient_to_physical(const MappingPointContext &context, const mfem::DenseMatrix &reference_vector_gradient, mfem::DenseMatrix &physical_vector_gradient) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(reference_vector_gradient.Width() == context.inverse_mapping_jacobian.Height(), "Reference vector gradient matrix has incompatible dimension with the inverse mapping jacobian ({} <=> {}).", reference_vector_gradient.Width(), context.inverse_mapping_jacobian.Height());

        physical_vector_gradient.SetSize(reference_vector_gradient.Height(), context.inverse_mapping_jacobian.Width());
        mfem::Mult(reference_vector_gradient, context.inverse_mapping_jacobian, physical_vector_gradient);
    }

    void map_physical_vector_gradient_to_reference(const MappingPointContext &context, const mfem::DenseMatrix &physical_vector_gradient, mfem::DenseMatrix &reference_vector_gradient) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(physical_vector_gradient.Width() == context.mapping_jacobian.Height(), "Physical vector gradient matrix has incompatible dimension with the mapping jacobian ({} <=> {}).", physical_vector_gradient.Width(), context.mapping_jacobian.Height());

        reference_vector_gradient.SetSize(physical_vector_gradient.Height(), context.mapping_jacobian.Width());
        mfem::Mult(physical_vector_gradient, context.mapping_jacobian, reference_vector_gradient);
    }

    double map_hdiv_divergence_to_physical(const MappingPointContext &context, double reference_divergence) {
        return reference_divergence / context.mapping_determinate;
    }

    void compute_hdiv_mass_tensor(const MappingPointContext &context, mfem::DenseMatrix &mass_tensor) {
        const int dimension = context.mapping_jacobian.Height();
        serif_assert<SERiFErrorCode::JACOBIAN_SHAPE_ERROR>((context.mapping_jacobian.Width() == dimension) && (context.mapping_jacobian.Height() == dimension), "Mapping jacobian is not square ({} x {}). It is expected to have shape {}x{}", context.mapping_jacobian.Height(), context.mapping_jacobian.Width(), dimension, dimension);
        serif_assert<SERiFErrorCode::INVERTED_JACOBIAN>(context.mapping_determinate > 0.0, "Mapping jacobian determinant is non-positive ({}). The mapping may be inverted or degenerate.", context.mapping_determinate);

        mass_tensor.SetSize(dimension, dimension);
        mfem::MultAtB(context.mapping_jacobian, context.mapping_jacobian, mass_tensor);
        mass_tensor *= 1/context.mapping_determinate;
    }

    void compute_scalar_diffusion_tensor(const MappingPointContext &context, mfem::DenseMatrix &diffusion_tensor) {
        const int dimension = context.mapping_jacobian.Height();
        serif_assert<SERiFErrorCode::JACOBIAN_SHAPE_ERROR>((context.inverse_mapping_jacobian.Width() == dimension) && (context.inverse_mapping_jacobian.Height() == dimension), "Inverse mapping jacobian is not square ({} x {}). It is expected to have shape {}x{}", context.inverse_mapping_jacobian.Height(), context.inverse_mapping_jacobian.Width(), dimension, dimension);
        serif_assert<SERiFErrorCode::INVERTED_JACOBIAN>(context.mapping_determinate > 0.0, "Mapping jacobian determinant is non-positive ({}). The mapping may be inverted or degenerate.", context.mapping_determinate);

        diffusion_tensor.SetSize(dimension, dimension);
        mfem::MultAtB(context.inverse_mapping_jacobian, context.inverse_mapping_jacobian, diffusion_tensor);
        diffusion_tensor *= context.mapping_determinate;
    }

    void map_hcurl_field_to_physical(const MappingPointContext &context, const mfem::Vector &reference_field, mfem::Vector &physical_field) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(reference_field.Size() == context.inverse_mapping_jacobian.Height(), "Reference field vector has incompatible dimension with the inverse mapping jacobian ({} <=> {}).", reference_field.Size(), context.inverse_mapping_jacobian.Height());

        physical_field.SetSize(reference_field.Size());
        context.inverse_mapping_jacobian.MultTranspose(reference_field, physical_field);
    }

    void map_physical_field_to_hcurl_reference(const MappingPointContext &context, const mfem::Vector &physical_field, mfem::Vector &reference_field) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(physical_field.Size() == context.mapping_jacobian.Height(), "Physical field vector has incompatible dimension with the mapping jacobian ({} <=> {}).", physical_field.Size(), context.mapping_jacobian.Height());

        reference_field.SetSize(physical_field.Size());
        context.mapping_jacobian.MultTranspose(physical_field, reference_field);
    }

    void map_hcurl_curl_to_physical(const MappingPointContext &context, const mfem::Vector &reference_curl, mfem::Vector &physical_curl) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(reference_curl.Size() == context.mapping_jacobian.Width(), "Reference curl vector has incompatible dimension with the mapping jacobian ({} <=> {}).", reference_curl.Size(), context.mapping_jacobian.Height());

        physical_curl.SetSize(reference_curl.Size());
        context.mapping_jacobian.Mult(reference_curl, physical_curl);
        physical_curl /= context.mapping_determinate;
    }

    void map_physical_curl_to_hcurl_reference(const MappingPointContext &context, const mfem::Vector &physical_curl, mfem::Vector &reference_curl) {
        serif_assert<SERiFErrorCode::INVALID_DIMENSION>(physical_curl.Size() == context.inverse_mapping_jacobian.Width(), "Physical curl vector has incompatible dimension with the inverse mapping jacobian ({} <=> {}).", physical_curl.Size(), context.inverse_mapping_jacobian.Height());

        reference_curl.SetSize(physical_curl.Size());
        context.inverse_mapping_jacobian.Mult(physical_curl, reference_curl);
        reference_curl *= context.mapping_determinate;
    }

    void compute_hdiv_mass_tensor_variation(const MappingPointContext &context, const MappingPointVariation &variation, mfem::DenseMatrix &mass_tensor_variation) {
        const double determinate = context.mapping_determinate;
        const double determinate_variation = variation.mapping_determinate_variation;
        const int dimension = context.inverse_mapping_jacobian.Width();

        mass_tensor_variation.SetSize(dimension, dimension);

        serif_assert<SERiFErrorCode::INVERTED_JACOBIAN>(determinate > 0.0, "Mapping jacobian determinant is non-positive ({}). The mapping may be inverted or degenerate.", determinate);
        serif_assert<SERiFErrorCode::NON_FINITE_DETERMINATE>(utils::checks::is_finite(determinate), "Mapping jacobian determinant is non-finite ({}). The mapping may be inverted or degenerate.", determinate);

        const mfem::DenseMatrix& jacobian = context.mapping_jacobian;
        const mfem::DenseMatrix& jacobian_variation = variation.mapping_jacobian_variation;
        const double inverse_determinate = 1.0 / determinate;
        const double determinate_scale = determinate_variation * inverse_determinate;

        for (int row = 0; row < dimension; ++row) {
            for (int col = 0; col < dimension; ++col) {
                double gram{0.0};
                double gram_variation{0.0};

                for (int inner = 0; inner < dimension; ++inner) {
                    gram += jacobian(inner, row) * jacobian(inner, col);
                    gram_variation += jacobian(inner, row) * jacobian_variation(inner, col) + jacobian_variation(inner, row) * jacobian(inner, col);
                }
                mass_tensor_variation(row, col) = inverse_determinate * (gram_variation - determinate_scale * gram);
            }
        }
    }
}