import numpy as np
import cv2
import sys

# Read in the two images
imgL = cv2.imread(sys.argv[1],1)
imgR = cv2.imread(sys.argv[2],1)
cv2.imshow('OriginalL', imgL)
cv2.imshow('OriginalR', imgR)

# Some parameters
levels = int(sys.argv[3]) if len(sys.argv) > 3 else 5
zoom = 3.0

# Compute the disparity image
stereo = cv2.StereoBM(1, 16, 15)
disparity = stereo.compute(cv2.cvtColor(imgR, cv2.COLOR_BGR2GRAY), cv2.cvtColor(imgL, cv2.COLOR_BGR2GRAY))
disparity = ((disparity-disparity.min()) * 255 / (disparity.max()-disparity.min())).astype(np.uint8)
cv2.imshow('Disparity', disparity)

# Fill in occluded pixels
nonZeroDisparityVals = disparity[disparity != 0]  # labels
nonZeroDisparityCoors = np.array([coor[0] for coor in cv2.findNonZero(disparity)])  # features
knn = cv2.KNearest()
knn.train(nonZeroDisparityCoors.astype(np.float32),nonZeroDisparityVals.astype(np.float32))

zeroDisparityCoors = np.array([coor[0] for coor in cv2.findNonZero(np.logical_not(disparity).astype(np.uint8))])
_, newDisparities, _, _ = knn.find_nearest(zeroDisparityCoors.astype(np.float32), 1)
disparity[disparity == 0] = newDisparities.flatten().astype(np.uint8)
cv2.imshow('Disparity Filled', disparity)

# Group disparity values into bins

# # Way 1: equally spaced bins
# disparityCutoffs = np.linspace(-1,255,levels+1)

# # Way 2: bins with equal number of pixels
# disparityValues = disparity.flatten()
# disparityValues.sort()
# disparityCutoffs = [disparityValues[i] for i in np.linspace(0,len(disparityValues)-1,levels+1)]
# disparityCutoffs[0] -= 1

# Way 3: K-means
data = np.float32(disparity.flatten())
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 50, 1.0)
_, labels, centers = cv2.kmeans(data, levels, criteria, 10, cv2.KMEANS_RANDOM_CENTERS)
disparity = centers[labels.flatten()].reshape(disparity.shape)
disparityCutoffs = [-1] + sorted(centers + 1)

# Resize sections of the image according to disparity cutoffs
masks = []
for i in range(len(disparityCutoffs)-1):
    masks.append(np.logical_and(disparity > disparityCutoffs[i], disparity <= disparityCutoffs[i+1]))
resizedImgs = []
resizedMasks = []
for zoom, mask in zip(np.linspace(1.0, zoom, levels), masks):
    resizedImgs.append(cv2.resize(imgL, (0,0), fx=zoom, fy=zoom))
    resizedMasks.append(cv2.resize(mask.astype(np.uint8), (0,0), fx=zoom, fy=zoom).astype(np.bool))

# Put together images
finalImg = np.zeros((int(imgL.shape[0] * zoom), int(imgL.shape[1] * zoom), 3), dtype=np.uint8)
for img, mask in zip(resizedImgs, resizedMasks):
    height, width, _ = img.shape
    startCol = len(finalImg[0]) / 2 - width / 2
    startRow = len(finalImg) / 2 - height / 2

    finalImgMask = np.zeros(finalImg.shape[:2], dtype=np.bool)
    finalImgMask[startRow:startRow+height,startCol:startCol+width] = mask
    finalImg[finalImgMask] = img[mask]
cv2.imshow('Final', finalImg)
cv2.imwrite('output/Final.jpg', finalImg)

cv2.waitKey(0)
cv2.destroyAllWindows()