#pragma once

#include <cmath>
#include <stdexcept>
#include <string_view>
#include <tuple>

#include "serif/discretization/blocks/fields.hpp"
#include "serif/equations/contributions.hpp"
#include "serif/equations/point.hpp"
#include "serif/utils/error/errors.hpp"

namespace serif::equations {
    /** @brief Solve Delta(Phi) = 4 pi G rho with one conforming scalar potential field. */
    class PrimalPoissonBalance final {
    public:
        using Test = discretization::blocks::gravity::potential;
        using Density = discretization::blocks::density::mass;
        using PotentialGradient = discretization::forms::Operand<Test, discretization::forms::Operation::Gradient>;
        using ResidualQuantity = dimensions::InverseTimeSquared;
        static constexpr std::string_view name = "primal Poisson balance";

        /** @brief Integrate grad(v) dot grad(Phi) throughout the star and vacuum. */
        struct PotentialFlux : Contribution<PotentialGradient> {
            using Operands = utils::types::TypeList<PotentialGradient>;

            /** @brief The potential flux and its derivative are grad(Phi) and identity. */
            [[nodiscard]] static PointLinearization<Operands, discretization::forms::VectorValue>
            evaluate(const EquationPoint<Operands>& point) {
                PointLinearization<Operands, discretization::forms::VectorValue> result;
                result.residual = point.fields.template get<PotentialGradient>();

                for (int axis = 0; axis < 3; ++axis) {
                    result.field_derivatives.template get<PotentialGradient>()[axis][axis] = 1.0;
                }

                return result;
            }
            /** @brief A linear flux introduces no additional coefficient degree. */
            [[nodiscard]] static int additional_order(PotentialGradient, int) { return 0; }
        };

        /** @brief Integrate +4 pi G v rho only where matter is supported. */
        struct DensitySource : Contribution<discretization::forms::Operand<Test>, discretization::domain::StellarDomains> {
            using Operands = utils::types::TypeList<discretization::forms::Operand<Density>>;

            double four_pi_G;

            /** @brief Store the already validated source coefficient. */
            explicit DensitySource(const double coefficient) : four_pi_G(coefficient) {}

            /** @brief Return the source with the positive sign of the primal weak form. */
            [[nodiscard]] PointLinearization<Operands> evaluate(const EquationPoint<Operands>& point) const {
                PointLinearization<Operands> result;
                result.residual = four_pi_G * point.fields.template get<Density>();
                result.field_derivatives.template get<Density>() = four_pi_G;
                return result;
            }

            /** @brief The source is linear in density. */
            [[nodiscard]] static int additional_order(discretization::forms::Operand<Density>, int) { return 0; }
        };

        /** @brief Set positive finite 4 pi G; impose the zero trace at infinity separately. */
        explicit PrimalPoissonBalance(const double four_pi_G) : m_four_pi_G(four_pi_G) {
            using namespace utils::errors;
            serif_ensure<SERiFErrorCode::INVALID_ARGUMENT, std::invalid_argument>(std::isfinite(four_pi_G) && four_pi_G > 0.0, "Primal Poisson balance requires positive finite 4 pi G.");
        }

        /** @brief Publish the diffusion term and stellar source to the existing volume assembler. */
        [[nodiscard]] auto contributions() const {
            return std::tuple{PotentialFlux{}, DensitySource{m_four_pi_G}};
        }
        /** @brief Inspect the physical source coefficient for serialization. */
        [[nodiscard]] double four_pi_G() const noexcept { return m_four_pi_G; }

    private:
        double m_four_pi_G;
    };
}