module;

#include <algorithm>
#include <cmath>
#include <numbers>
#include <optional>
#include <stdexcept>
#include <vector>

#include <mfem.hpp>

module mean_field;

import :seed.lane_emden;
import :utils.misc;

namespace {
    struct LaneEmdenPoint final {
        double coordinate{0.0};
        double value{0.0};
        double derivative{0.0};
    };

    struct LaneEmdenDerivative final {
        double value{0.0};
        double derivative{0.0};
    };

    [[nodiscard]] LaneEmdenDerivative evaluate_lane_emden_rhs(
        const double coordinate,
        const double value,
        const double derivative,
        const double polytropicIndex
    ) {
        const double nonnegativeValue = std::max(value, 0.0);
        return {
            .value      = derivative,
            .derivative = -2.0 * derivative / coordinate - std::pow(nonnegativeValue, polytropicIndex)
        };
    }

    [[nodiscard]] LaneEmdenPoint take_lane_emden_step(
        const LaneEmdenPoint &point,
        const double step,
        const double polytropicIndex
    ) {
        const LaneEmdenDerivative first =
            evaluate_lane_emden_rhs(point.coordinate, point.value, point.derivative, polytropicIndex);
        const LaneEmdenDerivative second = evaluate_lane_emden_rhs(
            point.coordinate + 0.5 * step, point.value + 0.5 * step * first.value,
            point.derivative + 0.5 * step * first.derivative, polytropicIndex
        );
        const LaneEmdenDerivative third = evaluate_lane_emden_rhs(
            point.coordinate + 0.5 * step, point.value + 0.5 * step * second.value,
            point.derivative + 0.5 * step * second.derivative, polytropicIndex
        );
        const LaneEmdenDerivative fourth = evaluate_lane_emden_rhs(
            point.coordinate + step, point.value + step * third.value, point.derivative + step * third.derivative,
            polytropicIndex
        );

        return {
            .coordinate = point.coordinate + step,
            .value = point.value + step / 6.0 * (first.value + 2.0 * second.value + 2.0 * third.value + fourth.value),
            .derivative =
                point.derivative +
                step / 6.0 * (first.derivative + 2.0 * second.derivative + 2.0 * third.derivative + fourth.derivative)
        };
    }

    [[nodiscard]] std::vector<LaneEmdenPoint> solve_lane_emden(
        const double polytropicIndex,
        const double coordinateLimit,
        const double integrationStep
    ) {
        if (!std::isfinite(polytropicIndex) || polytropicIndex < 0.0) {
            throw std::invalid_argument("Lane-Emden integration requires a finite, nonnegative polytropic index.");
        }
        if (!std::isfinite(coordinateLimit) || coordinateLimit <= 0.0) {
            throw std::invalid_argument("The Lane-Emden coordinate limit must be finite and positive.");
        }
        if (!std::isfinite(integrationStep) || integrationStep <= 0.0) {
            throw std::invalid_argument("The Lane-Emden integration step must be finite and positive.");
        }

        constexpr int maximumStepCount = 2'000'000;
        if (std::ceil(coordinateLimit / integrationStep) > static_cast<double>(maximumStepCount)) {
            throw std::invalid_argument("The requested Lane-Emden interval exceeds the integration step limit.");
        }

        const double initialCoordinate = std::min(1.0e-6, coordinateLimit);

        const double coordinateSquared = initialCoordinate * initialCoordinate;
        const double coordinateCubed   = coordinateSquared * initialCoordinate;
        const double coordinateFourth  = coordinateSquared * coordinateSquared;

        LaneEmdenPoint point{
            .coordinate = initialCoordinate,
            .value      = 1.0 - coordinateSquared / 6.0 + polytropicIndex * coordinateFourth / 120.0,
            .derivative = -initialCoordinate / 3.0 + polytropicIndex * coordinateCubed / 30.0
        };

        std::vector<LaneEmdenPoint> solution;
        solution.reserve(8192);
        solution.push_back({.coordinate = 0.0, .value = 1.0, .derivative = 0.0});
        solution.push_back(point);

        for (int stepIndex = 0; stepIndex < maximumStepCount && point.coordinate < coordinateLimit; ++stepIndex) {
            const double step        = std::min(integrationStep, coordinateLimit - point.coordinate);
            LaneEmdenPoint nextPoint = take_lane_emden_step(point, step, polytropicIndex);
            if (!std::isfinite(nextPoint.value)) {
                throw std::runtime_error(
                    "The Lane-Emden integration produced a non-finite solution before reaching its termination."
                );
            }
            if (nextPoint.value <= 0.0) {
                const double rootFraction = point.value / (point.value - nextPoint.value);
                solution.push_back(
                    {.coordinate = point.coordinate + rootFraction * (nextPoint.coordinate - point.coordinate),
                     .value      = 0.0,
                     .derivative = point.derivative + rootFraction * (nextPoint.derivative - point.derivative)}
                );
                return solution;
            }
            solution.push_back(nextPoint);
            point = nextPoint;
        }

        if (point.coordinate < coordinateLimit) {
            throw std::runtime_error("The Lane-Emden integration exceeded its step limit.");
        }
        return solution;
    }

    [[nodiscard]] double interpolate_lane_emden_value(
        const std::vector<LaneEmdenPoint> &solution,
        const double coordinate,
        std::size_t &lowerIndex
    ) {
        while (lowerIndex + 1 < solution.size() && solution[lowerIndex + 1].coordinate < coordinate) {
            ++lowerIndex;
        }
        if (lowerIndex + 1 >= solution.size()) {
            return 0.0;
        }

        const LaneEmdenPoint &lower = solution[lowerIndex];
        const LaneEmdenPoint &upper = solution[lowerIndex + 1];
        const double interval       = upper.coordinate - lower.coordinate;
        if (interval <= 0.0) {
            throw std::runtime_error("The Lane-Emden interpolation grid is not strictly increasing.");
        }
        const double fraction = (coordinate - lower.coordinate) / interval;
        return std::clamp(lower.value + fraction * (upper.value - lower.value), 0.0, 1.0);
    }
} // namespace

namespace mean_field::seed {
    DimensionlessLaneEmdenSolution integrateLaneEmden(
        const double polytropicIndex,
        const double coordinateLimit,
        const double integrationStep
    ) {
        const std::vector<LaneEmdenPoint> points = solve_lane_emden(polytropicIndex, coordinateLimit, integrationStep);

        DimensionlessLaneEmdenSolution solution{
            .coordinate          = mfem::Vector(static_cast<int>(points.size())),
            .theta               = mfem::Vector(static_cast<int>(points.size())),
            .thetaDerivative     = mfem::Vector(static_cast<int>(points.size())),
            .firstZeroCoordinate = std::nullopt
        };
        for (int index = 0; index < static_cast<int>(points.size()); ++index) {
            solution.coordinate(index)      = points[static_cast<std::size_t>(index)].coordinate;
            solution.theta(index)           = points[static_cast<std::size_t>(index)].value;
            solution.thetaDerivative(index) = points[static_cast<std::size_t>(index)].derivative;
        }
        if (points.back().value == 0.0) {
            solution.firstZeroCoordinate = points.back().coordinate;
        }
        return solution;
    }

    RadialProfile generateLaneEmdenProfile(
        const eos::Polytrope &equationOfState,
        const dimensions::DensityValue centralDensity,
        const int radialSampleCount
    ) {
        if (!std::isfinite(centralDensity.value()) || centralDensity.value() <= 0.0) {
            throw std::invalid_argument("A Lane-Emden seed central density must be finite and positive.");
        }
        if (radialSampleCount < 2) {
            throw std::invalid_argument("A Lane-Emden seed requires at least two radial samples.");
        }

        const double polytropicIndex = equationOfState.polytropic_index();
        if (!std::isfinite(polytropicIndex) || polytropicIndex < 1.0 || polytropicIndex >= 5.0) {
            throw std::invalid_argument("Lane-Emden seeds require a finite-radius polytrope with 1 <= n < 5.");
        }

        constexpr double seedCoordinateLimit = 2'000.0;
        constexpr double integrationStep     = 1.0e-3;
        const std::vector solution           = solve_lane_emden(polytropicIndex, seedCoordinateLimit, integrationStep);
        if (solution.back().value != 0.0) {
            throw std::runtime_error("The Lane-Emden integration did not reach its first zero within the step limit.");
        }
        const double surfaceCoordinate = solution.back().coordinate;
        const dimensions::SpecificEnthalpyValue centralEnthalpy =
            eos::evaluate<dimensions::quantity::SpecificEnthalpy>(equationOfState, centralDensity);
        const double radialScaleSquared = centralEnthalpy.value() / (4.0 * std::numbers::pi_v<double> *
                                                                     mean_field::utils::G * centralDensity.value());
        if (!std::isfinite(radialScaleSquared) || radialScaleSquared <= 0.0) {
            throw std::runtime_error("The polytropic Lane-Emden radial scale is not finite and positive.");
        }

        const double radialScale = std::sqrt(radialScaleSquared);
        RadialProfile profile{
            .radius                  = mfem::Vector(radialSampleCount),
            .density                 = mfem::Vector(radialSampleCount),
            .specificEnthalpy        = mfem::Vector(radialSampleCount),
            .stellarRadius           = dimensions::LengthValue{radialScale * surfaceCoordinate},
            .centralDensity          = centralDensity,
            .centralSpecificEnthalpy = centralEnthalpy
        };

        std::size_t interpolationIndex = 0;
        for (int sampleIndex = 0; sampleIndex < radialSampleCount; ++sampleIndex) {
            const double fraction = static_cast<double>(sampleIndex) / static_cast<double>(radialSampleCount - 1);
            const double dimensionlessRadius = fraction * surfaceCoordinate;
            const double laneEmdenValue =
                interpolate_lane_emden_value(solution, dimensionlessRadius, interpolationIndex);
            const dimensions::DensityValue density{centralDensity.value() * std::pow(laneEmdenValue, polytropicIndex)};

            profile.radius(sampleIndex)  = radialScale * dimensionlessRadius;
            profile.density(sampleIndex) = density.value();
            profile.specificEnthalpy(sampleIndex) =
                eos::evaluate<dimensions::quantity::SpecificEnthalpy>(equationOfState, density).value();
        }

        profile.radius(0)                      = 0.0;
        profile.density(0)                     = centralDensity.value();
        profile.specificEnthalpy(0)            = centralEnthalpy.value();
        const int surfaceIndex                 = radialSampleCount - 1;
        profile.radius(surfaceIndex)           = profile.stellarRadius.value();
        profile.density(surfaceIndex)          = 0.0;
        profile.specificEnthalpy(surfaceIndex) = 0.0;
        return profile;
    }
} // namespace mean_field::seed