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

#include <catch2/catch_approx.hpp>
#include <catch2/catch_test_macros.hpp>

import mean_field;
import test_helpers;

namespace {
    template <
        typename AnalyticValue,
        typename AnalyticDerivative>
    void check_dimensionless_solution(
        const mean_field::seed::DimensionlessLaneEmdenSolution &solution,
        AnalyticValue analyticValue,
        AnalyticDerivative analyticDerivative,
        const double tolerance
    ) {
        REQUIRE(solution.coordinate.Size() >= 2);
        REQUIRE(solution.theta.Size() == solution.coordinate.Size());
        REQUIRE(solution.thetaDerivative.Size() == solution.coordinate.Size());

        double maximumValueError      = 0.0;
        double maximumDerivativeError = 0.0;
        for (int index = 0; index < solution.coordinate.Size(); ++index) {
            const double coordinate = solution.coordinate(index);
            CHECK(std::isfinite(coordinate));
            CHECK(std::isfinite(solution.theta(index)));
            CHECK(std::isfinite(solution.thetaDerivative(index)));
            if (index > 0) {
                CHECK(coordinate > solution.coordinate(index - 1));
            }

            maximumValueError =
                std::max(maximumValueError, std::abs(solution.theta(index) - analyticValue(coordinate)));
            maximumDerivativeError = std::max(
                maximumDerivativeError, std::abs(solution.thetaDerivative(index) - analyticDerivative(coordinate))
            );
        }

        CHECK(maximumValueError < tolerance);
        CHECK(maximumDerivativeError < tolerance);
    }

    void check_profiles_are_identical(
        const mean_field::seed::RadialProfile &left,
        const mean_field::seed::RadialProfile &right
    ) {
        REQUIRE(left.radius.Size() == right.radius.Size());
        REQUIRE(left.density.Size() == right.density.Size());
        REQUIRE(left.specificEnthalpy.Size() == right.specificEnthalpy.Size());

        for (int index = 0; index < left.radius.Size(); ++index) {
            CHECK(left.radius(index) == right.radius(index));
            CHECK(left.density(index) == right.density(index));
            CHECK(left.specificEnthalpy(index) == right.specificEnthalpy(index));
        }

        CHECK(left.stellarRadius == right.stellarRadius);
        CHECK(left.centralDensity == right.centralDensity);
        CHECK(left.centralSpecificEnthalpy == right.centralSpecificEnthalpy);
    }
} // namespace

TEST_CASE(
    "Lane Emden Integration Matches The Analytic Incompressible Solution",
    tags::lane_emden_analytic
) {
    using Catch::Approx;

    const mean_field::seed::DimensionlessLaneEmdenSolution solution = mean_field::seed::integrateLaneEmden(0.0, 3.0);

    REQUIRE(solution.firstZeroCoordinate.has_value());
    CHECK(*solution.firstZeroCoordinate == Approx(std::sqrt(6.0)).margin(2.0e-7));
    CHECK(solution.theta(solution.theta.Size() - 1) == 0.0);
    check_dimensionless_solution(
        solution, [](const double coordinate) { return 1.0 - coordinate * coordinate / 6.0; },
        [](const double coordinate) { return -coordinate / 3.0; }, 2.0e-7
    );
}

TEST_CASE(
    "Lane Emden Integration Matches The Analytic Index One Solution",
    tags::lane_emden_analytic
) {
    using Catch::Approx;

    const mean_field::seed::DimensionlessLaneEmdenSolution solution = mean_field::seed::integrateLaneEmden(1.0, 4.0);

    REQUIRE(solution.firstZeroCoordinate.has_value());
    CHECK(*solution.firstZeroCoordinate == Approx(std::numbers::pi_v<double>).margin(2.0e-7));
    CHECK(solution.theta(solution.theta.Size() - 1) == 0.0);
    check_dimensionless_solution(
        solution, [](const double coordinate) { return coordinate == 0.0 ? 1.0 : std::sin(coordinate) / coordinate; },
        [](const double coordinate) {
            if (coordinate == 0.0) {
                return 0.0;
            }
            if (coordinate < 1.0e-4) {
                return -coordinate / 3.0 + coordinate * coordinate * coordinate / 30.0;
            }
            return (coordinate * std::cos(coordinate) - std::sin(coordinate)) / (coordinate * coordinate);
        },
        2.0e-7
    );
}

TEST_CASE(
    "Lane Emden Integration Matches The Analytic Index Five Infinite Solution",
    tags::lane_emden_analytic
) {
    using Catch::Approx;

    constexpr double coordinateLimit = 20.0;
    const mean_field::seed::DimensionlessLaneEmdenSolution solution =
        mean_field::seed::integrateLaneEmden(5.0, coordinateLimit);

    CHECK_FALSE(solution.firstZeroCoordinate.has_value());
    CHECK(solution.coordinate(solution.coordinate.Size() - 1) == Approx(coordinateLimit));
    CHECK(solution.theta(solution.theta.Size() - 1) > 0.0);
    check_dimensionless_solution(
        solution, [](const double coordinate) { return 1.0 / std::sqrt(1.0 + coordinate * coordinate / 3.0); },
        [](const double coordinate) { return -coordinate / 3.0 * std::pow(1.0 + coordinate * coordinate / 3.0, -1.5); },
        2.0e-7
    );
}

TEST_CASE(
    "Lane Emden Seed Uses The Stellar Model Central Density Phase Condition",
    tags::lane_emden_seed
) {
    using namespace mean_field;
    using Catch::Approx;

    const auto stellarModel = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}}),
        constraint::FixedCentralDensity({.RhoC = dimensions::DensityValue{2.0}})
    );
    const seed::LaneEmden strategy({.radialSampleCount = 64});

    STATIC_CHECK(seed::RadialSeedStrategyFor<seed::LaneEmden, decltype(stellarModel)>);

    const seed::RadialProfile profile = seed::generateRadialProfile(stellarModel, strategy);
    REQUIRE(profile.radius.Size() == 64);
    REQUIRE(profile.density.Size() == 64);
    REQUIRE(profile.specificEnthalpy.Size() == 64);
    CHECK(profile.centralDensity == dimensions::DensityValue{2.0});
    CHECK(profile.centralSpecificEnthalpy.value() == Approx(std::cbrt(2.0)));
    CHECK(profile.radius(0) == 0.0);
    CHECK(profile.radius(63) == profile.stellarRadius.value());
    CHECK(profile.density(0) == 2.0);
    CHECK(profile.density(63) == 0.0);
    CHECK(profile.specificEnthalpy(0) == profile.centralSpecificEnthalpy.value());
    CHECK(profile.specificEnthalpy(63) == 0.0);

    for (int index = 1; index < profile.radius.Size(); ++index) {
        CHECK(profile.radius(index) > profile.radius(index - 1));
        CHECK(profile.density(index) <= profile.density(index - 1));
        CHECK(profile.specificEnthalpy(index) <= profile.specificEnthalpy(index - 1));
        CHECK(profile.density(index) >= 0.0);
        CHECK(profile.specificEnthalpy(index) >= 0.0);
    }
}

TEST_CASE(
    "Explicit Lane Emden Seed Density Is Independent Of Model Invariants",
    tags::lane_emden_seed
) {
    using namespace mean_field;

    const auto unitMassModel = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{1.0}})
    );
    const auto largeMassModel = model::StellarModel(
        eos::Polytrope({.n = 3.0, .K = 0.25}), surface::Isobaric({.Psurf = dimensions::PressureValue{0.0}}),
        integral::FixedTotalMass({.Mtotal = dimensions::MassValue{9.0}})
    );
    const seed::LaneEmden strategy({.centralDensity = dimensions::DensityValue{1.7}, .radialSampleCount = 48});

    const seed::RadialProfile unitMassProfile  = seed::generateRadialProfile(unitMassModel, strategy);
    const seed::RadialProfile largeMassProfile = seed::generateRadialProfile(largeMassModel, strategy);
    check_profiles_are_identical(unitMassProfile, largeMassProfile);

    CHECK_THROWS_AS(
        seed::generateRadialProfile(unitMassModel, seed::LaneEmden({.radialSampleCount = 48})), std::invalid_argument
    );
}

TEST_CASE(
    "Legacy Polytropic Structure Seed Is An Exact Adapter Over Lane Emden Generation",
    tags::lane_emden_seed
) {
    using namespace mean_field;

    const eos::Polytrope equationOfState({.n = 3.0, .K = 0.25});
    const seed::RadialProfile profile =
        seed::generateLaneEmdenProfile(equationOfState, dimensions::DensityValue{1.25}, 40);
    const models::structure::StructureSeed legacySeed =
        models::structure::PolytropicStructure{equationOfState, 7.0}.makeInitialSeed(
            {.centralDensity = 1.25, .radialSampleCount = 40}
        );

    REQUIRE(legacySeed.radius.Size() == profile.radius.Size());
    for (int index = 0; index < profile.radius.Size(); ++index) {
        CHECK(legacySeed.radius(index) == profile.radius(index));
        CHECK(legacySeed.density(index) == profile.density(index));
        CHECK(legacySeed.enthalpy(index) == profile.specificEnthalpy(index));
    }
    CHECK(legacySeed.stellarRadius == profile.stellarRadius.value());
    CHECK(legacySeed.centralDensity == profile.centralDensity.value());
    CHECK(legacySeed.centralEnthalpy == profile.centralSpecificEnthalpy.value());
}

TEST_CASE(
    "Lane Emden Seed Rejects Invalid Numerical Prescriptions",
    tags::lane_emden_seed
) {
    using namespace mean_field;

    CHECK_THROWS_AS(seed::LaneEmden({.radialSampleCount = 1}), std::invalid_argument);
    CHECK_THROWS_AS(seed::LaneEmden({.centralDensity = dimensions::DensityValue{0.0}}), std::invalid_argument);
    CHECK_THROWS_AS(
        seed::LaneEmden({.centralDensity = dimensions::DensityValue{std::numeric_limits<double>::infinity()}}),
        std::invalid_argument
    );
    CHECK_THROWS_AS(
        seed::generateLaneEmdenProfile(eos::Polytrope({.n = 5.0, .K = 0.25}), dimensions::DensityValue{1.0}, 8),
        std::invalid_argument
    );
}