# SciPy Bundle Adjustment Example : https://scipy-cookbook.readthedocs.io/items/bundle_adjustment.html

from __future__ import print_function

import urllib.request
import bz2
import os
import numpy as np

BASE_URL = "http://grail.cs.washington.edu/projects/bal/data/ladybug/"
FILE_NAME = "problem-49-7776-pre.txt.bz2"
URL = BASE_URL + FILE_NAME

if not os.path.isfile(FILE_NAME):
    urllib.request.urlretrieve(URL, FILE_NAME)

def read_bal_data(file_name):
    with bz2.open(file_name, "rt") as file:
        n_cameras, n_points, n_observations = map(
            int, file.readline().split())

        camera_indices = np.empty(n_observations, dtype=int)
        point_indices = np.empty(n_observations, dtype=int)
        points_2d = np.empty((n_observations, 2))

        for i in range(n_observations):
            camera_index, point_index, x, y = file.readline().split()
            camera_indices[i] = int(camera_index)
            point_indices[i] = int(point_index)
            points_2d[i] = [float(x), float(y)]

        camera_params = np.empty(n_cameras * 9)
        for i in range(n_cameras * 9):
            camera_params[i] = float(file.readline())
        camera_params = camera_params.reshape((n_cameras, -1))

        points_3d = np.empty(n_points * 3)
        for i in range(n_points * 3):
            points_3d[i] = float(file.readline())
        points_3d = points_3d.reshape((n_points, -1))

    return camera_params, points_3d, camera_indices, point_indices, points_2d

camera_params, points_3d, camera_indices, point_indices, points_2d = read_bal_data(FILE_NAME)

'''
camera_params[0 ~ 2] = rotation vector (Rodrigue's rotation formula)
camera_params[3 ~ 5] = translation vector
camera_params[6] = focal distance
camera_params[7, 8] = distortion parameters
'''

n_cameras = camera_params.shape[0]
n_points = points_3d.shape[0]

n = 9 * n_cameras + 3 * n_points
m = 2 * points_2d.shape[0]

print("n_cameras: {}".format(n_cameras))
print("n_points: {}".format(n_points))
print("Total number of parameters: {}".format(n))
print("Total number of residuals: {}".format(m))

def rotate_rodrigues(points, rot_vecs):
    """
    Rotate points by given rotation vectors.
    
    Rodrigues' rotation formula is used.
    """
    theta = np.linalg.norm(rot_vecs, axis=1)[:, np.newaxis]

    with np.errstate(invalid='ignore'):
        v = rot_vecs / theta    # v = unit vector for describing an axis of rotation
        v = np.nan_to_num(v)

    dot = np.sum(points * v, axis=1)[:, np.newaxis]
    cos_theta = np.cos(theta)
    sin_theta = np.sin(theta)

    return cos_theta * points + sin_theta * np.cross(v, points) + dot * (1 - cos_theta) * v

def project(points, camera_params):
    """
    Convert 3-D points to 2-D by projecting onto images.
    Reprojection Step
    """

    ### Pinhole Camera Model : 3D world point --> 2D image plane ###
    # Rotate 3D world points by rotation vector
    points_proj = rotate_rodrigues(points, camera_params[:, :3])
    # Translate 3D world points by translation vector
    points_proj += camera_params[:, 3:6]
    # Divide by Z / 3D points are mapped on 2D image plane
    points_proj = -points_proj[:, :2] / points_proj[:, 2, np.newaxis]

    f = camera_params[:, 6]
    k1 = camera_params[:, 7]
    k2 = camera_params[:, 8]

    n = np.sum(points_proj**2, axis=1)
    r = 1 + k1 * n + k2 * n ** 2    # Rectify the points on 2D image plane based on lens distortion

    points_proj *= (r * f)[:, np.newaxis]   # Apply focal length as the final calculation variable

    return points_proj  # Return reprojected 2D points of 3D world points

def fun(params, n_cameras, n_points, camera_indices, point_indices, points_2d):
    """
    Compute residuals (Reprojection Error between reprojected 2D points and groundtruth 2D points)
    
    `params` contains camera parameters and 3-D coordinates.
    """
    camera_params = params[:n_cameras * 9].reshape((n_cameras, 9))
    points_3d = params[n_cameras * 9:].reshape((n_points, 3))
    points_proj = project(points_3d[point_indices], camera_params[camera_indices])

    return (points_proj - points_2d).ravel()    # Return the reprojection error (residual) between reprojected 2D points and groundtruth 2D points

from scipy.sparse import lil_matrix

def bundle_adjustment_sparsity(n_cameras, n_points, camera_indices, point_indices):
    '''
    Create initial Jacobian Matrix for non-linear optimization
    '''
    m = camera_indices.size * 2
    n = n_cameras * 9 + n_points * 3
    A = lil_matrix((m, n), dtype=int)

    i = np.arange(camera_indices.size)
    for s in range(9):
        A[2 * i, camera_indices * 9 + s] = 1
        A[2 * i + 1, camera_indices * 9 + s] = 1

    for s in range(3):
        A[2 * i, n_cameras * 9 + point_indices * 3 + s] = 1
        A[2 * i + 1, n_cameras * 9 + point_indices * 3 + s] = 1

    return A

# Visualize residuals (reprojection error) before non-linear optimization
import matplotlib.pyplot as plt

x0 = np.hstack((camera_params.ravel(), points_3d.ravel()))

f0 = fun(x0, n_cameras, n_points, camera_indices, point_indices, points_2d)

plt.figure(figsize=(20, 8))
plt.subplot(1, 2, 1)
plt.title('Residual (Reprojection Error) before BA')
plt.plot(f0)
plt.grid()

# Non-Linear Optimization Step of Bundle Adjustment
import time
from scipy.optimize import least_squares

A = bundle_adjustment_sparsity(n_cameras, n_points, camera_indices, point_indices)

t0 = time.time()
res = least_squares(fun, x0, jac_sparsity=A, verbose=2, x_scale='jac', ftol=1e-4, method='trf',
                    args=(n_cameras, n_points, camera_indices, point_indices, points_2d))
t1 = time.time()

print("Optimization took {0:.0f} seconds".format(t1 - t0))

plt.subplot(1, 2, 2)
plt.title('Residual (Reprojection Error) after BA')
plt.plot(res.fun)
plt.grid()

print(x0)

plt.show()
