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

// Adapted from https://github.com/ceres-solver/ceres-solver (internal/ceres/polynomial.cc,
// include/ceres/internal/sphere_manifold_functions.h, householder_vector.h)
// Copyright 2023 Google Inc. All rights reserved.
// BSD-3-Clause, see licenses/ceres-solver.txt

#include "LMSolver.h"

#include <Eigen/Eigenvalues>

namespace {
    void HouseholderVector3(const double x[3], double v[3], double &beta) {
        const double sigma = x[0] * x[0] + x[1] * x[1];
        v[0] = x[0];
        v[1] = x[1];
        v[2] = 1.0;
        beta = 0.0;
        const double x_pivot = x[2];
        if (sigma <= std::numeric_limits<double>::epsilon()) {
            if (x_pivot < 0.0)
                beta = 2.0;
            return;
        }
        const double mu = std::sqrt(x_pivot * x_pivot + sigma);
        const double v_pivot = (x_pivot <= 0.0) ? x_pivot - mu : -sigma / (x_pivot + mu);
        beta = 2.0 * v_pivot * v_pivot / (sigma + v_pivot * v_pivot);
        v[0] /= v_pivot;
        v[1] /= v_pivot;
    }

    double Norm3(const double x[3]) {
        return std::sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]);
    }

    using Vector = Eigen::VectorXd;
    using Matrix = Eigen::MatrixXd;

    double EvaluatePolynomial(const Vector &polynomial, double x) {
        double v = 0.0;
        for (int i = 0; i < polynomial.size(); ++i)
            v = v * x + polynomial(i);
        return v;
    }

    void BalanceCompanionMatrix(Matrix &companion_matrix) {
        Matrix offdiagonal = companion_matrix;
        offdiagonal.diagonal().setZero();
        const int degree = static_cast<int>(companion_matrix.rows());
        const double gamma = 0.9;
        bool scaling_has_changed;
        do {
            scaling_has_changed = false;
            for (int i = 0; i < degree; ++i) {
                const double col_norm = offdiagonal.col(i).lpNorm<1>();
                if (std::fpclassify(col_norm) != FP_ZERO) {
                    const double row_norm = offdiagonal.row(i).lpNorm<1>();
                    int exponent = 0;
                    std::frexp(row_norm / col_norm, &exponent);
                    exponent /= 2;
                    if (exponent != 0) {
                        const double scaled_col_norm = std::ldexp(col_norm, exponent);
                        const double scaled_row_norm = std::ldexp(row_norm, -exponent);
                        if (scaled_col_norm + scaled_row_norm < gamma * (col_norm + row_norm)) {
                            scaling_has_changed = true;
                            offdiagonal.row(i) *= std::ldexp(1.0, -exponent);
                            offdiagonal.col(i) *= std::ldexp(1.0, exponent);
                        }
                    }
                }
            }
        } while (scaling_has_changed);
        offdiagonal.diagonal() = companion_matrix.diagonal();
        companion_matrix = offdiagonal;
    }

    // Real parts of the roots, as Ceres' FindPolynomialRoots (the imaginary parts are not used here).
    bool FindPolynomialRoots(const Vector &polynomial_in, Vector &real) {
        if (polynomial_in.size() == 0)
            return false;
        int lead = 0;
        while (lead < polynomial_in.size() - 1 && polynomial_in(lead) == 0.0)
            ++lead;
        Vector polynomial = polynomial_in.tail(polynomial_in.size() - lead);
        const int degree = static_cast<int>(polynomial.size()) - 1;
        if (degree == 0) {
            real.resize(0);
            return true;
        }
        if (degree == 1) {
            real.resize(1);
            real(0) = -polynomial(1) / polynomial(0);
            return true;
        }
        if (degree == 2) {
            const double a = polynomial(0);
            const double b = polynomial(1);
            const double c = polynomial(2);
            const double D = b * b - 4 * a * c;
            const double sqrt_D = std::sqrt(std::fabs(D));
            real.setZero(2);
            if (D >= 0) {
                if (b >= 0) {
                    real(0) = (-b - sqrt_D) / (2.0 * a);
                    real(1) = (2.0 * c) / (-b - sqrt_D);
                } else {
                    real(0) = (2.0 * c) / (-b + sqrt_D);
                    real(1) = (-b + sqrt_D) / (2.0 * a);
                }
            } else {
                real(0) = -b / (2.0 * a);
                real(1) = -b / (2.0 * a);
            }
            return true;
        }
        polynomial /= polynomial(0);
        Matrix companion = Matrix::Zero(degree, degree);
        companion.diagonal(-1).setOnes();
        companion.col(degree - 1) = -polynomial.reverse().head(degree);
        BalanceCompanionMatrix(companion);
        Eigen::EigenSolver<Matrix> solver(companion, false);
        if (solver.info() != Eigen::Success)
            return false;
        real = solver.eigenvalues().real();
        return true;
    }

    void MinimizePolynomial(const Vector &polynomial, double x_min, double x_max,
                            double &optimal_x, double &optimal_value) {
        optimal_x = (x_min + x_max) / 2.0;
        optimal_value = EvaluatePolynomial(polynomial, optimal_x);
        const double x_min_value = EvaluatePolynomial(polynomial, x_min);
        if (x_min_value < optimal_value) {
            optimal_value = x_min_value;
            optimal_x = x_min;
        }
        const double x_max_value = EvaluatePolynomial(polynomial, x_max);
        if (x_max_value < optimal_value) {
            optimal_value = x_max_value;
            optimal_x = x_max;
        }
        if (polynomial.rows() <= 2)
            return;

        const int degree = static_cast<int>(polynomial.rows()) - 1;
        Vector derivative(degree);
        for (int i = 0; i < degree; ++i)
            derivative(i) = (degree - i) * polynomial(i);
        Vector roots_real;
        if (!FindPolynomialRoots(derivative, roots_real))
            return;
        for (int i = 0; i < roots_real.rows(); ++i) {
            const double root = roots_real(i);
            if (root < x_min || root > x_max)
                continue;
            const double value = EvaluatePolynomial(polynomial, root);
            if (value < optimal_value) {
                optimal_value = value;
                optimal_x = root;
            }
        }
    }

    Vector FindInterpolatingPolynomial(const std::vector<LMLineSample> &samples) {
        int num_constraints = 0;
        for (const auto &s: samples)
            num_constraints += (s.value_is_valid ? 1 : 0) + (s.gradient_is_valid ? 1 : 0);
        const int degree = num_constraints - 1;
        Matrix lhs = Matrix::Zero(num_constraints, num_constraints);
        Vector rhs = Vector::Zero(num_constraints);
        int row = 0;
        for (const auto &s: samples) {
            if (s.value_is_valid) {
                for (int j = 0; j <= degree; ++j)
                    lhs(row, j) = std::pow(s.x, degree - j);
                rhs(row) = s.value;
                ++row;
            }
            if (s.gradient_is_valid) {
                for (int j = 0; j < degree; ++j)
                    lhs(row, j) = (degree - j) * std::pow(s.x, degree - j - 1);
                rhs(row) = s.gradient;
                ++row;
            }
        }
        Eigen::FullPivLU<Matrix> lu(lhs);
        return lu.setThreshold(0.0).solve(rhs);
    }
}

void SpherePlus3(const double x[3], const double delta[2], double out[3]) {
    const double norm_delta = std::sqrt(delta[0] * delta[0] + delta[1] * delta[1]);
    if (norm_delta == 0.0) {
        out[0] = x[0];
        out[1] = x[1];
        out[2] = x[2];
        return;
    }
    double v[3], beta;
    HouseholderVector3(x, v, beta);
    const double sin_delta_by_delta = std::sin(norm_delta) / norm_delta;
    const double y[3] = {sin_delta_by_delta * delta[0], sin_delta_by_delta * delta[1], std::cos(norm_delta)};
    const double vy = v[0] * y[0] + v[1] * y[1] + v[2] * y[2];
    const double x_norm = Norm3(x);
    for (int i = 0; i < 3; i++)
        out[i] = x_norm * (y[i] - v[i] * (beta * vy));
}

void SpherePlusJacobian3(const double x[3], double jacobian[3][2]) {
    double v[3], beta;
    HouseholderVector3(x, v, beta);
    const double x_norm = Norm3(x);
    for (int i = 0; i < 2; ++i)
        for (int r = 0; r < 3; ++r)
            jacobian[r][i] = (-beta * v[i] * v[r] + (r == i ? 1.0 : 0.0)) * x_norm;
}

double LMInterpolatedStepSize(const LMLineSample &lowerbound, const LMLineSample &previous,
                              const LMLineSample &current, double min_step, double max_step) {
    if (!current.value_is_valid)
        return std::min(std::max(current.x * 0.5, min_step), max_step);

    std::vector<LMLineSample> samples{lowerbound, current};
    if (previous.value_is_valid)
        samples.push_back(previous);

    const Vector polynomial = FindInterpolatingPolynomial(samples);
    double step = 0.0, value = 0.0;
    MinimizePolynomial(polynomial, min_step, max_step, step, value);
    for (const auto &s: samples) {
        if (s.x < min_step || s.x > max_step)
            continue;
        const double v = EvaluatePolynomial(polynomial, s.x);
        if (v < value) {
            step = s.x;
            value = v;
        }
    }
    return step;
}