module;

#include <algorithm>
#include <cmath>
#include <limits>
#include <numbers>
#include <stdexcept>

#include <mfem.hpp>
#include <mpi.h>

module mean_field;

import :field.mfem;
import :seed.stellar_equilibrium_projection;
import :utils.domain;
import :utils.misc;

namespace {
    using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;

    void validate_profile(const mean_field::seed::RadialProfile &profile) {
        const int sampleCount = profile.radius.Size();
        if (sampleCount < 2 || profile.density.Size() != sampleCount ||
            profile.specificEnthalpy.Size() != sampleCount) {
            throw std::invalid_argument("A radial seed projection requires equally sized profiles with two samples.");
        }
        if (!std::isfinite(profile.stellarRadius.value()) || profile.stellarRadius.value() <= 0.0 ||
            !std::isfinite(profile.centralDensity.value()) || profile.centralDensity.value() <= 0.0 ||
            !std::isfinite(profile.centralSpecificEnthalpy.value()) || profile.centralSpecificEnthalpy.value() <= 0.0) {
            throw std::invalid_argument("A radial seed projection requires finite, positive physical scales.");
        }

        for (int index = 0; index < sampleCount; ++index) {
            if (!std::isfinite(profile.radius(index)) || !std::isfinite(profile.density(index)) ||
                !std::isfinite(profile.specificEnthalpy(index)) || profile.density(index) < 0.0 ||
                profile.specificEnthalpy(index) < 0.0) {
                throw std::invalid_argument("A radial seed projection received a non-finite or negative profile.");
            }
            if (index > 0 && profile.radius(index) <= profile.radius(index - 1)) {
                throw std::invalid_argument("A radial seed projection requires strictly increasing radii.");
            }
        }

        const int surfaceIndex   = sampleCount - 1;
        const double radialScale = std::max(profile.stellarRadius.value(), 1.0);
        if (std::abs(profile.radius(0)) > 64.0 * std::numeric_limits<double>::epsilon() * radialScale ||
            std::abs(profile.radius(surfaceIndex) - profile.stellarRadius.value()) >
                64.0 * std::numeric_limits<double>::epsilon() * radialScale ||
            profile.density(0) != profile.centralDensity.value() ||
            profile.specificEnthalpy(0) != profile.centralSpecificEnthalpy.value() ||
            profile.density(surfaceIndex) != 0.0 || profile.specificEnthalpy(surfaceIndex) != 0.0) {
            throw std::invalid_argument("A radial seed projection received inconsistent center or surface metadata.");
        }
    }

    [[nodiscard]] double interpolate_profile(
        const mfem::Vector &radius,
        const mfem::Vector &values,
        const double requestedRadius
    ) {
        if (requestedRadius <= radius(0)) {
            return values(0);
        }
        const int finalIndex = radius.Size() - 1;
        if (requestedRadius >= radius(finalIndex)) {
            return values(finalIndex);
        }

        int lowerIndex = 0;
        int upperIndex = finalIndex;
        while (upperIndex - lowerIndex > 1) {
            const int middleIndex = lowerIndex + (upperIndex - lowerIndex) / 2;
            if (radius(middleIndex) <= requestedRadius) {
                lowerIndex = middleIndex;
            } else {
                upperIndex = middleIndex;
            }
        }

        const double fraction = (requestedRadius - radius(lowerIndex)) / (radius(upperIndex) - radius(lowerIndex));
        return (1.0 - fraction) * values(lowerIndex) + fraction * values(upperIndex);
    }

    struct SurfaceRadiusRange final {
        double minimum;
        double maximum;
    };

    [[nodiscard]] SurfaceRadiusRange measure_surface_radius(const mean_field::fem::FEM &finiteElementModel) {
        if (finiteElementModel.surfaceDeformationFes == nullptr) {
            throw std::invalid_argument("Radial seed projection requires the surface-deformation space.");
        }

        mfem::ParFiniteElementSpace &surfaceSpace = *finiteElementModel.surfaceDeformationFes;
        const mean_field::field::ScalarBoundaryDofMap surfaceMap =
            mean_field::field::make_stellar_surface_scalar_dof_map<DomainSchema>(surfaceSpace);
        mfem::Vector radiusSquared(surfaceMap.local_size());
        radiusSquared = 0.0;

        mfem::ParGridFunction coordinateField(&surfaceSpace);
        for (int component = 0; component < surfaceSpace.GetMesh()->SpaceDimension(); ++component) {
            mfem::FunctionCoefficient coordinateCoefficient([component](const mfem::Vector &position) {
                return position(component);
            });
            coordinateField.ProjectCoefficient(coordinateCoefficient);
            mfem::Vector coordinateTrue;
            coordinateField.GetTrueDofs(coordinateTrue);
            const mfem::Vector surfaceCoordinate = surfaceMap.gather(coordinateTrue);
            for (int index = 0; index < radiusSquared.Size(); ++index) {
                radiusSquared(index) += surfaceCoordinate(index) * surfaceCoordinate(index);
            }
        }

        double localMinimum = std::numeric_limits<double>::infinity();
        double localMaximum = 0.0;
        for (int index = 0; index < radiusSquared.Size(); ++index) {
            const double radius = std::sqrt(radiusSquared(index));
            localMinimum        = std::min(localMinimum, radius);
            localMaximum        = std::max(localMaximum, radius);
        }

        double globalMinimum = 0.0;
        double globalMaximum = 0.0;
        MPI_Allreduce(&localMinimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN, surfaceSpace.GetComm());
        MPI_Allreduce(&localMaximum, &globalMaximum, 1, MPI_DOUBLE, MPI_MAX, surfaceSpace.GetComm());
        if (!std::isfinite(globalMinimum) || !std::isfinite(globalMaximum) || globalMinimum <= 0.0 ||
            globalMaximum < globalMinimum) {
            throw std::runtime_error("The stellar surface has no finite, positive radial extent.");
        }
        return {.minimum = globalMinimum, .maximum = globalMaximum};
    }
} // namespace

namespace mean_field::seed::detail {
    ProjectedRadialFields projectRadialFields(
        fem::FEM &finiteElementModel,
        const RadialProfile &profile,
        const dimensions::MassValue targetMass,
        const StellarEquilibriumProjectionOptions &options
    ) {
        validate_profile(profile);
        if (!std::isfinite(options.surfaceRadiusRelativeTolerance) || options.surfaceRadiusRelativeTolerance < 0.0) {
            throw std::invalid_argument("The surface-radius projection tolerance must be finite and nonnegative.");
        }
        const SurfaceRadiusRange surfaceRadius = measure_surface_radius(finiteElementModel);
        const double targetRadius              = profile.stellarRadius.value();
        const double comparisonScale           = std::max({targetRadius, surfaceRadius.maximum, 1.0e-300});
        const double relativeMismatch =
            std::max(std::abs(surfaceRadius.minimum - targetRadius), std::abs(surfaceRadius.maximum - targetRadius)) /
            comparisonScale;
        if (relativeMismatch > options.surfaceRadiusRelativeTolerance) {
            throw std::invalid_argument(
                "The radial seed surface does not coincide with the spherical reference discretization."
            );
        }

        if (finiteElementModel.densityFes == nullptr || finiteElementModel.enthalpyFes == nullptr ||
            finiteElementModel.displacementFes == nullptr || finiteElementModel.gravityFluxFes == nullptr ||
            finiteElementModel.gravityPotentialFes == nullptr) {
            throw std::invalid_argument("Radial seed projection requires the complete equilibrium discretization.");
        }

        mfem::FunctionCoefficient densityCoefficient([&profile](const mfem::Vector &position) {
            return interpolate_profile(profile.radius, profile.density, position.Norml2());
        });
        mfem::FunctionCoefficient enthalpyCoefficient([&profile](const mfem::Vector &position) {
            return interpolate_profile(profile.radius, profile.specificEnthalpy, position.Norml2());
        });

        mfem::ParGridFunction densityField(finiteElementModel.densityFes.get());
        mfem::ParGridFunction enthalpyField(finiteElementModel.enthalpyFes.get());
        mfem::ParGridFunction displacementField(finiteElementModel.displacementFes.get());
        densityField      = 0.0;
        enthalpyField     = 0.0;
        displacementField = 0.0;
        densityField.ProjectCoefficient(densityCoefficient);
        enthalpyField.ProjectCoefficient(enthalpyCoefficient);

        const physics::GravitySolution gravitySolution =
            physics::solve_gravity_field(finiteElementModel, options.gravity, densityField, displacementField);

        double radialMomentIntegral = 0.0;
        for (int index = 0; index + 1 < profile.radius.Size(); ++index) {
            const double leftRadius     = profile.radius(index);
            const double rightRadius    = profile.radius(index + 1);
            const double leftIntegrand  = profile.density(index) * std::pow(leftRadius, 4);
            const double rightIntegrand = profile.density(index + 1) * std::pow(rightRadius, 4);
            radialMomentIntegral += 0.5 * (rightRadius - leftRadius) * (leftIntegrand + rightIntegrand);
        }
        const double sphericalMomentOfInertia = (8.0 * std::numbers::pi / 3.0) * radialMomentIntegral;
        if (!std::isfinite(sphericalMomentOfInertia) || sphericalMomentOfInertia <= 0.0) {
            throw std::runtime_error("The radial profile has no finite, positive moment of inertia.");
        }

        const field::FieldDofGridFunctionAdapter densityAdapter =
            field::make_field_dof_grid_function_adapter<field::Density, DomainSchema>(*finiteElementModel.densityFes);
        const field::FieldDofGridFunctionAdapter enthalpyAdapter =
            field::make_field_dof_grid_function_adapter<field::Enthalpy, DomainSchema>(*finiteElementModel.enthalpyFes);
        const field::FieldDofGridFunctionAdapter gravityFluxAdapter =
            field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(
                *finiteElementModel.gravityFluxFes
            );
        const field::FieldDofGridFunctionAdapter gravityPotentialAdapter =
            field::make_field_dof_grid_function_adapter<field::Gravity, DomainSchema>(
                *finiteElementModel.gravityPotentialFes
            );

        return {
            .density                  = densityAdapter.gather(densityField),
            .gravityGradient          = gravityFluxAdapter.gather(gravitySolution.gradPhi),
            .gravityPotential         = gravityPotentialAdapter.gather(gravitySolution.phi),
            .specificEnthalpy         = enthalpyAdapter.gather(enthalpyField),
            .bernoulliConstant        = -utils::G * targetMass.value() / targetRadius,
            .sphericalMomentOfInertia = sphericalMomentOfInertia
        };
    }
} // namespace mean_field::seed::detail