"""
Utility function for image co-registration

Credits:
Copyright (c) 2017-2022 Matej Aleksandrov, Matej Batič, Grega Milčinski, Domagoj Korais, Matic Lubej (Sinergise)
Copyright (c) 2017-2022 Žiga Lukšič, Devis Peressutti, Nejc Vesel, Jovan Višnjić, Anže Zupanc (Sinergise)
Copyright (c) 2019-2020 Jernej Puc, Lojze Žust (Sinergise)
Copyright (c) 2017-2019 Blaž Sovdat, Andrej Burja (Sinergise)

This source code is licensed under the MIT license found in the LICENSE
file in the root directory of this source tree.
"""
# pylint: disable=invalid-name

import logging

import numpy as np

LOGGER = logging.getLogger(__name__)


def ransac(npts, model, n, k, t, d):
    """Fit model parameters to data using the RANSAC algorithm

    This implementation is written from pseudo-code found at
    http://en.wikipedia.org/w/index.php?title=RANSAC&oldid=116358182

    :param npts: A set of observed data points
    :param model: A model that can be fitted to data points
    :param n: The minimum number of data values required to fit the model
    :param k: The maximum number of iterations allowed in the algorithm
    :param t: A threshold value for determining when a data point fits a model
    :param d: The number of close data values required to assert that a model fits well to data
    :return: Model parameters which best fit the data (or None if no good model is found)
    """
    iterations = 0
    bestfit = None
    besterr = np.inf
    # best_inlier_idxs = None
    while iterations < k:
        maybe_idxs, test_idxs = random_partition(n, npts)
        maybemodel = model.fit(maybe_idxs)
        test_err = model.score(test_idxs, maybemodel)
        also_idxs = test_idxs[test_err < t]  # select indices of rows with accepted points

        LOGGER.debug("test_err.min() %f", test_err.min() if test_err.size else None)
        LOGGER.debug("test_err.max() %f", test_err.max() if test_err.size else None)
        LOGGER.debug("numpy.mean(test_err) %f", np.mean(test_err) if test_err.size else None)
        LOGGER.debug("iteration %d, len(alsoinliers) = %d", iterations, len(also_idxs))

        if len(also_idxs) > d:
            betteridxs = np.concatenate((maybe_idxs, also_idxs))
            bettermodel = model.fit(betteridxs)
            better_errs = model.score(betteridxs, bettermodel)
            thiserr = np.mean(better_errs)
            if thiserr < besterr:
                bestfit = bettermodel
                besterr = thiserr
                # best_inlier_idxs = np.concatenate((maybe_idxs, also_idxs))
        iterations += 1
    return bestfit


def random_partition(n, n_data):
    """return n random rows of data (and also the other len(data)-n rows)"""
    all_idxs = np.arange(n_data)
    np.random.shuffle(all_idxs)
    idxs1 = all_idxs[:n]
    idxs2 = all_idxs[n:]
    return idxs1, idxs2


class EstimateEulerTransformModel:
    """Estimate Euler transform linear system solved using linear least squares

    This class estimates an Euler 2D transformation between two cloud of 2D points using SVD decomposition
    """

    def __init__(self, src_pts, trg_pts):
        """Initialise target and source cloud points as Nx2 matrices. The transformation aligning source points to
        target points is estimated.

        :param src_pts: Array of source points
        :param trg_pts: Array of target points
        """
        self.src_pts = src_pts
        self.trg_pts = trg_pts

    def estimate_rigid_transformation(self, idx):
        """Estimate rigid transformation given a set of indices

        :param idx: Array of indices used to estimate the transformation
        :return: Estimated transformation matrix
        """
        # Look at points of given indices only
        src_pts = self.src_pts[idx, :]
        trg_pts = self.trg_pts[idx, :]
        # Get centroid location
        src_centroid, trg_centroid = np.mean(src_pts, axis=0), np.mean(trg_pts, axis=0)
        # Centre sets of points
        src_pts_centr = src_pts - src_centroid
        trg_pts_centr = trg_pts - trg_centroid
        # SVD decomposition
        u, _, v = np.linalg.svd(np.matmul(src_pts_centr.transpose(), trg_pts_centr))
        e = np.eye(2)
        e[1, 1] = np.linalg.det(np.matmul(v.transpose(), u.transpose()))
        # Estimate rotation matrix
        r = np.matmul(v.transpose(), np.matmul(e, u.transpose()))
        # Estimate translation
        t = trg_centroid - np.matmul(r, src_centroid.transpose()).transpose()
        # Fill in transformation matrix and return
        warp_matrix = np.zeros((2, 3))
        warp_matrix[:2, :2] = r
        warp_matrix[:, 2] = t
        return warp_matrix

    def fit(self, idx):
        """Estimate Euler transformation on points listed in `idx`

        :param idx: Indices used to estimate transformation
        :return: Transformation matrix
        """
        x = self.estimate_rigid_transformation(idx)
        return x

    def score(self, idx, warp_matrix):
        """Estimate the registration error of estimated transformation matrix

        :param idx: List of points used to estimate the transformation
        :param warp_matrix: Matrix estimating Euler transformation
        :return: Square root of Target Registration Error
        """
        # Transform source points with estimated transformation
        trg_fit = np.dot(warp_matrix, np.concatenate((self.src_pts[idx, :], np.ones((len(idx), 1))), axis=1).T).T
        # Compute error in transformation
        err_per_point = np.sqrt(
            np.sum((self.trg_pts[idx, :] - trg_fit[:, :2]) ** 2, axis=1)
        )  # sum squared error per row
        return err_per_point
