#pragma once

#include <optional>
#include <string_view>

#include "serif/discretization/blocks/fields.hpp"
#include "serif/equations/companion.hpp"
#include "serif/equations/point.hpp"
#include "serif/equations/rigid_rotation.hpp"

namespace serif::equations {
    /**
     * @brief Set the effective potential to one unknown level on the selected isobar.
     *
     * Integrate this point relation against the existing surface shape basis with
     * a fixed dimensionless reference measure. The constant mode is retained.
     * This is a boundary equation, not a volume displacement force.
     */
    class SurfacePotentialBalance final {
    public:
        using Potential = discretization::blocks::gravity::potential;
        using SurfaceLevel = discretization::blocks::barotropic_constant::mass_normalization;
        using AngularVelocity = discretization::blocks::rotation::angular_velocity;
        using Operands = utils::types::TypeList<discretization::forms::Operand<Potential>,
            discretization::forms::Operand<SurfaceLevel>, discretization::forms::Operand<AngularVelocity>>;
        using ResidualQuantity = dimensions::GravitationalPotential;
        static constexpr std::string_view name = "surface effective potential balance";

        /** @brief Use the same physical rotation axis, origin and companion as the interior equation. */
        explicit SurfacePotentialBalance(const RigidRotation::Vector& axis = {0.0, 0.0, 1.0},
            const RigidRotation::Vector& center = {}, std::optional<PointMassCompanion> companion = std::nullopt) :
            m_rotation(RigidRotation::unit_rotation(axis, center)), m_companion(companion) {}

        /** @brief Differentiate at fixed pulled-back potential; geometry moves only the centrifugal term. */
        [[nodiscard]] PointLinearization<Operands> evaluate(const EquationPoint<Operands>& point) const {
            const double omega = point.fields.get<AngularVelocity>();
            const double centrifugal = m_rotation.centrifugal_potential(point.position);
            PointLinearization<Operands> result;

            result.residual = point.fields.get<Potential>() - omega * omega * centrifugal - point.fields.get<SurfaceLevel>();
            result.field_derivatives.get<Potential>() = 1.0;
            result.field_derivatives.get<SurfaceLevel>() = -1.0;
            result.field_derivatives.get<AngularVelocity>() = -2.0 * omega * centrifugal;
            result.position_derivative = m_rotation.centrifugal_potential_gradient(point.position);

            for (double& component : result.position_derivative) {
                component *= -omega * omega;
            }

            if (m_companion) {
                result.residual += m_companion->potential(point.position);
                const auto gradient = m_companion->potential_gradient(point.position);
                for (int axis = 0; axis < 3; ++axis) { result.position_derivative[axis] += gradient[axis]; }
            }

            return result;
        }
        /** @brief Inspect the physical rotation convention. */
        [[nodiscard]] const RigidRotation& rotation() const noexcept { return m_rotation; }
        /** @brief Inspect the companion, if any. */
        [[nodiscard]] const std::optional<PointMassCompanion>& companion() const noexcept { return m_companion; }

    private:
        RigidRotation m_rotation;
        std::optional<PointMassCompanion> m_companion;
    };
}