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."
        );

        mfem::DenseMatrix determinant_correction(dimension, dimension);
        ComputeHDivMassTensor(context, determinant_correction);
        determinant_correction *= determinant_variation / determinant;

        mfem::DenseMatrix right_jacobian_variation(dimension, dimension);
        mfem::MultAtB(
            context.mapping_jacobian, variation.mapping_jacobian_variation,
            right_jacobian_variation
        );
        mfem::MultAtB(
            variation.mapping_jacobian_variation, context.mapping_jacobian,
            mass_tensor_variation
        );

        mass_tensor_variation += right_jacobian_variation;
        mass_tensor_variation *= 1 / determinant;
        mass_tensor_variation -= determinant_correction;
    }
} // namespace mean_field::mapping