import numpy as np
from matplotlib import pyplot as plt
import cv2
plt.figure(figsize = (10, 10))

def disparity_ssd(l_image, r_image):
    l_gray = cv2.cvtColor(l_image, cv2.COLOR_RGB2GRAY)
    r_gray = cv2.cvtColor(r_image, cv2.COLOR_RGB2GRAY)
    result = np.zeros(l_image.shape)
    win = 10
    
    for row in range(l_gray.shape[0]):
        for col in range(l_gray.shape[1]):
            row_u = row-win if row-win > 0 else 0
            row_d = row+win if row+win < l_gray.shape[0] else l_gray.shape[0]
            col_l = col-win if col-win > 0 else 0
            col_r = col+win if col+win < l_gray.shape[1] else l_gray.shape[1]
            
            patch = l_gray[row_u:row_d,col_l:col_r]
            candi = r_gray[row_u:row_d]
            
            res = cv2.matchTemplate(candi, patch, cv2.TM_CCORR_NORMED)
            res_row, res_col = np.unravel_index(res.argmax(), res.shape)
            res_col = res_col + win
            result[row,col] = res_col - col
            
    return result
            

l_image = cv2.imread('/home/andy/Desktop/computer_vision/ps2_python_template/input/pair2-L.png')
r_image = cv2.imread('/home/andy/Desktop/computer_vision/ps2_python_template/input/pair2-R.png')

plt.subplot(221), plt.imshow(l_image)
plt.subplot(222), plt.imshow(r_image)

d_l = disparity_ssd(l_image, r_image)
d_r = disparity_ssd(r_image, l_image)

plt.subplot(223), plt.imshow(d_l)
plt.subplot(224), plt.imshow(d_r)

plt.show()