module;

#include <algorithm>
#include <cmath>
#include <limits>
#include <numbers>
#include <stdexcept>
#include <utility>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

module mean_field;

import :deformation.spherical_harmonic_surface;

namespace mean_field::deformation {
    namespace {
        [[nodiscard]] SurfaceDeformationDescriptor harmonicDescriptor() noexcept {
            return {
                .name                               = "SphericalHarmonicSurface",
                .spatialDimension                   = 3,
                .motionKind                         = SurfaceMotionKind::Radial,
                .linearOnReferenceGeometry          = true,
                .requiresStarShapedReferenceSurface = true,
                .hasExactDerivativeTranspose        = true,
                .hasExactPullbackDerivative         = true,
                // A spectral cutoff alone is not a proof that either physical
                // gauge is removed; retain the conservative nodal descriptor.
                .translationTreatment               = GeometricGaugeTreatment::Retained,
                .orientationTreatment               = GeometricGaugeTreatment::Retained
            };
        }

        [[nodiscard]] long long countModes(const SphericalHarmonicSurfaceOptions &options) noexcept {
            const long long degree = options.maximumDegree;
            const long long order  = options.maximumAzimuthalOrder.value_or(options.maximumDegree);
            return (order + 1) * (order + 1) + (degree - order) * (2 * order + 1);
        }

        [[nodiscard]] mfem::Vector zeroCenter() {
            mfem::Vector center(3);
            center = 0.;
            return center;
        }

        void requireFinite(
            const mfem::Vector &values,
            const char *message
        ) {
            for (int i = 0; i < values.Size(); ++i)
                if (!std::isfinite(values(i)))
                    throw std::invalid_argument(message);
        }

        [[nodiscard]] double checkedDouble(long double value) {
            const double result = static_cast<double>(value);
            if (!std::isfinite(result))
                throw std::invalid_argument("Spherical harmonic surface action produced a nonfinite value.");
            return result;
        }

        // Recurrence on normalized associated Legendre functions avoids the
        // factorial growth of an unnormalized polynomial followed by scaling.
        [[nodiscard]] std::vector<double> evaluateModes(
            const double x,
            const double y,
            const double z,
            const int maximumDegree,
            const int maximumOrder,
            const std::vector<SphericalHarmonicMode> &modes
        ) {
            const std::size_t stride = static_cast<std::size_t>(maximumOrder) + 1;
            std::vector<double> legendre((static_cast<std::size_t>(maximumDegree) + 1) * stride, 0.);
            const auto p = [&](int l, int m) -> double & { return legendre[static_cast<std::size_t>(l) * stride + m]; };
            const double azimuth = std::atan2(y, x), cylindricalRadius = std::hypot(x, y);
            p(0, 0) = 1. / std::sqrt(4. * std::numbers::pi);
            for (int m = 0; m <= maximumOrder; ++m) {
                if (m > 0)
                    p(m, m) = -std::sqrt((2. * m + 1.) / (2. * m)) * cylindricalRadius * p(m - 1, m - 1);
                if (m < maximumDegree)
                    p(m + 1, m) = std::sqrt(2. * m + 3.) * z * p(m, m);
                for (int l = m + 2; l <= maximumDegree; ++l) {
                    const double ll = l, mm = m;
                    const double a = std::sqrt((4. * ll * ll - 1.) / (ll * ll - mm * mm));
                    const double b = std::sqrt(
                        ((2. * ll + 1.) * ((ll - 1.) * (ll - 1.) - mm * mm)) / ((2. * ll - 3.) * (ll * ll - mm * mm))
                    );
                    p(l, m) = a * z * p(l - 1, m) - b * p(l - 2, m);
                }
            }
            std::vector<double> values;
            values.reserve(modes.size());
            for (const auto mode : modes) {
                double value = p(mode.degree, std::abs(mode.order));
                if (mode.order < 0)
                    value *= std::sqrt(2.) * std::sin(-mode.order * azimuth);
                if (mode.order > 0)
                    value *= std::sqrt(2.) * std::cos(mode.order * azimuth);
                if (!std::isfinite(value))
                    throw std::invalid_argument("Nonfinite spherical harmonic reference basis.");
                values.push_back(value);
            }
            return values;
        }

        [[nodiscard]] long double columnDot(
            const mfem::DenseMatrix &matrix,
            int first,
            int second
        ) {
            long double result = 0.;
            for (int row = 0; row < matrix.Height(); ++row)
                result += static_cast<long double>(matrix(row, first)) * matrix(row, second);
            return result;
        }

        void requireIndependentModes(const mfem::DenseMatrix &basis) {
            mfem::DenseMatrix orthogonal(basis);
            // This detects unresolved/aliased nodal modes without changing
            // their physical normalization or silently dropping coordinates.
            constexpr double minimumRelativePivot = 1e-12;
            for (int column = 0; column < basis.Width(); ++column) {
                const double original = std::sqrt(checkedDouble(columnDot(basis, column, column)));
                for (int pass = 0; pass < 2; ++pass)
                    for (int previous = 0; previous < column; ++previous) {
                        const double projection = checkedDouble(columnDot(orthogonal, previous, column));
                        for (int row = 0; row < basis.Height(); ++row)
                            orthogonal(row, column) -= projection * orthogonal(row, previous);
                    }
                const double pivot = std::sqrt(checkedDouble(columnDot(orthogonal, column, column)));
                if (!(original > 0.) || !std::isfinite(pivot) || !(pivot > minimumRelativePivot * original))
                    throw std::invalid_argument(
                        "Selected spherical harmonics are numerically dependent on the surface nodes; "
                        "reduce their degree/order or refine the surface discretization."
                    );
                for (int row = 0; row < basis.Height(); ++row)
                    orthogonal(row, column) /= pivot;
            }
        }
    } // namespace

    void SphericalHarmonicSurfaceOptions::validate() const {
        if (maximumDegree < 0 ||
            (maximumAzimuthalOrder && (*maximumAzimuthalOrder < 0 || *maximumAzimuthalOrder > maximumDegree)))
            throw std::invalid_argument(
                "Spherical harmonic degree must be nonnegative and the azimuthal cutoff must lie in [0,L]."
            );
        if (countModes(*this) > std::numeric_limits<int>::max())
            throw std::invalid_argument("The spherical harmonic mode count exceeds supported vector indexing.");
    }

    SphericalHarmonicSurface::SphericalHarmonicSurface(SphericalHarmonicSurfaceOptions options)
        : SphericalHarmonicSurface(
              zeroCenter(),
              std::move(options)
          ) {
    }

    SphericalHarmonicSurface::SphericalHarmonicSurface(
        mfem::Vector referenceCenter,
        SphericalHarmonicSurfaceOptions options
    )
        : m_referenceCenter(std::move(referenceCenter)),
          m_options(std::move(options)) {
        validate();
    }

    const mfem::Vector &SphericalHarmonicSurface::referenceCenter() const noexcept {
        return m_referenceCenter;
    }
    const SphericalHarmonicSurfaceOptions &SphericalHarmonicSurface::options() const noexcept {
        return m_options;
    }
    int SphericalHarmonicSurface::modeCount() const noexcept {
        return static_cast<int>(countModes(m_options));
    }
    SurfaceDeformationDescriptor SphericalHarmonicSurface::descriptor() const noexcept {
        return harmonicDescriptor();
    }

    void SphericalHarmonicSurface::validate() const {
        m_options.validate();
        if (m_referenceCenter.Size() != 3)
            throw std::invalid_argument("SphericalHarmonicSurface requires a three-dimensional reference center.");
        requireFinite(m_referenceCenter, "Spherical harmonic reference-center coordinates must be finite.");
    }

    PreparedSphericalHarmonicSurface::PreparedSphericalHarmonicSurface(
        PreparedNodalRadialSurface radialSurface,
        std::vector<SphericalHarmonicMode> modes,
        mfem::DenseMatrix nodalBasis
    )
        : m_radialSurface(std::move(radialSurface)),
          m_modes(std::move(modes)),
          m_nodalBasis(std::move(nodalBasis)) {
    }

    SurfaceDeformationDescriptor PreparedSphericalHarmonicSurface::descriptor() const noexcept {
        return harmonicDescriptor();
    }
    int PreparedSphericalHarmonicSurface::parameterCount() const noexcept {
        return static_cast<int>(m_modes.size());
    }
    long long PreparedSphericalHarmonicSurface::globalParameterCount() const noexcept {
        return parameterCount();
    }
    long long PreparedSphericalHarmonicSurface::globalParameterOffset() const noexcept {
        return 0;
    }
    int PreparedSphericalHarmonicSurface::spatialDimension() const noexcept {
        return 3;
    }
    int PreparedSphericalHarmonicSurface::surfaceNodeCount() const noexcept {
        return m_radialSurface.parameterCount();
    }
    int PreparedSphericalHarmonicSurface::surfaceDisplacementSize() const noexcept {
        return m_radialSurface.surfaceDisplacementSize();
    }
    long long PreparedSphericalHarmonicSurface::globalSurfaceDisplacementSize() const noexcept {
        return m_radialSurface.globalSurfaceDisplacementSize();
    }
    long long PreparedSphericalHarmonicSurface::globalSurfaceDisplacementOffset() const noexcept {
        return m_radialSurface.globalSurfaceDisplacementOffset();
    }
    int PreparedSphericalHarmonicSurface::surfaceDisplacementDof(
        int node,
        int component
    ) const {
        return m_radialSurface.surfaceDisplacementDof(node, component);
    }
    double PreparedSphericalHarmonicSurface::radialDirection(
        int node,
        int component
    ) const {
        return m_radialSurface.radialDirection(node, component);
    }
    double PreparedSphericalHarmonicSurface::referenceRadius(int node) const {
        return m_radialSurface.referenceRadius(node);
    }
    const mfem::Vector &PreparedSphericalHarmonicSurface::referenceCenter() const noexcept {
        return m_radialSurface.referenceCenter();
    }
    const field::ScalarBoundaryDofMap &PreparedSphericalHarmonicSurface::surfaceDofMap() const noexcept {
        return m_radialSurface.surfaceDofMap();
    }
    std::span<const SphericalHarmonicMode> PreparedSphericalHarmonicSurface::modes() const noexcept {
        return m_modes;
    }
    const mfem::DenseMatrix &PreparedSphericalHarmonicSurface::nodalBasis() const noexcept {
        return m_nodalBasis;
    }

    void PreparedSphericalHarmonicSurface::requireParameters(const mfem::Vector &parameters) const {
        if (parameters.Size() != parameterCount())
            throw std::invalid_argument("Spherical harmonic parameter vector has an incompatible size.");
        requireFinite(parameters, "Spherical harmonic parameters and directions must be finite.");
    }

    void PreparedSphericalHarmonicSurface::requireSurfaceSize(const mfem::Vector &surfaceDisplacement) const {
        if (surfaceDisplacement.Size() != surfaceDisplacementSize())
            throw std::invalid_argument("Spherical harmonic nodal surface vector has an incompatible size.");
    }

    void PreparedSphericalHarmonicSurface::applyEmbedding(
        const mfem::Vector &parameters,
        mfem::Vector &surfaceDisplacement
    ) const {
        for (int node = 0; node < surfaceNodeCount(); ++node) {
            long double amplitude = 0.;
            for (int mode = 0; mode < parameterCount(); ++mode)
                amplitude += static_cast<long double>(m_nodalBasis(node, mode)) * parameters(mode);
            const double radialAmplitude = checkedDouble(amplitude);
            for (int component = 0; component < 3; ++component)
                surfaceDisplacement(surfaceDisplacementDof(node, component)) =
                    checkedDouble(static_cast<long double>(radialAmplitude) * radialDirection(node, component));
        }
    }

    void PreparedSphericalHarmonicSurface::buildSurfaceDisplacement(
        const mfem::Vector &parameters,
        mfem::Vector &surfaceDisplacement
    ) const {
        requireParameters(parameters);
        requireSurfaceSize(surfaceDisplacement);
        applyEmbedding(parameters, surfaceDisplacement);
    }

    void PreparedSphericalHarmonicSurface::applyJacobian(
        const mfem::Vector &parameters,
        const mfem::Vector &parameterDirection,
        mfem::Vector &surfaceDisplacementDirection
    ) const {
        requireParameters(parameters);
        requireParameters(parameterDirection);
        requireSurfaceSize(surfaceDisplacementDirection);
        applyEmbedding(parameterDirection, surfaceDisplacementDirection);
    }

    void PreparedSphericalHarmonicSurface::applyJacobianTranspose(
        const mfem::Vector &parameters,
        const mfem::Vector &surfaceDisplacementDual,
        mfem::Vector &parameterDual
    ) const {
        requireParameters(parameters);
        requireSurfaceSize(surfaceDisplacementDual);
        requireFinite(surfaceDisplacementDual, "Spherical harmonic surface dual must be finite.");
        if (parameterDual.Size() != parameterCount())
            throw std::invalid_argument("Spherical harmonic parameter dual has an incompatible size.");
        for (int mode = 0; mode < parameterCount(); ++mode) {
            long double work = 0.;
            for (int node = 0; node < surfaceNodeCount(); ++node) {
                long double radialWork = 0.;
                for (int component = 0; component < 3; ++component)
                    radialWork += static_cast<long double>(radialDirection(node, component)) *
                                  surfaceDisplacementDual(surfaceDisplacementDof(node, component));
                work += static_cast<long double>(m_nodalBasis(node, mode)) * radialWork;
            }
            parameterDual(mode) = checkedDouble(work);
        }
    }

    void PreparedSphericalHarmonicSurface::applyPullbackDerivative(
        const mfem::Vector &parameters,
        const mfem::Vector &parameterDirection,
        const mfem::Vector &surfaceDisplacementDual,
        mfem::Vector &parameterDualAction
    ) const {
        requireParameters(parameters);
        requireParameters(parameterDirection);
        requireSurfaceSize(surfaceDisplacementDual);
        requireFinite(surfaceDisplacementDual, "Spherical harmonic surface dual must be finite.");
        if (parameterDualAction.Size() != parameterCount())
            throw std::invalid_argument("Spherical harmonic pullback action has an incompatible size.");
        // The basis is compiled on the reference surface and never reweighted
        // as the star moves, so the derivative of its transpose is exactly zero.
        parameterDualAction = 0.;
    }

    PreparedSphericalHarmonicSurface compileSurfaceDeformationPrescription(
        const SphericalHarmonicSurface &prescription,
        const SurfaceDeformationCompilationContext &context
    ) {
        int ranks = 0;
        if (MPI_Comm_size(context.scalarFiniteElementSpace().GetComm(), &ranks) != MPI_SUCCESS)
            throw std::runtime_error("Could not inspect the spherical harmonic surface communicator.");
        // A replicated K-vector would be double-counted by distributed Newton
        // inner products. Reject it until modal ownership is implemented.
        if (ranks != 1)
            throw std::invalid_argument(
                "Spherical harmonic surface compilation currently supports one MPI rank; distributed modal ownership "
                "is not implemented."
            );
        prescription.validate();
        auto radial =
            compileSurfaceDeformationPrescription(NodalRadialSurface{prescription.referenceCenter()}, context);
        if (radial.spatialDimension() != 3)
            throw std::invalid_argument("Spherical harmonics require a three-dimensional reference mesh.");
        if (prescription.modeCount() > radial.parameterCount())
            throw std::invalid_argument("Selected spherical harmonic mode count exceeds the number of surface nodes.");
        const int degree = prescription.options().maximumDegree;
        const int cutoff = prescription.options().maximumAzimuthalOrder.value_or(degree);
        std::vector<SphericalHarmonicMode> modes;
        modes.reserve(prescription.modeCount());
        for (int l = 0; l <= degree; ++l)
            for (int m = -std::min(l, cutoff); m <= std::min(l, cutoff); ++m)
                modes.push_back({l, m});
        mfem::DenseMatrix basis(radial.parameterCount(), prescription.modeCount());
        for (int node = 0; node < radial.parameterCount(); ++node) {
            const auto values = evaluateModes(
                radial.radialDirection(node, 0), radial.radialDirection(node, 1), radial.radialDirection(node, 2),
                degree, cutoff, modes
            );
            for (int mode = 0; mode < prescription.modeCount(); ++mode)
                basis(node, mode) = values[mode];
        }
        requireIndependentModes(basis);
        return PreparedSphericalHarmonicSurface(std::move(radial), std::move(modes), std::move(basis));
    }
} // namespace mean_field::deformation