module;
#include <mfem.hpp>

module mean_field;
import :mapping.types;

namespace mean_field::mapping {
    ///////////////////////////////
    /// MappedScalarCoefficient ///
    //////////////////////////////
    MappedScalarCoefficient::MappedScalarCoefficient(
        const DomainMapper &map,
        Coefficient &coeff,
        const COORDINATE_SPACE coord_space
    )
        : m_map(map),
          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);
            const double detJ = m_map.ComputeDetJ(T, ip);
            return f_val * fabs(detJ);
        }
        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 &map,
        mfem::Coefficient &sigma,
        const int dim
    )
        : MatrixCoefficient(dim),
          m_map(map),
          m_scalar(&sigma),
          m_tensor(nullptr) { };

    MappedDiffusionCoefficient::MappedDiffusionCoefficient(
        const DomainMapper &map,
        MatrixCoefficient &sigma
    )
        : MatrixCoefficient(sigma.GetHeight()),
          m_map(map),
          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);

        mfem::DenseMatrix J(dim, dim), JInv(dim, dim);
        m_map.ComputeJacobian(T, J);
        const double detJ = J.Det();
        mfem::CalcInverse(J, JInv);

        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 &map,
        VectorCoefficient &coeff
    )
        : VectorCoefficient(coeff.GetVDim()),
          m_map(map),
          m_coeff(coeff) { };

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

        mfem::DenseMatrix JInv(dim, dim);
        m_map.ComputeInverseJacobian(T, JInv);
        double detJ = m_map.ComputeDetJ(T, ip);

        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 &map,
        Func f // std::function<double(const mfem::Vector&)>
    )
        : m_f(std::move(f)),
          m_map(map) { };

    double PhysicalPositionFunctionCoefficient::Eval(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) {
        T.SetIntPoint(&ip);
        mfem::Vector x;
        m_map.GetPhysicalPoint(T, ip, x);
        return m_f(x);
    }

    MappedHDivMassCoefficient::MappedHDivMassCoefficient(
        const DomainMapper &map,
        const int dim
    )
        : MatrixCoefficient(dim),
          m_map(map) {
    }

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

        mfem::DenseMatrix map_jacobian(height, height);
        m_map.ComputeJacobian(transformation, map_jacobian);

        const double map_determinant = map_jacobian.Det();

        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