module;
#include <mfem.hpp>

module mean_field;
import :mapping.types;

namespace mean_field::mapping {
    void MapHDivFluxToPhysical(
        const MappingPointContext &context,
        const mfem::Vector &reference_flux,
        mfem::Vector &physical_flux
    ) {
        MFEM_VERIFY(
            reference_flux.Size() == context.mapping_jacobian.Width(),
            "The reference H(div) flux has the wrong dimension."
        );

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

    void MapPhysicalFluxToHDivReference(
        const MappingPointContext &context,
        const mfem::Vector &physical_flux,
        mfem::Vector &reference_flux
    ) {
        MFEM_VERIFY(
            physical_flux.Size() == context.inverse_mapping_jacobian.Width(),
            "The physical H(div) flux has the wrong dimension."
        );

        reference_flux.SetSize(physical_flux.Size());
        context.inverse_mapping_jacobian.Mult(physical_flux, reference_flux);
        reference_flux *= context.mapping_determinant;
    }

    void MapReferenceGradientToPhysical(
        const MappingPointContext &context,
        const mfem::Vector &reference_gradient,
        mfem::Vector &physical_gradient
    ) {
        MFEM_VERIFY(
            reference_gradient.Size() == context.inverse_mapping_jacobian.Height(),
            "The reference scalar gradient has the wrong dimension."
        );

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

    void MapPhysicalGradientToReference(
        const MappingPointContext &context,
        const mfem::Vector &physical_gradient,
        mfem::Vector &reference_gradient
    ) {
        MFEM_VERIFY(
            physical_gradient.Size() == context.mapping_jacobian.Height(),
            "The physical scalar gradient has the wrong dimension."
        );

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

    void MapReferenceVectorGradientToPhysical(
        const MappingPointContext &context,
        const mfem::DenseMatrix &reference_gradient,
        mfem::DenseMatrix &physical_gradient
    ) {
        MFEM_VERIFY(
            reference_gradient.Width() == context.inverse_mapping_jacobian.Height(),
            "The reference vector gradient has the wrong dimension."
        );

        physical_gradient.SetSize(reference_gradient.Height(), context.inverse_mapping_jacobian.Width());
        mfem::Mult(reference_gradient, context.inverse_mapping_jacobian, physical_gradient);
    }

    void MapPhysicalVectorGradientToReference(
        const MappingPointContext &context,
        const mfem::DenseMatrix &physical_gradient,
        mfem::DenseMatrix &reference_gradient
    ) {
        MFEM_VERIFY(
            physical_gradient.Width() == context.mapping_jacobian.Height(),
            "The physical vector gradient has the wrong dimension."
        );

        reference_gradient.SetSize(physical_gradient.Height(), context.mapping_jacobian.Width());
        mfem::Mult(physical_gradient, context.mapping_jacobian, reference_gradient);
    }

    double MapHDivDivergenceToPhysical(
        const MappingPointContext &context,
        const double reference_divergence
    ) {
        return reference_divergence / context.mapping_determinant;
    }

    void ComputeHDivMassTensor(
        const MappingPointContext &context,
        mfem::DenseMatrix &mass_tensor
    ) {
        const int dimension = context.mapping_jacobian.Height();

        MFEM_VERIFY(context.mapping_jacobian.Width() == dimension, "The mapping Jacobian must be square.");
        MFEM_VERIFY(context.mapping_determinant > 0.0, "The mapping determinant must be positive.");

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

    void ComputeScalarDiffusionTensor(
        const MappingPointContext &context,
        mfem::DenseMatrix &diffusion_tensor
    ) {
        const int dimension = context.inverse_mapping_jacobian.Height();

        MFEM_VERIFY(
            context.inverse_mapping_jacobian.Width() == dimension, "The inverse mapping Jacobian must be square."
        );
        MFEM_VERIFY(context.mapping_determinant > 0.0, "The mapping determinant must be positive.");

        diffusion_tensor.SetSize(dimension, dimension);
        mfem::MultABt(context.inverse_mapping_jacobian, context.inverse_mapping_jacobian, diffusion_tensor);
        diffusion_tensor *= context.mapping_determinant;
    }

    void MapHCurlFieldToPhysical(
        const MappingPointContext &context,
        const mfem::Vector &reference_field,
        mfem::Vector &physical_field
    ) {
        MFEM_VERIFY(
            reference_field.Size() == context.inverse_mapping_jacobian.Height(),
            "The reference H(curl) field has the wrong dimension."
        );

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

    void MapPhysicalFieldToHCurlReference(
        const MappingPointContext &context,
        const mfem::Vector &physical_field,
        mfem::Vector &reference_field
    ) {
        MFEM_VERIFY(
            physical_field.Size() == context.mapping_jacobian.Height(),
            "The physical H(curl) field has the wrong dimension."
        );

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

    void MapHCurlCurlToPhysical(
        const MappingPointContext &context,
        const mfem::Vector &reference_curl,
        mfem::Vector &physical_curl
    ) {
        MFEM_VERIFY(
            reference_curl.Size() == context.mapping_jacobian.Width(),
            "The reference H(curl) curl has the wrong dimension."
        );

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

    void MapPhysicalCurlToHCurlReference(
        const MappingPointContext &context,
        const mfem::Vector &physical_curl,
        mfem::Vector &reference_curl
    ) {
        MFEM_VERIFY(
            physical_curl.Size() == context.inverse_mapping_jacobian.Width(),
            "The physical H(curl) curl has the wrong dimension."
        );

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

    // TODO: Investigate these
    void ComputeHCurlMassTensor(
        const MappingPointContext &context,
        mfem::DenseMatrix &mass_tensor
    ) {
        ComputeScalarDiffusionTensor(context, mass_tensor);
    }

    void ComputeHCurlCurlTensor(
        const MappingPointContext &context,
        mfem::DenseMatrix &curl_tensor
    ) {
        ComputeHDivMassTensor(context, curl_tensor);
    }

    void ComputeHDivMassTensorVariation(
        const MappingPointContext &context,
        const MappingPointVariation &variation,
        mfem::DenseMatrix &mass_tensor_variation
    ) {
        const double determinant           = context.mapping_determinant;
        const double determinant_variation = variation.mapping_determinant_variation;
        const int dimension                = context.inverse_mapping_jacobian.Width();

        mass_tensor_variation.SetSize(dimension, dimension);

        MFEM_VERIFY(
            std::isfinite(determinant) && determinant > 0.0, "The mapping determinant must be positive and finite."
        );
        MFEM_VERIFY(std::isfinite(determinant_variation), "The mapping determinant variation must be finite.");

        const mfem::DenseMatrix &jacobian          = context.mapping_jacobian;
        const mfem::DenseMatrix &jacobianVariation = variation.mapping_jacobian_variation;
        const double inverseDeterminant            = 1.0 / determinant;
        const double determinantScale              = determinant_variation * inverseDeterminant;

        for (int row = 0; row < dimension; ++row) {
            for (int column = 0; column < dimension; ++column) {
                double gram{0.0};
                double gramVariation{0.0};
                for (int inner = 0; inner < dimension; ++inner) {
                    gram += jacobian(inner, row) * jacobian(inner, column);
                    gramVariation += jacobian(inner, row) * jacobianVariation(inner, column) +
                                     jacobianVariation(inner, row) * jacobian(inner, column);
                }
                mass_tensor_variation(row, column) = inverseDeterminant * (gramVariation - determinantScale * gram);
            }
        }
    }
} // namespace mean_field::mapping