module;
#include <mfem.hpp>

module mean_field;
import :mapping.types;

namespace mean_field::mapping {
    ///////////////////////////////
    /// MappedScalarCoefficient ///
    //////////////////////////////
    MappedScalarCoefficient::MappedScalarCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        Coefficient &coeff,
        const COORDINATE_SPACE coord_space
    )
        : m_mapping(mapper, displacement, compactification_coordinate),
          m_coeff(coeff),
          m_coord_space(coord_space) { };

    double MappedScalarCoefficient::Eval(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        T.SetIntPoint(&ip);
        double f_val = 0.0;

        switch (m_coord_space) {
        case COORDINATE_SPACE::PHYSICAL: {
            f_val             = eval_at_point(m_coeff, T, ip);
            VolumeMappingContext context;
            MFEM_VERIFY(
                m_mapping.EvaluateVolume(T, ip, context) == MappingStatus::valid,
                "Mapped scalar coefficient encountered an invalid mapping."
            );
            return f_val * std::abs(context.mapping.mapping_determinant);
        }
        case COORDINATE_SPACE::REFERENCE: {
            f_val = m_coeff.Eval(T, ip);
            return f_val;
        }
        }
    }

    double MappedScalarCoefficient::eval_at_point(
        Coefficient &c,
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        return c.Eval(T, ip);
    }

    //////////////////////////////////
    /// MappedDiffusionCoefficient ///
    //////////////////////////////////

    MappedDiffusionCoefficient::MappedDiffusionCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        mfem::Coefficient &sigma,
        const int dim
    )
        : MatrixCoefficient(dim),
          m_mapping(mapper, displacement, compactification_coordinate),
          m_scalar(&sigma),
          m_tensor(nullptr) { };

    MappedDiffusionCoefficient::MappedDiffusionCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        MatrixCoefficient &sigma
    )
        : MatrixCoefficient(sigma.GetHeight()),
          m_mapping(mapper, displacement, compactification_coordinate),
          m_scalar(nullptr),
          m_tensor(&sigma) { };

    void MappedDiffusionCoefficient::Eval(
        mfem::DenseMatrix &K,
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        const int dim = height;
        T.SetIntPoint(&ip);

        VolumeMappingContext context;
        MFEM_VERIFY(
            m_mapping.EvaluateVolume(T, ip, context) == MappingStatus::valid,
            "Mapped diffusion coefficient encountered an invalid mapping."
        );
        const mfem::DenseMatrix &JInv = context.mapping.inverse_mapping_jacobian;
        const double detJ = context.mapping.mapping_determinant;

        if (m_scalar) {
            const double sig_val = m_scalar->Eval(T, ip);
            mfem::MultABt(JInv, JInv, K);
            K *= sig_val * fabs(detJ);
        } else {
            mfem::DenseMatrix sig_mat(dim, dim);
            m_tensor->Eval(sig_mat, T, ip);

            mfem::DenseMatrix temp(dim, dim);
            Mult(JInv, sig_mat, temp);

            MultABt(temp, JInv, K);
            K *= fabs(detJ);
        }
    }

    ///////////////////////////////
    /// MappedVectorCoefficient ///
    ///////////////////////////////
    MappedVectorCoefficient::MappedVectorCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        VectorCoefficient &coeff
    )
        : VectorCoefficient(coeff.GetVDim()),
          m_mapping(mapper, displacement, compactification_coordinate),
          m_coeff(coeff) { };

    void MappedVectorCoefficient::Eval(
        mfem::Vector &V,
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        const int dim = vdim;
        T.SetIntPoint(&ip);

        VolumeMappingContext context;
        MFEM_VERIFY(
            m_mapping.EvaluateVolume(T, ip, context) == MappingStatus::valid,
            "Mapped vector coefficient encountered an invalid mapping."
        );
        const mfem::DenseMatrix &JInv = context.mapping.inverse_mapping_jacobian;
        const double detJ = context.mapping.mapping_determinant;

        mfem::Vector C_phys(dim);
        m_coeff.Eval(C_phys, T, ip);

        V.SetSize(dim);
        JInv.MultTranspose(C_phys, V);
        V *= fabs(detJ);
    }

    ///////////////////////////////////////////
    /// PhysicalPositionFunctionCoefficient ///
    ///////////////////////////////////////////
    PhysicalPositionFunctionCoefficient::PhysicalPositionFunctionCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        Func f // std::function<double(const mfem::Vector&)>
    )
        : m_f(std::move(f)),
          m_mapping(mapper, displacement, compactification_coordinate) { };

    double PhysicalPositionFunctionCoefficient::Eval(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        T.SetIntPoint(&ip);
        MappingPointContext context;
        MFEM_VERIFY(
            m_mapping.EvaluatePoint(T, ip, context) == MappingStatus::valid,
            "Physical-position coefficient encountered an invalid mapping."
        );
        return m_f(context.physical_position);
    }

    MappedHDivMassCoefficient::MappedHDivMassCoefficient(
        const DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        const int dim
    )
        : MatrixCoefficient(dim),
          m_mapping(mapper, displacement, compactification_coordinate) {
    }

    void MappedHDivMassCoefficient::Eval(
        mfem::DenseMatrix &matrix,
        mfem::ElementTransformation &transformation,
        const mfem::IntegrationPoint &integration_point
    ) {
        transformation.SetIntPoint(&integration_point);

        VolumeMappingContext context;
        MFEM_VERIFY(
            m_mapping.EvaluateVolume(transformation, integration_point, context) == MappingStatus::valid,
            "Mapped H(div) coefficient encountered an invalid mapping."
        );
        const mfem::DenseMatrix &map_jacobian = context.mapping.mapping_jacobian;
        const double map_determinant = context.mapping.mapping_determinant;

        MFEM_VERIFY(map_determinant > 0.0, "Domain mapping has a non-positive Jacobian determinant.");

        mfem::MultAtB(map_jacobian, map_jacobian, matrix);
        matrix *= 1.0 / std::abs(map_determinant);
    }
} // namespace mean_field::mapping