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

#include "AndersonMixing.h"

std::vector<double> AndersonMixing::Next(const std::vector<double> &x, const std::vector<double> &fx,
                                         const std::vector<double> &w) {
    // Anderson acceleration, type II, after Walker & Ni (2011) SIAM J. Numer. Anal. 49, 1715-1735
    // (Anderson (1965) J. ACM 12, 547-560): with dX and dR the differences of the states and of the
    // residuals of consecutive rounds, gamma = argmin |r - dR gamma|_w and the next state is
    // F(x) - (dX + dR) gamma.
    const int n = static_cast<int>(x.size());
    const Eigen::Map<const Eigen::VectorXd> xv(x.data(), n), fxv(fx.data(), n), wv(w.data(), n);
    const Eigen::VectorXd r = fxv - xv;
    std::vector<uint8_t> support(n);
    for (int i = 0; i < n; ++i) support[i] = w[i] > 0.0 ? 1 : 0;
    const double r_norm = (wv.array() * r.array().square()).sum();

    if (stopped || (x_last.size() == n && r_norm > r_norm_last)) {
        stopped = true;
        return fx;
    }
    if (x_last.size() == n && support == support_last) {
        dx.push_back(xv - x_last);
        dr.push_back(r - r_last);
        if (static_cast<int>(dx.size()) > depth) {
            dx.pop_front();
            dr.pop_front();
        }
    } else if (!dx.empty()) {
        dx.clear();
        dr.clear();
        ++restarts;
    }
    x_last = xv;
    r_last = r;
    support_last = std::move(support);
    r_norm_last = r_norm;

    std::vector<double> out = fx;
    const int m = static_cast<int>(dr.size());
    if (m == 0) return out;
    const Eigen::VectorXd sqrt_w = wv.cwiseMax(0.0).cwiseSqrt();
    Eigen::MatrixXd M(n, m);
    for (int j = 0; j < m; ++j) M.col(j) = sqrt_w.cwiseProduct(dr[j]);
    const Eigen::VectorXd gamma = M.colPivHouseholderQr().solve(sqrt_w.cwiseProduct(r));
    for (int j = 0; j < m; ++j)
        for (int i = 0; i < n; ++i)
            if (w[i] > 0.0) out[i] -= gamma[j] * (dx[j][i] + dr[j][i]);
    return out;
}