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

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

using namespace poselib;

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

namespace test::homography {

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, Eigen::Matrix3d &H, std::vector<Point2D> &x1, std::vector<Point2D> &x2, Camera &cam1,
                 Camera &cam2, const std::string &case_name = "homography_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);
    CameraPose p = p2.compose(p1.inverse());

    Eigen::Vector3d normal = test_rng::symmetric_vec3(rng);
    normal.normalize();
    Eigen::Vector3d X0 = test_rng::symmetric_vec3(rng);

    X0 = p1.apply_inverse(X0);
    normal = p1.derotate(normal);

    // normal.dot(X0) = alpha
    normal = normal / normal.dot(X0);

    // X = lambda * x1
    // 1 = n'*X = lambda * n'*x1
    // lambda = 1 / n'*x1

    // (1/*n'*x1) * R * x1 + t =
    // (R + t*n') * x1

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

        // n'*(lambda * Xi) = 1
        // lambda = 1 / n'*Xi
        Xi = (1.0 / normal.dot(Xi)) * Xi;

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

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

    H = p.R() + p.t * normal.transpose();
}

void setup_scene_w_lines(int N_pts, int N_lines, Eigen::Matrix3d &H, std::vector<Point2D> &x1, std::vector<Point2D> &x2,
                         std::vector<Line2D> &lines1, std::vector<Line2D> &lines2, Camera &cam1, Camera &cam2) {

    std::vector<Point2D> x1_all, x2_all;
    setup_scene(N_pts + 2 * N_lines, H, x1_all, x2_all, cam1, cam2);

    for (int i = 0; i < N_pts; ++i) {
        x1.push_back(x1_all[i]);
        x2.push_back(x2_all[i]);
    }
    for (int i = N_pts; i < N_pts + 2 * N_lines; i += 2) {
        Line2D l1, l2;
        l1.x1 = x1_all[i];
        l1.x2 = x1_all[i + 1];
        l2.x1 = x2_all[i];
        l2.x2 = x2_all[i + 1];
        lines1.push_back(l1);
        lines2.push_back(l2);
    }
}

// Point Homographies

bool test_homography_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);

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

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

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

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

    return true;
}

bool test_homography_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);

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

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

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), Eigen::Matrix3d>(refiner, H, 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, H);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, H);
    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_homography_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);

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

    // Add some noise
    test_rng::Rng noise_rng = test_rng::make_rng("homography_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);
    }

    PinholeHomographyRefiner refiner(x1, x2);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &H, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-6, "test_homography_refinement"));

    return true;
}

//////////////////////////////////////
// Line Homographies

bool test_line_homography_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);

    Eigen::Matrix3d H;
    std::vector<Eigen::Vector2d> x1, x2;
    std::vector<Line2D> lines1, lines2;
    setup_scene_w_lines(N, N, H, x1, x2, lines1, lines2, camera, camera);

    NormalAccumulator acc;
    PinholeLineHomographyRefiner refiner(lines1, lines2);
    acc.initialize(refiner.num_params);

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

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

    return true;
}

bool test_line_homography_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);

    Eigen::Matrix3d H;
    std::vector<Eigen::Vector2d> x1, x2;
    std::vector<Line2D> lines1, lines2;
    setup_scene_w_lines(N, N, H, x1, x2, lines1, lines2, camera, camera);

    PinholeLineHomographyRefiner<UniformWeightVector, TestAccumulator> refiner(lines1, lines2);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), Eigen::Matrix3d>(refiner, H, 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, H);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, H);
    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_line_homography_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);

    Eigen::Matrix3d H;
    std::vector<Eigen::Vector2d> x1, x2;
    std::vector<Line2D> lines1, lines2;
    setup_scene_w_lines(N, N, H, x1, x2, lines1, lines2, camera, camera);

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

    PinholeLineHomographyRefiner refiner(lines1, lines2);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &H, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-6, "test_line_homography_refinement"));

    return true;
}

///////////////////////////
// Point + Line homography refinement

bool test_point_line_homography_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);

    Eigen::Matrix3d H;
    std::vector<Eigen::Vector2d> x1, x2;
    std::vector<Line2D> lines1, lines2;
    setup_scene_w_lines(N, N, H, x1, x2, lines1, lines2, camera, camera);

    PinholeHomographyRefiner<UniformWeightVector, TestAccumulator> point_refiner(x1, x2);
    PinholeLineHomographyRefiner<UniformWeightVector, TestAccumulator> line_refiner(lines1, lines2);
    HybridRefiner<Eigen::Matrix3d, TestAccumulator> refiner;
    refiner.register_refiner(&point_refiner);
    refiner.register_refiner(&line_refiner);

    const double delta = 1e-6;
    double jac_err = verify_jacobian<decltype(refiner), Eigen::Matrix3d>(refiner, H, 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, H);
    acc.reset_jacobian();
    refiner.compute_jacobian(acc, H);
    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_point_line_homography_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);

    Eigen::Matrix3d H;
    std::vector<Eigen::Vector2d> x1, x2;
    std::vector<Line2D> lines1, lines2;
    setup_scene_w_lines(N, N, H, x1, x2, lines1, lines2, camera, camera);

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

    PinholeHomographyRefiner point_refiner(x1, x2);
    PinholeLineHomographyRefiner line_refiner(lines1, lines2);

    HybridRefiner<Eigen::Matrix3d> refiner;
    refiner.register_refiner(&point_refiner);
    refiner.register_refiner(&line_refiner);

    BundleOptions bundle_opt;
    bundle_opt.step_tol = 1e-12;
    BundleStats stats = lm_impl(refiner, &H, bundle_opt, print_iteration);
    REQUIRE(check_bundle_cost_and_gradient(stats, 1e-6, "test_point_line_homography_refinement"));

    return true;
}

} // namespace test::homography

using namespace test::homography;
std::vector<Test> register_optim_homography_test() {
    return {TEST(test_homography_normal_acc),          TEST(test_homography_jacobian),
            TEST(test_homography_refinement),          TEST(test_line_homography_normal_acc),
            TEST(test_line_homography_jacobian),       TEST(test_line_homography_refinement),
            TEST(test_point_line_homography_jacobian), TEST(test_point_line_homography_refinement)};
}
