module;
#include <cmath>
#include <format>
#include <stdexcept>
export module mean_field:eos.polytrope;
export import :eos.evaluation;

export namespace mean_field::eos {
    class Polytrope final {
    public:
        using Relations = RelationCatalog<
            PressureFromDensity,
            PressureFromSpecificEnthalpy,
            SpecificEnthalpyFromDensity,
            SpecificEnthalpyFromPressure,
            DensityFromSpecificEnthalpy>;

        Polytrope(
            const double polytropic_index,
            const double polytropic_constant
        )
            : m_polytropic_index(polytropic_index),
              m_polytropic_constant(polytropic_constant),
              m_enthalpy_scale((polytropic_index + 1.0) * polytropic_constant) {
            if (!std::isfinite(polytropic_index) || polytropic_index < 1.0) {
                throw std::invalid_argument(
                    std::format(
                        "The differentiable polytropic closure requires a "
                        "finite polytropic index greater than or equal to one. "
                        "Instead a value of {} has been provided",
                        polytropic_index
                    )
                );
            }

            if (!std::isfinite(polytropic_constant) || polytropic_constant <= 0.0) {
                throw std::invalid_argument(
                    std::format(
                        "The polytropic constant must be finite and positive. "
                        "Instead a value of {} has been provided",
                        polytropic_constant
                    )
                );
            }
        };

        [[nodiscard]] double polytropic_index() const noexcept {
            return m_polytropic_index;
        }

        [[nodiscard]] double polytropic_constant() const noexcept {
            return m_polytropic_constant;
        }

        [[nodiscard]] double enthalpy_scale() const noexcept {
            return m_enthalpy_scale;
        }

        [[nodiscard]] PressureValue evaluate(
            PressureFromDensity,
            const DensityValue density
        ) const {
            validate_nonnegativity(density.value(), "density");
            if (density.value() == 0.0) {
                return PressureValue{0.0};
            }

            return PressureValue{m_polytropic_constant * std::pow(density.value(), 1.0 + 1.0 / m_polytropic_index)};
        }

        [[nodiscard]] SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromDensity,
            const DensityValue density
        ) const {
            validate_nonnegativity(density.value(), "density");
            if (density.value() == 0.0) {
                return SpecificEnthalpyValue{0.0};
            }

            return SpecificEnthalpyValue{m_enthalpy_scale * std::pow(density.value(), 1.0 / m_polytropic_index)};
        }

        [[nodiscard]] DensityValue evaluate(
            DensityFromSpecificEnthalpy,
            const SpecificEnthalpyValue specificEnthalpy
        ) const {
            validate_finite(specificEnthalpy.value(), "specific enthalpy");

            if (specificEnthalpy.value() <= 0.0) {
                return DensityValue{0.0};
            }

            return DensityValue{std::pow(specificEnthalpy.value() / m_enthalpy_scale, m_polytropic_index)};
        }

        [[nodiscard]] PressureValue evaluate(
            PressureFromSpecificEnthalpy,
            const SpecificEnthalpyValue specificEnthalpy
        ) const {
            const DensityValue density = evaluate(DensityFromSpecificEnthalpy{}, specificEnthalpy);

            if (specificEnthalpy.value() <= 0.0) {
                return PressureValue{0.0};
            }

            return PressureValue{density.value() * specificEnthalpy.value() / (m_polytropic_index + 1.0)};
        }

        [[nodiscard]] SpecificEnthalpyValue evaluate(
            SpecificEnthalpyFromPressure,
            const PressureValue pressure
        ) const {
            validate_nonnegativity(pressure.value(), "pressure");
            if (pressure.value() == 0.0) {
                return SpecificEnthalpyValue{0.0};
            }

            const double indexPlusOne = m_polytropic_index + 1.0;

            return SpecificEnthalpyValue{
                indexPlusOne * std::pow(m_polytropic_constant, m_polytropic_index / indexPlusOne) *
                std::pow(pressure.value(), 1.0 / indexPlusOne)
            };
        }

        [[nodiscard]] PartialDerivative<
            quantity::Density,
            quantity::SpecificEnthalpy>
        partialDerivative(
            DensityFromSpecificEnthalpy,
            WithRespectTo<quantity::SpecificEnthalpy>,
            const SpecificEnthalpyValue specificEnthalpy
        ) const {
            validate_finite(specificEnthalpy.value(), "specific enthalpy");
            if (specificEnthalpy.value() < 0.0) {
                return PartialDerivative<quantity::Density, quantity::SpecificEnthalpy>{0.0};
            }

            if (specificEnthalpy.value() == 0.0) {
                return PartialDerivative<quantity::Density, quantity::SpecificEnthalpy>{
                    m_polytropic_index == 1.0 ? 1.0 / m_enthalpy_scale : 0.0
                };
            }

            return PartialDerivative<quantity::Density, quantity::SpecificEnthalpy>{
                m_polytropic_index / m_enthalpy_scale *
                std::pow(specificEnthalpy.value() / m_enthalpy_scale, m_polytropic_index - 1.0)
            };
        }

        [[nodiscard]] PartialDerivative<
            quantity::Pressure,
            quantity::SpecificEnthalpy>
        partialDerivative(
            PressureFromSpecificEnthalpy,
            WithRespectTo<quantity::SpecificEnthalpy>,
            const SpecificEnthalpyValue specificEnthalpy
        ) const {
            const DensityValue density = evaluate(DensityFromSpecificEnthalpy{}, specificEnthalpy);

            return PartialDerivative<quantity::Pressure, quantity::SpecificEnthalpy>{density.value()};
        }

        [[nodiscard]] PartialDerivative<
            quantity::Pressure,
            quantity::Density>
        partialDerivative(
            PressureFromDensity,
            WithRespectTo<quantity::Density>,
            const DensityValue density
        ) const {
            validate_nonnegativity(density.value(), "density");
            if (density.value() == 0.0) {
                return PartialDerivative<quantity::Pressure, quantity::Density>{0.0};
            }

            return PartialDerivative<quantity::Pressure, quantity::Density>{
                m_polytropic_constant * (1.0 + 1.0 / m_polytropic_index) *
                std::pow(density.value(), 1.0 / m_polytropic_index)
            };
        }

    private:
        static void validate_finite(
            const double value,
            const char *quantity
        ) {
            if (!std::isfinite(value)) {
                throw EvaluationError(
                    EvaluationErrorCode::nonfinite_input, std::format(
                                                              "The {} must be finite. Instead a value of {} has been "
                                                              "provided",
                                                              quantity, value
                                                          )
                );
            }
        }

        static void validate_nonnegativity(
            const double value,
            const char *quantity
        ) {
            validate_finite(value, quantity);
            if (value < 0.0) {
                throw EvaluationError(
                    EvaluationErrorCode::outside_domain, std::format(
                                                             "The {} must be non-negative. Instead a value of {} "
                                                             "has been "
                                                             "provided",
                                                             quantity, value
                                                         )
                );
            }
        }
        double m_polytropic_index;
        double m_polytropic_constant;
        double m_enthalpy_scale;
    };
} // namespace mean_field::eos