import math
import cv2
import os
import numpy as np
from .read_write_model import read_images_binary


def align_model(model, rot, trans, scale):
    return (np.matmul(rot, model) + trans) * scale


def align(model, data):
    '''
    Source: https://vision.in.tum.de/data/datasets/rgbd-dataset/tools
            #absolute_trajectory_error_ate
    Align two trajectories using the method of Horn (closed-form).
    
    Input:
    model -- first trajectory (3xn)
    data -- second trajectory (3xn)

    Output:
    rot -- rotation matrix (3x3)
    trans -- translation vector (3x1)
    trans_error -- translational error per point (1xn)

    '''

    if model.shape[1] < 3:
        print('Need at least 3 points for ATE: {}'.format(model))
        return np.identity(3), np.zeros((3, 1)), 1

    # Get zero centered point cloud
    model_zerocentered = model - model.mean(1, keepdims=True)
    data_zerocentered = data - data.mean(1, keepdims=True)

    # constructed covariance matrix
    W = np.zeros((3, 3))
    for column in range(model.shape[1]):
        W += np.outer(model_zerocentered[:, column],
                      data_zerocentered[:, column])

    # SVD
    U, d, Vh = np.linalg.linalg.svd(W.transpose())
    S = np.identity(3)
    if (np.linalg.det(U) * np.linalg.det(Vh) < 0):
        S[2, 2] = -1
    rot = np.matmul(np.matmul(U, S), Vh)
    trans = data.mean(1, keepdims=True) - np.matmul(
        rot, model.mean(1, keepdims=True))

    # apply rot and trans to point cloud
    model_aligned = align_model(model, rot, trans, 1.0)
    model_aligned_zerocentered = model_aligned - model_aligned.mean(
        1, keepdims=True)

    # calc scale based on distance to point cloud center
    data_dist = np.sqrt((data_zerocentered * data_zerocentered).sum(axis=0))
    model_aligned_dist = np.sqrt(
        (model_aligned_zerocentered * model_aligned_zerocentered).sum(axis=0))
    scale_array = data_dist / model_aligned_dist
    scale = np.median(scale_array)

    return rot, trans, scale


def quaternion_matrix(quaternion):
    '''Return homogeneous rotation matrix from quaternion.

    >>> M = quaternion_matrix([0.99810947, 0.06146124, 0, 0])
    >>> numpy.allclose(M, rotation_matrix(0.123, [1, 0, 0]))
    True
    >>> M = quaternion_matrix([1, 0, 0, 0])
    >>> numpy.allclose(M, numpy.identity(4))
    True
    >>> M = quaternion_matrix([0, 1, 0, 0])
    >>> numpy.allclose(M, numpy.diag([1, -1, -1, 1]))
    True
    '''

    q = np.array(quaternion, dtype=np.float64, copy=True)
    n = np.dot(q, q)
    if n < _EPS:
        return np.identity(4)

    q *= math.sqrt(2.0 / n)
    q = np.outer(q, q)

    return np.array(
        [[1.0 - q[2, 2] - q[3, 3], q[1, 2] - q[3, 0], q[1, 3] + q[2, 0], 0.0],
         [q[1, 2] + q[3, 0], 1.0 - q[1, 1] - q[3, 3], q[2, 3] - q[1, 0], 0.0],
         [q[1, 3] - q[2, 0], q[2, 3] + q[1, 0], 1.0 - q[1, 1] - q[2, 2], 0.0],
         [0.0, 0.0, 0.0, 1.0]])


def quaternion_from_matrix(matrix, isprecise=False):
    '''Return quaternion from rotation matrix.

    If isprecise is True, the input matrix is assumed to be a precise rotation
    matrix and a faster algorithm is used.

    >>> q = quaternion_from_matrix(numpy.identity(4), True)
    >>> numpy.allclose(q, [1, 0, 0, 0])
    True
    >>> q = quaternion_from_matrix(numpy.diag([1, -1, -1, 1]))
    >>> numpy.allclose(q, [0, 1, 0, 0]) or numpy.allclose(q, [0, -1, 0, 0])
    True
    >>> R = rotation_matrix(0.123, (1, 2, 3))
    >>> q = quaternion_from_matrix(R, True)
    >>> numpy.allclose(q, [0.9981095, 0.0164262, 0.0328524, 0.0492786])
    True
    >>> R = [[-0.545, 0.797, 0.260, 0], [0.733, 0.603, -0.313, 0],
    ...      [-0.407, 0.021, -0.913, 0], [0, 0, 0, 1]]
    >>> q = quaternion_from_matrix(R)
    >>> numpy.allclose(q, [0.19069, 0.43736, 0.87485, -0.083611])
    True
    >>> R = [[0.395, 0.362, 0.843, 0], [-0.626, 0.796, -0.056, 0],
    ...      [-0.677, -0.498, 0.529, 0], [0, 0, 0, 1]]
    >>> q = quaternion_from_matrix(R)
    >>> numpy.allclose(q, [0.82336615, -0.13610694, 0.46344705, -0.29792603])
    True
    >>> R = random_rotation_matrix()
    >>> q = quaternion_from_matrix(R)
    >>> is_same_transform(R, quaternion_matrix(q))
    True
    >>> R = euler_matrix(0.0, 0.0, numpy.pi/2.0)
    >>> numpy.allclose(quaternion_from_matrix(R, isprecise=False),
    ...                quaternion_from_matrix(R, isprecise=True))
    True

    '''

    M = np.array(matrix, dtype=np.float64, copy=False)[:4, :4]
    if isprecise:
        q = np.empty((4, ))
        t = np.trace(M)
        if t > M[3, 3]:
            q[0] = t
            q[3] = M[1, 0] - M[0, 1]
            q[2] = M[0, 2] - M[2, 0]
            q[1] = M[2, 1] - M[1, 2]
        else:
            i, j, k = 1, 2, 3
            if M[1, 1] > M[0, 0]:
                i, j, k = 2, 3, 1
            if M[2, 2] > M[i, i]:
                i, j, k = 3, 1, 2
            t = M[i, i] - (M[j, j] + M[k, k]) + M[3, 3]
            q[i] = t
            q[j] = M[i, j] + M[j, i]
            q[k] = M[k, i] + M[i, k]
            q[3] = M[k, j] - M[j, k]
        q *= 0.5 / math.sqrt(t * M[3, 3])
    else:
        m00 = M[0, 0]
        m01 = M[0, 1]
        m02 = M[0, 2]
        m10 = M[1, 0]
        m11 = M[1, 1]
        m12 = M[1, 2]
        m20 = M[2, 0]
        m21 = M[2, 1]
        m22 = M[2, 2]

        # symmetric matrix K
        K = np.array([[m00 - m11 - m22, 0.0, 0.0, 0.0],
                      [m01 + m10, m11 - m00 - m22, 0.0, 0.0],
                      [m02 + m20, m12 + m21, m22 - m00 - m11, 0.0],
                      [m21 - m12, m02 - m20, m10 - m01, m00 + m11 + m22]])
        K /= 3.0

        # quaternion is eigenvector of K that corresponds to largest eigenvalue
        w, V = np.linalg.eigh(K)
        q = V[[3, 0, 1, 2], np.argmax(w)]

    if q[0] < 0.0:
        np.negative(q, q)

    return q

def is_colmap_img_valid(colmap_img_file):
    '''Return validity of a colmap reconstruction'''

    images_bin = read_images_binary(colmap_img_file)
    # Check if everything is finite for this subset
    for key in images_bin.keys():
        q = np.asarray(images_bin[key].qvec).flatten()
        t = np.asarray(images_bin[key].tvec).flatten()

        is_cur_valid = True
        is_cur_valid = is_cur_valid and q.shape == (4, )
        is_cur_valid = is_cur_valid and t.shape == (3, )
        is_cur_valid = is_cur_valid and np.all(np.isfinite(q))
        is_cur_valid = is_cur_valid and np.all(np.isfinite(t))

        # If any is invalid, immediately return
        if not is_cur_valid:
            return False

    return True

def get_best_colmap_index(colmap_output_path):
    '''
    Determines the colmap model with the most images if there is more than one.
    '''

    # First find the colmap reconstruction with the most number of images.
    best_index, best_num_images = -1, 0

    # Check all valid sub reconstructions.
    if os.path.exists(colmap_output_path):
        idx_list = [
            _d for _d in os.listdir(colmap_output_path)
            if os.path.isdir(os.path.join(colmap_output_path, _d))
        ]
    else:
        idx_list = []

    for cur_index in idx_list:
        cur_output_path = os.path.join(colmap_output_path, cur_index)
        if os.path.isdir(cur_output_path):
            colmap_img_file = os.path.join(cur_output_path, 'images.bin')
            images_bin = read_images_binary(colmap_img_file)
            # Check validity
            if not is_colmap_img_valid(colmap_img_file):
                continue
            # Find the reconstruction with most number of images
            if len(images_bin) > best_num_images:
                best_index = int(cur_index)
                best_num_images = len(images_bin)

    return str(best_index)