MeanField_SERiF/src/include/serif/solver/nonlinear/observers.hpp

#pragma once

#include <chrono>
#include <functional>
#include <type_traits>

#include "serif/solver/nonlinear/result.hpp"

namespace serif::solver {
    /** @brief Optional elapsed wall time, kept separately from numerical restart history. */
    struct SolveTiming {
        double seconds{0.0};
    };

    /**
     * @brief Measure solves explicitly while retaining an ordinary progress or stop callback.
     * @note Reusing the same timing record accumulates resumed solves. Timing starts before initial
     * evaluation and ends on any returned numerical result, including a rejected initial state.
     * Exceptions propagate normally; an interrupted call contributes no completed-solve duration.
     */
    template <typename Observer = NoObserver>
    class TimedObserver {
    public:
        /** @brief Borrow the timing destination and own the user's callback. */
        explicit TimedObserver(SolveTiming& timing, Observer observer = {}) :
            m_timing(timing), m_observer(std::move(observer)) {}

        /** @brief Begin measurement when the solver starts evaluating its initial state. */
        void started() {
            m_start = std::chrono::steady_clock::now();
            if constexpr (requires { m_observer.started(); }) {
                m_observer.started();
            }
        }

        /** @brief Accumulate a normally completed call independently of its convergence status. */
        void finished(SolveStatus status, const SolveDiagnostics& diagnostics) {
            if constexpr (requires { m_observer.finished(status, diagnostics); }) {
                m_observer.finished(status, diagnostics);
            }
            m_timing.seconds += std::chrono::duration<double>(std::chrono::steady_clock::now() - m_start).count();
        }

        /** @brief Preserve the callback's void or stop-request return type. */
        decltype(auto) operator()(SolveEvent event, const auto& state, const auto& system) {
            return std::invoke(m_observer, std::move(event), state, system);
        }

    private:
        SolveTiming& m_timing;
        Observer m_observer;
        std::chrono::steady_clock::time_point m_start;
    };

    namespace detail {
        /** @brief Adapt callbacks and latch requests; only the Newton driver decides when stopping is safe. */
        template <typename Observer, typename Strategies>
        class SolveObservation {
        public:
            /** @brief Borrow the single diagnostics record and the selected numerical execution. */
            SolveObservation(Observer& observer, Strategies& strategies, SolveDiagnostics& diagnostics) :
                m_observer(observer), m_strategies(strategies), m_diagnostics(diagnostics) {
                if constexpr (requires { observer.started(); }) {
                    observer.started();
                }
            }

            /** @brief Publish an immutable snapshot and remember a stop request until an accepted boundary. */
            void notify(EventKind kind, const auto& state, const auto& system, int iteration = 0,
                int trial = 0, double step = 0.0, std::optional<MetricEvaluation> metric = {}, std::string message = {}) {
                m_strategies.contribute_diagnostics(m_diagnostics);
                SolveEvent event{.kind = kind, .iteration = iteration, .trial = trial, .step = step,
                    .metric = metric, .message = std::move(message), .diagnostics = m_diagnostics};
                if constexpr (std::same_as<std::invoke_result_t<Observer&, SolveEvent, decltype(state), decltype(system)>, SolveControl>) {
                    m_stop_requested |= std::invoke(m_observer, std::move(event), state, system) == SolveControl::stop;
                } else {
                    std::invoke(m_observer, std::move(event), state, system);
                }
            }

            /** @brief Check the latched request without letting callbacks accept or reject a trial. */
            [[nodiscard]] bool stop_requested() const noexcept {
                return m_stop_requested;
            }

            /** @brief Finalize reporting even when the initial state was rejected and no snapshot exists. */
            [[nodiscard]] auto finish(auto& result, SolveStatus status, std::string message, MetricEvaluation metric = {}) {
                m_strategies.contribute_diagnostics(m_diagnostics);
                result.finish(status, std::move(message), m_diagnostics);
                if (result.has_state()) {
                    notify(EventKind::finished, result.state(), result.system(), m_diagnostics.attempted_iterations,
                        0, 0.0, metric, result.message());
                }
                if constexpr (requires { m_observer.finished(status, m_diagnostics); }) {
                    m_observer.finished(status, m_diagnostics);
                }
                return std::move(result);
            }

        private:
            Observer& m_observer;
            Strategies& m_strategies;
            SolveDiagnostics& m_diagnostics;
            bool m_stop_requested{false};
        };
    }
}