module;
#include <mfem.hpp>

module mean_field;
import :mapping.types;

namespace {
    double get_positive_map_jacobian(
        const mean_field::mapping::DomainMapper &domain_mapper,
        mfem::ElementTransformation &transformation,
        const mfem::IntegrationPoint &integration_point,
        mfem::DenseMatrix &map_jacobian
    ) {
        transformation.SetIntPoint(&integration_point);
        domain_mapper.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.");
        return map_determinant;
    }
} // namespace

namespace mean_field::mapping {
    DomainMapper::DomainMapper(
        const double r_star_ref,
        const double r_inf_ref
    )
        : m_d(nullptr),
          m_r_star_ref(r_star_ref),
          m_r_inf_ref(r_inf_ref) {
        InitAllScratchSpaces();
        CalcIsIdentity() ? m_displacement_is_identity = true : m_displacement_is_identity = false;
    }

    DomainMapper::DomainMapper(
        const mfem::GridFunction &d,
        const double r_star_ref,
        const double r_inf_ref
    )
        : m_d(&d),
          m_dim(d.FESpace()->GetMesh()->Dimension()),
          m_r_star_ref(r_star_ref),
          m_r_inf_ref(r_inf_ref) {
        InitAllScratchSpaces();
        CalcIsIdentity() ? m_displacement_is_identity = true : m_displacement_is_identity = false;
    }

    bool DomainMapper::is_vacuum(const mfem::ElementTransformation &T) const {
        if (T.ElementType == mfem::ElementTransformation::ELEMENT) {
            return T.Attribute == m_vacuum_attr;
        } else if (T.ElementType == mfem::ElementTransformation::BDR_ELEMENT) {
            return T.Attribute == m_vacuum_attr - 1;
            // TODO: In a more robust code this should really be read from the
            // stroid API to ensure that the vacuum boundary is really 1 - the
            // vacuum material attribute
        }
        return false;
    }

    void DomainMapper::SetDisplacement(const mfem::GridFunction &d) {
        if (m_dim != d.FESpace()->GetMesh()->Dimension()) {
            const std::string err_msg = std::format(
                "Dimension mismatch: DomainMapper is initialized for dimension "
                "{}, "
                "but provided displacement field has "
                "dimension {}.",
                m_dim, d.FESpace()->GetMesh()->Dimension()
            );
            throw std::invalid_argument(err_msg);
        }
        m_d = &d;
        InvalidateCache();

        CalcIsIdentity() ? m_displacement_is_identity = true : m_displacement_is_identity = false;
    }

    bool DomainMapper::HasCompactification() const noexcept {
        return std::isfinite(m_r_star_ref) && std::isfinite(m_r_inf_ref) && m_r_star_ref > 0.0 &&
               m_r_inf_ref > m_r_star_ref && m_xi_clamp > 0.0 && m_xi_clamp < 1.0;
    }

    bool DomainMapper::HasDisplacementField() const noexcept {
        return m_d != nullptr;
    }

    bool DomainMapper::CalcIsIdentity() const {
        if (m_d == nullptr) {
            return true;
        }

        const int local_identity          = m_d->Normlinf() == 0.0 ? 1 : 0;

        const auto *parallel_displacement = dynamic_cast<const mfem::ParGridFunction *>(m_d);

        if (parallel_displacement == nullptr) {
            return local_identity == 1;
        }

        int global_identity = 0;
        MPI_Allreduce(
            &local_identity, &global_identity, 1, MPI_INT, MPI_MIN, parallel_displacement->ParFESpace()->GetComm()
        );

        return global_identity == 1;
    }

    void DomainMapper::ResetDisplacement() {
        m_d = nullptr;
        InvalidateCache();
        CalcIsIdentity() ? m_displacement_is_identity = true : m_displacement_is_identity = false;
    }

    void DomainMapper::ComputeJacobian(
        mfem::ElementTransformation &T,
        mfem::DenseMatrix &J
    ) const {
        J.SetSize(m_dim, m_dim);
        J     = 0.0;
        m_J_D = 0.0;
        if (!HasDisplacementField()) {
            for (int i = 0; i < m_dim; ++i) {
                m_J_D(i, i) = 1.0; // Identity mapping
            }
        } else {
            UpdateElementCache(T);
            m_dshape.SetSize(m_fe->GetDof(), m_dim);
            m_fe->CalcPhysDShape(T, m_dshape);
            mfem::MultAtB(m_dof_mat, m_dshape, m_J_D);

            for (int i = 0; i < m_dim; ++i) {
                m_J_D(i, i) += 1.0;
            }
        }

        if (is_vacuum(T)) {
            T.Transform(T.GetIntPoint(), m_x_ref);

            if (!HasDisplacementField()) {
                m_x_disp = m_x_ref;
            } else {
                m_shape.SetSize(m_fe->GetDof());
                m_fe->CalcShape(T.GetIntPoint(), m_shape);
                m_dof_mat.MultTranspose(m_shape, m_d_val);
                add(m_x_ref, m_d_val, m_x_disp);
            }

            ComputeKelvinJacobian(m_x_ref, m_x_disp, m_J_D, J);
        } else {
            J = m_J_D;
        }
    }

    double DomainMapper::ComputeDetJ(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) const {
        if (!HasDisplacementField() && !is_vacuum(T))
            return 1.0; // If no mapping, the determinant of the Jacobian is 1
        T.SetIntPoint(&ip);
        mfem::DenseMatrix J;
        ComputeJacobian(T, J);
        return J.Det();
    }

    void DomainMapper::ComputeMappedDiffusionTensor(
        mfem::ElementTransformation &T,
        mfem::DenseMatrix &D
    ) const {
        ComputeJacobian(T, m_J_temp);
        const double detJ = m_J_temp.Det();
        mfem::CalcInverse(m_J_temp, m_JInv_temp);
        D.SetSize(m_dim, m_dim);
        mfem::MultABt(m_JInv_temp, m_JInv_temp, D);
        D *= fabs(detJ);
    }

    void DomainMapper::ComputeInverseJacobian(
        mfem::ElementTransformation &T,
        mfem::DenseMatrix &JInv
    ) const {
        ComputeJacobian(T, m_J_temp);
        JInv.SetSize(m_dim, m_dim);
        mfem::CalcInverse(m_J_temp, JInv);
    }

    VolumeQuadratureContext DomainMapper::GetQuadratureContext(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip
    ) const {
        const int dim = T.GetSpaceDim();
        mfem::DenseMatrix J_map(dim, dim), J_inv(dim, dim);
        ComputeJacobian(T, J_map);
        mfem::DenseMatrix J_full(dim, dim);
        mfem::Mult(J_map, T.Jacobian(), J_full);
        mfem::CalcInverse(J_full, J_inv);
        const double detJ   = std::fabs(ComputeDetJ(T, ip));
        const double weight = ip.weight * T.Weight() * detJ;
        return {.J_inv = J_inv, .detJ = detJ, .weight = weight};
    }

    FaceQuadratureContext DomainMapper::GetFaceQuadratureContext(
        mfem::FaceElementTransformations &T,
        const mfem::IntegrationPoint &ip
    ) const {
        const int dim = T.GetSpaceDim();
        T.SetAllIntPoints(&ip);

        mfem::Vector n_raw(dim);
        mfem::CalcOrtho(T.Jacobian(), n_raw);

        if (!HasDisplacementField() && !is_vacuum(T)) {
            const double n_raw_mag = n_raw.Norml2();
            mfem::Vector n_unit(dim);
            n_unit = n_raw;
            n_unit /= n_raw_mag;
            return FaceQuadratureContext{.normal = n_unit, .ds = ip.weight * n_raw_mag, .v_dot_n_scale = 1.0};
        }

        // Nanson's Formula
        // (https://en.wikiversity.org/wiki/Continuum_mechanics/Volume_change_and_area_change)
        // Since the displacement field lives in H1 it should be irrelevant if
        // we pick Elem1 or Elem2
        mfem::DenseMatrix J_map(dim, dim);
        ComputeJacobian(*T.Elem1, J_map);
        const double detJ_map = J_map.Det();

        mfem::DenseMatrix J_map_inv(dim, dim);
        mfem::CalcInverse(J_map, J_map_inv);

        mfem::Vector n_phys(dim);
        J_map_inv.MultTranspose(n_raw, n_phys);
        n_phys *= detJ_map;

        const double n_phys_mag = n_phys.Norml2();
        mfem::Vector n_unit(dim);
        n_unit = n_phys;
        n_unit /= n_phys_mag;

        const double n_raw_mag = n_raw.Norml2();

        return FaceQuadratureContext{
            .normal = n_unit, .ds = ip.weight * n_raw_mag, .v_dot_n_scale = n_phys_mag / n_raw_mag
        };
    }

    void DomainMapper::GetPhysicalPoint(
        mfem::ElementTransformation &T,
        const mfem::IntegrationPoint &ip,
        mfem::Vector &x_phys
    ) const {
        x_phys.SetSize(m_dim);
        T.Transform(ip, m_x_ref);

        if (!HasDisplacementField()) {
            x_phys = m_x_ref;
        } else {
            UpdateElementCache(T);

            m_shape.SetSize(m_fe->GetDof());
            m_fe->CalcShape(ip, m_shape);

            m_dof_mat.MultTranspose(m_shape, m_d_val);
            add(m_x_ref, m_d_val, x_phys);
        }
        if (is_vacuum(T)) {
            ApplyKelvinMapping(m_x_ref, x_phys);
        }
    }

    void DomainMapper::GetVectorValue(
        const int i,
        const mfem::IntegrationPoint &ip,
        mfem::Vector &val
    ) const {
        m_d->GetVectorValue(i, ip, val);
    }

    void DomainMapper::MapHDivFluxToPhysical(
        mfem::ElementTransformation &transformation,
        const mfem::IntegrationPoint &integration_point,
        const mfem::Vector &reference_flux,
        mfem::Vector &physical_flux
    ) const {
        MFEM_VERIFY(reference_flux.Size() == m_dim, "The reference H(div) flux has the wrong dimension.");

        mfem::DenseMatrix map_jacobian(m_dim, m_dim);
        const double map_determinant =
            get_positive_map_jacobian(*this, transformation, integration_point, map_jacobian);

        mfem::Vector mapped_flux(m_dim);
        map_jacobian.Mult(reference_flux, mapped_flux);
        mapped_flux /= map_determinant;

        physical_flux = mapped_flux;
    }

    void DomainMapper::MapPhysicalFluxToHDivReference(
        mfem::ElementTransformation &transformation,
        const mfem::IntegrationPoint &integration_point,
        const mfem::Vector &physical_flux,
        mfem::Vector &reference_flux
    ) const {
        MFEM_VERIFY(physical_flux.Size() == m_dim, "The physical flux has the wrong dimension.");

        mfem::DenseMatrix map_jacobian(m_dim, m_dim);
        const double map_determinant =
            get_positive_map_jacobian(*this, transformation, integration_point, map_jacobian);

        mfem::DenseMatrix inverse_map_jacobian(m_dim, m_dim);
        mfem::CalcInverse(map_jacobian, inverse_map_jacobian);

        mfem::Vector mapped_flux(m_dim);
        inverse_map_jacobian.Mult(physical_flux, mapped_flux);
        mapped_flux *= map_determinant;

        reference_flux = mapped_flux;
    }

    void DomainMapper::MapReferenceGradientToPhysical(
        mfem::ElementTransformation &transformation,
        const mfem::IntegrationPoint &integration_point,
        const mfem::Vector &reference_gradient,
        mfem::Vector &physical_gradient
    ) const {
        MFEM_VERIFY(reference_gradient.Size() == m_dim, "The reference gradient has the wrong dimension.");

        mfem::DenseMatrix map_jacobian(m_dim, m_dim);
        get_positive_map_jacobian(*this, transformation, integration_point, map_jacobian);

        mfem::DenseMatrix inverse_map_jacobian(m_dim, m_dim);
        mfem::CalcInverse(map_jacobian, inverse_map_jacobian);

        mfem::Vector mapped_gradient(m_dim);
        inverse_map_jacobian.MultTranspose(reference_gradient, mapped_gradient);

        physical_gradient = mapped_gradient;
    }

    const mfem::GridFunction *DomainMapper::GetDisplacement() const {
        return m_d;
    }

    double DomainMapper::GetPhysInfRadius() const {
        return 1.0 - m_xi_clamp;
    }

    size_t DomainMapper::GetCacheHits() const {
        return m_cache_hits;
    }

    size_t DomainMapper::GetCacheMisses() const {
        return m_cache_misses;
    }

    double DomainMapper::GetCacheHitRate() const {
        return (static_cast<double>(m_cache_hits)) / static_cast<double>(m_cache_misses + m_cache_hits);
    }

    void DomainMapper::ResetCacheStats() const {
        m_cache_hits   = 0;
        m_cache_misses = 0;
    }

    void DomainMapper::InitAllScratchSpaces() const {
        m_J_D.SetSize(m_dim, m_dim);
        m_J_temp.SetSize(m_dim, m_dim);
        m_JInv_temp.SetSize(m_dim, m_dim);
        m_x_ref.SetSize(m_dim);
        m_x_disp.SetSize(m_dim);
        m_d_val.SetSize(m_dim);
    }

    void DomainMapper::ApplyKelvinMapping(
        const mfem::Vector &x_ref,
        mfem::Vector &x_phys
    ) const {
        const double r_ref  = x_ref.Norml2();
        double xi           = (r_ref - m_r_star_ref) / (m_r_inf_ref - m_r_star_ref);
        xi                  = std::clamp(xi, 0.0, m_xi_clamp);
        const double factor = m_r_star_ref / (r_ref * (1 - xi));
        x_phys *= factor;
    }

    void DomainMapper::ComputeKelvinJacobian(
        const mfem::Vector &x_ref,
        const mfem::Vector &x_disp,
        const mfem::DenseMatrix &J_D,
        mfem::DenseMatrix &J
    ) const {
        const double r_ref   = x_ref.Norml2();
        const double delta_R = m_r_inf_ref - m_r_star_ref;

        double xi            = (r_ref - m_r_star_ref) / delta_R;
        xi                   = std::clamp(xi, 0.0, m_xi_clamp);

        const double denom   = 1.0 - xi;

        const double k       = m_r_star_ref / (r_ref * denom);

        const double dk_dr =
            m_r_star_ref * ((1.0 / (delta_R * r_ref * denom * denom)) - (1.0 / (r_ref * r_ref * denom)));

        J.SetSize(m_dim, m_dim);
        const double outer_factor = dk_dr / r_ref;

        for (int i = 0; i < m_dim; ++i) {
            for (int j = 0; j < m_dim; ++j) {
                J(i, j) = outer_factor * x_disp(i) * x_ref(j) + k * J_D(i, j);
            }
        }
    }

    void DomainMapper::InvalidateCache() const {
        m_cached_elem_id = -1;
    }

    void DomainMapper::UpdateElementCache(const mfem::ElementTransformation &T) const {
        if (!HasDisplacementField())
            return;

        if (T.ElementNo != m_cached_elem_id || T.ElementType != m_cached_elem_type) {
            m_cache_misses++;
            m_cached_elem_id                    = T.ElementNo;
            m_cached_elem_type                  = T.ElementType;

            const mfem::FiniteElementSpace *fes = m_d->FESpace();
            mfem::Array<int> vdofs;

            if (T.ElementType == mfem::ElementTransformation::ELEMENT) {
                m_fe = fes->GetFE(m_cached_elem_id);
                fes->GetElementVDofs(m_cached_elem_id, vdofs);
            } else {
                m_fe = fes->GetBE(m_cached_elem_id);
                fes->GetBdrElementVDofs(m_cached_elem_id, vdofs);
            }

            m_d->GetSubVector(vdofs, m_elem_dofs);

            const int nd = m_fe->GetDof();
            const int vd = fes->GetVDim();

            m_dof_mat.UseExternalData(m_elem_dofs.GetData(), nd, vd);
        } else {
            m_cache_hits++;
        }
    }
} // namespace mean_field::mapping