module;

#include <cmath>
#include <mfem.hpp>

export module mean_field:physics.rigid_rotation;

export namespace mean_field::physics {
    class RigidRotation final {
    public:
        RigidRotation(
            const mfem::Vector &angularVelocity,
            const mfem::Vector &center
        )
            : m_angularVelocity(angularVelocity),
              m_center(center) {
            MFEM_VERIFY(
                m_angularVelocity.Size() == 3, "RigidRotation requires a three-dimensional "
                                               "angular-velocity vector."
            );

            MFEM_VERIFY(m_center.Size() == 3, "RigidRotation requires a three-dimensional center.");

            for (int component = 0; component < 3; ++component) {
                MFEM_VERIFY(
                    std::isfinite(m_angularVelocity(component)), "RigidRotation received a non-finite "
                                                                 "angular-velocity component."
                );

                MFEM_VERIFY(
                    std::isfinite(m_center(component)), "RigidRotation received a non-finite center component."
                );
            }
        }

        [[nodiscard]] double potential(const mfem::Vector &physicalPosition) const {
            MFEM_VERIFY(
                physicalPosition.Size() == 3, "RigidRotation::potential requires a "
                                              "three-dimensional position."
            );

            const double relativeX = physicalPosition(0) - m_center(0);

            const double relativeY = physicalPosition(1) - m_center(1);

            const double relativeZ = physicalPosition(2) - m_center(2);

            const double crossX    = m_angularVelocity(1) * relativeZ - m_angularVelocity(2) * relativeY;

            const double crossY    = m_angularVelocity(2) * relativeX - m_angularVelocity(0) * relativeZ;

            const double crossZ    = m_angularVelocity(0) * relativeY - m_angularVelocity(1) * relativeX;

            return 0.5 * (crossX * crossX + crossY * crossY + crossZ * crossZ);
        }

        [[nodiscard]] double potential_directional_derivative(
            const mfem::Vector &physicalPosition,
            const mfem::Vector &physicalPositionVariation
        ) const {
            MFEM_VERIFY(
                physicalPosition.Size() == 3, "RigidRotation derivative requires a "
                                              "three-dimensional position."
            );

            MFEM_VERIFY(
                physicalPositionVariation.Size() == 3, "RigidRotation derivative requires a "
                                                       "three-dimensional direction."
            );

            double angularVelocitySquared     = 0.0;
            double angularVelocityDotPosition = 0.0;

            for (int component = 0; component < 3; ++component) {
                const double relativePosition = physicalPosition(component) - m_center(component);

                angularVelocitySquared += m_angularVelocity(component) * m_angularVelocity(component);

                angularVelocityDotPosition += m_angularVelocity(component) * relativePosition;
            }

            double derivative = 0.0;

            for (int component = 0; component < 3; ++component) {
                const double relativePosition  = physicalPosition(component) - m_center(component);

                const double gradientComponent = angularVelocitySquared * relativePosition -
                                                 angularVelocityDotPosition * m_angularVelocity(component);

                derivative += gradientComponent * physicalPositionVariation(component);
            }

            return derivative;
        }

        /*
         * Gradient of the positive rigid-rotation potential
         *
         *   Psi = 0.5 |Omega x (x - x_0)|^2.
         *
         * This points away from the rotation axis. The rotational
         * displacement residual uses its negative.
         */
        void potential_gradient(
            const mfem::Vector &physicalPosition,
            mfem::Vector &gradient
        ) const {
            MFEM_VERIFY(
                physicalPosition.Size() == 3, "RigidRotation::potential_gradient requires a "
                                              "three-dimensional position."
            );

            double angularVelocitySquared     = 0.0;
            double angularVelocityDotPosition = 0.0;

            for (int component = 0; component < 3; ++component) {
                const double relativePosition = physicalPosition(component) - m_center(component);

                angularVelocitySquared += m_angularVelocity(component) * m_angularVelocity(component);

                angularVelocityDotPosition += m_angularVelocity(component) * relativePosition;
            }

            gradient.SetSize(3);

            for (int component = 0; component < 3; ++component) {
                const double relativePosition = physicalPosition(component) - m_center(component);

                gradient(component)           = angularVelocitySquared * relativePosition -
                                                angularVelocityDotPosition * m_angularVelocity(component);
            }
        }

        /*
         * Hessian action of Psi. The Hessian is constant for rigid
         * rotation, so only the physical-position direction is required.
         */
        void potential_gradient_directional_derivative(
            const mfem::Vector &physicalPositionVariation,
            mfem::Vector &gradientVariation
        ) const {
            MFEM_VERIFY(
                physicalPositionVariation.Size() == 3, "RigidRotation gradient derivative requires a "
                                                       "three-dimensional direction."
            );

            const double angularVelocitySquared      = m_angularVelocity * m_angularVelocity;

            const double angularVelocityDotVariation = m_angularVelocity * physicalPositionVariation;

            gradientVariation.SetSize(3);
            gradientVariation = physicalPositionVariation;
            gradientVariation *= angularVelocitySquared;
            gradientVariation.Add(-angularVelocityDotVariation, m_angularVelocity);
        }

        [[nodiscard]] const mfem::Vector &angular_velocity() const noexcept {
            return m_angularVelocity;
        }

        [[nodiscard]] const mfem::Vector &center() const noexcept {
            return m_center;
        }

    private:
        mfem::Vector m_angularVelocity;
        mfem::Vector m_center;
    };
} // namespace mean_field::physics