# Copyright 2020-present NAVER Corp. Under BSD 3-clause license

"""
Colmap database import as basic kapture objects functions
"""

from kapture.io.tar import TarCollection
import logging
import numpy as np
from tqdm import tqdm
from typing import Tuple, Optional

# kapture
import kapture
import kapture.io.features
# local
from .database import COLMAPDatabase, blob_to_array, pair_id_to_image_ids
from .database_extra import exists_table
from .cameras import CAMERA_MODEL_NAMES, get_camera_kapture_id_from_colmap_id

logger = logging.getLogger('colmap')


def get_cameras_from_database(database: COLMAPDatabase) -> kapture.Sensors:
    """
    Creates kapture sensors from the colmap database.

    :param database: colmap database
    :return: kapture sensors
    """
    logger.info('parsing cameras  ...')
    kapture_cameras = kapture.Sensors()

    for camera_id, model_id, width, height, params, prior_focal_length in database.execute(
            'SELECT camera_id, model, width, height, params, prior_focal_length FROM cameras;'):
        if model_id not in CAMERA_MODEL_NAMES:
            logger.warning(f'unable to convert colmap camera model ({model_id}) for camera {camera_id}.')
            # use 0 as default
            model_id = 0

        camera_id = get_camera_kapture_id_from_colmap_id(camera_id)
        model_name = CAMERA_MODEL_NAMES[model_id]

        #  By setting the prior_focal_length flag to 0 or 1,
        #  you can give a hint whether the reconstruction algorithm should trust the focal length value.
        params = blob_to_array(params, np.float64)
        params = [width, height] + params.tolist()

        kapture_camera = kapture.Camera(model_name, params)
        kapture_cameras[camera_id] = kapture_camera
    return kapture_cameras


def get_images_and_trajectories_from_database(database: COLMAPDatabase
                                              ) -> Tuple[kapture.RecordsCamera, kapture.Trajectories]:
    """
    Creates records_camera and trajectories from colmap images table
    In trajectories, timestamps are made up from colmap image id.

    :param database: colmap database
    :return: kapture records_camera and trajectories
    """
    logging.info('parsing images ...')
    kapture_images = kapture.RecordsCamera()
    kapture_trajectories = kapture.Trajectories()
    hide_progressbar = logger.getEffectiveLevel() > logging.INFO
    for image_id, name, camera_id, prior_qw, prior_qx, prior_qy, prior_qz, prior_tx, prior_ty, prior_tz \
            in tqdm(database.execute('SELECT image_id, name, camera_id, '
                                     'prior_qw, prior_qx, prior_qy, prior_qz, '
                                     'prior_tx, prior_ty, prior_tz  FROM images;'
                                     ), disable=hide_progressbar):
        # images
        timestamp = int(image_id)
        camera_id = get_camera_kapture_id_from_colmap_id(camera_id)
        kapture_images[timestamp, camera_id] = name
        # trajectories
        prior_q = [prior_qw, prior_qx, prior_qy, prior_qz]
        prior_t = [prior_tx, prior_ty, prior_tz]
        # do not register the pose part if its invalid.
        is_undefined = all(v is None for v in prior_q + prior_t)
        if is_undefined:
            # just ignore
            continue
        prior_pose = kapture.PoseTransform(prior_q, prior_t)
        kapture_trajectories[timestamp, camera_id] = prior_pose

    if len(kapture_trajectories) == 0:
        # if there is no pose at all, just don't bother.
        kapture_trajectories = None

    return kapture_images, kapture_trajectories


def get_keypoints_from_database(database: COLMAPDatabase,
                                records_camera: kapture.RecordsCamera,
                                kapture_dirpath: str,
                                tar_handlers: Optional[TarCollection] = None,
                                keypoints_type: str = 'SIFT'
                                ) -> Optional[kapture.Keypoints]:
    """
    Writes keypoints files and return the kapture keypoints from the colmap database.
    Requires records_camera timestamp == colmap_image_id

    :param database: colmap database.
    :param records_camera: input images.
    :param kapture_dirpath: input root path to kapture.
    :param tar_handlers: collection of preloaded tar archives
    :param keypoints_type: type of keypoints, name of the keypoints subfolder
     (by default, in colmap, it is SIFT, but can be imported)
    :return: kapture keypoints
    """
    image_filenames = set()
    dtype = np.float32
    dsize = None  # usually 6, will be retrieved on first keypoints of DB

    empty_keypoints = []

    # DB query
    colmap_keypoints_request = (
        (image_id, blob_to_array(data, dtype, (rows, cols)) if (rows > 0 and cols > 0) else None)
        for image_id, rows, cols, data in database.execute("SELECT image_id, rows, cols, data FROM keypoints"))

    hide_progressbar = logger.getEffectiveLevel() > logging.INFO
    for colmap_image_id, image_keypoints in tqdm(colmap_keypoints_request, disable=hide_progressbar):
        # retrieve image path from image_id
        assert len(records_camera[colmap_image_id]) == 1
        timestamp = colmap_image_id
        image_filename = next((v for v in records_camera[timestamp].values()), None)
        assert image_filename
        keypoints_filepath = kapture.io.features.get_keypoints_fullpath(keypoints_type,
                                                                        kapture_dirpath,
                                                                        image_filename,
                                                                        tar_handlers)

        # handle no keypoints (delay)
        if image_keypoints is None or image_keypoints.shape[0] == 0:
            logger.warning(f'image={image_filename} has 0 keypoints')
            empty_keypoints.append((image_filename, keypoints_filepath))
            continue

        if image_keypoints.dtype != dtype:
            raise ValueError('inconsistent keypoints type.')
        if dsize is None:
            dsize = int(image_keypoints.shape[1])
        elif dsize != image_keypoints.shape[1]:
            raise ValueError('inconsistent keypoints size.')

        # save the actual file
        kapture.io.features.image_keypoints_to_file(keypoints_filepath, image_keypoints)
        # register it into kapture
        image_filenames.add(image_filename)

    # do all the empty keypoints at the end when we know for sure dsize
    for image_filename, keypoints_filepath in empty_keypoints:
        if dsize is None:
            dsize = 6
        image_keypoints = np.zeros((0, dsize), dtype=dtype)
        # save the actual file
        kapture.io.features.image_keypoints_to_file(keypoints_filepath, image_keypoints)
        # register it into kapture
        image_filenames.add(image_filename)

    if image_filenames:
        return kapture.Keypoints(keypoints_type, dtype, dsize, image_filenames)
    else:
        return None


def get_descriptors_from_database(database: COLMAPDatabase,
                                  images: kapture.RecordsCamera,
                                  kapture_dirpath: str,
                                  tar_handlers: Optional[TarCollection] = None,
                                  keypoints_type: str = 'SIFT',
                                  descriptors_type: str = 'SIFT'
                                  ) -> Optional[kapture.Descriptors]:
    """
    Writes descriptors files and return the list in kapture format from the colmap database.

    :param database: colmap database.
    :param images: list of images (as RecordsCamera).
    :param kapture_dirpath: input root path to kapture.
    :param tar_handlers: collection of preloaded tar archives
    :param keypoints_type: type of keypoints, name of the keypoints subfolder
     (by default, in colmap, it is SIFT, but can be imported)
    :param descriptors_type: type of descriptors to export, name of the descriptors subfolder
     (by default, in colmap, it is SIFT, but can be imported)
    :return: kapture descriptors
    """
    image_filenames = set()
    dtype = np.uint8  # values in the range 0…255
    # see https://colmap.github.io/tutorial.html#feature-detection-and-extraction
    dsize = None  # usually uint8, 128, will be retrieved on first descriptor of DB

    empty_descriptors = []

    colmap_descriptors = (
        (image_id, blob_to_array(data, dtype, (rows, cols)) if (rows > 0 and cols > 0) else None)
        for image_id, rows, cols, data in database.execute("SELECT image_id, rows, cols, data FROM descriptors"))
    hide_progressbar = logger.getEffectiveLevel() > logging.INFO
    for image_id, image_descriptors in tqdm(colmap_descriptors, disable=hide_progressbar):
        # retrieve image path from image_id (actually the timestamp)
        image_filename = next((v for v in images[image_id].values()), None)
        assert image_filename
        descriptors_filepath = kapture.io.features.get_descriptors_fullpath(descriptors_type,
                                                                            kapture_dirpath,
                                                                            image_filename,
                                                                            tar_handlers)

        if image_descriptors is None or image_descriptors.shape[0] == 0:
            logger.warning(f'image={image_id}:{image_filename} has 0 descriptors.')
            empty_descriptors.append((image_filename, descriptors_filepath))
            continue

        if image_descriptors.dtype != dtype:
            raise ValueError('inconsistent descriptors type.')
        if dsize is None:
            dsize = int(image_descriptors.shape[1])
        elif dsize != image_descriptors.shape[1]:
            raise ValueError('inconsistent descriptors size.')

        # save the actual file
        kapture.io.features.image_descriptors_to_file(descriptors_filepath, image_descriptors)
        # register it into to kapture
        image_filenames.add(image_filename)

    # do all the empty descriptors at the end when we know for sure dsize
    for image_filename, descriptors_filepath in empty_descriptors:
        if dsize is None:
            dsize = 128
        image_descriptors = np.zeros((0, dsize), dtype=dtype)
        # save the actual file
        kapture.io.features.image_descriptors_to_file(descriptors_filepath, image_descriptors)
        # register it into kapture
        image_filenames.add(image_filename)

    if image_filenames:
        return kapture.Descriptors(descriptors_type, dtype, dsize, keypoints_type, "L2", image_filenames)
    else:
        return None


def get_matches_from_database(database: COLMAPDatabase,
                              images: kapture.RecordsCamera,
                              kapture_dirpath: str,
                              tar_handlers: Optional[TarCollection] = None,
                              keypoints_type: str = 'SIFT',
                              no_geometric_filtering: bool = False) -> kapture.Matches:
    """
    Writes Matches files and return the list in kapture format from the colmap database.

    :param database: input colmap database.
    :param images: input list of images (as RecordsCamera).
    :param kapture_dirpath: input root path to kapture.
    :param tar_handlers: collection of preloaded tar archives
    :param keypoints_type: type of keypoints, name of the keypoints subfolder
    (by default, in colmap, it is SIFT, but can be imported)
    :param no_geometric_filtering: only retrieve matches with geometric consistency.
    :return: kapture matches
    """
    kapture_matches = kapture.Matches()
    # check there is geometric matches available
    matches_table_name = 'matches'
    if not no_geometric_filtering:
        if not exists_table('two_view_geometries', database):
            logger.warning("No table: two_view_geometries: skipping geometric filtering")
        else:
            request = database.execute(
                'SELECT COUNT (*) FROM two_view_geometries')
            nb_verified_matches = next(request)[0]
            if nb_verified_matches > 0:
                request = database.execute('SELECT COUNT (*) FROM matches')
                nb_total_matches = next(request)[0]
                logger.info('keeps {}% of verified matches ({}/{}) ...'.format(
                    nb_verified_matches / nb_total_matches * 100, nb_verified_matches, nb_total_matches))
                matches_table_name = 'two_view_geometries'

    colmap_matches = [(pair_id_to_image_ids(pair_id), blob_to_array(data, np.uint32, (rows, cols)))
                      if rows > 0 else (pair_id_to_image_ids(pair_id), np.empty((rows, cols), dtype=np.uint32))
                      for pair_id, rows, cols, data
                      in database.execute('SELECT pair_id, rows, cols, data FROM {}'.format(matches_table_name))]

    logger.debug('matches: {}'.format(len(colmap_matches)))
    hide_progressbar = logger.getEffectiveLevel() > logging.INFO
    for (image_id1, image_id2), image_matches in tqdm(colmap_matches, disable=hide_progressbar):
        if image_id1 not in images or image_id2 not in images:
            logger.critical('inconsistent image ID {} or {}'.format(image_id1, image_id2))
            continue
        filename1 = next((v for v in images[image_id1].values()), None)
        filename2 = next((v for v in images[image_id2].values()), None)
        assert filename1 and filename2
        if (filename1, filename2) != kapture.Matches.lexical_order(filename1, filename2):
            # have to swap matches (keypoint image1, keypoint image2) become (keypoint image2, keypoint image1)
            image_matches = image_matches[:, ::-1]
            filename1, filename2 = kapture.Matches.lexical_order(filename1, filename2)

        # actually write the file
        # convert colmap image matches into kapture (cast to float and add a score column)
        image_matches = image_matches.astype(np.float)
        image_matches = np.hstack([image_matches, np.zeros((image_matches.shape[0], 1))])
        image_matches_filepath = kapture.io.features.get_matches_fullpath((filename1, filename2),
                                                                          keypoints_type,
                                                                          kapture_dirpath,
                                                                          tar_handlers)
        kapture.io.features.image_matches_to_file(image_matches_filepath, image_matches)
        # register the matching in kapture
        kapture_matches.add(filename1, filename2)

    return kapture_matches
