#include "optim_test_utils.h"
#include "test.h"

#include <PoseLib/misc/camera_models.h>
#include <PoseLib/robust/optim/jacobian_accumulator.h>
#include <PoseLib/robust/optim/lm_impl.h>
#include <PoseLib/robust/optim/relative.h>
#include <PoseLib/robust/robust_loss.h>

using namespace poselib;

//////////////////////////////
// Relative pose

namespace test::relative {

CameraPose random_camera(test_rng::Rng &rng) {
    Eigen::Vector3d cc = test_rng::symmetric_vec3(rng);
    cc.normalize();
    cc *= 2.0;

    // Lookat point
    Eigen::Vector3d p = test_rng::symmetric_vec3(rng, 0.1);

    Eigen::Vector3d r3 = p - cc;
    r3.normalize();

    Eigen::Vector3d r2 = r3.cross(Eigen::Vector3d::UnitX());
    r2.normalize();
    Eigen::Vector3d r1 = r2.cross(r3);

    Eigen::Matrix3d R;
    R.row(0) = r1;
    R.row(1) = r2;
    R.row(2) = r3;

    return CameraPose(R, -R * cc);
}

void setup_scene(int N, CameraPose &pose, std::vector<Point2D> &x1, std::vector<Point2D> &x2, Camera &cam1,
                 Camera &cam2, const std::string &case_name = "relative_scene", size_t case_index = 0) {

    test_rng::Rng rng = test_rng::make_rng(case_name, case_index);
    CameraPose p1 = random_camera(rng);
    CameraPose p2 = random_camera(rng);

    for (size_t i = 0; i < N; ++i) {
        Eigen::Vector3d Xi = test_rng::symmetric_vec3(rng);

        Eigen::Vector2d xi;
        cam1.project(p1.apply(Xi), &xi);
        x1.push_back(xi);

        cam2.project(p2.apply(Xi), &xi);
        x2.push_back(xi);
    }

    Eigen::Matrix3d R = p2.R() * p1.R().transpose();
    Eigen::Vector3d t = p2.t - p2.R() * p1.R().transpose() * p1.t;

    pose = CameraPose(R, t);
}

bool test_relative_pose_normal_acc() {

    const size_t N = 10;
    std::string camera_str = "0 PINHOLE 1 1 1.0 1.0 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    NormalAccumulator acc;
    PinholeRelativePoseRefiner refiner(x1, x2);
    acc.initialize(refiner.num_params);

    // Check that residual is zero
    acc.reset_residual();
    refiner.compute_residual(acc, pose);
    double residual = acc.get_residual();
    REQUIRE_SMALL(residual, 1e-6);

    // Check the gradient is zero
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, pose);
    REQUIRE_SMALL(acc.Jtr.norm(), 1e-6);

    return true;
}

bool test_relative_pose_jacobian() {
    const size_t N = 10;
    std::string camera_str = "0 PINHOLE 1 1 1.0 1.0 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    PinholeRelativePoseRefiner<UniformWeightVector, TestAccumulator> refiner(x1, x2);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), CameraPose>(refiner, pose, delta);
    REQUIRE_SMALL(jac_err, 1e-6)

    // Test that compute_residual and compute_jacobian are compatible
    TestAccumulator acc;
    acc.reset_residual();
    double r1 = refiner.compute_residual(acc, pose);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, pose);
    double r2 = 0.0;
    for (int i = 0; i < acc.rs.size(); ++i) {
        r2 += acc.weights[i] * acc.rs[i].squaredNorm();
    }
    REQUIRE_SMALL(std::abs(r1 - r2), 1e-10);

    return true;
}

bool test_relative_pose_refinement() {
    const size_t N = 10;
    std::string camera_str = "0 PINHOLE 1 1 1.0 1.0 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("relative_pose_refinement_noise");
    for (size_t i = 0; i < N; ++i) {
        x1[i] += test_rng::symmetric_vec2(noise_rng, 0.001);
        x2[i] += test_rng::symmetric_vec2(noise_rng, 0.001);
    }

    PinholeRelativePoseRefiner refiner(x1, x2);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &pose, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_relative_pose_refinement"));

    return true;
}

bool test_shared_focal_relative_pose_normal_acc() {
    const size_t N = 10;
    std::string camera_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);
    ImagePair image_pair = ImagePair(pose, camera, camera);

    NormalAccumulator acc;
    SharedFocalRelativePoseRefiner refiner(x1, x2);
    acc.initialize(refiner.num_params);

    // Check that residual is zero
    acc.reset_residual();
    refiner.compute_residual(acc, image_pair);
    double residual = acc.get_residual();
    REQUIRE_SMALL(residual, 1e-6);

    // Check the gradient is zero
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, image_pair);
    REQUIRE_SMALL(acc.Jtr.norm(), 1e-6);

    return true;
}

bool test_shared_focal_relative_pose_jacobian() {
    const size_t N = 10;
    std::string camera_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);
    ImagePair image_pair = ImagePair(pose, camera, camera);

    SharedFocalRelativePoseRefiner<UniformWeightVector, TestAccumulator> refiner(x1, x2);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), ImagePair>(refiner, image_pair, delta);
    REQUIRE_SMALL(jac_err, 1e-6)

    // Test that compute_residual and compute_jacobian are compatible
    TestAccumulator acc;
    acc.reset_residual();
    double r1 = refiner.compute_residual(acc, image_pair);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, image_pair);
    double r2 = 0.0;
    for (int i = 0; i < acc.rs.size(); ++i) {
        r2 += acc.weights[i] * acc.rs[i].squaredNorm();
    }
    REQUIRE_SMALL(std::abs(r1 - r2), 1e-10);

    return true;
}
//
bool test_shared_focal_relative_pose_refinement() {
    const size_t N = 10;
    std::string camera_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);
    ImagePair image_pair = ImagePair(pose, camera, camera);

    // Add some noise
    test_rng::Rng rng = test_rng::make_rng("shared_focal_relative_pose_refinement_noise");
    for (int i = 0; i < N; ++i) {
        x1[i] += 5e-4 * test_rng::symmetric_vec2(rng);
        x2[i] += 5e-4 * test_rng::symmetric_vec2(rng);
    }

    SharedFocalRelativePoseRefiner refiner(x1, x2);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &image_pair, bundle_opt, print_iteration);
    log_bundle_stats(stats, "test_shared_focal_relative_pose_refinement");
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_shared_focal_relative_pose_refinement"));

    return true;
}

bool test_varying_focal_relative_pose_normal_acc() {
    const size_t N = 10;
    std::string camera1_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    std::string camera2_str = "0 SIMPLE_PINHOLE 1 1 0.8 0.0 0.0";
    Camera camera1;
    camera1.initialize_from_txt(camera1_str);
    Camera camera2;
    camera2.initialize_from_txt(camera2_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera1, camera2);
    ImagePair image_pair = ImagePair(pose, camera1, camera2);

    NormalAccumulator acc;
    VaryingFocalRelativePoseRefiner refiner(x1, x2);
    acc.initialize(refiner.num_params);

    // Check that residual is zero
    acc.reset_residual();
    refiner.compute_residual(acc, image_pair);
    double residual = acc.get_residual();
    REQUIRE_SMALL(residual, 1e-6);

    // Check the gradient is zero
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, image_pair);
    REQUIRE_SMALL(acc.Jtr.norm(), 1e-6);

    return true;
}

bool test_varying_focal_relative_pose_jacobian() {
    const size_t N = 10;
    std::string camera1_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    std::string camera2_str = "0 SIMPLE_PINHOLE 1 1 0.8 0.0 0.0";
    Camera camera1;
    camera1.initialize_from_txt(camera1_str);
    Camera camera2;
    camera2.initialize_from_txt(camera2_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera1, camera2);
    ImagePair image_pair = ImagePair(pose, camera1, camera2);
    VaryingFocalRelativePoseRefiner<UniformWeightVector, TestAccumulator> refiner(x1, x2);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), ImagePair>(refiner, image_pair, delta);
    REQUIRE_SMALL(jac_err, 1e-6)

    // Test that compute_residual and compute_jacobian are compatible
    TestAccumulator acc;
    acc.reset_residual();
    double r1 = refiner.compute_residual(acc, image_pair);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, image_pair);
    double r2 = 0.0;
    for (int i = 0; i < acc.rs.size(); ++i) {
        r2 += acc.weights[i] * acc.rs[i].squaredNorm();
    }
    REQUIRE_SMALL(std::abs(r1 - r2), 1e-10);

    return true;
}
//
bool test_varying_focal_relative_pose_refinement() {
    const size_t N = 10;
    std::string camera1_str = "0 SIMPLE_PINHOLE 1 1 1.2 0.0 0.0";
    std::string camera2_str = "0 SIMPLE_PINHOLE 1 1 0.8 0.0 0.0";
    Camera camera1;
    camera1.initialize_from_txt(camera1_str);
    Camera camera2;
    camera2.initialize_from_txt(camera2_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera1, camera2);
    ImagePair image_pair = ImagePair(pose, camera1, camera2);

    // Add some noise
    test_rng::Rng rng = test_rng::make_rng("varying_focal_relative_pose_refinement_noise");
    for (int i = 0; i < N; ++i) {
        x1[i] += 5e-4 * test_rng::symmetric_vec2(rng);
        x2[i] += 5e-4 * test_rng::symmetric_vec2(rng);
    }

    VaryingFocalRelativePoseRefiner refiner(x1, x2);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &image_pair, bundle_opt, print_iteration);
    log_bundle_stats(stats, "test_varying_focal_relative_pose_refinement");
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_varying_focal_relative_pose_refinement"));

    return true;
}

bool test_tangent_sampson_fix_camera_relative_pose_jacobian() {
    const size_t N = 10;
    std::string camera_str = "3 RADIAL 1936 1296 2425.85 932.38 629.325 -0.04012 0.00123";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    std::vector<Eigen::Vector3d> d1, d2;
    std::vector<Eigen::Matrix<double, 3, 2>> J1inv, J2inv;
    camera.unproject_with_jac(x1, &d1, &J1inv);
    camera.unproject_with_jac(x2, &d2, &J2inv);

    FixCameraRelativePoseRefiner<UniformWeightVector, TestAccumulator> refiner(d1, d2, J1inv, J2inv);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), CameraPose>(refiner, pose, delta);
    REQUIRE_SMALL(jac_err, 1e-6)

    // Test that compute_residual and compute_jacobian are compatible
    TestAccumulator acc;
    acc.reset_residual();
    double r1 = refiner.compute_residual(acc, pose);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, pose);
    double r2 = 0.0;
    for (int i = 0; i < acc.rs.size(); ++i) {
        r2 += acc.weights[i] * acc.rs[i].squaredNorm();
    }
    REQUIRE_SMALL(std::abs(r1 - r2), 1e-10);

    return true;
}

bool test_tangent_sampson_fix_camera_relative_pose_refinement() {
    const size_t N = 25;
    std::string camera_str = "3 RADIAL 1936 1296 2425.85 932.38 629.325 -0.04012 0.00123";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("fix_camera_relative_pose_refinement_noise");
    for (size_t i = 0; i < N; ++i) {
        x1[i] += test_rng::symmetric_vec2(noise_rng, 2.0);
        x2[i] += test_rng::symmetric_vec2(noise_rng, 2.0);
    }

    double f = camera.focal();
    for (int i = 0; i < N; ++i) {
        x1[i] /= f;
        x2[i] /= f;
    }
    camera.rescale(1.0 / f);

    std::vector<Eigen::Vector3d> d1, d2;
    std::vector<Eigen::Matrix<double, 3, 2>> J1inv, J2inv;
    camera.unproject_with_jac(x1, &d1, &J1inv);
    camera.unproject_with_jac(x2, &d2, &J2inv);
    FixCameraRelativePoseRefiner refiner(d1, d2, J1inv, J2inv);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &pose, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_tangent_sampson_fix_camera_relative_pose_refinement"));

    return true;
}

bool test_tangent_sampson_camera_relative_pose_jacobian() {
    const size_t N = 50;
    std::string camera1_str = "3 RADIAL 1936 1296 2425.85 932.38 629.325 -0.04012 0.00123";
    std::string camera2_str = "0 SIMPLE_RADIAL 1936 1296 2425.85 932.383 628.265 -0.0397695";
    Camera camera1, camera2;
    camera1.initialize_from_txt(camera1_str);
    camera2.initialize_from_txt(camera2_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera1, camera2);

    BundleOptions opt;
    opt.refine_principal_point = false;
    opt.refine_focal_length = true;
    opt.refine_extra_params = true;
    std::vector<size_t> cam1_ref_idx = camera1.get_param_refinement_idx(opt);
    std::vector<size_t> cam2_ref_idx = camera2.get_param_refinement_idx(opt);

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("camera_relative_pose_jacobian_noise");
    for (size_t i = 0; i < N; ++i) {
        x1[i] += test_rng::symmetric_vec2(noise_rng, 2.0);
        x2[i] += test_rng::symmetric_vec2(noise_rng, 2.0);
    }

    CameraRelativePoseRefiner<UniformWeightVector, TestAccumulator> refiner(x1, x2, cam1_ref_idx, cam2_ref_idx, false);

    ImagePair pair(pose, camera1, camera2);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), ImagePair>(refiner, pair, delta);
    REQUIRE_SMALL(jac_err, 1e-6)

    // Test that compute_residual and compute_jacobian are compatible
    TestAccumulator acc;
    acc.reset_residual();
    double r1 = refiner.compute_residual(acc, pair);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, pair);
    double r2 = 0.0;
    for (int i = 0; i < acc.rs.size(); ++i) {
        r2 += acc.weights[i] * acc.rs[i].squaredNorm();
    }
    REQUIRE_SMALL(std::abs(r1 - r2), 1e-10);

    return true;
}

bool test_tangent_sampson_camera_relative_pose_refinement() {
    const size_t N = 50;
    std::string camera1_str = "3 RADIAL 1936 1296 2425.85 932.38 629.325 -0.04012 0.00123";
    std::string camera2_str = "0 SIMPLE_RADIAL 1936 1296 2425.85 932.383 628.265 -0.0397695";
    camera2_str = camera1_str;
    Camera camera1, camera2;
    camera1.initialize_from_txt(camera1_str);
    camera2.initialize_from_txt(camera2_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera1, camera2);

    BundleOptions opt;
    opt.refine_principal_point = false;
    opt.refine_focal_length = true;
    opt.refine_extra_params = true;
    std::vector<size_t> cam1_ref_idx = camera1.get_param_refinement_idx(opt);
    std::vector<size_t> cam2_ref_idx = camera2.get_param_refinement_idx(opt);

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("camera_relative_pose_refinement_noise");
    for (size_t i = 0; i < N; ++i) {
        x1[i] += test_rng::symmetric_vec2(noise_rng, 0.1);
        x2[i] += test_rng::symmetric_vec2(noise_rng, 0.1);
    }

    double f1 = camera1.focal();
    double f2 = camera2.focal();
    for (int i = 0; i < N; ++i) {
        x1[i] /= f1;
        x2[i] /= f2;
    }
    camera1.rescale(1.0 / f1);
    camera2.rescale(1.0 / f2);

    ImagePair pair(pose, camera1, camera2);

    CameraRelativePoseRefiner refiner(x1, x2, cam1_ref_idx, cam2_ref_idx, false);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &pair, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_tangent_sampson_camera_relative_pose_refinement"));

    return true;
}

bool test_tangent_sampson_shared_camera_relative_pose_refinement() {
    const size_t N = 50;
    std::string camera_str = "0 SIMPLE_RADIAL 1936 1296 2425.85 932.383 628.265 -0.0397695";
    Camera camera;
    camera.initialize_from_txt(camera_str);

    CameraPose pose;
    std::vector<Eigen::Vector2d> x1, x2;
    setup_scene(N, pose, x1, x2, camera, camera);

    BundleOptions opt;
    opt.refine_principal_point = false;
    opt.refine_focal_length = true;
    opt.refine_extra_params = true;
    std::vector<size_t> cam_ref_idx = camera.get_param_refinement_idx(opt);

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("shared_camera_relative_pose_refinement_noise");
    for (size_t i = 0; i < N; ++i) {
        x1[i] += test_rng::symmetric_vec2(noise_rng, 0.1);
        x2[i] += test_rng::symmetric_vec2(noise_rng, 0.1);
    }

    double f = camera.focal();
    for (int i = 0; i < N; ++i) {
        x1[i] /= f;
        x2[i] /= f;
    }
    camera.rescale(1.0 / f);

    ImagePair pair(pose, camera, camera);

    CameraRelativePoseRefiner refiner(x1, x2, cam_ref_idx, cam_ref_idx, true);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &pair, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-8, "test_tangent_sampson_shared_camera_relative_pose_refinement"));

    return true;
}

} // namespace test::relative

using namespace test::relative;
std::vector<Test> register_optim_relative_test() {
    return {TEST(test_relative_pose_normal_acc),
            TEST(test_relative_pose_jacobian),
            TEST(test_relative_pose_refinement),
            TEST(test_shared_focal_relative_pose_normal_acc),
            TEST(test_shared_focal_relative_pose_jacobian),
            TEST(test_shared_focal_relative_pose_refinement),
            TEST(test_varying_focal_relative_pose_normal_acc),
            TEST(test_varying_focal_relative_pose_jacobian),
            TEST(test_varying_focal_relative_pose_refinement),
            TEST(test_tangent_sampson_fix_camera_relative_pose_jacobian),
            TEST(test_tangent_sampson_fix_camera_relative_pose_refinement),
            TEST(test_tangent_sampson_camera_relative_pose_jacobian),
            TEST(test_tangent_sampson_camera_relative_pose_refinement),
            TEST(test_tangent_sampson_shared_camera_relative_pose_refinement)};
}
