// Copyright (c) 2023, Viktor Larsson
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//
//     * Neither the name of the copyright holder nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
// ARE DISCLAIMED. IN NO EVENT SHALL COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#include "test.h"
#include "test_rng.h"

#include <Eigen/Dense>
#include <PoseLib/misc/camera_models.h>
#include <PoseLib/robust/optim/jacobian_accumulator.h>
#include <PoseLib/robust/robust_loss.h>
#include <PoseLib/types.h>
#include <algorithm>

#ifndef POSELIB_OPTIM_TEST_UTILS_H_
#define POSELIB_OPTIM_TEST_UTILS_H_

using namespace poselib;

// Accumulator used in Jacobian tests. Stores every (residual, Jacobian, weight)
// triple so that finite-difference checks can compare them against the
// analytically computed Jacobians.
class TestAccumulator {
  public:
    TestAccumulator() {}

    void reset_residual() { residual_acc = 0; }
    void add_residual(const double res, const double w = 1.0) { residual_acc += w * res * res; }
    template <typename Derived> void add_residual(const Eigen::MatrixBase<Derived> &res, const double w = 1.0) {
        residual_acc += w * res.squaredNorm();
    }
    double get_residual() const { return residual_acc; }

    void reset_jacobian() {
        Js.clear();
        rs.clear();
        weights.clear();
    }
    template <typename Derived>
    void add_jacobian(const double res, const Eigen::MatrixBase<Derived> &jac, const double w = 1.0) {
        Js.push_back(jac);
        rs.push_back(Eigen::Matrix<double, 1, 1>(res));
        weights.push_back(w);
    }
    template <typename Derived1, typename Derived2>
    void add_jacobian(const Eigen::MatrixBase<Derived1> &res, const Eigen::MatrixBase<Derived2> &jac,
                      const double w = 1.0) {
        Js.push_back(jac);
        rs.push_back(res);
        weights.push_back(w);
    }

    double residual_acc;
    std::vector<Eigen::MatrixXd> Js;
    std::vector<Eigen::VectorXd> rs;
    std::vector<double> weights;

    // Debug stuff
    Eigen::MatrixXd JtJ;
    Eigen::VectorXd Jtr;
    std::shared_ptr<RobustLoss> loss_fcn;
};

// Verify the analytical Jacobian against a central finite-difference estimate.
// Returns the maximum relative error over all residual blocks; the caller
// typically asserts this is below ~1e-3. delta is the finite-difference step.
template <typename Refiner, typename Model> double verify_jacobian(Refiner &refiner, const Model &m, double delta) {
    int num_params = refiner.num_params;
    TestAccumulator acc;
    std::vector<TestAccumulator> acc_backward(num_params);
    std::vector<TestAccumulator> acc_forward(num_params);

    // Compute residual and jacobian
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, m);

    int num_res = acc.rs.size();

    // Compute finite differences
    std::vector<Eigen::MatrixXd> J_est;
    J_est.resize(num_res);
    for (int k = 0; k < num_res; ++k) {
        J_est[k].resize(acc.Js[k].rows(), acc.Js[k].cols());
    }

    for (int i = 0; i < num_params; ++i) {
        Eigen::Matrix<double, Eigen::Dynamic, 1> dp(num_params, 1);
        dp.setZero();
        dp(i) = delta;

        Model m_forward = refiner.step(dp, m);
        Model m_backward = refiner.step(-dp, m);

        acc_backward[i].reset_jacobian();
        acc_forward[i].reset_jacobian();
        refiner.compute_jacobian(acc_forward[i], m_forward);
        refiner.compute_jacobian(acc_backward[i], m_backward);

        for (int k = 0; k < num_res; ++k) {
            J_est[k].col(i) = (acc_forward[i].rs[k] - acc_backward[i].rs[k]) / (2.0 * delta);
        }
    }

    // std::cout << "J_est=\n" << J_est[0] << "\nJ=\n" << acc.Js[0] << "\n";
    // std::cout << "------\n";

    // Compare the jacobian with the residual
    double max_err = 0.0;
    for (int k = 0; k < num_res; ++k) {
        double err = (J_est[k] - acc.Js[k]).norm() / std::max(1.0, J_est[k].norm());
        max_err = std::max(err, max_err);
        if (max_err > 1e-3) {
            std::cout << "Jacobian failure! \n J=\n" << acc.Js[k] << "\nJ (finite) = \n" << J_est[k] << "\n";
            break;
        }
    }
    return max_err;
}

// Returns true if project(x) ≈ xp and unproject(xp) ≈ x/|x|.
inline bool check_valid_camera_projection(const Camera &camera, const Eigen::Vector2d &xp, const Eigen::Vector3d &x) {
    Eigen::Vector3d x0;
    Eigen::Vector2d xp0;

    camera.unproject(xp, &x0);
    camera.project(x, &xp0);

    return ((xp - xp0).norm() / (1e-6 + xp.norm()) < 1e-6) && ((x.normalized() - x0.normalized()).norm() < 1e-6);
}

// Divide all points and the camera's intrinsics by max_dim so that the
// optimization operates in roughly unit-scale coordinates. Returns the scale
// factor applied (useful if the caller needs to undo the scaling later).
inline double normalize_camera_points(std::vector<Point2D> &points, Camera *camera) {
    const double scale = std::max(1.0, camera->max_dim());
    for (Point2D &point : points) {
        point /= scale;
    }
    camera->rescale(1.0 / scale);
    return scale;
}

// Overload for multi-camera rigs: normalizes each camera independently.
inline void normalize_camera_points(std::vector<std::vector<Point2D>> &points, std::vector<Camera> *cameras) {
    for (size_t camera_idx = 0; camera_idx < cameras->size(); ++camera_idx) {
        normalize_camera_points(points[camera_idx], &(*cameras)[camera_idx]);
    }
}

// Logging helpers. Output is captured by ScopedStreamCapture during test runs
// and only printed to the terminal when a test fails, so feel free to log
// verbosely — passing tests stay silent.
inline std::string bundle_stats_message(const BundleStats &stats, const std::string &label) {
    std::ostringstream ss;
    ss << label << ", initial_cost=" << stats.initial_cost << ", cost=" << stats.cost
       << ", grad_norm=" << stats.grad_norm << ", step_norm=" << stats.step_norm << ", iterations=" << stats.iterations
       << ", invalid_steps=" << stats.invalid_steps << ", lambda=" << stats.lambda;
    return ss.str();
}

inline void log_test_message(const std::string &message) { std::cout << message << "\n"; }

inline void log_test_case(const std::string &label, const std::string &value) {
    std::cout << label << "=" << value << "\n";
}

inline void log_bundle_stats(const BundleStats &stats, const std::string &label) {
    std::cout << "bundle_stats: " << bundle_stats_message(stats, label) << "\n";
}

// Check that the optimizer reduced the cost and converged (small gradient).
// Prints a structured failure message to stdout (captured on failure) and
// returns false so callers can use: REQUIRE(check_bundle_cost_and_gradient(...))
inline bool check_bundle_cost_and_gradient(const BundleStats &stats, double grad_threshold, const std::string &label) {
    if (!(stats.cost < stats.initial_cost)) {
        std::cout << "Failure: refinement did not reduce cost. (" << bundle_stats_message(stats, label) << ")\n";
        return false;
    }
    if (std::isnan(stats.grad_norm) || std::abs(stats.grad_norm) > grad_threshold) {
        std::cout << "Failure: gradient norm too large. (" << bundle_stats_message(stats, label)
                  << ", grad_threshold=" << grad_threshold << ")\n";
        return false;
    }
    return true;
}

// Extended convergence check that also asserts the final step size is small.
// Useful when a near-zero gradient is necessary but not sufficient (e.g. when
// the optimizer may have stalled at a large step due to a poor line-search).
inline bool check_bundle_cost_gradient_and_step(const BundleStats &stats, double grad_threshold, double step_threshold,
                                                const std::string &label) {
    if (!check_bundle_cost_and_gradient(stats, grad_threshold, label)) {
        return false;
    }
    if (std::isnan(stats.step_norm) || std::abs(stats.step_norm) > step_threshold) {
        std::cout << "Failure: step norm too large. (" << bundle_stats_message(stats, label)
                  << ", step_threshold=" << step_threshold << ")\n";
        return false;
    }
    return true;
}

// Sample image points from a jittered 5-column grid in the stable interior of the image.
inline Eigen::Vector2d image_sample(const Camera &cam, size_t idx, size_t count, test_rng::Rng &rng) {
    const size_t cols = 5;
    const size_t rows = std::max<size_t>(1, (count + cols - 1) / cols);
    const size_t row = idx / cols;
    const size_t col = idx % cols;
    const Eigen::Vector2d jitter = test_rng::symmetric_vec2(rng, 0.12);
    const double u = std::clamp((static_cast<double>(col) + 0.5 + jitter(0)) / static_cast<double>(cols), 0.2, 0.8);
    const double v = std::clamp((static_cast<double>(row) + 0.5 + jitter(1)) / static_cast<double>(rows), 0.2, 0.8);
    return Eigen::Vector2d(u * cam.width, v * cam.height);
}

#endif
