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

export namespace mean_field::eos {
    class Polytrope final : public EquationOfState {
    public:
        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]] double pressure_from_density(const double density) const override {
            validate_nonnegativity(density, "density");
            if (density == 0.0) {
                return 0.0;
            }

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

        [[nodiscard]] double enthalpy_from_density(const double density) const override {
            validate_nonnegativity(density, "density");
            if (density == 0.0) {
                return 0.0;
            }

            return m_enthalpy_scale * std::pow(density, 1.0 / m_polytropic_index);
        }

        [[nodiscard]] double density_from_enthalpy(const double enthalpy) const override {
            validate_finite(enthalpy, "enthalpy");

            if (enthalpy <= 0.0) {
                return 0.0;
            }

            return std::pow(enthalpy / m_enthalpy_scale, m_polytropic_index);
        }

        [[nodiscard]] double pressure_from_enthalpy(const double enthalpy) const override {
            validate_finite(enthalpy, "enthalpy");

            if (enthalpy <= 0.0) {
                return 0.0;
            }

            return density_from_enthalpy(enthalpy) * enthalpy / (m_polytropic_index + 1.0);
        }

        [[nodiscard]] double density_derivative_from_enthalpy(const double enthalpy) const override {
            validate_finite(enthalpy, "enthalpy");
            if (enthalpy < 0.0) {
                return 0.0;
            }

            if (enthalpy == 0.0) {
                return m_polytropic_index == 1.0 ? 1.0 / m_enthalpy_scale : 0.0;
            }

            return m_polytropic_index / m_enthalpy_scale *
                   std::pow(enthalpy / m_enthalpy_scale, m_polytropic_index - 1.0);
        }

        [[nodiscard]] double pressure_derivative_from_enthalpy(const double enthalpy) const override {
            validate_finite(enthalpy, "enthalpy");

            if (enthalpy <= 0.0) {
                return 0.0;
            }

            return density_from_enthalpy(enthalpy);
        }

        [[nodiscard]] double pressure_derivative_from_density(const double density) const override {
            validate_nonnegativity(density, "density");
            if (density == 0.0) {
                return 0.0;
            }

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

        [[nodiscard]] double enthalpy_from_pressure(double pressure) const override {
            validate_nonnegativity(pressure, "pressure");
            const double np1 = m_polytropic_index + 1;
            return np1 * std::pow(m_polytropic_constant, m_polytropic_index / np1) * std::pow(pressure, 1.0 / np1);
        }

    private:
        static void validate_finite(
            const double value,
            const char *quantity
        ) {
            if (!std::isfinite(value)) {
                throw std::domain_error(
                    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 std::domain_error(
                    std::format(
                        "The {} must be non-negative. Instead a value of {} "
                        "has been "
                        "provided",
                        quantity, value
                    )
                );
            }
        }

    public:

    private:
        double m_polytropic_index;
        double m_polytropic_constant;
        double m_enthalpy_scale;
    };
} // namespace mean_field::eos