import os
import shutil
import sys
import tkinter
from tkinter import filedialog
import cv2
import numpy as np
import numpy.matlib
import time
from open3d.cpu.pybind.geometry import PointCloud
from open3d.cpu.pybind.utility import Vector3dVector
from open3d.cpu.pybind.visualization import draw_geometries
from open3d.cpu.pybind.io import write_point_cloud
from Caliborate_helper import Calibrator


tkinter.Tk().withdraw()
def preprocess(img1, img2):
    if (img1.ndim == 3):
        img1 = cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY)
    if (img2.ndim == 3):
        img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
    img1 = cv2.equalizeHist(img1)
    img2 = cv2.equalizeHist(img2)
    return img1, img2
def stereo_calibration(x_nums,y_nums,unit_size,left_photo_path,right_photo_path,mtx_left,dist_left,mtx_right,dist_right):
    world_position = []
    left_image_position = []
    right_image_position =[]
    objp=np.zeros((x_nums*y_nums,3),np.float32)
    objp[:,:2]=np.mgrid[:x_nums,:y_nums].T.reshape(-1, 2)
    objp*=unit_size
    world_position.append(objp)
    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,30,0.001)
    left_image = cv2.imread(left_photo_path)
    left_gray_image = cv2.cvtColor(left_image,cv2.COLOR_RGB2GRAY)
    right_image = cv2.imread(right_photo_path)
    right_gray_image = cv2.cvtColor(right_image,cv2.COLOR_RGB2GRAY)
    _,left_corners = cv2.findChessboardCorners(left_gray_image,(x_nums,y_nums),None)
    left_exact_corners=cv2.cornerSubPix(left_gray_image,left_corners,(11,11),(-1,-1),criteria)
    _,right_corners = cv2.findChessboardCorners(left_gray_image,(x_nums,y_nums),None)
    right_exact_corners=cv2.cornerSubPix(right_gray_image,right_corners,(11,11),(-1,-1),criteria)
    left_image_position.append(left_exact_corners)
    right_image_position.append(right_exact_corners)
    flags=0
    flags|=cv2.CALIB_FIX_INTRINSIC
    flags|=cv2.CALIB_USE_INTRINSIC_GUESS
    img_shape=left_gray_image.shape[::-1]
    ret,_,_,_,_,R,T,E,F=cv2.stereoCalibrate(world_position,left_image_position,right_image_position,mtx_left,dist_left,mtx_right,dist_right,flags=flags,imageSize=img_shape)
    return R,T
def undistortion(image, camera_matrix, dist_coeff):
    undistortion_image = cv2.undistort(image, camera_matrix, dist_coeff)
    return undistortion_image
def stereo(left, right, left_rectified, right_rectified,height, width):
    global file_path
    file_path3 = filedialog.askdirectory(title=u'选择左图文件夹')
    img_dir1 = file_path3
    file_path4 = filedialog.askdirectory(title=u'选择右图文件夹')
    img_dir2 = file_path4
    img_dir11 = img_dir1+"\\1left.png"
    img_dir22 = img_dir2+"\\1right.png"
    shape_inner_corner = (9,6)
    size_grid = 0.025
    calibrator1 =Calibrator(img_dir1, shape_inner_corner, size_grid)
    mtxL, distL = calibrator1.calibrate_camera()
    calibrator2 =Calibrator(img_dir2, shape_inner_corner, size_grid)
    mtxR, distR = calibrator2.calibrate_camera()
    R,T=stereo_calibration(9, 6,25, img_dir11,img_dir22,mtxL,distL,mtxR,distR)
    with open(r"StereoCamera_1.txt", 'w', encoding="utf8") as f:
        f.write('mtxL:'+str(mtxL)+'\n \n ')
        f.write('distL:'+str(distL)+'\n \n ')
        f.write('mtxR:'+str(mtxR)+'\n \n ')
        f.write('distR:'+str(distR)+'\n \n ')
        f.write('旋转矩阵:'+str(R)+'\n \n ')
        f.write('平移矩阵:'+str(T)+'\n \n ')
        f.write('基线距离:'+str(T[0])+'\n \n ')
    print(mtxL)
    print(distL)
    print(mtxR)
    print(distR)
    print(R)
    print(T)
    print(T[0])
    R1, R2, P1, P2, Q, roi1, roi2 = cv2.stereoRectify(mtxL, distL, mtxR, distR,
                                                      (width, height), R, T, alpha=0)
    map1x, map1y = cv2.initUndistortRectifyMap(mtxL, distL, R1, P1, (width, height), cv2.CV_32FC1)
    map2x, map2y = cv2.initUndistortRectifyMap(mtxR, distR, R2, P2, (width, height), cv2.CV_32FC1)
    iml_1 = cv2.imread(left, 1)
    imr_1 = cv2.imread(right, 1)
    rectifyed_img1 = cv2.remap(iml_1, map1x, map1y, cv2.INTER_AREA)
    rectifyed_img2 = cv2.remap(imr_1, map2x, map2y, cv2.INTER_AREA)
    cv2.imwrite(left_rectified,rectifyed_img1)
    cv2.imwrite(right_rectified,rectifyed_img2)
    return Q
def draw_line(image1, image2):
    height = max(image1.shape[0], image2.shape[0])
    width = image1.shape[1] + image2.shape[1]
    output = np.zeros((height, width, 3), dtype=np.uint8)
    output[0:image1.shape[0], 0:image1.shape[1]] = image1
    output[0:image2.shape[0], image1.shape[1]:] = image2
    line_interval = 50
    for k in range(height // line_interval):
        cv2.line(output, (0, line_interval * (k + 1)), (2 * width, line_interval * (k + 1)), (0, 255, 0), thickness=2,
                 lineType=cv2.LINE_AA)
    return output
def stereoMatchSGBM(left_image, right_image, down_scale=True):
    if left_image.ndim == 2:
        img_channels = 1
    else:
        img_channels = 3
    blockSize = 3
    paraml = {'minDisparity': 6,
              'numDisparities': 80,
              'blockSize': blockSize,
              'P1': 8 * img_channels * blockSize ** 2,
              'P2': 32 * img_channels * blockSize ** 2,
              'disp12MaxDiff': 1,
              'preFilterCap': 63,
              'uniquenessRatio': 15,
              'speckleWindowSize': 100,
              'speckleRange': 1,
              'mode': cv2.STEREO_SGBM_MODE_HH
              }
    left_matcher = cv2.StereoSGBM_create(**paraml)
    paramr = paraml
    paramr['minDisparity'] = -paraml['numDisparities']
    right_matcher = cv2.StereoSGBM_create(**paramr)
    size = (left_image.shape[1], left_image.shape[0])
    if down_scale == False:
        disparity_left = left_matcher.compute(left_image, right_image)
        disparity_right = right_matcher.compute(right_image, left_image)
    else:
        left_image_down = cv2.pyrDown(left_image)
        right_image_down = cv2.pyrDown(right_image)
        factor = left_image.shape[1] / left_image_down.shape[1]

        disparity_left_half = left_matcher.compute(left_image_down, right_image_down)
        disparity_right_half = right_matcher.compute(right_image_down, left_image_down)
        disparity_left = cv2.resize(disparity_left_half, size, interpolation=cv2.INTER_AREA)
        disparity_right = cv2.resize(disparity_right_half, size, interpolation=cv2.INTER_AREA)
        disparity_left = factor * disparity_left
        disparity_right = factor * disparity_right
    trueDisp_left = disparity_left.astype(np.float32) / 16.
    trueDisp_right = disparity_right.astype(np.float32) / 16.
    return trueDisp_left, trueDisp_right
def shengchengopen(left_image,disparity_image,pointcloud):
    scale_factor = 1
    img_left = cv2.imread(left_image)
    img_disparity = cv2.imread(disparity_image, cv2.IMREAD_GRAYSCALE) / scale_factor
    h, w = img_disparity.shape
    xyz_points = np.zeros((h * w, 6))
    Y = np.matlib.repmat(np.array(range(h)).reshape(-1, 1), 1, w).flatten().astype(int)
    X = np.matlib.repmat(np.array(range(w)), h, 1).flatten().astype(int)
    Z = img_disparity.flatten().astype(int)
    B = img_left[:, :, 0].flatten() / 255
    G = img_left[:, :, 1].flatten() / 255
    R = img_left[:, :, 2].flatten() / 255
    xyz_points[:, 0] = X
    xyz_points[:, 1] = Y
    xyz_points[:, 2] = Z
    xyz_points[:, 3] = R
    xyz_points[:, 4] = G
    xyz_points[:, 5] = B
    pcd = PointCloud()
    pcd.points = Vector3dVector(xyz_points[:, :3])
    pcd.colors = Vector3dVector(xyz_points[:, 3:])
    pcd.transform([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    draw_geometries([pcd])
    write_point_cloud(pointcloud, pcd)
if __name__ == '__main__':
    file_path1 = filedialog.askopenfilename(title=u'选择左图')
    iml = file_path1
    file_path2 = filedialog.askopenfilename(title=u'选择右图')
    imr = file_path2
    iml_1=cv2.imread(iml,1)
    imr_1=cv2.imread(imr,1)
    height, width = iml_1.shape[0:2]
    if (iml is None) or (imr is None):
        print("Error: Images are empty, please check your image's path!")
        sys.exit(0)
    iml_rectified = './double_camera/1_test_imagesleft_rec.jpg'
    imr_rectified = './double_camera/1_test_imagesright_rec.jpg'
    Q=stereo(iml,imr,iml_rectified,imr_rectified,height, width)
    iml_rectified_1=cv2.imread(iml_rectified,1)
    imr_rectified_1=cv2.imread(imr_rectified,1)
    line = draw_line(iml_rectified_1, imr_rectified_1)
    cv2.imwrite('./double_camera/1_check_rectification4.png', line)
    o_path = "./double_camera/1_1.jpg"
    iml_, imr_ = preprocess(iml_rectified_1,imr_rectified_1)
    disp,_=stereoMatchSGBM(iml_rectified_1, imr_rectified_1, True)
    cv2.imwrite(o_path, disp * 4)
    points_3d = cv2.reprojectImageTo3D(disp, Q)
    def onMouse(event, x, y, flags, param):
        if event == cv2.EVENT_LBUTTONDOWN:
            print('点 (%d, %d) 的三维坐标 (%f, %f, %f)' % (x, y, points_3d[y, x, 0], points_3d[y, x, 1], points_3d[y, x, 2]))
            dis = ((points_3d[y, x, 0] ** 2 + points_3d[y, x, 1] ** 2 + points_3d[y, x, 2] ** 2) ** 0.5) / 1000
            print('点 (%d, %d) 距离左摄像头的相对距离为 %0.3f m' % (x, y, dis))
            output_path = 'out.txt'
            with open(output_path, 'w', encoding='utf-8') as file1:
                print('点 (%d, %d) 的三维坐标 (%f, %f, %f)' % (
                x, y, points_3d[y, x, 0], points_3d[y, x, 1], points_3d[y, x, 2]), file=file1)
                print('点 (%d, %d) 距离左摄像头的相对距离为 %0.3f m' % (x, y, dis), file=file1)
    disp = cv2.imread(o_path, 1)
    cv2.namedWindow("disparity", 0)
    cv2.resizeWindow("disparity", 640, 480)
    cv2.imshow("disparity", disp)
    cv2.setMouseCallback("disparity", onMouse, 0)
    cv2.waitKey(0)
    shengchengopen(iml, o_path, './double_camera/1_16.ply')
    shutil.copy(os.path.join('./double_camera/1_1.jpg'), os.path.join('./double_camera/4_1.jpg'))
    shutil.copy(os.path.join('./double_camera/1_check_rectification4.png'), os.path.join('./double_camera/4_check_rectification4.png'))
    shutil.copy(os.path.join('./double_camera/1_16.ply'), os.path.join('./double_camera/4_16.ply'))