module;
#include <mfem.hpp>
module mean_field;

namespace mean_field::integrators {
    ViscosityIntegrator::ViscosityIntegrator(
        const mapping::DomainMapper &map,
        const double mu,
        const int quad_boost
    )
        : m_map(map),
          m_mu(mu),
          m_quad_boost(quad_boost) {
    }

    void ViscosityIntegrator::SetMu(const double mu) {
        m_mu = mu;
    }

    void ViscosityIntegrator::AssembleElementVector(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array<mfem::Vector *> &elvec
    ) {
        if (utils::is_vacuum(Tr, elvec)) {
            return;
        }

        void *data_before                 = (void *)elvec[0]->GetData();
        int size_before                   = elvec[0]->Size();

        const mfem::FiniteElement *fe_v   = el[0];
        const mfem::FiniteElement *fe_rho = el[1];

        const int dof_v                   = fe_v->GetDof();
        const int dof_rho                 = fe_rho->GetDof();
        const int dim                     = Tr.GetSpaceDim();

        const mfem::Vector &v_dofs        = *elfun[0];

        mfem::Vector &r_v                 = *elvec[0];
        r_v.SetSize(dof_v * dim);
        r_v = 0.0;
        if (elvec[1]) {
            elvec[1]->SetSize(dof_rho);
            *elvec[1] = 0.0;
        }

        mfem::DenseMatrix dshape_v_ref(dof_v, dim), dshape_v_phys(dof_v, dim);

        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(fe_v->GetGeomType(), 2 * fe_v->GetOrder() + m_quad_boost);

        for (int q = 0; q < ir->GetNPoints(); ++q) {

            const mfem::IntegrationPoint &ip = ir->IntPoint(q);
            Tr.SetIntPoint(&ip);

            auto [J_inv, detJ, weight] = m_map.GetQuadratureContext(Tr, ip);

            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);

            // ∇v(c,j) = δj v_c
            mfem::DenseMatrix grad_v(dim, dim);
            grad_v = 0.0;
            for (int n = 0; n < dof_v; ++n) {
                for (int c = 0; c < dim; ++c) {
                    const double vn_c = v_dofs(n + c * dof_v);
                    for (int j = 0; j < dim; ++j) {
                        grad_v(c, j) += vn_c * dshape_v_phys(n, j);
                    }
                }
            }

            double div_v = 0.0;
            for (int c = 0; c < dim; ++c) {
                div_v += grad_v(c, c);
            }

            // D_cj = dj v_c + dc v_j - (2/3) δcj ∇v
            // R_v_i^c += μ * weight * ∑_j (δj φi) D_cj
            const double mu_w = m_mu * weight;
            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    double acc = 0.0;
                    for (int j = 0; j < dim; ++j) {
                        double D_cj = grad_v(c, j) + grad_v(j, c);
                        if (c == j)
                            D_cj -= (2.0 / 3.0) * div_v;
                        acc += dshape_v_phys(i, j) * D_cj;
                    }
                    r_v(i + c * dof_v) += mu_w * acc;
                }
            }
        }
    }

    void ViscosityIntegrator::AssembleElementGrad(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        const mfem::FiniteElement *fe_v   = el[0];
        const mfem::FiniteElement *fe_rho = el[1];

        const int dof_v                   = fe_v->GetDof();
        const int dof_rho                 = fe_rho->GetDof();
        const int dim                     = Tr.GetSpaceDim();

        mfem::DenseMatrix *dv_dv          = elmats(0, 0);
        mfem::DenseMatrix *dv_drho        = elmats(0, 1);

        if (dv_drho)
            *dv_drho = 0.0;
        if (dv_dv)
            *dv_dv = 0.0;
        if (elmats(1, 0))
            *elmats(1, 0) = 0.0;
        if (elmats(1, 1))
            *elmats(1, 1) = 0.0;
        if (!dv_dv)
            return;

        mfem::DenseMatrix dshape_v_ref(dof_v, dim), dshape_v_phys(dof_v, dim);

        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(fe_v->GetGeomType(), 2 * fe_v->GetOrder());
        for (int q = 0; q < ir->GetNPoints(); ++q) {
            const mfem::IntegrationPoint &ip = ir->IntPoint(q);
            Tr.SetIntPoint(&ip);

            auto [J_inv, detJ, weight] = m_map.GetQuadratureContext(Tr, ip);

            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);

            const double mu_w = m_mu * weight;

            for (int i = 0; i < dof_v; ++i) {
                for (int n = 0; n < dof_v; ++n) {
                    double dot_grad = 0.0;
                    for (int j = 0; j < dim; ++j) {
                        dot_grad += dshape_v_phys(i, j) * dshape_v_phys(n, j);
                    }

                    for (int c = 0; c < dim; ++c) {
                        const int row = i + c * dof_v;
                        for (int d = 0; d < dim; ++d) {
                            const int col = n + d * dof_v;

                            double val    = 0.0;
                            if (c == d)
                                val += dot_grad;

                            val += dshape_v_phys(i, d) * dshape_v_phys(n, c);

                            val -= (2.0 / 3.0) * dshape_v_phys(i, c) * dshape_v_phys(n, d);
                            (*dv_dv)(row, col) += mu_w * val;
                        }
                    }
                }
            }
        }
    }

} // namespace mean_field::integrators