#include "serif/tests/equations/fixture.hpp"
#include "serif/tests/mapping/volume_variation.hpp"

#include <algorithm>
#include <cmath>

namespace serif::tests::equations {
    ::mfem::Vector manual_closure(
        Fixture& f,
        const eos::models::Polytrope& eos,
        const ::mfem::Vector& density,
        const ::mfem::Vector& enthalpy,
        const ::mfem::Vector& displacement,
        const ::mfem::Vector* density_direction,
        const ::mfem::Vector* enthalpy_direction,
        const ::mfem::Vector* displacement_direction
    ) {
        ::mfem::GridFunction rho(&f.density_space), h(&f.enthalpy_space), u(&f.displacement_space);
        ::mfem::GridFunction drho(&f.density_space), dh(&f.enthalpy_space), du(&f.displacement_space);
        rho = f.density_binding.scatter(density);
        h = f.enthalpy_binding.scatter(enthalpy);
        u = f.displacement_binding.scatter(displacement);
        drho = 0.0;
        dh = 0.0;
        du = 0.0;
        if (density_direction != nullptr) {
            drho = f.density_binding.scatter(*density_direction);
        }
        if (enthalpy_direction != nullptr) {
            dh = f.enthalpy_binding.scatter(*enthalpy_direction);
        }
        if (displacement_direction != nullptr) {
            du = f.displacement_binding.scatter(*displacement_direction);
        }
        const bool action =
            density_direction != nullptr || enthalpy_direction != nullptr || displacement_direction != nullptr;
        const double n = eos.polytropic_index();
        const double scale = (n + 1.0) * eos.polytropic_constant();
        ::mfem::Vector result(f.density_space.GetVSize());
        result = 0.0;
        mapping::DMWorkspace workspace;
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (!Schema::domain_id_belongs_to<discretization::domain::StellarDomains>(f.mesh.GetAttribute(element))) {
                continue;
            }
            auto& transformation = *f.mesh.GetElementTransformation(element);
            const auto& test = *f.density_space.GetFE(element);
            const auto& displacement_element = *f.displacement_space.GetFE(element);
            const int order = std::max(
                                  2 * test.GetOrder(), test.GetOrder() + f.enthalpy_space.GetFE(element)->GetOrder() +
                                                           static_cast<int>(std::ceil(std::max(0.0, n - 1.0) * 3))
                              ) +
                              transformation.OrderW() + 3 * (displacement_element.GetOrder() - 1);
            const auto& rule = ::mfem::IntRules.Get(transformation.GetGeometryType(), order);
            ::mfem::Array<int> dofs;
            auto* dof_transformation = f.displacement_space.GetElementVDofs(element, dofs);
            ::mfem::Vector u_dofs, du_dofs;
            u.GetSubVector(dofs, u_dofs);
            du.GetSubVector(dofs, du_dofs);
            if (dof_transformation != nullptr) {
                dof_transformation->InvTransformPrimal(u_dofs);
                dof_transformation->InvTransformPrimal(du_dofs);
            }
            const auto u_data = mapping::element_displacement_data_from_element_vdofs(displacement_element, u_dofs);
            const auto du_data = mapping::element_displacement_data_from_element_vdofs(displacement_element, du_dofs);
            ::mfem::Vector compactification_dofs(f.enthalpy_space.GetFE(element)->GetDof());
            compactification_dofs = 0.0;
            const mapping::ElementCompactificationData compactification(
                *f.enthalpy_space.GetFE(element), compactification_dofs
            );
            const mapping::ElementMappingData mapping_data{u_data, compactification};
            ::mfem::Vector shape(test.GetDof()), local(test.GetDof());
            local = 0.0;
            for (int point = 0; point < rule.GetNPoints(); ++point) {
                const auto& ip = rule.IntPoint(point);
                const double enthalpy_value = h.GetValue(element, ip);
                const double eos_density = std::pow(enthalpy_value / scale, n);
                const double eos_derivative = n / scale * std::pow(enthalpy_value / scale, n - 1.0);
                const double closure = rho.GetValue(element, ip) - eos_density;
                const double variation = drho.GetValue(element, ip) - eos_derivative * dh.GetValue(element, ip);
                mapping::VolumeMappingContext context;
                const auto status = f.mapper.evaluate_volume(mapping_data, transformation, ip, workspace, context);
                utils::errors::serif_ensure<utils::errors::SERiFErrorCode::INVALID_MAPPING>(
                    status == mapping::MappingStatus::VALID, "Invalid manual-reference geometry."
                );
                double integrand = context.quadrature.physical_integration_weight * closure;
                if (action) {
                    mapping::VolumeMappingVariation geometry_variation;
                    const auto variation_status = serif::tests::evaluate_volume_variation(
                        f.mapper, mapping_data, du_data, transformation, ip, context, workspace, geometry_variation
                    );
                    utils::errors::serif_ensure<utils::errors::SERiFErrorCode::INVALID_MAPPING>(
                        variation_status == mapping::MappingStatus::VALID, "Invalid manual-reference variation."
                    );
                    integrand = context.quadrature.physical_integration_weight * variation +
                                geometry_variation.weight_variation * closure;
                }
                test.CalcShape(ip, shape);
                local.Add(integrand, shape);
            }
            auto* test_transformation = f.density_space.GetElementVDofs(element, dofs);
            if (test_transformation != nullptr) {
                test_transformation->TransformDual(local);
            }
            result.AddElementVector(dofs, local);
        }
        return f.density_binding.gather(result);
    }
} // namespace serif::tests::equations