# python -m pip install opencv-python
import cv2 as cv
# python -m pip install numpy==1.19.3
import numpy as np
# python -m pip install matplotlib
from matplotlib import pyplot as plt

# We already have individual calibrations for each "camera" from Calib.py
# The stereoCalibrate function should thus skip finding them
stereoflags = cv.CALIB_FIX_INTRINSIC
# To do this we need to import the previously found camera properties
# Left
retL = np.load('./calibParams/retL.npy', None, True)
CmatrixL = np.load('./calibParams/CmatrixL.npy', None, True)
distCoeffsL = np.load('./calibParams/distCoeffsL.npy', None, True)
rvecsL = np.load('./calibParams/rvecsL.npy', None, True)
tvecsL = np.load('./calibParams/tvecsL.npy', None, True)
img_pointsL = np.load('./calibParams/img_pointsL.npy', None, True)
# Right
retR = np.load('./calibParams/retR.npy', None, True)
CmatrixR = np.load('./calibParams/CmatrixR.npy', None, True)
distCoeffsR = np.load('./calibParams/distCoeffsR.npy', None, True)
rvecsR = np.load('./calibParams/rvecsR.npy', None, True)
tvecsR = np.load('./calibParams/tvecsR.npy', None, True)
img_pointsR = np.load('./calibParams/img_pointsR.npy', None, True)
# Shared
obj_points = np.load('./calibParams/obj_points.npy', None, True)

# Set the criteria
criteria_stereo = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001)

# Need an image to calculate the translation and rotation of the two cameras
image = cv.imread('../IMAGES/calib3/IMGP1882.jpg')
# Make sure it is grayscale
image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)
Y, X = image.shape
# Split the image
imageL = image[0:int(Y), 0:int(X/2)-100]
imageR = image[0:int(Y), int(X/2)+100:int(X)]

print(image.shape[::-1])

# Commented out to save runtime, parameters are saved and loaded in after the first run
# Run the cv.stereoCalibrate function to find the Essential and Fundamental matrix
# Note that ROTmatrix and TRANSmatrix are the position of the first (left) camera with respect to the second (right) camera coordinate system
# retS, NMatrixL, distCoeffsL, NMatrixR, distCoeffsR, ROTmatrix, TRANSmatrix, Essmatrix, Fundmatrix = cv.stereoCalibrate(
#     obj_points, img_pointsL, img_pointsR, CmatrixL, distCoeffsL, CmatrixR, distCoeffsR, image.shape[::-1], stereoflags, criteria_stereo)

# Show some output
# print("Finished stereoCalibrate")
# print(Essmatrix)
# print(Fundmatrix)


Fundmatrix = np.load('./epiParams/fundmatrix.npy')
Essmatrix = np.load('./epiParams/essmatrix.npy')
TRANSmatrix = np.load('./epiParams/transmatrix.npy')
ROTmatrix = np.load('./epiParams/rotmatrix.npy')
distCoeffsR = np.load('./epiParams/distR.npy')
NMatrixR = np.load('./epiParams/nmatrixR.npy')
distCoeffsL = np.load('./epiParams/distL.npy')
NMatrixL = np.load('./epiParams/nmatrixL.npy')

# Save outputs, this script takes some time to run
np.save('./epiParams/fundmatrix', Fundmatrix)
np.save('./epiParams/essmatrix', Essmatrix)
np.save('./epiParams/transmatrix', TRANSmatrix)
np.save('./epiParams/rotmatrix', ROTmatrix)
np.save('./epiParams/distR', distCoeffsR)
np.save('./epiParams/nmatrixR', NMatrixR)
np.save('./epiParams/distL', distCoeffsL)
np.save('./epiParams/nmatrixL', NMatrixL)
print("Saved params")


# Another way of doing it: https://docs.opencv.org/master/da/de9/tutorial_py_epipolar_geometry.html
# Point matching to create the fundamental matrix
img1 = imageL  #queryimage # left image
img2 = imageR #trainimage # right image

sift = cv.SIFT_create()
# find the keypoints and descriptors with SIFT
kp1, des1 = sift.detectAndCompute(img1,None)
kp2, des2 = sift.detectAndCompute(img2,None)
# FLANN parameters
FLANN_INDEX_KDTREE = 1
index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
search_params = dict(checks=50)
flann = cv.FlannBasedMatcher(index_params,search_params)
matches = flann.knnMatch(des1,des2,k=2)
matchesMask = [[0,0] for i in range(len(matches))]
pts1 = []
pts2 = []
# ratio test as per Lowe's paper
for i,(m,n) in enumerate(matches):
    if m.distance < 0.8*n.distance:
        pts2.append(kp2[m.trainIdx].pt)
        pts1.append(kp1[m.queryIdx].pt)

# Draw the keypoint matches between both pictures
draw_params = dict(matchColor=(0, 255, 0),
    singlePointColor=(255, 0, 0),
    matchesMask=matchesMask[300:500],
    flags=cv.DrawMatchesFlags_DEFAULT)

keypoint_matches = cv.drawMatchesKnn(img1, kp1, img2, kp2, matches[300:500], None, **draw_params)
keypoint_matches = cv.resize(keypoint_matches, (1000,1000))
cv.imshow("Keypoint matches", keypoint_matches)


pts1 = np.int32(pts1)
pts2 = np.int32(pts2)
F, mask = cv.findFundamentalMat(pts1,pts2,cv.FM_LMEDS)
# We select only inlier points
pts1 = pts1[mask.ravel()==1]
pts2 = pts2[mask.ravel()==1]

print(F)

def drawlines(img1,img2,lines,pts1,pts2):
    ''' img1 - image on which we draw the epilines for the points in img2
        lines - corresponding epilines '''
    r,c = img1.shape
    img1 = cv.cvtColor(img1,cv.COLOR_GRAY2BGR)
    img2 = cv.cvtColor(img2,cv.COLOR_GRAY2BGR)
    np.random.seed(0)
    for r,pt1,pt2 in zip(lines,pts1,pts2):
        color = tuple(np.random.randint(0,255,3).tolist())
        x0,y0 = map(int, [0, -r[2]/r[1] ])
        x1,y1 = map(int, [c, -(r[2]+r[0]*c)/r[1] ])
        img1 = cv.line(img1, (x0,y0), (x1,y1), color,1)
        img1 = cv.circle(img1,tuple(pt1),5,color,-1)
        img2 = cv.circle(img2,tuple(pt2),5,color,-1)
    return img1,img2

# Find epilines corresponding to points in right image (second image) and
# drawing its lines on left image
lines1 = cv.computeCorrespondEpilines(pts2.reshape(-1,1,2), 2,F)
lines1 = lines1.reshape(-1,3)
img5,img6 = drawlines(img1,img2,lines1,pts1,pts2)
# Find epilines corresponding to points in left image (first image) and
# drawing its lines on right image
lines2 = cv.computeCorrespondEpilines(pts1.reshape(-1,1,2), 1,F)
lines2 = lines2.reshape(-1,3)
img3,img4 = drawlines(img2,img1,lines2,pts2,pts1)
plt.subplot(121),plt.imshow(img5)
plt.subplot(122),plt.imshow(img3)
plt.show()

h1, w1 = img1.shape
h2, w2 = img2.shape
_, H1, H2 = cv.stereoRectifyUncalibrated(np.float32(pts1), np.float32(pts2), Fundmatrix, imgSize=(w1, h2))

img1_rect = cv.warpPerspective(img1, H1, (w1, h1))
img2_rect = cv.warpPerspective(img2, H2, (w2, h2))
img1_rectg = cv.resize(img1_rect, (1000,1000))
img2_rectg = cv.resize(img2_rect, (1000,1000))
cv.imshow("rectified_1",img1_rectg)
cv.imshow("rectified_2",img2_rectg)
cv.waitKey(0)

# Draw the rectified images
fig, axes = plt.subplots(1, 2, figsize=(15, 10))
axes[0].imshow(img1_rect, cmap="gray")
axes[1].imshow(img2_rect, cmap="gray")
axes[0].axhline(250)
axes[1].axhline(250)
axes[0].axhline(450)
axes[1].axhline(450)
axes[0].axhline(900)
axes[1].axhline(900)
axes[0].axhline(1250)
axes[1].axhline(1250)
axes[0].axhline(1800)
axes[1].axhline(1800)

plt.suptitle("Rectified images")
plt.savefig("rectified_images.png")
plt.show()

imgL = img1_rect        # The rectified left image
imgR = img2_rect        # The rectified right image

# Settings for cv.StereoSGBM_create
minDisparity = 1
numDisparities = 64
blockSize = 1
disp12MaxDiff = 1
uniquenessRatio = 10
speckleWindowSize = 0
speckleRange = 8
stereo = cv.StereoSGBM_create(minDisparity=minDisparity, numDisparities=numDisparities, blockSize=blockSize, disp12MaxDiff=disp12MaxDiff, uniquenessRatio=uniquenessRatio,
    speckleWindowSize=speckleWindowSize, speckleRange=speckleRange)

# Calculate the disparity map
disp = stereo.compute(imgL, imgR).astype(np.float32)
# Normalize the values to spread them across the viewable range
disp = cv.normalize(disp,0,255,cv.NORM_MINMAX)

# Resize for display
disp = cv.resize(disp, (1000,1000))
cv.imshow("disparity",disp)
cv.waitKey(0)