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;
        }

        [[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