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_force, quadratureOptions.gravity_force);

        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