module;

#include <cmath>
#include <stdexcept>
#include <string>

export module mean_field_extension_example.ideal_gas_radiation;

import mean_field;

/*
 * This file is intended to be read from top to bottom by a physicist who is
 * adding an equation of state (EOS).  The comments explain the small amount
 * of type-system vocabulary required by MeanField; the thermodynamics remain
 * visible as ordinary equations.
 */
export namespace mean_field::extension_example {
    namespace eos_quantity = mean_field::dimensions::quantity;

    /*
     * A relation is only a compile-time sentence:
     *
     *     output = f(input 1, input 2, ...).
     *
     * Input order is significant.  These declarations say that density is
     * the first argument and temperature is the second argument.  They do not
     * allocate data and have no runtime cost.
     */
    using PressureFromDensityAndTemperature = mean_field::eos::Relation<
        eos_quantity::Pressure,
        eos_quantity::Density,
        eos_quantity::Temperature>;

    using SpecificInternalEnergyFromDensityAndTemperature = mean_field::eos::Relation<
        eos_quantity::SpecificInternalEnergy,
        eos_quantity::Density,
        eos_quantity::Temperature>;

    using SpecificEnthalpyFromDensityAndTemperature = mean_field::eos::Relation<
        eos_quantity::SpecificEnthalpy,
        eos_quantity::Density,
        eos_quantity::Temperature>;

    /*
     * A monatomic ideal gas plus equilibrium radiation:
     *
     *   R       = k_B / (mu m_u)
     *   P_gas   = rho R T
     *   P_rad   = a T^4 / 3
     *   u       = (3/2) R T + a T^4 / rho
     *   h       = u + P/rho
     *           = (5/2) R T + 4 a T^4 / (3 rho)
     *
     * The scalar QuantityValue wrappers identify what a number means.  They
     * intentionally do not perform unit conversion.  Every number supplied
     * here must therefore use one coherent unit system.
     */
    class IdealGasRadiation final {
    public:
        struct Parameters final {
            /* Mean particle mass in atomic-mass units. */
            double meanMolecularWeight{0.61};

            /* CGS defaults: erg K^-1, g, and erg cm^-3 K^-4. */
            double boltzmannConstant{1.380649e-16};
            double atomicMassUnit{1.66053906660e-24};
            double radiationConstant{7.5657e-15};
        };

        /*
         * This one alias makes the EOS a constitutive-law specification that
         * can be placed directly in model::StellarModel(...).  There is no
         * registry edit and no central list of EOS combinations to maintain.
         */
        using ModelDefinition = mean_field::eos::ConstitutiveLaw<IdealGasRadiation,"IdealGasRadiation">;

        /*
         * The catalog is the complete public claim made by this EOS.  If an
         * evaluate overload below is missing or has the wrong argument order,
         * eos::EquationOfStateModel<IdealGasRadiation> becomes false at
         * compile time.
         */
        using Relations = mean_field::eos::RelationCatalog<
            PressureFromDensityAndTemperature,
            SpecificInternalEnergyFromDensityAndTemperature,
            SpecificEnthalpyFromDensityAndTemperature
        >;

        struct PressureContributions final {
            mean_field::dimensions::PressureValue gas;
            mean_field::dimensions::PressureValue radiation;

            [[nodiscard]] mean_field::dimensions::PressureValue total() const noexcept {
                return gas + radiation;
            }
        };

        explicit IdealGasRadiation(const Parameters parameters)
            : m_parameters(validatedParameters(parameters)),
              m_specificGasConstant(
                  m_parameters.boltzmannConstant /(m_parameters.meanMolecularWeight * m_parameters.atomicMassUnit)
              ) {}

        [[nodiscard]] const Parameters &parameters() const noexcept {
            return m_parameters;
        }

        [[nodiscard]] double specificGasConstant() const noexcept {
            return m_specificGasConstant;
        }

        /*
         * Named component functions are not required by the EOS protocol.
         * They are provided because they make diagnostics and physics tests
         * easier to read than repeated algebra in client code.
         */
        [[nodiscard]] PressureContributions pressureContributions(
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);

            const double rho = density.value();
            const double T   = temperature.value();
            return PressureContributions{
                .gas = mean_field::dimensions::PressureValue{rho * m_specificGasConstant * T},
                .radiation = mean_field::dimensions::PressureValue{
                    m_parameters.radiationConstant * fourthPower(T) / 3.0
                }
            };
        }

        [[nodiscard]] mean_field::dimensions::SpecificInternalEnergyValue gasSpecificInternalEnergy(
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateTemperature(temperature);
            return mean_field::dimensions::SpecificInternalEnergyValue{
                1.5 * m_specificGasConstant * temperature.value()
            };
        }

        [[nodiscard]] mean_field::dimensions::SpecificInternalEnergyValue radiationSpecificInternalEnergy(
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::dimensions::SpecificInternalEnergyValue{
                m_parameters.radiationConstant * fourthPower(temperature.value()) / density.value()
            };
        }

        /* The evaluate overloads implement the three declared relations. */
        [[nodiscard]] mean_field::dimensions::PressureValue evaluate(
            PressureFromDensityAndTemperature,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            return pressureContributions(density, temperature).total();
        }

        [[nodiscard]] mean_field::dimensions::SpecificInternalEnergyValue evaluate(
            SpecificInternalEnergyFromDensityAndTemperature,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            const auto gas       = gasSpecificInternalEnergy(temperature);
            const auto radiation = radiationSpecificInternalEnergy(density, temperature);
            return gas + radiation;
        }

        [[nodiscard]] mean_field::dimensions::SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromDensityAndTemperature,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);

            const double rho = density.value();
            const double T   = temperature.value();
            return mean_field::dimensions::SpecificEnthalpyValue{
                2.5 * m_specificGasConstant * T +
                4.0 * m_parameters.radiationConstant * fourthPower(T) / (3.0 * rho)
            };
        }

        /*
         * Jacobian entries are ordinary analytic partial derivatives.  The
         * WithRespectTo tag prevents accidentally returning dP/dT from the
         * overload that promised dP/drho.
         */
        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::Pressure,
            eos_quantity::Density>
        partialDerivative(
            PressureFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Density>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::eos::PartialDerivative<
                eos_quantity::Pressure,
                eos_quantity::Density>{m_specificGasConstant * temperature.value()};
        }

        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::Pressure,
            eos_quantity::Temperature>
        partialDerivative(
            PressureFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Temperature>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            const double T = temperature.value();
            return mean_field::eos::PartialDerivative<
                eos_quantity::Pressure,
                eos_quantity::Temperature>{
                density.value() * m_specificGasConstant +
                4.0 * m_parameters.radiationConstant * cube(T) / 3.0
            };
        }

        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::SpecificInternalEnergy,
            eos_quantity::Density>
        partialDerivative(
            SpecificInternalEnergyFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Density>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::eos::PartialDerivative<
                eos_quantity::SpecificInternalEnergy,
                eos_quantity::Density>{
                -m_parameters.radiationConstant * fourthPower(temperature.value()) /
                square(density.value())
            };
        }

        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::SpecificInternalEnergy,
            eos_quantity::Temperature>
        partialDerivative(
            SpecificInternalEnergyFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Temperature>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::eos::PartialDerivative<
                eos_quantity::SpecificInternalEnergy,
                eos_quantity::Temperature>{
                1.5 * m_specificGasConstant +
                4.0 * m_parameters.radiationConstant * cube(temperature.value()) / density.value()
            };
        }

        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::SpecificEnthalpy,
            eos_quantity::Density>
        partialDerivative(
            SpecificEnthalpyFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Density>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::eos::PartialDerivative<
                eos_quantity::SpecificEnthalpy,
                eos_quantity::Density>{
                -4.0 * m_parameters.radiationConstant * fourthPower(temperature.value()) /
                (3.0 * square(density.value()))
            };
        }

        [[nodiscard]] mean_field::eos::PartialDerivative<
            eos_quantity::SpecificEnthalpy,
            eos_quantity::Temperature>
        partialDerivative(
            SpecificEnthalpyFromDensityAndTemperature,
            mean_field::eos::WithRespectTo<eos_quantity::Temperature>,
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) const {
            validateMaterialState(density, temperature);
            return mean_field::eos::PartialDerivative<
                eos_quantity::SpecificEnthalpy,
                eos_quantity::Temperature>{
                2.5 * m_specificGasConstant +
                16.0 * m_parameters.radiationConstant * cube(temperature.value()) /
                    (3.0 * density.value())
            };
        }

    private:
        [[nodiscard]] static Parameters validatedParameters(const Parameters parameters) {
            requirePositiveFinite(parameters.meanMolecularWeight, "mean molecular weight");
            requirePositiveFinite(parameters.boltzmannConstant, "Boltzmann constant");
            requirePositiveFinite(parameters.atomicMassUnit, "atomic mass unit");
            requireNonnegativeFinite(parameters.radiationConstant, "radiation constant");
            return parameters;
        }

        static void requirePositiveFinite(const double value, const char *name) {
            if (!std::isfinite(value) || value <= 0.0) {
                throw std::invalid_argument(
                    std::string{"IdealGasRadiation requires a finite, positive "} + name + "."
                );
            }
        }

        static void requireNonnegativeFinite(const double value, const char *name) {
            if (!std::isfinite(value) || value < 0.0) {
                throw std::invalid_argument(
                    std::string{"IdealGasRadiation requires a finite, nonnegative "} + name + "."
                );
            }
        }

        static void validateMaterialState(
            const mean_field::dimensions::DensityValue density,
            const mean_field::dimensions::TemperatureValue temperature
        ) {
            if (!std::isfinite(density.value()) || !std::isfinite(temperature.value())) {
                throw mean_field::eos::EvaluationError{
                    mean_field::eos::EvaluationErrorCode::nonfinite_input,
                    "IdealGasRadiation requires finite density and temperature."
                };
            }
            if (density.value() <= 0.0 || temperature.value() < 0.0) {
                throw mean_field::eos::EvaluationError{
                    mean_field::eos::EvaluationErrorCode::outside_domain,
                    "IdealGasRadiation requires rho > 0 and T >= 0."
                };
            }
        }

        static void validateTemperature(const mean_field::dimensions::TemperatureValue temperature) {
            if (!std::isfinite(temperature.value())) {
                throw mean_field::eos::EvaluationError{
                    mean_field::eos::EvaluationErrorCode::nonfinite_input,
                    "IdealGasRadiation requires finite temperature."
                };
            }
            if (temperature.value() < 0.0) {
                throw mean_field::eos::EvaluationError{
                    mean_field::eos::EvaluationErrorCode::outside_domain,
                    "IdealGasRadiation requires T >= 0."
                };
            }
        }

        [[nodiscard]] static double square(const double value) noexcept {
            return value * value;
        }

        [[nodiscard]] static double cube(const double value) noexcept {
            return value * value * value;
        }

        [[nodiscard]] static double fourthPower(const double value) noexcept {
            const double squared = square(value);
            return squared * squared;
        }

        Parameters m_parameters;
        double m_specificGasConstant;
    };

    /*
     * These assertions are executable documentation.  They prove that the
     * class and every derivative satisfy the public extension protocol.
     */
    static_assert(mean_field::models::SelfDescribingModelSpecification<IdealGasRadiation>);
    static_assert(mean_field::eos::EquationOfStateModel<IdealGasRadiation>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  PressureFromDensityAndTemperature,
                  eos_quantity::Density>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  PressureFromDensityAndTemperature,
                  eos_quantity::Temperature>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  SpecificInternalEnergyFromDensityAndTemperature,
                  eos_quantity::Density>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  SpecificInternalEnergyFromDensityAndTemperature,
                  eos_quantity::Temperature>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  SpecificEnthalpyFromDensityAndTemperature,
                  eos_quantity::Density>);
    static_assert(mean_field::eos::SupportsPartialDerivative<
                  IdealGasRadiation,
                  SpecificEnthalpyFromDensityAndTemperature,
                  eos_quantity::Temperature>);
} // namespace mean_field::extension_example