module;

#include <algorithm>
#include <cmath>
#include <expected>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>

#include <mfem.hpp>
#include <stroid/stroid.h>

module mean_field;

import :boundary.contexts;
import :field.mfem;
import :mapping.coefficients;
import :utils.misc;
import :utils.user;

namespace mean_field::fem {
    FEM setup_fem(
        const std::string &filename,
        const utils::Args &args,
        const int extraRefine
    ) {
        FEM fem;

        using GravityPotential   = field::Gravity::Potential;
        using GravityFlux        = field::Gravity::Flux;
        using DisplacementVector = field::Displacement::Vector;
        using DensityScalar      = field::Density::Scalar;
        using EnthalpyScalar     = field::Enthalpy::Scalar;

        // =====================================================================
        // Section 1: Mesh construction
        // =====================================================================

        fem.smesh                = stroid::IO::LoadStroidMesh(filename).value();

        if (extraRefine > 0) {
            stroid::refinement::UniformRefinement(fem.smesh, extraRefine);
        }

        int mpiSize = 1;
        MPI_Comm_size(MPI_COMM_WORLD, &mpiSize);

        const std::unique_ptr<int[]> meshPartitioning(
            fem.smesh.mesh->GeneratePartitioning(mpiSize, 1)
        );

        fem.mesh = std::make_unique<mfem::ParMesh>(
            MPI_COMM_WORLD, *fem.smesh.mesh, meshPartitioning.get(), 1
        );

        fem.mesh->EnsureNodes();

        // =====================================================================
        // Section 2: Exterior compactification coordinate
        // =====================================================================

        if (fem.smesh.exterior_coordinate == nullptr) {
            throw std::runtime_error("Exterior coordinate not set.");
        }

        if (fem.smesh.exterior_coordinate->space == nullptr) {
            throw std::runtime_error("Space for exterior coordinate not set.");
        }

        if (fem.smesh.exterior_coordinate->values == nullptr) {
            throw std::runtime_error("Values for exterior coordinate not set.");
        }

        const mfem::FiniteElementSpace &serialCoordinateSpace =
            *fem.smesh.exterior_coordinate->space;

        const mfem::GridFunction &serialCoordinate =
            *fem.smesh.exterior_coordinate->values;

        if (serialCoordinate.FESpace() != &serialCoordinateSpace) {
            throw std::runtime_error(
                "Exterior coordinate values are not associated with the "
                "supplied finite-element space."
            );
        }

        if (serialCoordinateSpace.GetMesh() != fem.smesh.mesh.get()) {
            throw std::runtime_error(
                "Exterior coordinate space is not associated with the "
                "loaded STROID mesh."
            );
        }

        if (serialCoordinateSpace.GetVDim() != 1) {
            throw std::runtime_error(
                "Exterior coordinate must be a scalar field."
            );
        }

        if (serialCoordinate.Size() != serialCoordinateSpace.GetVSize()) {
            throw std::runtime_error(
                "Exterior coordinate value count does not match its "
                "finite-element space."
            );
        }

        const int compactificationOrder =
            serialCoordinateSpace.GetMaxElementOrder();

        const int dimension     = fem.mesh->Dimension();

        fem.compactificationFec = std::make_unique<mfem::H1_FECollection>(
            compactificationOrder, dimension
        );

        fem.compactificationFes = std::make_unique<mfem::ParFiniteElementSpace>(
            fem.mesh.get(), fem.compactificationFec.get()
        );

        mfem::ParGridFunction distributedCoordinate(
            fem.mesh.get(), &serialCoordinate, meshPartitioning.get()
        );

        if (distributedCoordinate.Size() !=
            fem.compactificationFes->GetVSize()) {
            throw std::runtime_error(
                "Distributed exterior coordinate does not match the "
                "constructed parallel finite-element space."
            );
        }

        fem.compactificationCoordinate =
            std::make_unique<mfem::ParGridFunction>(
                fem.compactificationFes.get()
            );

        *fem.compactificationCoordinate = distributedCoordinate;

        double localMinimum = std::numeric_limits<double>::infinity();

        double localMaximum = -std::numeric_limits<double>::infinity();

        for (int index = 0; index < fem.compactificationCoordinate->Size();
             ++index) {
            const double value = (*fem.compactificationCoordinate)(index);

            if (!std::isfinite(value)) {
                throw std::runtime_error(
                    "Exterior coordinate contains a non-finite value."
                );
            }

            localMinimum = std::min(localMinimum, value);

            localMaximum = std::max(localMaximum, value);
        }

        double globalMinimum = 0.0;
        double globalMaximum = 0.0;

        MPI_Allreduce(
            &localMinimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN,
            MPI_COMM_WORLD
        );

        MPI_Allreduce(
            &localMaximum, &globalMaximum, 1, MPI_DOUBLE, MPI_MAX,
            MPI_COMM_WORLD
        );

        constexpr double coordinateTolerance = 1.0e-12;

        if (globalMinimum < -coordinateTolerance ||
            globalMaximum > 1.0 + coordinateTolerance) {
            throw std::runtime_error(
                "Exterior coordinate lies outside the expected "
                "interval [0, 1]."
            );
        }

        // =====================================================================
        // Section 3: Compile-time field realization
        // =====================================================================

        // ---------------------------------------------------------------------
        // Gravity potential: scalar L2
        // ---------------------------------------------------------------------

        fem.gravityPotentialFec =
            GravityField::make_fec<GravityPotential>(dimension);

        fem.gravityPotentialFes = GravityField::make_fespace<GravityPotential>(
            *fem.mesh, *fem.gravityPotentialFec
        );

        // ---------------------------------------------------------------------
        // Gravity flux: H(div)/RT. Basis choices are encoded by field.mfem.
        // ---------------------------------------------------------------------

        fem.gravityFluxFec = GravityField::make_fec<GravityFlux>(dimension);

        fem.gravityFluxFes = GravityField::make_fespace<GravityFlux>(
            *fem.mesh, *fem.gravityFluxFec
        );

        // ---------------------------------------------------------------------
        // Displacement: vector H1. Ordering is encoded by field.mfem.
        // ---------------------------------------------------------------------

        fem.displacementFec =
            DisplacementField::make_fec<DisplacementVector>(dimension);

        fem.displacementFes =
            DisplacementField::make_fespace<DisplacementVector>(
                *fem.mesh, *fem.displacementFec
            );

        fem.displacement =
            std::make_unique<mfem::ParGridFunction>(fem.displacementFes.get());

        *fem.displacement = 0.0;

        // ---------------------------------------------------------------------
        // Density: scalar discontinuous L2
        // ---------------------------------------------------------------------

        fem.densityFec    = DensityField::make_fec<DensityScalar>(dimension);

        fem.densityFes    = DensityField::make_fespace<DensityScalar>(
            *fem.mesh, *fem.densityFec
        );

        // ---------------------------------------------------------------------
        // Specific enthalpy: scalar continuous H1
        // ---------------------------------------------------------------------

        fem.enthalpyFec = EnthalpyField::make_fec<EnthalpyScalar>(dimension);

        fem.enthalpyFes = EnthalpyField::make_fespace<EnthalpyScalar>(
            *fem.mesh, *fem.enthalpyFec
        );

        // =====================================================================
        // Section 4: Domain mapping
        // =====================================================================

        auto [stellarRadiusReference, infinityRadiusReference] =
            utils::discover_bounds(fem.mesh.get(), 3)
                .or_else(
                    [](const boundary::BoundsError &)
                        -> std::expected<
                            boundary::Bounds, boundary::BoundsError> {
                        throw std::runtime_error(
                            "Unable to determine vacuum-domain reference "
                            "boundaries."
                        );
                    }
                )
                .value();

        fem.mapping = std::make_unique<mapping::DomainMapper>(
            *fem.displacement, stellarRadiusReference, infinityRadiusReference
        );

        // =====================================================================
        // Section 5: Block offsets
        //
        // Legacy layouts only. New coupled operators use :utils.blocks forms.
        //
        // Main system:    [Displacement | Density]
        // Gravity system: [Flux | Potential]
        // =====================================================================

        fem.blockTrueOffsets.SetSize(3);
        fem.blockTrueOffsets[0] = 0;

        fem.blockTrueOffsets[1] = fem.displacementFes->GetTrueVSize();

        fem.blockTrueOffsets[2] =
            fem.blockTrueOffsets[1] + fem.densityFes->GetTrueVSize();

        fem.gravityBlockTrueOffsets.SetSize(3);
        fem.gravityBlockTrueOffsets[0] = 0;

        fem.gravityBlockTrueOffsets[1] = fem.gravityFluxFes->GetTrueVSize();

        fem.gravityBlockTrueOffsets[2] =
            fem.gravityBlockTrueOffsets[1] +
            fem.gravityPotentialFes->GetTrueVSize();

        // =====================================================================
        // Section 6: Multipole data
        // =====================================================================

        fem.com.SetSize(dimension);
        fem.com = 0.0;

        fem.Q.SetSize(dimension, dimension);
        fem.Q = 0.0;

        // =====================================================================
        // Section 7: Essential boundaries and domain masks
        // =====================================================================

        fem.essentialDisplacementTdofs.SetSize(0);

        populate_element_mask(
            fem.mesh.get(), utils::DOMAINS::STELLAR,
            fem.gravityContext.stellar_mask
        );

        const int boundaryAttributeCount = fem.mesh->bdr_attributes.Max();

        fem.boundaryContext.inf_bounds.SetSize(boundaryAttributeCount);

        fem.boundaryContext.stellar_bounds.SetSize(boundaryAttributeCount);

        fem.boundaryContext.inf_bounds     = 0;
        fem.boundaryContext.stellar_bounds = 0;

        fem.boundaryContext.inf_bounds
            [static_cast<int>(boundary::Boundaries::INF_SURFACE) - 1] = 1;

        fem.boundaryContext.stellar_bounds
            [static_cast<int>(boundary::Boundaries::STELLAR_SURFACE) - 1] = 1;

        // =====================================================================
        // Section 8: Gravity solver context
        // =====================================================================

        fem.gravityContext.minres =
            std::make_unique<mfem::MINRESSolver>(fem.mesh->GetComm());

        fem.gravityContext.minres->SetRelTol(1.0e-12);
        fem.gravityContext.minres->SetAbsTol(1.0e-12);
        fem.gravityContext.minres->SetMaxIter(1000);
        fem.gravityContext.minres->SetPrintLevel(0);

        fem.gravityContext.prec_Phi = std::make_unique<mfem::HypreBoomerAMG>();

        fem.gravityContext.prec_Phi->SetPrintLevel(0);

        fem.gravityContext.block_prec =
            std::make_unique<mfem::BlockDiagonalPreconditioner>(
                fem.gravityBlockTrueOffsets
            );

        fem.gravityContext.minres->SetPreconditioner(
            *fem.gravityContext.block_prec
        );

        // =====================================================================
        // Section 9: Vacuum true-DOF masks
        // =====================================================================

        {
            mfem::Array<int> vacuumMask;

            utils::populate_element_mask(
                fem.mesh.get(), utils::DOMAINS::VACUUM, vacuumMask
            );

            utils::populate_domain_tdofs(
                fem.displacementFes.get(), vacuumMask,
                fem.vacuumDisplacementTdofs
            );

            utils::populate_domain_tdofs(
                fem.densityFes.get(), vacuumMask, fem.vacuumDensityTdofs
            );

            utils::populate_domain_tdofs(
                fem.enthalpyFes.get(), vacuumMask, fem.vacuumEnthalpyTdofs
            );
        }

        // =====================================================================
        // Section 10: Quadrature policy
        // =====================================================================

        const quadrature::QuadratureOptions &quadratureOptions =
            args.quadrature;

        if (quadratureOptions.validation.reject_negative_boosts &&
            quadratureOptions.global_boost < 0) {
            throw std::invalid_argument(
                "Global quadrature boost cannot be negative."
            );
        }

        quadrature::RuleSet quadratureRuleSet = quadrature::make_rule_set(
            quadratureOptions.mode, quadratureOptions.global_boost
        );

        if (quadratureOptions.fallback_fixed_order.has_value()) {
            if (*quadratureOptions.fallback_fixed_order < 0) {
                throw std::invalid_argument(
                    "Fallback quadrature order cannot be negative."
                );
            }

            quadratureRuleSet.fallback.fixed_order =
                quadratureOptions.fallback_fixed_order;
        }

        auto apply_quadrature_options =
            [&quadratureOptions](
                quadrature::RuleControl &ruleControl,
                const quadrature::QuadratureTermOptions &termOptions
            ) {
                if (termOptions.fixed_order.has_value() &&
                    *termOptions.fixed_order < 0) {
                    throw std::invalid_argument(
                        "Fixed quadrature order cannot be negative."
                    );
                }

                if (quadratureOptions.validation.reject_negative_boosts &&
                    termOptions.additional_boost < 0) {
                    throw std::invalid_argument(
                        "Term quadrature boost cannot be negative."
                    );
                }

                ruleControl.boost += termOptions.additional_boost;

                if (termOptions.fixed_order.has_value()) {
                    ruleControl.fixed_order = termOptions.fixed_order;
                }
            };

        apply_quadrature_options(
            quadratureRuleSet.gravity_hdiv_mass,
            quadratureOptions.gravity_hdiv_mass
        );

        apply_quadrature_options(
            quadratureRuleSet.gravity_divergence,
            quadratureOptions.gravity_divergence
        );

        apply_quadrature_options(
            quadratureRuleSet.gravity_source, quadratureOptions.gravity_source
        );

        apply_quadrature_options(
            quadratureRuleSet.gravity_boundary,
            quadratureOptions.gravity_boundary
        );

        apply_quadrature_options(
            quadratureRuleSet.centrifugal, quadratureOptions.centrifugal
        );

        apply_quadrature_options(
            quadratureRuleSet.density_projection,
            quadratureOptions.density_projection
        );

        apply_quadrature_options(
            quadratureRuleSet.eos_closure, quadratureOptions.eos_closure
        );

        apply_quadrature_options(
            quadratureRuleSet.hydrostatic_equilibrium,
            quadratureOptions.hydrostatic_equilibrium
        );

        apply_quadrature_options(
            quadratureRuleSet.isobaric_surface,
            quadratureOptions.isobaric_surface
        );

        apply_quadrature_options(
            quadratureRuleSet.mesh_extension, quadratureOptions.mesh_extension
        );

        apply_quadrature_options(
            quadratureRuleSet.mass_conservation,
            quadratureOptions.mass_conservation
        );

        apply_quadrature_options(
            quadratureRuleSet.mass_normalization,
            quadratureOptions.mass_normalization
        );

        apply_quadrature_options(
            quadratureRuleSet.center_of_mass, quadratureOptions.center_of_mass
        );

        apply_quadrature_options(
            quadratureRuleSet.quadrupole, quadratureOptions.quadrupole
        );

        apply_quadrature_options(
            quadratureRuleSet.gravitational_energy,
            quadratureOptions.gravitational_energy
        );

        apply_quadrature_options(
            quadratureRuleSet.pressure_integral,
            quadratureOptions.pressure_integral
        );

        apply_quadrature_options(
            quadratureRuleSet.pressure_force, quadratureOptions.pressure_force
        );

        apply_quadrature_options(
            quadratureRuleSet.virial, quadratureOptions.virial
        );

        apply_quadrature_options(
            quadratureRuleSet.error_norm, quadratureOptions.error_norm
        );

        apply_quadrature_options(
            quadratureRuleSet.roles.discretization,
            quadratureOptions.roles.discretization
        );

        apply_quadrature_options(
            quadratureRuleSet.roles.preconditioner,
            quadratureOptions.roles.preconditioner
        );

        apply_quadrature_options(
            quadratureRuleSet.roles.diagnostic,
            quadratureOptions.roles.diagnostic
        );

        apply_quadrature_options(
            quadratureRuleSet.roles.projection,
            quadratureOptions.roles.projection
        );

        fem.quadratureFactory = std::make_unique<quadrature::RuleFactory>(
            quadrature::Policy(std::move(quadratureRuleSet))
        );

        // =====================================================================
        // Section 11: Stateless domain mapper
        // =====================================================================

        auto exteriorDomain = std::make_unique<
            const mapping::compactification::KelvinCompactification>(
            args.kelvin_options
        );

        fem.domainMapperStateless =
            std::make_unique<mapping::DomainMapperStateless>(
                args.domain_mapper_options, std::move(exteriorDomain)
            );

        return fem;
    }
} // namespace mean_field::fem