module;
#include <mfem.hpp>

module mean_field;

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();
    }

    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();
    }

    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();
    }

    bool DomainMapper::IsIdentity() const {
        return (m_d == nullptr);
    }

    void DomainMapper::ResetDisplacement() {
        m_d = nullptr;
        InvalidateCache();
    }

    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 (IsIdentity()) {
            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 (IsIdentity()) {
                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 (IsIdentity() && !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);
    }

    DomainMapper::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};
    }

    DomainMapper::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 (IsIdentity()) {
            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 (IsIdentity()) {
            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);
    }

    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 (IsIdentity()) 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++;
        }
    }
}