module;

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

module mean_field;

namespace {
    bool vector_is_finite(const mfem::Vector &vector) {
        for (int i = 0; i < vector.Size(); ++i) {
            if (!std::isfinite(vector(i)))
                return false;
        }
        return true;
    }

    bool matrix_is_finite(const mfem::DenseMatrix &matrix) {
        for (int i = 0; i < matrix.Height(); ++i) {
            for (int j = 0; j < matrix.Width(); ++j) {
                if (!std::isfinite(matrix(i, j)))
                    return false;
            }
        }
        return true;
    }
} // namespace

namespace mean_field::mapping::compactification {
    KelvinCompactification::KelvinCompactification(options::KelvinCompactificationOptions options)
        : m_options(options) {
        if (!std::isfinite(m_options.r_star_ref) || !std::isfinite(m_options.r_inf_ref)) {
            throw std::invalid_argument("Kelvin compactification radii must be finite.");
        }

        if (m_options.r_star_ref <= 0.0 || m_options.r_inf_ref <= m_options.r_star_ref) {
            throw std::invalid_argument("Kelvin compactification requires 0 < r_star_ref < r_inf_ref.");
        }

        if (!std::isfinite(m_options.coordinate_tolerance) || m_options.coordinate_tolerance < 0.0 ||
            m_options.coordinate_tolerance >= 1.0) {
            throw std::invalid_argument(
                "Kelvin compactification coordinate tolerance must be finite "
                "and lie "
                "in [0, 1)."
            );
        }
    }

    MappingStatus KelvinCompactification::ComputeRadialFactors(
        const double compactification_coordinate,
        RadialFactors &factors
    ) const {
        if (!std::isfinite(compactification_coordinate))
            return MappingStatus::non_finite_input;

        const double tolerance = m_options.coordinate_tolerance;

        if (compactification_coordinate < -tolerance || compactification_coordinate > 1.0 + tolerance) {
            return MappingStatus::outside_reference_domain;
        }

        double coordinate = compactification_coordinate;
        if (coordinate < 0.0)
            coordinate = 0.0;

        if (coordinate >= 1.0 - tolerance) {
            return MappingStatus::at_compactified_infinity;
        }

        const double radial_extent        = m_options.r_inf_ref - m_options.r_star_ref;
        const double computational_radius = m_options.r_star_ref + coordinate * radial_extent;

        if (!std::isfinite(computational_radius) || computational_radius <= 0.0) {
            return MappingStatus::invalid_reference_radius;
        }

        const double one_minus_coordinate = 1.0 - coordinate;
        const double denominator          = computational_radius * one_minus_coordinate;

        if (!std::isfinite(denominator) || denominator <= 0.0) {
            return MappingStatus::non_finite_result;
        }

        const double scale            = m_options.r_star_ref / denominator;
        const double scale_derivative = scale * (1.0 / one_minus_coordinate - radial_extent / computational_radius);

        if (!std::isfinite(scale) || !std::isfinite(scale_derivative)) {
            return MappingStatus::non_finite_result;
        }

        factors.coordinate           = coordinate;
        factors.computational_radius = computational_radius;
        factors.scale                = scale;
        factors.scale_derivative     = scale_derivative;

        return MappingStatus::valid;
    }

    MappingStatus KelvinCompactification::Evaluate(
        const ExteriorMapInput &input,
        ExteriorMapResult &result
    ) const {
        const int dimension = input.reference_position.Size();

        if (dimension <= 0 || input.displaced_position.Size() != dimension ||
            input.compactification_coordinate_gradient.Size() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (input.displacement_jacobian.Height() != dimension || input.displacement_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (!vector_is_finite(input.reference_position) || !vector_is_finite(input.displaced_position) ||
            !vector_is_finite(input.compactification_coordinate_gradient) ||
            !matrix_is_finite(input.displacement_jacobian)) {
            return MappingStatus::non_finite_input;
        }

        RadialFactors factors;
        const MappingStatus factor_status = ComputeRadialFactors(input.compactification_coordinate, factors);
        if (factor_status != MappingStatus::valid)
            return factor_status;

        result.physical_position.SetSize(dimension);
        result.mapping_jacobian.SetSize(dimension, dimension);

        for (int i = 0; i < dimension; ++i) {
            result.physical_position(i) = factors.scale * input.displaced_position(i);

            for (int j = 0; j < dimension; ++j) {
                const double scale_gradient = factors.scale_derivative * input.compactification_coordinate_gradient(j);
                result.mapping_jacobian(i, j) =
                    factors.scale * input.displacement_jacobian(i, j) + input.displaced_position(i) * scale_gradient;
            }
        }

        if (!vector_is_finite(result.physical_position) || !matrix_is_finite(result.mapping_jacobian)) {
            return MappingStatus::non_finite_result;
        }

        const double mapping_determinant = result.mapping_jacobian.Det();

        if (!std::isfinite(mapping_determinant))
            return MappingStatus::non_finite_result;
        if (mapping_determinant <= 0.0)
            return MappingStatus::non_positive_determinant;

        return MappingStatus::valid;
    }

    MappingStatus KelvinCompactification::EvaluateVariation(
        const ExteriorMapInput &input,
        const ExteriorMapResult &result,
        const ExteriorMapDirection &direction,
        ExteriorMapVariation &variation
    ) const {
        const int dimension = input.reference_position.Size();

        if (dimension <= 0 || input.displaced_position.Size() != dimension ||
            input.compactification_coordinate_gradient.Size() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (input.displacement_jacobian.Height() != dimension || input.displacement_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (result.physical_position.Size() != dimension || result.mapping_jacobian.Height() != dimension ||
            result.mapping_jacobian.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (direction.displaced_position_variation.Size() != dimension ||
            direction.displacement_jacobian_variation.Height() != dimension ||
            direction.displacement_jacobian_variation.Width() != dimension) {
            return MappingStatus::invalid_dimension;
        }

        if (!vector_is_finite(input.reference_position) || !vector_is_finite(input.displaced_position) ||
            !vector_is_finite(input.compactification_coordinate_gradient) ||
            !matrix_is_finite(input.displacement_jacobian)) {
            return MappingStatus::non_finite_input;
        }

        if (!vector_is_finite(result.physical_position) || !matrix_is_finite(result.mapping_jacobian) ||
            !vector_is_finite(direction.displaced_position_variation) ||
            !matrix_is_finite(direction.displacement_jacobian_variation)) {
            return MappingStatus::non_finite_input;
        }

        RadialFactors factors;
        const MappingStatus factor_status = ComputeRadialFactors(input.compactification_coordinate, factors);
        if (factor_status != MappingStatus::valid)
            return factor_status;

        variation.physical_position_variation.SetSize(dimension);
        variation.mapping_jacobian_variation.SetSize(dimension, dimension);

        for (int i = 0; i < dimension; ++i) {
            variation.physical_position_variation(i) = factors.scale * direction.displaced_position_variation(i);

            for (int j = 0; j < dimension; ++j) {
                const double scale_gradient = factors.scale_derivative * input.compactification_coordinate_gradient(j);
                variation.mapping_jacobian_variation(i, j) =
                    factors.scale * direction.displacement_jacobian_variation(i, j) +
                    direction.displaced_position_variation(i) * scale_gradient;
            }
        }

        if (!vector_is_finite(variation.physical_position_variation) ||
            !matrix_is_finite(variation.mapping_jacobian_variation)) {
            return MappingStatus::non_finite_result;
        }

        return MappingStatus::valid;
    }

    std::string_view KelvinCompactification::GetName() const noexcept {
        return "KelvinCompactification";
    }

    double KelvinCompactification::GetReferenceStellarRadius() const noexcept {
        return m_options.r_star_ref;
    }

    double KelvinCompactification::GetReferenceInfinityRadius() const noexcept {
        return m_options.r_inf_ref;
    }

    double KelvinCompactification::GetCoordinateTolerance() const noexcept {
        return m_options.coordinate_tolerance;
    }
} // namespace mean_field::mapping::compactification