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 &map,
            mfem::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(
            mfem::Coefficient &c,
            mfem::ElementTransformation &T,
            const mfem::IntegrationPoint &ip
        );

    private:
        const DomainMapper &m_map;
        mfem::Coefficient &m_coeff;
        COORDINATE_SPACE m_coord_space;
    };

    class MappedDiffusionCoefficient : public mfem::MatrixCoefficient {
    public:
        MappedDiffusionCoefficient(
            const DomainMapper &map,
            mfem::Coefficient &sigma,
            int dim
        );

        MappedDiffusionCoefficient(
            const DomainMapper &map,
            mfem::MatrixCoefficient &sigma
        );

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

    private:
        const DomainMapper &m_map;
        mfem::Coefficient *m_scalar;
        mfem::MatrixCoefficient *m_tensor;
    };

    class MappedVectorCoefficient : public mfem::VectorCoefficient {
    public:
        MappedVectorCoefficient(
            const DomainMapper &map,
            mfem::VectorCoefficient &coeff
        );

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

    private:
        const DomainMapper &m_map;
        mfem::VectorCoefficient &m_coeff;
    };

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

        PhysicalPositionFunctionCoefficient(
            const DomainMapper &map,
            Func f
        );

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

    private:
        Func m_f;
        const DomainMapper &m_map;
    };

    class MappedHDivMassCoefficient final : public mfem::MatrixCoefficient {
    public:
        MappedHDivMassCoefficient(const DomainMapper& map, const int dim);

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

    };
}