#pragma once
/**
 * @file rotational_force.hpp
 * @author Emily M. Boudreaux
 *
 * @brief File containing the RotationalForce contribution class, which computes the rotational force contribution to the displacement balance equation.
 */

#include <string_view>

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

namespace serif::equations {
    /**
     * @class RotationalForce
     *
     * @brief The rotational force contribution to the displacement balance equation. This equation computes the rotational force contribution
     * to the displacement balance equation, which is given by \f$-\int_{\Omega_{\text{stellar}}}\rho\nabla\Psi \cdot w \, dV\f$ where \f$\Psi\f$ is the centrifugal potential
     * and \f$w\f$ is the test function.
     *
     * @contribution Residuals to the displacement vector residual block where only on the parts of that block which are part of the stellar domain.
     *
     */
    class RotationalForce final : public Contribution<discretization::forms::Operand<discretization::blocks::displacement::vector>, discretization::domain::StellarDomains> {
    public:
        using Density = discretization::blocks::density::mass;                          ///< Alias for mass density
        using DensityValue = discretization::forms::Operand<Density>;                   ///< Operand type for mass density

        using AngularVelocity = discretization::blocks::rotation::angular_velocity;     ///< Alias for angular velocity
        using AngularVelocityValue = discretization::forms::Operand<AngularVelocity>;   ///< Operand type for angular velocity

        using Operands = utils::types::TypeList<DensityValue, AngularVelocityValue>;    ///< The rotational force depends on the mass density and the angular velocity operands.
        using Result = discretization::forms::VectorValue;                              ///< The result type will be a vector value.

        static constexpr std::string_view name = "rotational force";                    ///< Compile time name for the rotational force equation, used for debugging and logging.

        /**
         * @brief Construct a rotational force contribution with a given axis and center of rotation.
         *
         * @param axis The axis of rotation, which must be finite. Defaults to the z-axis.
         * @param center The center of rotation, which must be finite. Defaults to the origin.
         */
        explicit RotationalForce(
            const Result &axis = {0.0, 0.0, 1.0},
            const Result &center = {}
        ) :
        m_rotation(RigidRotation::unit_rotation(axis, center)) {}


        /**
         * @brief Evaluate the rotational force contribution at a given point.
         *
         * @param point The point at which to evaluate the rotational force contribution, which must contain the mass density and angular velocity operands.
         * @return A PointLinearization containing the residual and derivatives of the rotational force contribution at the given point.
         *
         * Here the point-wise integrand is computed as a three-vector, where each component \f$r_{i}\f$ is given by \f$ r_{i} = \rho\left(-\omega^{2}\nabla\Psi_{i}\right)
         *
         * There are then derivatives of this given with respect to mass density, angular velocity, and position
         *
         * - \frac{\partial r_{i}}{\partial \rho} = -\omega^{2}\nabla\Psi_{i}
         * - \frac{\partial r_{i}}{\partial \omega} = -2\omega\rho\nabla\Psi_{i}
         * - \frac{\partial r_{i}}{\partial x_{j}} = -\rho \omega^{2} \frac{\partial^{2} \Psi_{i}}{\partial x_{i} \partial x_{j}}
         */
        [[nodiscard]] PointLinearization<Operands, Result> evaluate(
            const EquationPoint<Operands>& point
        ) const {
            const double density = point.fields.get<Density>();
            const double omega = point.fields.get<AngularVelocity>();
            const double omega_squared = omega * omega;
            const auto gradient = m_rotation.centrifugal_potential_gradient(point.position);

            PointLinearization<Operands, Result> result;

            for (int row = 0; row < 3; ++row) {
                result.residual[row] = density * (-omega_squared * gradient[row]);
                result.field_derivatives.get<Density>()[row] = -omega_squared * gradient[row];
                result.field_derivatives.get<AngularVelocity>()[row] = -2.0 * omega * density * gradient[row];
                for (int column = 0; column < 3; ++column) {
                    result.position_derivative[row][column] = -density * (omega_squared * m_rotation.centrifugal_potential_hessian()[row][column]);
                }
            }

            return result;
        }

        /**
         * @brief Get the additional quadrature order required for the operands of this equation when integrating density.
         *
         * @param operand-type The operand type for which to determine the additional quadrature order. Note this is only used to select the correct overload.
         * @param position_order The quadrature order associated with the position.
         * @return The additional quadrature order required for the given operand type. For density this just returns the position order;
         */
        [[nodiscard]] static int additional_order(DensityValue /*operand-type*/, const int position_order) {
            return position_order;
        }

        /**
         * @brief Get the additional quadrature order required for the operands of this equation when integrating angular velocity.
         *
         * @param operand-type The operand type for which to determine the additional quadrature order. Note this is only used to select the correct overload.
         * @param position_order The quadrature order associated with the position.
         * @return The additional quadrature order required for the given operand type. For angular velocity this returns the sum of the density functional space order and the position order.
         */
        [[nodiscard]] static int additional_order(AngularVelocityValue /*operand-type*/, const int position_order) {
            return Density::FunctionalSpace::order + position_order;
        }

    private:
        RigidRotation m_rotation;  ///< The rigid rotation object used to compute the centrifugal potential and its derivatives, which is constructed from the axis and center of rotation provided to the constructor.
    };
} // namespace serif::equations