from cv2_45 import cv2
import numpy as np
import time
from scipy.interpolate import interp1d


def fillMissingValue(img_in):
    img_out = np.copy(img_in)

    #x = np.arange(img_in.shape[0])

    #Proceed column by column
    for i in range(img_in.shape[1]):
        col = img_in[:,i]

        col_r = col[:,0]
        col_g = col[:,1]
        col_b = col[:,2]
        r = interpolate(col_r)
        g = interpolate(col_g)
        b = interpolate(col_b)
        img_out[:,i,0] = r
        img_out[:,i,1] = g
        img_out[:,i,2] = b

    return img_out

def interpolate(y):
    x = np.arange(len(y))
    idx = np.nonzero(y)
    interp = interp1d(x[idx],y[idx],fill_value="extrapolate" )

    return interp(x)






def fillMissingValue2(img_in):
    rows, cols, _ = img_in.shape

    img_out = np.copy(img_in)

    x = np.arange(rows)
    img2d = img.reshape(rows, cols * 3)
    idx = np.nonzero(np.sum(img2d, axis=1))

    #Proceed column by column
    for i in range(cols):
        col = img_in[:,i]
        rgb = interpolate2(x, col, idx)
        img_out[:,i] = rgb

    return img_out

def interpolate2(x, y, idx):
    interp = interp1d(x[idx],y[idx,:],axis=1, fill_value="extrapolate" )
    return interp(x)



def fillMissingValue3(img_in):
    rows, cols, _ = img_in.shape

    x = np.arange(rows)
    img2d = img.reshape(rows, cols * 3)
    idx = np.nonzero(np.sum(img2d, axis=1))

    interp = interp1d(x[idx],img2d[idx,:],axis=1, fill_value="extrapolate" )
    return np.uint8(interp(x)).reshape(rows, cols, 3)




if __name__ == "__main__":
    img = cv2.imread("lena.jpg")
    img = cv2.resize(img, (1024,1024))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    for row in [5,24,27,40,90,110,120,200,300,400,500,700,930]:
        img[row,:,:] = 0

    start1 = time.time()
    img2 = fillMissingValue(img)
    end1 = time.time()

    start2 = time.time()
    img3 = fillMissingValue2(img)
    end2 = time.time()

    start3 = time.time()
    img4 = fillMissingValue3(img)
    end3 = time.time()

    print (img2 == img3).all()
    print (img2 == img4).all()

    print("Process time (1): {}".format(np.round(end1-start1,3)))
    print("Process time (2): {}".format(np.round(end2-start2,3)))
    print("Process time (3): {}".format(np.round(end3-start3,3)))