// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <algorithm>
#include <cmath>
#include <deque>
#include <functional>
#include <vector>

// Limited-memory BFGS whose memory outlives one call, so that a refinement's macro-cycles - each a few dozen
// iterations on a target whose scale and likelihood parameters were refitted in between - go on from what the
// earlier ones learnt of the curvature instead of starting each from a steepest-descent step. The initial inverse
// Hessian is a diagonal the caller gives (SetDiagonal: an estimate of the Hessian's diagonal, so that a parameter
// with ten times the curvature of another takes a tenth of its step), scaled each iteration by the latest pair
// (Oren-Shanno in that metric). The line search backtracks from the full step to sufficient decrease.
// Following Nocedal (1980) Math. Comp. 35, 773-782 and Liu & Nocedal (1989) Math. Program. 45, 503-528
class RefinementMinimiser {
public:
    // f(x, g): the target at x, and its gradient in g (sized as x); false where it cannot be evaluated there.
    using Target = std::function<bool(const std::vector<double> &x, double &f, std::vector<double> &g)>;

    explicit RefinementMinimiser(size_t memory = 20) : memory_(memory) {}

    // Forget the pairs (the target changed beyond what the memory still describes, e.g. another weight).
    void Reset() {
        s_.clear();
        y_.clear();
        sy_.clear();
    }
    // The estimated diagonal of the Hessian; every element > 0.
    void SetDiagonal(std::vector<double> h) { h_ = std::move(h); }

    struct Outcome {
        int iterations = 0;
        int evaluations = 0;
        double f_start = 0, f_end = 0;
    };

    // At most `iterations` iterations from x (x is where it ended); stops early where an iteration lowers f by
    // less than `f_tolerance` relative to |f|, or `stop` says so (checked before every iteration). `max_step`, in
    // the parameters' units, is the most any parameter moves in the first trial of a line search.
    Outcome Minimise(const Target &target, std::vector<double> &x, int iterations, double f_tolerance, double max_step,
                     const std::function<bool()> &stop = nullptr) {
        const size_t n = x.size();
        if (h_.size() != n)
            h_.assign(n, 1.0);
        Outcome out;
        double f = 0;
        std::vector<double> g(n), d(n), x_new(n), g_new(n);
        if (!target(x, f, g))
            return out;
        out.evaluations++;
        out.f_start = out.f_end = f;
        for (int it = 0; it < iterations; it++) {
            if (stop && stop())
                break;
            Direction(g, d);
            double gd = Dot(g, d);
            if (!(gd < 0)) {   // not a descent direction: the memory no longer fits, start again from the diagonal
                Reset();
                Direction(g, d);
                gd = Dot(g, d);
                if (!(gd < 0))
                    break;
            }
            double d_max = 0;
            for (double v : d)
                d_max = std::max(d_max, std::fabs(v));
            double alpha = d_max > max_step ? max_step / d_max : 1.0;
            double f_new = 0;
            bool accepted = false;
            for (int ls = 0; ls < 10; ls++) {
                for (size_t k = 0; k < n; k++)
                    x_new[k] = x[k] + alpha * d[k];
                const bool ok = target(x_new, f_new, g_new);
                out.evaluations++;
                if (ok && f_new <= f + 1e-4 * alpha * gd) {
                    accepted = true;
                    break;
                }
                // the minimum of the quadratic through f, its slope and f_new, kept within [0.1, 0.5] of the step
                const double q = ok ? -gd * alpha * alpha / (2.0 * (f_new - f - gd * alpha)) : 0.1 * alpha;
                alpha = std::clamp(q, 0.1 * alpha, 0.5 * alpha);
            }
            if (!accepted)
                break;
            std::vector<double> s(n), y(n);
            for (size_t k = 0; k < n; k++) {
                s[k] = x_new[k] - x[k];
                y[k] = g_new[k] - g[k];
            }
            const double sy = Dot(s, y);
            if (sy > 1e-10 * std::sqrt(Dot(s, s) * Dot(y, y))) {
                s_.push_back(std::move(s));
                y_.push_back(std::move(y));
                sy_.push_back(sy);
                if (s_.size() > memory_) {
                    s_.pop_front();
                    y_.pop_front();
                    sy_.pop_front();
                }
            }
            const double decrease = f - f_new;
            x.swap(x_new);
            g.swap(g_new);
            f = f_new;
            out.iterations++;
            out.f_end = f;
            if (decrease <= f_tolerance * std::max(std::fabs(f), 1.0))
                break;
        }
        return out;
    }

private:
    static double Dot(const std::vector<double> &a, const std::vector<double> &b) {
        double s = 0;
        for (size_t k = 0; k < a.size(); k++)
            s += a[k] * b[k];
        return s;
    }

    // d = -H g by the two-loop recursion, H0 = gamma diag(1/h).
    void Direction(const std::vector<double> &g, std::vector<double> &d) const {
        const size_t n = g.size(), m = s_.size();
        std::vector<double> q = g, a(m), rho(m);
        for (size_t i = m; i-- > 0;) {
            rho[i] = 1.0 / sy_[i];
            a[i] = rho[i] * Dot(s_[i], q);
            for (size_t k = 0; k < n; k++)
                q[k] -= a[i] * y_[i][k];
        }
        double gamma = 1.0;
        if (m > 0) {
            double yhy = 0;
            for (size_t k = 0; k < n; k++)
                yhy += y_[m - 1][k] * y_[m - 1][k] / h_[k];
            gamma = sy_[m - 1] / yhy;
        }
        for (size_t k = 0; k < n; k++)
            q[k] *= gamma / h_[k];
        for (size_t i = 0; i < m; i++) {
            const double b = rho[i] * Dot(y_[i], q);
            for (size_t k = 0; k < n; k++)
                q[k] += s_[i][k] * (a[i] - b);
        }
        for (size_t k = 0; k < n; k++)
            d[k] = -q[k];
    }

    size_t memory_;
    std::deque<std::vector<double>> s_, y_;
    std::deque<double> sy_;   // s.y of each pair, as taken when it was kept
    std::vector<double> h_;
};