import cv2
import numpy as np
from sklearn.neighbors import KernelDensity

__all__ = [
    'build_transformation_matrix',
    'update_transformation_matrix',
    'estimate_partial_transform',
    'removeOutliers',
    'guidedfilter'
]


def build_transformation_matrix(transform):
    """Convert transform list to transformation matrix

    :param transform: transform list as [dx, dy, da]
    :return: transform matrix as 2d (2, 3) numpy array
    """
    transform_matrix = np.zeros((2, 3))

    transform_matrix[0, 0] = np.cos(transform[2])
    transform_matrix[0, 1] = -np.sin(transform[2])
    transform_matrix[1, 0] = np.sin(transform[2])
    transform_matrix[1, 1] = np.cos(transform[2])
    transform_matrix[0, 2] = transform[0]
    transform_matrix[1, 2] = transform[1]

    return transform_matrix


def update_transformation_matrix(M, m):

    # extend M and m to 3x3 by adding an [0,0,1] to their 3rd row
    M_ = np.concatenate([M, np.zeros([1, 3])], axis=0)
    M_[-1, -1] = 1
    m_ = np.concatenate([m, np.zeros([1, 3])], axis=0)
    m_[-1, -1] = 1

    M_new = np.matmul(m_, M_)
    return M_new[0:2, :]


def estimate_partial_transform(matched_keypoints):
    """Wrapper of cv2.estimateRigidTransform for convenience in vidstab process

    :param matched_keypoints: output of match_keypoints util function; tuple of (cur_matched_kp, prev_matched_kp)
    :return: transform as list of [dx, dy, da]
    """
    prev_matched_kp, cur_matched_kp = matched_keypoints
    transform = cv2.estimateAffinePartial2D(np.array(prev_matched_kp),
                                            np.array(cur_matched_kp))[0]

    if transform is not None:
        # translation x
        dx = transform[0, 2]
        # translation y
        dy = transform[1, 2]
        # rotation
        da = np.arctan2(transform[1, 0], transform[0, 0])
    else:
        dx = dy = da = 0

    return [dx, dy, da]


def removeOutliers(prev_pts, curr_pts):

    d = np.sum((prev_pts - curr_pts)**2, axis=-1)**0.5

    d_ = np.array(d).reshape(-1, 1)
    kde = KernelDensity(kernel='gaussian', bandwidth=0.5).fit(d_)
    density = np.exp(kde.score_samples(d_))

    prev_pts = prev_pts[np.where((density >= 0.1))]
    curr_pts = curr_pts[np.where((density >= 0.1))]

    return prev_pts, curr_pts


def boxfilter(img, r):
    (rows, cols) = img.shape
    imDst = np.zeros_like(img)

    imCum = np.cumsum(img, 0)
    imDst[0: r+1, :] = imCum[r: 2*r+1, :]
    imDst[r+1: rows-r, :] = imCum[2*r+1: rows, :] - imCum[0: rows-2*r-1, :]
    imDst[rows-r: rows, :] = np.tile(imCum[rows-1, :],
                                     [r, 1]) - imCum[rows-2*r-1: rows-r-1, :]

    imCum = np.cumsum(imDst, 1)
    imDst[:, 0: r+1] = imCum[:, r: 2*r+1]
    imDst[:, r+1: cols-r] = imCum[:, 2*r+1: cols] - imCum[:, 0: cols-2*r-1]
    imDst[:, cols-r: cols] = np.tile(imCum[:, cols-1],
                                     [r, 1]).T - imCum[:, cols-2*r-1: cols-r-1]

    return imDst


def guidedfilter(img, p, r, eps):
    (rows, cols) = img.shape
    N = boxfilter(np.ones([rows, cols]), r)

    meanI = boxfilter(img, r) / N
    meanP = boxfilter(p, r) / N
    meanIp = boxfilter(img * p, r) / N
    covIp = meanIp - meanI * meanP

    meanII = boxfilter(img * img, r) / N
    varI = meanII - meanI * meanI

    a = covIp / (varI + eps)
    b = meanP - a * meanI

    meanA = boxfilter(a, r) / N
    meanB = boxfilter(b, r) / N

    q = meanA * img + meanB
    return q