module;
#include <mfem.hpp>

module mean_field;

namespace mean_field::integrators {
    ContinuityVolumeIntegrator::ContinuityVolumeIntegrator(
        const mapping::DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate
    )
        : m_mapping(
              mapper,
              displacement,
              compactification_coordinate
          ) { };

    void ContinuityVolumeIntegrator::AssembleElementVector(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array<mfem::Vector *> &elvec
    ) {
        m_mapping.InvalidateCache();

        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];

        void *data_rho_before             = elvec[1] ? (void *)elvec[1]->GetData() : nullptr;
        int size_rho_before               = elvec[1] ? elvec[1]->Size() : -1;

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

        mfem::Vector shape_v(dof_v), shape_rho(dof_rho);
        mfem::DenseMatrix dshape_rho_ref(dof_rho, dim), dshape_rho_phys(dof_rho, 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_mapping.GetQuadratureContext(Tr, ip);

            fe_v->CalcShape(ip, shape_v);
            fe_rho->CalcShape(ip, shape_rho);
            fe_rho->CalcDShape(ip, dshape_rho_ref);
            mfem::Mult(dshape_rho_ref, J_inv, dshape_rho_phys);

            mfem::Vector v_val(dim);
            v_val = 0.0;
            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    const int row = i + c * dof_v;
                    v_val(c) += v_dofs(row) * shape_v(i);
                }
            }
            double rho_val = 0.0;
            for (int i = 0; i < dof_rho; ++i) {
                rho_val += rho_dofs(i) * shape_rho(i);
            }

            for (int i = 0; i < dof_rho; ++i) {
                double grad_dot_rhov = 0.0;
                for (int c = 0; c < dim; ++c) {
                    grad_dot_rhov += dshape_rho_phys(i, c) * rho_val * v_val(c);
                }
                r_rho(i) -= grad_dot_rhov * weight;
            }
        }
    }

    void ContinuityVolumeIntegrator::AssembleElementGrad(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        m_mapping.InvalidateCache();

        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 *drho_dv        = elmats(1, 0);
        mfem::DenseMatrix *drho_drho      = elmats(1, 1);

        if (elmats(0, 0))
            *elmats(0, 0) = 0.0;
        if (elmats(0, 1))
            *elmats(0, 1) = 0.0;

        if (drho_dv)
            *drho_dv = 0.0;
        if (drho_drho)
            *drho_drho = 0.0;

        mfem::Vector shape_v(dof_v), shape_rho(dof_rho);
        mfem::DenseMatrix dshape_rho_ref(dof_rho, dim), dshape_rho_phys(dof_rho, 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_mapping.GetQuadratureContext(Tr, ip);

            fe_v->CalcShape(ip, shape_v);
            fe_rho->CalcShape(ip, shape_rho);
            fe_rho->CalcDShape(ip, dshape_rho_ref);
            mfem::Mult(dshape_rho_ref, J_inv, dshape_rho_phys);

            mfem::Vector v_val(dim);
            v_val = 0.0;
            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    const int row = i + c * dof_v;
                    v_val(c) += v_dofs(row) * shape_v(i);
                }
            }

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

            if (drho_dv) {
                for (int i = 0; i < dof_rho; ++i) {
                    for (int j = 0; j < dof_v; ++j) {
                        for (int d = 0; d < dim; ++d) {
                            const int col = j + d * dof_v;
                            (*drho_dv)(i, col) -= dshape_rho_phys(i, d) * rho_val * shape_v(j) * weight;
                        }
                    }
                }
            }

            if (drho_drho) {
                for (int i = 0; i < dof_rho; ++i) {
                    double grad_psi_dot_v = 0.0;
                    for (int c = 0; c < dim; ++c) {
                        grad_psi_dot_v += dshape_rho_phys(i, c) * v_val(c);
                    }
                    for (int j = 0; j < dof_rho; ++j) {
                        (*drho_drho)(i, j) -= grad_psi_dot_v * shape_rho(j) * weight;
                    }
                }
            }
        }
    }

    ContinuityFaceIntegrator::ContinuityFaceIntegrator(
        const mapping::DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate
    )
        : m_mapping(
              mapper,
              displacement,
              compactification_coordinate
          ) {
    }

    void ContinuityFaceIntegrator::AssembleFaceVector(
        const mfem::Array<const mfem::FiniteElement *> &el1,
        const mfem::Array<const mfem::FiniteElement *> &el2,
        mfem::FaceElementTransformations &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array<mfem::Vector *> &elvect
    ) {
        m_mapping.InvalidateCache();

        const mfem::FiniteElement *fe_v_minus   = el1[0];
        const mfem::FiniteElement *fe_v_plus    = el2[0];

        const mfem::FiniteElement *fe_rho_minus = el1[1];
        const mfem::FiniteElement *fe_rho_plus  = el2[1];

        const int dof_v_minus                   = fe_v_minus->GetDof();
        const int dof_v_plus                    = fe_v_plus->GetDof();

        const int dof_rho_minus                 = fe_rho_minus->GetDof();
        const int dof_rho_plus                  = fe_rho_plus->GetDof();

        const int dim                           = Tr.GetSpaceDim();

        if (elvect[0]) {
            elvect[0]->SetSize(dim * dof_v_minus + dim * dof_v_plus);
            *elvect[0] = 0.0;
        }
        mfem::Vector &r_rho = *elvect[1];
        r_rho.SetSize(dof_rho_minus + dof_rho_plus);
        r_rho                = 0.0;

        const int attr_minus = Tr.Elem1->Attribute;
        const int attr_plus  = (Tr.Elem2 != nullptr) ? Tr.Elem2->Attribute : -1;
        using DomainSchema   = utils::domain::CoreEnvelopeVacuumDomainSchema;
        if (DomainSchema::template attribute_belongs_to<utils::domain::Vacuum>(attr_minus) ||
            DomainSchema::template attribute_belongs_to<utils::domain::Vacuum>(attr_plus)) {
            return; // No flux contribution for vacuum faces
        }

        if (Tr.Elem2 == nullptr) {
            return; // Boundary face,
        }

        const mfem::Vector &v_dofs      = *elfun[0]; // Size: dim * dof_v_minus + dim*dof_v_plus
        const mfem::Vector &rho_dofs    = *elfun[1]; // Size: dof_rho_minus + dof_rho_plus

        // Helpers to auto offset to the correct point in the dof array
        auto rho_minus_dof              = [&](const int i) { return rho_dofs(i); };
        auto rho_plus_dof               = [&](const int i) { return rho_dofs(i + dof_rho_minus); };
        auto v_minus_dof                = [&](const int k, const int c) { return v_dofs(k + c * dof_v_minus); };

        const int p_v                   = fe_v_minus->GetOrder();
        const int p_rho                 = fe_rho_minus->GetOrder();
        const int int_order             = 2 * std::max(p_v, p_rho) + 1;

        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(Tr.GetGeometryType(), int_order);

        mfem::Vector shape_v_minus(dof_v_minus), shape_rho_minus(dof_rho_minus), shape_rho_plus(dof_rho_plus);

        for (int q = 0; q < ir->GetNPoints(); ++q) {
            const mfem::IntegrationPoint &face_ip = ir->IntPoint(q);

            Tr.SetAllIntPoints(&face_ip);
            const mfem::IntegrationPoint &ip_minus = Tr.GetElement1IntPoint();
            const mfem::IntegrationPoint &ip_plus  = Tr.GetElement2IntPoint();

            auto [n_unit, ds, v_dot_n_scale]       = m_mapping.GetFaceQuadratureContext(Tr, face_ip);

            fe_v_minus->CalcShape(ip_minus, shape_v_minus);
            fe_rho_minus->CalcShape(ip_minus, shape_rho_minus);
            fe_rho_plus->CalcShape(ip_plus, shape_rho_plus);

            // v dot n
            // u_n = ∑ n_c * ∑ v_kc * φ_k
            double u_n = 0.0;
            for (int c = 0; c < dim; ++c) {
                double v_c = 0.0;
                for (int k = 0; k < dof_v_minus; ++k) {
                    v_c += v_minus_dof(k, c) * shape_v_minus(k);
                }
                u_n += v_c * n_unit(c);
            }

            double rho_minus_val = 0.0;
            for (int i = 0; i < dof_rho_minus; ++i) {
                rho_minus_val += shape_rho_minus(i) * rho_minus_dof(i);
            }

            double rho_plus_val = 0.0;
            for (int i = 0; i < dof_rho_plus; ++i) {
                rho_plus_val += shape_rho_plus(i) * rho_plus_dof(i);
            }

            // Upwind density
            // I use the convention that the flow is positive when moving from
            // minus to plus
            const double rho_up        = (u_n >= 0) ? rho_minus_val : rho_plus_val;

            const double flux_weighted = u_n * rho_up * ds;

            // Note the normals need to be in opposite directions for these two
            // fluxes
            for (int i = 0; i < dof_rho_minus; ++i) {
                r_rho(i) += shape_rho_minus(i) * flux_weighted;
            }

            for (int i = 0; i < dof_rho_plus; ++i) {
                r_rho(dof_rho_minus + i) -= shape_rho_plus(i) * flux_weighted;
            }
        }
    }

    void ContinuityFaceIntegrator::AssembleFaceGrad(
        const mfem::Array<const mfem::FiniteElement *> &el1,
        const mfem::Array<const mfem::FiniteElement *> &el2,
        mfem::FaceElementTransformations &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        m_mapping.InvalidateCache();

        const mfem::FiniteElement *fe_v_minus   = el1[0];
        const mfem::FiniteElement *fe_v_plus    = el2[0];
        const mfem::FiniteElement *fe_rho_minus = el1[1];
        const mfem::FiniteElement *fe_rho_plus  = el2[1];

        const int dof_v_minus                   = fe_v_minus->GetDof();
        const int dof_v_plus                    = fe_v_plus->GetDof();
        const int dof_rho_minus                 = fe_rho_minus->GetDof();
        const int dof_rho_plus                  = fe_rho_plus->GetDof();
        const int dim                           = Tr.GetSpaceDim();

        const int N_v_total                     = dim * (dof_v_minus + dof_v_plus);
        const int N_rho_total                   = dof_rho_minus + dof_rho_plus;

        auto size_and_zero_mat                  = [&](mfem::DenseMatrix *mat, const int r_size, const int c_size) {
            if (mat) {
                mat->SetSize(r_size, c_size);
                *mat = 0.0;
            }
        };

        size_and_zero_mat(elmats(0, 0), N_v_total, N_v_total);
        size_and_zero_mat(elmats(0, 1), N_v_total, N_rho_total);
        size_and_zero_mat(elmats(1, 0), N_rho_total, N_v_total);
        size_and_zero_mat(elmats(1, 1), N_rho_total, N_rho_total);

        if (skip_face(Tr))
            return;

        mfem::DenseMatrix *drho_dv   = elmats(1, 0);
        mfem::DenseMatrix *drho_drho = elmats(1, 1);

        if (!drho_dv && !drho_drho)
            return;

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

        const int int_order             = 2 * std::max(fe_v_minus->GetOrder(), fe_rho_minus->GetOrder()) + 1;
        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(Tr.GetGeometryType(), int_order);

        mfem::Vector shape_v_minus(dof_v_minus), shape_rho_minus(dof_rho_minus), shape_rho_plus(dof_rho_plus);

        for (int q = 0; q < ir->GetNPoints(); ++q) {
            const mfem::IntegrationPoint &face_ip = ir->IntPoint(q);
            Tr.SetAllIntPoints(&face_ip);
            const mfem::IntegrationPoint &ip_minus = Tr.GetElement1IntPoint();
            const mfem::IntegrationPoint &ip_plus  = Tr.GetElement2IntPoint();

            auto [n_unit, ds, v_dot_n_scale]       = m_mapping.GetFaceQuadratureContext(Tr, face_ip);

            fe_v_minus->CalcShape(ip_minus, shape_v_minus);
            fe_rho_minus->CalcShape(ip_minus, shape_rho_minus);
            fe_rho_plus->CalcShape(ip_plus, shape_rho_plus);

            const double u_n     = compute_u_n(v_dofs, shape_v_minus, n_unit, dof_v_minus, dim);

            double rho_minus_val = 0.0;
            for (int i = 0; i < dof_rho_minus; ++i) {
                rho_minus_val += shape_rho_minus(i) * rho_dofs(i);
            }

            double rho_plus_val = 0.0;
            for (int i = 0; i < dof_rho_plus; ++i) {
                rho_plus_val += shape_rho_plus(i) * rho_dofs(dof_rho_minus + i);
            }

            const bool upwind_minus = (u_n >= 0.0);
            const double rho_up     = upwind_minus ? rho_minus_val : rho_plus_val;

            // (1, 1)
            if (drho_drho) {
                const double u_w = u_n * ds;

                if (upwind_minus) {
                    for (int ip = 0; ip < dof_rho_minus; ++ip) {
                        const double col_w = u_w * shape_rho_minus(ip);
                        for (int i = 0; i < dof_rho_minus; ++i) {
                            (*drho_drho)(i, ip) += shape_rho_minus(i) * col_w;
                        }
                        for (int j = 0; j < dof_rho_plus; ++j) {
                            (*drho_drho)(dof_rho_minus + j, ip) -= shape_rho_plus(j) * col_w;
                        }
                    }
                } else {
                    for (int jp = 0; jp < dof_rho_plus; ++jp) {
                        const double col_w = u_w * shape_rho_plus(jp);
                        const int col_idx  = dof_rho_minus + jp;
                        for (int i = 0; i < dof_rho_minus; ++i) {
                            (*drho_drho)(i, col_idx) += shape_rho_minus(i) * col_w;
                        }
                        for (int j = 0; j < dof_rho_plus; ++j) {
                            (*drho_drho)(dof_rho_minus + j, col_idx) -= shape_rho_plus(j) * col_w;
                        }
                    }
                }
            }
            // (1, 0)
            if (drho_dv) {
                const double rho_w = rho_up * ds;
                for (int c = 0; c < dim; ++c) {
                    const double n_c_rho_w = n_unit(c) * rho_w;
                    for (int k = 0; k < dof_v_minus; ++k) {
                        const int col_idx  = k + c * dof_v_minus;
                        const double col_w = n_c_rho_w * shape_v_minus(k);
                        for (int i = 0; i < dof_rho_minus; ++i) {
                            (*drho_dv)(i, col_idx) += shape_rho_minus(i) * col_w;
                        }
                        for (int j = 0; j < dof_rho_plus; ++j) {
                            (*drho_dv)(dof_rho_minus + j, col_idx) -= shape_rho_plus(j) * col_w;
                        }
                    }
                }
            }
        }
    }

    bool ContinuityFaceIntegrator::skip_face(const mfem::FaceElementTransformations &Tr) {
        const int attr_minus = Tr.Elem1->Attribute;
        const int attr_plus  = (Tr.Elem2 != nullptr) ? Tr.Elem2->Attribute : -1;
        using DomainSchema   = utils::domain::CoreEnvelopeVacuumDomainSchema;
        if (DomainSchema::template attribute_belongs_to<utils::domain::Vacuum>(attr_minus) ||
            DomainSchema::template attribute_belongs_to<utils::domain::Vacuum>(attr_plus)) {
            return true; // No flux contribution for vacuum faces
        }
        if (Tr.Elem2 == nullptr) {
            return true; // Boundary face,
        }
        return false;
    }

    double ContinuityFaceIntegrator::compute_u_n(
        const mfem::Vector &v_dofs,
        const mfem::Vector &shape_v_minus,
        const mfem::Vector &n_unit,
        int dof_v_minus,
        int dim
    ) {
        double u_n = 0.0;
        for (int c = 0; c < dim; ++c) {
            double v_c = 0.0;
            for (int k = 0; k < dof_v_minus; ++k) {
                v_c += v_dofs(k + c * dof_v_minus) * shape_v_minus(k);
            }
            u_n += v_c * n_unit(c);
        }
        return u_n;
    }
} // namespace mean_field::integrators