module;
#include <cmath>
#include <format>
#include <numbers>
module mean_field;
import :model.structure.polytropic;

namespace mean_field::models::structure {
    PolytropicStructure::PolytropicStructure(
        eos::Polytrope equationOfState,
        const double targetMass
    )
        : m_equationOfState(std::move(equationOfState)),
          m_targetMass(targetMass) {
        validate();
    }

    const eos::EquationOfState &PolytropicStructure::equationOfState() const noexcept {
        return m_equationOfState;
    }

    double PolytropicStructure::targetMass() const noexcept {
        return m_targetMass;
    }

    StructureSeed PolytropicStructure::makeInitialSeed(const StructureSeedRequest &request) const {
        validateSeedRequest(request);

        const double polytropicIndex                        = m_equationOfState.polytropic_index();
        const std::vector<LaneEmdenPoint> laneEmdenSolution = solveLaneEmden(polytropicIndex);
        const double surfaceCoordinate                      = laneEmdenSolution.back().coordinate;
        const double centralEnthalpy = m_equationOfState.enthalpy_from_density(request.centralDensity);
        const double radialScaleSquared =
            centralEnthalpy / (4.0 * std::numbers::pi_v<double> * mean_field::utils::G * request.centralDensity);

        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);

        StructureSeed seed;

        seed.radius.SetSize(request.radialSampleCount);
        seed.density.SetSize(request.radialSampleCount);
        seed.enthalpy.SetSize(request.radialSampleCount);

        seed.stellarRadius             = radialScale * surfaceCoordinate;
        seed.centralDensity            = request.centralDensity;
        seed.centralEnthalpy           = centralEnthalpy;

        std::size_t interpolationIndex = 0;

        for (int sampleIndex = 0; sampleIndex < request.radialSampleCount; ++sampleIndex) {
            const double sampleFraction =
                static_cast<double>(sampleIndex) / static_cast<double>(request.radialSampleCount - 1);

            const double dimensionlessRadius = sampleFraction * surfaceCoordinate;
            const double laneEmdenValue =
                interpolateLaneEmdenValue(laneEmdenSolution, dimensionlessRadius, interpolationIndex);
            const double density       = request.centralDensity * std::pow(laneEmdenValue, polytropicIndex);

            seed.radius(sampleIndex)   = radialScale * dimensionlessRadius;
            seed.density(sampleIndex)  = density;
            seed.enthalpy(sampleIndex) = m_equationOfState.enthalpy_from_density(density);
        }

        seed.radius(0)              = 0.0;
        seed.density(0)             = request.centralDensity;
        seed.enthalpy(0)            = centralEnthalpy;

        const int surfaceIndex      = request.radialSampleCount - 1;

        seed.radius(surfaceIndex)   = seed.stellarRadius;
        seed.density(surfaceIndex)  = 0.0;
        seed.enthalpy(surfaceIndex) = 0.0;

        return seed;
    }

    void PolytropicStructure::validate() const {
        const double polytropicIndex = m_equationOfState.polytropic_index();

        if (!std::isfinite(polytropicIndex) || polytropicIndex < 1.0 || polytropicIndex >= 5.0) {
            throw std::invalid_argument(
                std::format(
                    "PolytropicStructure requires a finite-radius "
                    "polytrope with 1 <= n < 5. Instead n = {} was "
                    "provided.",
                    polytropicIndex
                )
            );
        }

        if (!std::isfinite(m_targetMass) || m_targetMass <= 0.0) {
            throw std::invalid_argument(
                std::format(
                    "The target stellar mass must be finite and "
                    "positive. Instead a value of {} was provided.",
                    m_targetMass
                )
            );
        }
    }

    void PolytropicStructure::validateSeedRequest(const StructureSeedRequest &request) {
        if (!std::isfinite(request.centralDensity) || request.centralDensity <= 0.0) {
            throw std::invalid_argument(
                std::format(
                    "The seed central density must be finite and "
                    "positive. Instead a value of {} was provided.",
                    request.centralDensity
                )
            );
        }

        if (request.radialSampleCount < 2) {
            throw std::invalid_argument(
                std::format(
                    "A polytropic seed requires at least two radial "
                    "samples. Instead {} samples were requested.",
                    request.radialSampleCount
                )
            );
        }
    }

    PolytropicStructure::LaneEmdenDerivative PolytropicStructure::evaluateLaneEmdenRhs(
        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)
        };
    }

    PolytropicStructure::LaneEmdenPoint PolytropicStructure::takeLaneEmdenStep(
        const LaneEmdenPoint &point,
        const double step,
        const double polytropicIndex
    ) {
        const LaneEmdenDerivative first =
            evaluateLaneEmdenRhs(point.coordinate, point.value, point.derivative, polytropicIndex);

        const LaneEmdenDerivative second = evaluateLaneEmdenRhs(
            point.coordinate + 0.5 * step, point.value + 0.5 * step * first.value,
            point.derivative + 0.5 * step * first.derivative, polytropicIndex
        );

        const LaneEmdenDerivative third = evaluateLaneEmdenRhs(
            point.coordinate + 0.5 * step, point.value + 0.5 * step * second.value,
            point.derivative + 0.5 * step * second.derivative, polytropicIndex
        );

        const LaneEmdenDerivative fourth = evaluateLaneEmdenRhs(
            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)
        };
    }

    std::vector<PolytropicStructure::LaneEmdenPoint> PolytropicStructure::solveLaneEmden(const double polytropicIndex) {
        constexpr double initialCoordinate = 1.0e-6;
        constexpr double integrationStep   = 1.0e-3;
        constexpr int maximumStepCount     = 2'000'000;

        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; ++stepIndex) {
            LaneEmdenPoint nextPoint = takeLaneEmdenStep(point, integrationStep, polytropicIndex);

            if (!std::isfinite(nextPoint.value)) {
                throw std::runtime_error(
                    "The Lane-Emden integration produced a non-finite "
                    "solution before reaching the stellar surface."
                );
            }

            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;
        }

        throw std::runtime_error(
            "The Lane-Emden integration did not reach its first zero "
            "within the configured step limit."
        );
    }

    double PolytropicStructure::interpolateLaneEmdenValue(
        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 mean_field::models::structure