// 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>

#include "../../common/ParallelFor.h"

// 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)>;

    // `nthreads`: the threads the sweeps over the parameters run on (the dot products stay serial chains).
    explicit RefinementMinimiser(size_t memory = 20, size_t nthreads = 1) : memory_(memory), nthreads_(nthreads) {}

    // 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++) {
                Sweep(n, [&](size_t 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;
            // the pair's vectors are the ones the memory last let go, where it did: made anew, their pages fault in
            std::vector<double> s = std::move(spare_s_), y = std::move(spare_y_);
            s.resize(n);
            y.resize(n);
            Sweep(n, [&](size_t 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_) {
                    spare_s_ = std::move(s_.front());
                    spare_y_ = std::move(y_.front());
                    s_.pop_front();
                    y_.pop_front();
                    sy_.pop_front();
                }
            } else {
                spare_s_ = std::move(s);
                spare_y_ = std::move(y);
            }
            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;
    }

    // fn(k) for every k below n on the threads, each element on its own: the serial loop's bits.
    template <typename Fn>
    void Sweep(size_t n, Fn fn) const {
        ParallelChunks(static_cast<int>(n), ThreadsForWork(n, nthreads_), [&](int lo, int hi) {
            for (int k = lo; k < hi; k++)
                fn(static_cast<size_t>(k));
        });
    }

    // d = -H g by the two-loop recursion, H0 = gamma diag(1/h). Every dot product is a serial chain of n
    // additions, so each sweep over the parameters takes the next pair's dot product while it applies the
    // previous pair's update to q - the same operations on each element in the same order as a sweep for
    // each would do, with the update's memory traffic hidden under the chain.
    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 = q_;   // kept from call to call: made anew, its pages fault in every time
        q.assign(g.begin(), g.end());
        std::vector<double> a(m), rho(m);
        for (size_t i = 0; i < m; i++)
            rho[i] = 1.0 / sy_[i];
        double dot = m > 0 ? Dot(s_[m - 1], q) : 0.0, yhy = 0.0;
        for (size_t i = m; i-- > 0;) {
            a[i] = rho[i] * dot;
            const double ai = a[i];
            const double *y = y_[i].data();
            dot = 0.0;
            if (i > 0) {
                const double *s = s_[i - 1].data();
                for (size_t k = 0; k < n; k++) {
                    q[k] -= ai * y[k];
                    dot += s[k] * q[k];
                }
            } else {
                const double *yl = y_[m - 1].data();
                for (size_t k = 0; k < n; k++) {
                    q[k] -= ai * y[k];
                    yhy += yl[k] * yl[k] / h_[k];
                }
            }
        }
        const double gamma = m > 0 ? sy_[m - 1] / yhy : 1.0;
        if (m > 0) {
            const double *y = y_[0].data();
            for (size_t k = 0; k < n; k++) {
                q[k] *= gamma / h_[k];
                dot += y[k] * q[k];
            }
        } else
            for (size_t k = 0; k < n; k++)
                q[k] *= gamma / h_[k];
        for (size_t i = 0; i < m; i++) {
            // b on its own and (a[i] - b) in the loop, as before: written as a[i] - rho[i] * dot the compiler
            // contracts it into a fused multiply-add, and the bits change
            const double b = rho[i] * dot;
            const double *s = s_[i].data();
            dot = 0.0;
            if (i + 1 < m) {
                const double *y = y_[i + 1].data();
                for (size_t k = 0; k < n; k++) {
                    q[k] += s[k] * (a[i] - b);
                    dot += y[k] * q[k];
                }
            } else
                for (size_t k = 0; k < n; k++)
                    q[k] += s[k] * (a[i] - b);
        }
        for (size_t k = 0; k < n; k++)
            d[k] = -q[k];
    }

    size_t memory_, nthreads_;
    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> spare_s_, spare_y_;   // the pair the memory last let go, for the next one to reuse
    mutable std::vector<double> q_;           // Direction()'s working vector
    std::vector<double> h_;
};