module;
#include "mean_field.h"

export module mean_field:mapping.coefficients;
export import :mapping.domain_mapper;
export import :mapping.types;

export namespace mean_field::mapping {
    class MappedScalarCoefficient : public mfem::Coefficient {
    public:
        MappedScalarCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            Coefficient &coeff,
            COORDINATE_SPACE coord_space = COORDINATE_SPACE::PHYSICAL
        );

        double Eval(
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        ) override;

    private:
        static double eval_at_point(
            Coefficient &c,
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        );

    private:
        GridFunctionMappingEvaluator m_mapping;
        Coefficient &m_coeff;
        COORDINATE_SPACE m_coord_space;
    };

    class MappedDiffusionCoefficient : public mfem::MatrixCoefficient {
    public:
        MappedDiffusionCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            mfem::Coefficient &sigma,
            int dim
        );

        MappedDiffusionCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            MatrixCoefficient &sigma
        );

        void Eval(
            mfem::DenseMatrix &K,
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        ) override;

    private:
        GridFunctionMappingEvaluator m_mapping;
        mfem::Coefficient *m_scalar;
        MatrixCoefficient *m_tensor;
    };

    class MappedVectorCoefficient : public mfem::VectorCoefficient {
    public:
        MappedVectorCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            VectorCoefficient &coeff
        );

        void Eval(
            mfem::Vector &V,
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        ) override;

    private:
        GridFunctionMappingEvaluator m_mapping;
        VectorCoefficient &m_coeff;
    };

    class PhysicalPositionFunctionCoefficient : public mfem::Coefficient {
    public:
        using Func = std::function<double(const mfem::Vector &x)>;

        PhysicalPositionFunctionCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            Func f
        );

        double Eval(
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        ) override;

    private:
        Func m_f;
        GridFunctionMappingEvaluator m_mapping;
    };

    class MappedHDivMassCoefficient final : public mfem::MatrixCoefficient {
    public:
        MappedHDivMassCoefficient(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            const int dim
        );

        void Eval(
            mfem::DenseMatrix &matrix,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point
        ) override;

    private:
        GridFunctionMappingEvaluator m_mapping;
    };
} // namespace mean_field::mapping