module;
#include <mfem.hpp>

module mean_field;

namespace mean_field::integrators {
    AdvectionIntegrator::AdvectionIntegrator(const mapping::DomainMapper &map) : m_map(map) {
    }

    void AdvectionIntegrator::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;
        }

        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];
        const mfem::Vector &rho_dofs      = *elfun[1];

        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::Vector shape_v(dof_v), shape_rho(dof_rho);
        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() + 1);

        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->CalcShape(ip, shape_v);
            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);
            fe_rho->CalcShape(ip, shape_rho);

            double rho_val = 0.0;
            for (int i = 0; i < dof_rho; ++i) {
                rho_val += rho_dofs(i) * shape_rho(i);
            }

            mfem::Vector v_val(dim);
            v_val = 0.0;
            mfem::DenseMatrix grad_v(dim, dim);
            grad_v = 0.0;

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    const double v_ic = v_dofs(i + c * dof_v);
                    v_val(c) += v_ic * shape_v(i);
                    for (int d = 0; d < dim; ++d) {
                        grad_v(c, d) += v_ic * dshape_v_phys(i, d);
                    }
                }
            }

            mfem::Vector adv_val(dim);
            adv_val = 0.0;
            for (int c = 0; c < dim; ++c) {
                for (int d = 0; d < dim; ++d) {
                    adv_val(c) += v_val(d) * grad_v(c, d);
                }
            }

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    r_v(i + c * dof_v) += shape_v(i) * rho_val * adv_val(c) * weight;
                }
            }
        }
    }

    void AdvectionIntegrator::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();

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

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

        if (dv_dv)
            *dv_dv = 0.0;
        if (dv_drho)
            *dv_drho = 0.0;

        mfem::Vector shape_v(dof_v), shape_rho(dof_rho);
        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() + 1);

        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->CalcShape(ip, shape_v);
            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);
            fe_rho->CalcShape(ip, shape_rho);

            double rho_val = 0.0;
            for (int i = 0; i < dof_rho; ++i) {
                rho_val += rho_dofs(i) * shape_rho(i);
            }

            mfem::Vector v_val(dim);
            v_val = 0.0;
            mfem::DenseMatrix grad_v(dim, dim);
            grad_v = 0.0;

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    double v_ic = v_dofs(i + c * dof_v);
                    v_val(c) += v_ic * shape_v(i);
                    for (int d = 0; d < dim; ++d) {
                        grad_v(c, d) += v_ic * dshape_v_phys(i, d);
                    }
                }
            }

            mfem::Vector adv_val(dim);
            adv_val = 0.0;
            for (int c = 0; c < dim; ++c) {
                for (int d = 0; d < dim; ++d) {
                    adv_val(c) += v_val(d) * grad_v(c, d);
                }
            }

            // Jacobian wrt. Velocity: dR_v/dv
            if (dv_dv) {
                for (int i = 0; i < dof_v; ++i) {
                    // Test function index
                    for (int c = 0; c < dim; ++c) {
                        // Test function component
                        int row = i + c * dof_v;
                        for (int j = 0; j < dof_v; ++j) {
                            // Trial function index
                            double v_dot_grad_phi_j = 0.0;

                            for (int k = 0; k < dim; ++k) {
                                v_dot_grad_phi_j += v_val(k) * dshape_v_phys(j, k);
                            }

                            for (int d = 0; d < dim; ++d) {
                                // Trial function component
                                int col      = j + d * dof_v;

                                // \rho (\delta \vec{v} \cdot \nabla \vec{v})
                                // \delta v is along direction 'd' for the cth
                                // component of advection
                                double termA = shape_v(j) * grad_v(c, d);

                                // \rho(\vec{v} \cdot \nabla \delta \vec{v})
                                // Only non-zero when the advected component
                                // matches the test component
                                double termB = (c == d) ? v_dot_grad_phi_j : 0.0;

                                (*dv_dv)(row, col) += shape_v(i) * rho_val * (termA + termB) * weight;
                            }
                        }
                    }
                }
            }

            // Jacobian wrt. Density: dR_v / drho
            if (dv_drho) {
                for (int i = 0; i < dof_v; ++i) {
                    for (int c = 0; c < dim; ++c) {
                        int row = i + c * dof_v;
                        for (int j = 0; j < dof_rho; ++j) {
                            int col     = j;

                            // \delta \rho * (\vec{v} \cdot \nabla \vec{v})
                            double term = shape_rho(j) * adv_val(c);
                            (*dv_drho)(row, col) += shape_v(i) * term * weight;
                        }
                    }
                }
            }
        }
    }
} // namespace mean_field::integrators