module;

#include <cmath>
#include <concepts>
#include <cstdint>
#include <exception>
#include <functional>
#include <optional>
#include <span>
#include <stdexcept>
#include <string_view>
#include <type_traits>
#include <utility>

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

export module mean_field:solver.newton;

export import :solver.linear_backend;

export namespace mean_field::solver::nonlinear {
    /*
     * Backtracking counts the full Newton trial as its first trial.  A
     * contraction is applied only after that candidate has been rejected.
     */
    struct BacktrackingOptions final {
        double initialStepLength{1.0};
        double contractionFactor{0.5};
        double fractionToBoundarySafety{0.9};
        double sufficientDecrease{1.0e-4};
        double minimumStepLength{1.0e-8};
        int maximumTrials{20};

        void Validate() const {
            if (!std::isfinite(initialStepLength) || initialStepLength <= 0.0) {
                throw std::invalid_argument("Newton backtracking requires a finite, positive initial step length.");
            }
            if (!std::isfinite(contractionFactor) || contractionFactor <= 0.0 || contractionFactor >= 1.0) {
                throw std::invalid_argument(
                    "Newton backtracking requires a finite contraction factor strictly between zero and one."
                );
            }
            if (!std::isfinite(fractionToBoundarySafety) || fractionToBoundarySafety <= 0.0 ||
                fractionToBoundarySafety >= 1.0) {
                throw std::invalid_argument(
                    "Newton backtracking requires a finite fraction-to-boundary safety factor strictly between zero "
                    "and one."
                );
            }
            if (!std::isfinite(sufficientDecrease) || sufficientDecrease <= 0.0 || sufficientDecrease >= 1.0) {
                throw std::invalid_argument(
                    "Newton backtracking requires a finite sufficient-decrease factor strictly between zero and one."
                );
            }
            if (!std::isfinite(minimumStepLength) || minimumStepLength <= 0.0 ||
                minimumStepLength > initialStepLength) {
                throw std::invalid_argument(
                    "Newton backtracking requires a finite, positive minimum step no larger than the initial step."
                );
            }
            if (maximumTrials <= 0) {
                throw std::invalid_argument("Newton backtracking requires at least one permitted trial.");
            }
        }
    };

    struct NewtonOptions final {
        double relativeTolerance{1.0e-8};
        double absoluteTolerance{0.0};
        int maximumIterations{50};
        LinearSolveControl linearSolve{};
        BacktrackingOptions backtracking{};

        void Validate() const {
            if (!std::isfinite(relativeTolerance) || relativeTolerance < 0.0) {
                throw std::invalid_argument("A Newton solve requires a finite, non-negative relative tolerance.");
            }
            if (!std::isfinite(absoluteTolerance) || absoluteTolerance < 0.0) {
                throw std::invalid_argument("A Newton solve requires a finite, non-negative absolute tolerance.");
            }
            if (maximumIterations <= 0) {
                throw std::invalid_argument("A Newton solve requires at least one permitted iteration.");
            }
            linearSolve.Validate();
            backtracking.Validate();
        }

        [[nodiscard]] double ConvergenceThreshold(const double initialResidualNorm) const {
            Validate();
            if (!std::isfinite(initialResidualNorm) || initialResidualNorm < 0.0) {
                throw std::invalid_argument(
                    "A Newton convergence threshold requires a finite, non-negative initial residual norm."
                );
            }
            const double relativeThreshold = relativeTolerance * initialResidualNorm;
            if (!std::isfinite(relativeThreshold)) {
                throw std::invalid_argument("The Newton relative convergence threshold must be finite.");
            }
            return absoluteTolerance > relativeThreshold ? absoluteTolerance : relativeThreshold;
        }
    };

    /*
     * The MVP globalization merit is phi(x) = 0.5 ||F_normalized(x)||^2.
     * A metric customization receives the solve communicator and must return
     * communicator-consistent values or throw collectively. The Newton engine
     * reduces every predicate that drives control flow, but it cannot make a
     * rank-local exception inside an arbitrary callback collective-safe.
     */
    struct NormalizedResidualMetric final { };

    struct MetricEvaluation final {
        double residualNorm{0.0};
        double merit{0.0};
    };

    [[nodiscard]] inline MetricEvaluation getMetric(
        const NormalizedResidualMetric &,
        const mfem::Vector &normalizedResidual,
        const MPI_Comm communicator
    ) {
        if (communicator == MPI_COMM_NULL) {
            throw std::invalid_argument("A nonlinear metric requires a valid communicator.");
        }

        double localSquaredNorm = 0.0;
        for (int index = 0; index < normalizedResidual.Size(); ++index) {
            const double value = static_cast<double>(normalizedResidual(index));
            localSquaredNorm += value * value;
        }

        double globalSquaredNorm = 0.0;
        if (MPI_Allreduce(&localSquaredNorm, &globalSquaredNorm, 1, MPI_DOUBLE, MPI_SUM, communicator) != MPI_SUCCESS) {
            throw std::runtime_error("The nonlinear metric could not reduce the normalized residual norm.");
        }

        const double residualNorm = std::sqrt(globalSquaredNorm);
        return {.residualNorm = residualNorm, .merit = 0.5 * residualNorm * residualNorm};
    }

    template <typename Metric = NormalizedResidualMetric>
        requires std::move_constructible<std::remove_cvref_t<Metric>>
    class Newton final {
    public:
        using MetricType = std::remove_cvref_t<Metric>;

        Newton()
            requires std::default_initializable<MetricType>
            : Newton(
                  NewtonOptions{},
                  MetricType{}
              ) {
        }

        explicit Newton(NewtonOptions options)
            requires std::default_initializable<MetricType>
            : Newton(
                  std::move(options),
                  MetricType{}
              ) {
        }

        Newton(
            NewtonOptions options,
            MetricType metric
        )
            : m_options(std::move(options)),
              m_metric(std::move(metric)) {
            m_options.Validate();
        }

        [[nodiscard]] const NewtonOptions &options() const noexcept {
            return m_options;
        }

        [[nodiscard]] const MetricType &metric() const noexcept {
            return m_metric;
        }

    private:
        NewtonOptions m_options;
        [[no_unique_address]] MetricType m_metric;
    };

    Newton() -> Newton<NormalizedResidualMetric>;
    Newton(NewtonOptions) -> Newton<NormalizedResidualMetric>;

    template <typename Metric>
    Newton(
        NewtonOptions,
        Metric
    ) -> Newton<std::remove_cvref_t<Metric>>;

    template <typename Candidate> struct IsNewtonConfiguration : std::false_type { };

    template <typename Metric> struct IsNewtonConfiguration<Newton<Metric>> : std::true_type { };

    template <typename Candidate>
    concept NewtonConfiguration = IsNewtonConfiguration<std::remove_cvref_t<Candidate>>::value;

    enum class IterationDisposition : std::uint8_t {
        unspecified,
        accepted,
        converged,
        inadmissible_state,
        non_finite_state,
        non_finite_residual,
        linear_solve_failure,
        globalization_failure,
        stagnation,
        iteration_limit
    };

    /*
     * Event spans borrow solver workspaces and are valid only for the duration
     * of the callback. Copy values that must outlive the callback.
     */
    struct BeforeIteration final {
        int iteration{0};
        double initialResidualNorm{0.0};
        double residualNorm{0.0};
        double relativeResidualNorm{0.0};
        double merit{0.0};
        MPI_Comm communicator{MPI_COMM_NULL};
        std::span<const mfem::real_t> physicalState{};
        std::span<const mfem::real_t> normalizedState{};
        std::span<const mfem::real_t> normalizedResidual{};
    };

    enum class LineSearchTrialDisposition : std::uint8_t {
        accepted,
        inadmissible_state,
        non_finite_state,
        non_finite_residual,
        insufficient_decrease
    };

    struct AfterLineSearchTrial final {
        int iteration{0};
        int trial{0};
        double stepLength{0.0};
        LineSearchTrialDisposition disposition{LineSearchTrialDisposition::insufficient_decrease};
        std::string_view rejectionSource{};
        std::optional<MetricEvaluation> metric{};
        std::optional<double> minimumJacobianDeterminant{};
        double preparationSeconds{0.0};
        double metricSeconds{0.0};
        MPI_Comm communicator{MPI_COMM_NULL};
        std::span<const mfem::real_t> candidatePhysicalState{};
        std::span<const mfem::real_t> candidateNormalizedState{};
        std::span<const mfem::real_t> candidateNormalizedResidual{};
    };

    struct AfterIteration final {
        int iteration{0};
        IterationDisposition disposition{IterationDisposition::unspecified};
        bool stepAccepted{false};
        double acceptedStepLength{0.0};
        int lineSearchTrials{0};
        double initialResidualNorm{0.0};
        double previousResidualNorm{0.0};
        double residualNorm{0.0};
        double relativeResidualNorm{0.0};
        double merit{0.0};
        double iterationSeconds{0.0};
        double lineSearchSeconds{0.0};
        double trialPreparationSeconds{0.0};
        double metricEvaluationSeconds{0.0};
        double preconditionerRefreshSeconds{0.0};
        double rollbackSeconds{0.0};
        std::optional<LinearSolveReport> linearSolve{};
        MPI_Comm communicator{MPI_COMM_NULL};
        std::span<const mfem::real_t> physicalState{};
        std::span<const mfem::real_t> normalizedState{};
        std::span<const mfem::real_t> normalizedResidual{};
    };

    struct NoObserver final { };

    template <typename Candidate>
    concept BeforeIterationCallback = std::invocable<Candidate &, const BeforeIteration &> &&
                                      std::same_as<std::invoke_result_t<Candidate &, const BeforeIteration &>, void>;

    template <typename Candidate>
    concept AfterIterationCallback = std::invocable<Candidate &, const AfterIteration &> &&
                                     std::same_as<std::invoke_result_t<Candidate &, const AfterIteration &>, void>;

    template <typename Candidate>
    concept LineSearchTrialCallback =
        std::invocable<Candidate &, const AfterLineSearchTrial &> &&
        std::same_as<std::invoke_result_t<Candidate &, const AfterLineSearchTrial &>, void>;

    template <BeforeIterationCallback BeforeCallback, AfterIterationCallback AfterCallback>
    class CallbackObserver final {
    public:
        CallbackObserver(
            BeforeCallback before,
            AfterCallback after
        )
            : m_before(std::move(before)),
              m_after(std::move(after)) {
        }

        void beforeIteration(const BeforeIteration &event) noexcept(std::is_nothrow_invocable_v<
                                                                    BeforeCallback &,
                                                                    const BeforeIteration &>) {
            std::invoke(m_before, event);
        }

        void afterIteration(const AfterIteration &event) noexcept(std::is_nothrow_invocable_v<
                                                                  AfterCallback &,
                                                                  const AfterIteration &>) {
            std::invoke(m_after, event);
        }

    private:
        [[no_unique_address]] BeforeCallback m_before;
        [[no_unique_address]] AfterCallback m_after;
    };

    template <
        typename BeforeCallback,
        typename AfterCallback>
        requires BeforeIterationCallback<std::decay_t<BeforeCallback>> &&
                 AfterIterationCallback<std::decay_t<AfterCallback>> &&
                 std::constructible_from<
                     std::decay_t<BeforeCallback>,
                     BeforeCallback> &&
                 std::constructible_from<
                     std::decay_t<AfterCallback>,
                     AfterCallback>
    [[nodiscard]] auto makeObserver(
        BeforeCallback &&before,
        AfterCallback &&after
    ) {
        return CallbackObserver<std::decay_t<BeforeCallback>, std::decay_t<AfterCallback>>{
            std::forward<BeforeCallback>(before), std::forward<AfterCallback>(after)
        };
    }

    template <
        BeforeIterationCallback BeforeCallback,
        LineSearchTrialCallback TrialCallback,
        AfterIterationCallback AfterCallback>
    class DetailedCallbackObserver final {
    public:
        DetailedCallbackObserver(
            BeforeCallback before,
            TrialCallback trial,
            AfterCallback after
        )
            : m_before(std::move(before)),
              m_trial(std::move(trial)),
              m_after(std::move(after)) {
        }

        void beforeIteration(const BeforeIteration &event) noexcept(std::is_nothrow_invocable_v<
                                                                    BeforeCallback &,
                                                                    const BeforeIteration &>) {
            std::invoke(m_before, event);
        }

        void afterLineSearchTrial(const AfterLineSearchTrial &event) noexcept(std::is_nothrow_invocable_v<
                                                                              TrialCallback &,
                                                                              const AfterLineSearchTrial &>) {
            std::invoke(m_trial, event);
        }

        void afterIteration(const AfterIteration &event) noexcept(std::is_nothrow_invocable_v<
                                                                  AfterCallback &,
                                                                  const AfterIteration &>) {
            std::invoke(m_after, event);
        }

    private:
        [[no_unique_address]] BeforeCallback m_before;
        [[no_unique_address]] TrialCallback m_trial;
        [[no_unique_address]] AfterCallback m_after;
    };

    template <
        typename BeforeCallback,
        typename TrialCallback,
        typename AfterCallback>
        requires BeforeIterationCallback<std::decay_t<BeforeCallback>> &&
                 LineSearchTrialCallback<std::decay_t<TrialCallback>> &&
                 AfterIterationCallback<std::decay_t<AfterCallback>> &&
                 std::constructible_from<
                     std::decay_t<BeforeCallback>,
                     BeforeCallback> &&
                 std::constructible_from<
                     std::decay_t<TrialCallback>,
                     TrialCallback> &&
                 std::constructible_from<
                     std::decay_t<AfterCallback>,
                     AfterCallback>
    [[nodiscard]] auto makeObserver(
        BeforeCallback &&before,
        TrialCallback &&trial,
        AfterCallback &&after
    ) {
        return DetailedCallbackObserver<
            std::decay_t<BeforeCallback>, std::decay_t<TrialCallback>, std::decay_t<AfterCallback>>{
            std::forward<BeforeCallback>(before), std::forward<TrialCallback>(trial), std::forward<AfterCallback>(after)
        };
    }

    namespace detail {
        [[nodiscard]] inline double NextBacktrackingStepLength(
            const double rejectedStepLength,
            const double acceptedMinimumJacobianDeterminant,
            const std::optional<double> rejectedMinimumJacobianDeterminant,
            const bool rejectedByInvertedGeometry,
            const BacktrackingOptions &options
        ) noexcept {
            const double contractedStepLength = rejectedStepLength * options.contractionFactor;
            if (!rejectedByInvertedGeometry || !rejectedMinimumJacobianDeterminant.has_value() ||
                !std::isfinite(acceptedMinimumJacobianDeterminant) || acceptedMinimumJacobianDeterminant <= 0.0 ||
                !std::isfinite(*rejectedMinimumJacobianDeterminant) || *rejectedMinimumJacobianDeterminant > 0.0) {
                return contractedStepLength;
            }

            const double determinantChange = acceptedMinimumJacobianDeterminant - *rejectedMinimumJacobianDeterminant;
            if (!std::isfinite(determinantChange) || determinantChange <= 0.0) {
                return contractedStepLength;
            }

            const double estimatedBoundaryStep =
                rejectedStepLength * acceptedMinimumJacobianDeterminant / determinantChange;
            const double safeguardedStep = options.fractionToBoundarySafety * estimatedBoundaryStep;
            if (!std::isfinite(safeguardedStep) || safeguardedStep <= 0.0 || safeguardedStep >= rejectedStepLength) {
                return contractedStepLength;
            }

            /*
             * Keep the configured backtracking ladder intact.  The geometry
             * certificate is used only to skip rungs that its local boundary
             * estimate says are unsafe; it does not introduce a new trial
             * length between two rungs.  This preserves the candidates that
             * ordinary backtracking would eventually test while avoiding the
             * expensive preparation of the skipped, inverted geometries.
             */
            if (contractedStepLength <= safeguardedStep) {
                return contractedStepLength;
            }

            const double rung =
                std::ceil(std::log(safeguardedStep / rejectedStepLength) / std::log(options.contractionFactor));
            double skippedStep = rejectedStepLength * std::pow(options.contractionFactor, rung);
            if (!std::isfinite(skippedStep) || skippedStep <= 0.0 || skippedStep >= rejectedStepLength) {
                return contractedStepLength;
            }
            if (skippedStep > safeguardedStep) {
                skippedStep *= options.contractionFactor;
            }
            return skippedStep;
        }

        template <typename Observer>
        inline constexpr bool isNoObserver = std::same_as<std::remove_cvref_t<Observer>, NoObserver>;

        template <typename Observer>
        concept ObservesBeforeIteration =
            !isNoObserver<Observer> &&
            requires(std::remove_reference_t<Observer> &observer, const BeforeIteration &event) {
                { observer.beforeIteration(event) } -> std::same_as<void>;
            };

        template <typename Observer>
        concept ObservesAfterIteration =
            !isNoObserver<Observer> &&
            requires(std::remove_reference_t<Observer> &observer, const AfterIteration &event) {
                { observer.afterIteration(event) } -> std::same_as<void>;
            };

        template <typename Observer>
        concept ObservesLineSearchTrial =
            !isNoObserver<Observer> &&
            requires(std::remove_reference_t<Observer> &observer, const AfterLineSearchTrial &event) {
                { observer.afterLineSearchTrial(event) } -> std::same_as<void>;
            };

        template <typename Callback>
        void InvokeObserverHookCollectively(
            const MPI_Comm communicator,
            const char *remoteFailureMessage,
            Callback &&callback
        ) {
            std::exception_ptr localFailure;
            try {
                std::invoke(std::forward<Callback>(callback));
            } catch (...) {
                localFailure = std::current_exception();
            }

            const int localFailureFlag = localFailure != nullptr ? 1 : 0;
            int globalFailureFlag      = 0;
            if (MPI_Allreduce(&localFailureFlag, &globalFailureFlag, 1, MPI_INT, MPI_MAX, communicator) !=
                MPI_SUCCESS) {
                if (localFailure != nullptr) {
                    std::rethrow_exception(localFailure);
                }
                throw std::runtime_error("The nonlinear solver could not synchronize an observer callback.");
            }

            if (globalFailureFlag != 0) {
                if (localFailure != nullptr) {
                    std::rethrow_exception(localFailure);
                }
                throw std::runtime_error(remoteFailureMessage);
            }
        }

        template <typename Observer>
        void InvokeBeforeIteration(
            Observer &observer,
            const BeforeIteration &event
        ) {
            if constexpr (ObservesBeforeIteration<Observer>) {
                if constexpr (noexcept(observer.beforeIteration(event))) {
                    observer.beforeIteration(event);
                } else {
                    InvokeObserverHookCollectively(
                        event.communicator, "An observer before-iteration callback failed on another rank.",
                        [&observer, &event] { observer.beforeIteration(event); }
                    );
                }
            }
        }

        template <typename Observer>
        void InvokeAfterIteration(
            Observer &observer,
            const AfterIteration &event
        ) {
            if constexpr (ObservesAfterIteration<Observer>) {
                if constexpr (noexcept(observer.afterIteration(event))) {
                    observer.afterIteration(event);
                } else {
                    InvokeObserverHookCollectively(
                        event.communicator, "An observer after-iteration callback failed on another rank.",
                        [&observer, &event] { observer.afterIteration(event); }
                    );
                }
            }
        }

        template <typename Observer>
        void InvokeAfterLineSearchTrial(
            Observer &observer,
            const AfterLineSearchTrial &event
        ) {
            if constexpr (ObservesLineSearchTrial<Observer>) {
                if constexpr (noexcept(observer.afterLineSearchTrial(event))) {
                    observer.afterLineSearchTrial(event);
                } else {
                    InvokeObserverHookCollectively(
                        event.communicator, "An observer line-search callback failed on another rank.",
                        [&observer, &event] { observer.afterLineSearchTrial(event); }
                    );
                }
            }
        }
    } // namespace detail

    /*
     * Observers run synchronously on every solve rank. Ordinary callback
     * exceptions are synchronized before the solver proceeds, so all ranks can
     * unwind together; explicitly noexcept callbacks bypass that synchronization.
     * A callback must still not enter an MPI collective on only a subset of
     * ranks. A before/after pair is guaranteed for iterations that finish by
     * returning an evaluation report. Infrastructure exceptions unwind
     * immediately and do not promise an after callback.
     */
    template <typename Candidate>
    concept NewtonObserver = detail::isNoObserver<Candidate> || detail::ObservesBeforeIteration<Candidate> ||
                             detail::ObservesLineSearchTrial<Candidate> || detail::ObservesAfterIteration<Candidate>;
} // namespace mean_field::solver::nonlinear