module;

#include <cmath>
#include <format>
#include <stdexcept>
#include <utility>

#include <mfem.hpp>

module mean_field;

import :deformation.nodal_radial_surface;

namespace mean_field::deformation {
    namespace {
        [[nodiscard]] SurfaceDeformationDescriptor nodalRadialDescriptor(const int spatialDimension) noexcept {
            return {
                .name                               = "NodalRadialSurface",
                .spatialDimension                   = spatialDimension,
                .motionKind                         = SurfaceMotionKind::Radial,
                .linearOnReferenceGeometry          = true,
                .requiresStarShapedReferenceSurface = true,
                .hasExactDerivativeTranspose        = true,
                .hasExactPullbackDerivative         = true,
                .translationTreatment               = GeometricGaugeTreatment::Retained,
                .orientationTreatment               = GeometricGaugeTreatment::Retained
            };
        }

        void requireFiniteVector(
            const mfem::Vector &vector,
            const char *message
        ) {
            for (int index = 0; index < vector.Size(); ++index) {
                if (!std::isfinite(vector(index))) {
                    throw std::invalid_argument(message);
                }
            }
        }
    } // namespace

    SurfaceDeformationCompilationContext::SurfaceDeformationCompilationContext(
        mfem::ParFiniteElementSpace &scalarFiniteElementSpace,
        field::ScalarBoundaryDofMap surfaceDofMap
    )
        : m_scalarFiniteElementSpace(&scalarFiniteElementSpace),
          m_surfaceDofMap(std::move(surfaceDofMap)) {
        if (scalarFiniteElementSpace.Nonconforming()) {
            throw std::invalid_argument(
                "Surface deformation compilation currently requires a conforming scalar finite-element space."
            );
        }
        if (scalarFiniteElementSpace.GetVDim() != 1) {
            throw std::invalid_argument("Surface deformation compilation requires a scalar finite-element space.");
        }
        if (scalarFiniteElementSpace.GetMesh() == nullptr) {
            throw std::invalid_argument("Surface deformation compilation requires a finite-element mesh.");
        }
        if (m_surfaceDofMap.volume_true_dof_size() != scalarFiniteElementSpace.GetTrueVSize()) {
            throw std::invalid_argument(
                "The surface DOF map and scalar finite-element space have incompatible true-DOF sizes."
            );
        }
        if (m_surfaceDofMap.global_size() <= 0) {
            throw std::invalid_argument("Surface deformation compilation requires at least one surface coordinate.");
        }
    }

    mfem::ParFiniteElementSpace &SurfaceDeformationCompilationContext::scalarFiniteElementSpace() const noexcept {
        return *m_scalarFiniteElementSpace;
    }

    const field::ScalarBoundaryDofMap &SurfaceDeformationCompilationContext::surfaceDofMap() const noexcept {
        return m_surfaceDofMap;
    }

    NodalRadialSurface::NodalRadialSurface(mfem::Vector referenceCenter)
        : m_referenceCenter(std::move(referenceCenter)) {
        validate();
    }

    const mfem::Vector &NodalRadialSurface::referenceCenter() const noexcept {
        return m_referenceCenter;
    }

    SurfaceDeformationDescriptor NodalRadialSurface::descriptor() const noexcept {
        return nodalRadialDescriptor(m_referenceCenter.Size());
    }

    void NodalRadialSurface::validate() const {
        if (m_referenceCenter.Size() <= 0) {
            throw std::invalid_argument("NodalRadialSurface requires a non-empty reference center.");
        }
        requireFiniteVector(m_referenceCenter, "NodalRadialSurface reference-center coordinates must be finite.");
    }

    PreparedNodalRadialSurface::PreparedNodalRadialSurface(
        const SurfaceDeformationDescriptor descriptor,
        mfem::Vector referenceCenter,
        field::ScalarBoundaryDofMap surfaceDofMap,
        mfem::Vector radialDirections,
        mfem::Vector referenceRadii
    )
        : m_descriptor(descriptor),
          m_referenceCenter(std::move(referenceCenter)),
          m_surfaceDofMap(std::move(surfaceDofMap)),
          m_radialDirections(std::move(radialDirections)),
          m_referenceRadii(std::move(referenceRadii)) {
    }

    SurfaceDeformationDescriptor PreparedNodalRadialSurface::descriptor() const noexcept {
        return m_descriptor;
    }

    int PreparedNodalRadialSurface::parameterCount() const noexcept {
        return m_surfaceDofMap.local_size();
    }

    long long PreparedNodalRadialSurface::globalParameterCount() const noexcept {
        return m_surfaceDofMap.global_size();
    }

    long long PreparedNodalRadialSurface::globalParameterOffset() const noexcept {
        return m_surfaceDofMap.global_offset();
    }

    int PreparedNodalRadialSurface::spatialDimension() const noexcept {
        return m_descriptor.spatialDimension;
    }

    int PreparedNodalRadialSurface::surfaceDisplacementSize() const noexcept {
        return spatialDimension() * parameterCount();
    }

    long long PreparedNodalRadialSurface::globalSurfaceDisplacementSize() const noexcept {
        return static_cast<long long>(spatialDimension()) * globalParameterCount();
    }

    long long PreparedNodalRadialSurface::globalSurfaceDisplacementOffset() const noexcept {
        return static_cast<long long>(spatialDimension()) * globalParameterOffset();
    }

    int PreparedNodalRadialSurface::surfaceDisplacementDof(
        const int parameterDof,
        const int component
    ) const {
        if (parameterDof < 0 || parameterDof >= parameterCount()) {
            throw std::out_of_range("Parameter DOF is outside PreparedNodalRadialSurface.");
        }
        if (component < 0 || component >= spatialDimension()) {
            throw std::out_of_range("Surface-displacement component is outside PreparedNodalRadialSurface.");
        }
        return spatialDimension() * parameterDof + component;
    }

    double PreparedNodalRadialSurface::radialDirection(
        const int parameterDof,
        const int component
    ) const {
        return m_radialDirections(surfaceDisplacementDof(parameterDof, component));
    }

    double PreparedNodalRadialSurface::referenceRadius(const int parameterDof) const {
        if (parameterDof < 0 || parameterDof >= parameterCount()) {
            throw std::out_of_range("Parameter DOF is outside PreparedNodalRadialSurface.");
        }
        return m_referenceRadii(parameterDof);
    }

    const mfem::Vector &PreparedNodalRadialSurface::referenceCenter() const noexcept {
        return m_referenceCenter;
    }

    const field::ScalarBoundaryDofMap &PreparedNodalRadialSurface::surfaceDofMap() const noexcept {
        return m_surfaceDofMap;
    }

    void PreparedNodalRadialSurface::buildSurfaceDisplacement(
        const mfem::Vector &parameters,
        mfem::Vector &surfaceDisplacement
    ) const {
        requireParameterSize(parameters);
        requireSurfaceDisplacementSize(surfaceDisplacement);

        for (int parameterDof = 0; parameterDof < parameterCount(); ++parameterDof) {
            for (int component = 0; component < spatialDimension(); ++component) {
                const int surfaceDof            = spatialDimension() * parameterDof + component;
                surfaceDisplacement(surfaceDof) = parameters(parameterDof) * m_radialDirections(surfaceDof);
            }
        }
    }

    void PreparedNodalRadialSurface::applyJacobian(
        const mfem::Vector &parameters,
        const mfem::Vector &parameterDirection,
        mfem::Vector &surfaceDisplacementDirection
    ) const {
        requireParameterSize(parameters);
        requireParameterSize(parameterDirection);
        requireSurfaceDisplacementSize(surfaceDisplacementDirection);

        for (int parameterDof = 0; parameterDof < parameterCount(); ++parameterDof) {
            for (int component = 0; component < spatialDimension(); ++component) {
                const int surfaceDof = spatialDimension() * parameterDof + component;
                surfaceDisplacementDirection(surfaceDof) =
                    parameterDirection(parameterDof) * m_radialDirections(surfaceDof);
            }
        }
    }

    void PreparedNodalRadialSurface::applyJacobianTranspose(
        const mfem::Vector &parameters,
        const mfem::Vector &surfaceDisplacementDual,
        mfem::Vector &parameterDual
    ) const {
        requireParameterSize(parameters);
        requireSurfaceDisplacementSize(surfaceDisplacementDual);
        requireParameterSize(parameterDual);

        for (int parameterDof = 0; parameterDof < parameterCount(); ++parameterDof) {
            double radialWork = 0.0;
            for (int component = 0; component < spatialDimension(); ++component) {
                const int surfaceDof = spatialDimension() * parameterDof + component;
                radialWork += m_radialDirections(surfaceDof) * surfaceDisplacementDual(surfaceDof);
            }
            parameterDual(parameterDof) = radialWork;
        }
    }

    void PreparedNodalRadialSurface::applyPullbackDerivative(
        const mfem::Vector &parameters,
        const mfem::Vector &parameterDirection,
        const mfem::Vector &surfaceDisplacementDual,
        mfem::Vector &parameterDualAction
    ) const {
        requireParameterSize(parameters);
        requireParameterSize(parameterDirection);
        requireSurfaceDisplacementSize(surfaceDisplacementDual);
        requireParameterSize(parameterDualAction);

        parameterDualAction = 0.0;
    }

    void PreparedNodalRadialSurface::requireParameterSize(const mfem::Vector &parameters) const {
        if (parameters.Size() != parameterCount()) {
            throw std::invalid_argument(
                std::format(
                    "Nodal radial parameter vector has size {}, but the prepared surface requires {}.",
                    parameters.Size(), parameterCount()
                )
            );
        }
    }

    void PreparedNodalRadialSurface::requireSurfaceDisplacementSize(const mfem::Vector &surfaceDisplacement) const {
        if (surfaceDisplacement.Size() != surfaceDisplacementSize()) {
            throw std::invalid_argument(
                std::format(
                    "Surface displacement vector has size {}, but the prepared nodal radial surface requires {}.",
                    surfaceDisplacement.Size(), surfaceDisplacementSize()
                )
            );
        }
    }

    PreparedNodalRadialSurface compileSurfaceDeformationPrescription(
        const NodalRadialSurface &prescription,
        const SurfaceDeformationCompilationContext &context
    ) {
        prescription.validate();

        mfem::ParFiniteElementSpace &scalarSpace = context.scalarFiniteElementSpace();
        const mfem::Mesh *mesh                   = scalarSpace.GetMesh();

        if (mesh == nullptr) {
            throw std::invalid_argument("Nodal radial surface compilation requires a reference mesh.");
        }
        if (prescription.referenceCenter().Size() != mesh->SpaceDimension()) {
            throw std::invalid_argument(
                std::format(
                    "NodalRadialSurface reference center has dimension {}, but the reference mesh has spatial "
                    "dimension {}.",
                    prescription.referenceCenter().Size(), mesh->SpaceDimension()
                )
            );
        }

        const field::ScalarBoundaryDofMap &surfaceDofMap = context.surfaceDofMap();
        const int parameterCount                         = surfaceDofMap.local_size();
        const int spatialDimension                       = mesh->SpaceDimension();

        mfem::Vector referencePositions(spatialDimension * parameterCount);
        mfem::ParGridFunction coordinateField(&scalarSpace);

        for (int component = 0; component < spatialDimension; ++component) {
            mfem::FunctionCoefficient coordinateCoefficient([component](const mfem::Vector &position) {
                return position(component);
            });

            coordinateField.ProjectCoefficient(coordinateCoefficient);

            mfem::Vector coordinateTrueDofs;
            coordinateField.GetTrueDofs(coordinateTrueDofs);
            const mfem::Vector surfaceCoordinates = surfaceDofMap.gather(coordinateTrueDofs);

            for (int parameterDof = 0; parameterDof < parameterCount; ++parameterDof) {
                referencePositions(spatialDimension * parameterDof + component) = surfaceCoordinates(parameterDof);
            }
        }

        mfem::Vector radialDirections(referencePositions.Size());
        mfem::Vector referenceRadii(parameterCount);

        for (int parameterDof = 0; parameterDof < parameterCount; ++parameterDof) {
            double radiusSquared = 0.0;

            for (int component = 0; component < spatialDimension; ++component) {
                const int surfaceDof = spatialDimension * parameterDof + component;
                const double radialCoordinate =
                    referencePositions(surfaceDof) - prescription.referenceCenter()(component);

                radialDirections(surfaceDof) = radialCoordinate;
                radiusSquared += radialCoordinate * radialCoordinate;
            }

            const double radius = std::sqrt(radiusSquared);
            if (!std::isfinite(radius) || radius <= 0.0) {
                throw std::invalid_argument(
                    "Every nodal radial surface coordinate must have a finite positive distance from the reference "
                    "center."
                );
            }

            referenceRadii(parameterDof) = radius;
            for (int component = 0; component < spatialDimension; ++component) {
                radialDirections(spatialDimension * parameterDof + component) /= radius;
            }
        }

        return PreparedNodalRadialSurface(
            nodalRadialDescriptor(spatialDimension), prescription.referenceCenter(), surfaceDofMap,
            std::move(radialDirections), std::move(referenceRadii)
        );
    }
} // namespace mean_field::deformation