import numpy as np
import sklearn.neighbors
import scipy.sparse
import warnings
import imageio

nn = 10


def knn_matte(img, trimap, mylambda=100):
    [m, n, c] = img.shape
    img, trimap = img/255.0, trimap/255.0
    foreground = (trimap > 0.99).astype(int)
    background = (trimap < 0.01).astype(int)
    all_constraints = foreground + background

    print('Finding nearest neighbors')
    a, b = np.unravel_index(np.arange(m*n), (m, n))
    feature_vec = np.append(np.transpose(img.reshape(m*n,c)), [ a, b]/np.sqrt(m*m + n*n), axis=0).T
    nbrs = sklearn.neighbors.NearestNeighbors(n_neighbors=10, n_jobs=4).fit(feature_vec)
    knns = nbrs.kneighbors(feature_vec)[1]

    # Compute Sparse A
    print('Computing sparse A')
    row_inds = np.repeat(np.arange(m*n), 10)
    col_inds = knns.reshape(m*n*10)
    vals = 1 - np.linalg.norm(feature_vec[row_inds] - feature_vec[col_inds], axis=1)/(c+2)
    A = scipy.sparse.coo_matrix((vals, (row_inds, col_inds)),shape=(m*n, m*n))

    D_script = scipy.sparse.diags(np.ravel(A.sum(axis=1)))
    L = D_script-A
    # D = scipy.sparse.diags(np.ravel(all_constraints[:,:, 0]))
    D = scipy.sparse.diags(np.ravel(all_constraints[:,:]))
    # v = np.ravel(foreground[:,:,0])
    v = np.ravel(foreground[:,:])
    c = 2*mylambda*np.transpose(v)
    H = 2*(L + mylambda*D)

    print('Solving linear system for alpha')
    warnings.filterwarnings('error')
    alpha = []
    try:
        alpha = np.minimum(np.maximum(scipy.sparse.linalg.spsolve(H, c), 0), 1).reshape(m, n)
    except Warning:
        x = scipy.sparse.linalg.lsqr(H, c)
        alpha = np.minimum(np.maximum(x[0], 0), 1).reshape(m, n)
    return alpha


def main():
    # img = scipy.misc.imread('donkey.png')[:,:,:3]
    # trimap = scipy.misc.imread('donkeyTrimap.png')[:,:,:3]

    img = imageio.imread('donkey.png')[:,:,:3]
    trimap = imageio.imread('donkeyTrimap.png')#[:,:,:3]
    alpha = knn_matte(img, trimap)
    imageio.imsave('my_donkeyAlpha.png', alpha)
    plt.title('Alpha Matte')
    plt.imshow(alpha, cmap='gray')
    plt.show()

def my_main():
    # img = scipy.misc.imread('donkey.png')[:,:,:3]
    # trimap = scipy.misc.imread('donkeyTrimap.png')[:,:,:3]
    import cv2

    files = sorted(os.listdir('./examples/images/'))
    for f in files:
        img = imageio.imread(os.path.join('./examples/images', f))[:,:,:3]
        trimap = imageio.imread(os.path.join('./examples/trimaps', f))#[:,:,:3
        #print(img.shape)
        #exit(0)
        img = cv2.resize(img, (1280, 1280), cv2.INTER_LINEAR)
        img_hat = cv2.resize(img, (320, 320), cv2.INTER_LINEAR)
        trimap = cv2.resize(trimap, (320, 320), cv2.INTER_NEAREST)

        alpha = knn_matte(img_hat, trimap)

        alpha = cv2.resize(alpha, (1280, 1280), cv2.INTER_NEAREST)
        # imageio.imsave(os.path.join('./examples/mattes', f), alpha)
        print(img.shape, alpha.shape)
        final = np.concatenate((img, alpha[...,np.newaxis]), axis=2)
        imageio.imsave(os.path.join('./examples/colors', f), final)
    #plt.title('Alpha Matte')
    #plt.imshow(alpha, cmap='gray')
    #plt.show()

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import scipy.misc
    import os
    my_main()
    
