import os
import os.path as osp
import glob
import cv2 as cv
import numpy as np
import json
import argparse


def calibrate_camera(imgFolder, chessboardSize, squareSize):
    # 设置输出目录
    outputFolder = osp.join(imgFolder, "output")
    if not osp.exists(outputFolder):
        os.makedirs(outputFolder)

    # 图片路径
    imgPaths = []
    for extension in ["jpg", "png", "jpeg"]:
        imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
    if len(imgPaths) == 0:
        print("No images found!")
        return

    # 存储世界坐标和像素坐标
    # 计算出棋盘格中每个网格角点的坐标，之后当成世界坐标
    board_w, board_h = chessboardSize
    board_grid = np.zeros((board_w * board_h, 3), np.float32)
    board_grid[:, :2] = np.mgrid[0:board_w, 0:board_h].T.reshape(-1, 2) * squareSize

    pointsWorld = []
    pointsPixel = []

    # 遍历图片
    for imgPath in imgPaths:
        img = cv.imread(imgPath)
        gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)

        # 查找角点
        ret, corners = cv.findChessboardCorners(gray, (board_w, board_h), None)
        if ret:
            cv.cornerSubPix(gray, corners, (11, 11), (-1, -1),
                            (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 30, 0.001))
            pointsWorld.append(board_grid)
            pointsPixel.append(corners)

            cv.drawChessboardCorners(img, (board_w, board_h), corners, ret)
            cv.imwrite(osp.join(outputFolder, osp.basename(imgPath)), img)

    # 标定相机
    ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld, pointsPixel, gray.shape[::-1], None, None)
    print("Intrinsic matrix:\n", mtx.astype(np.float32))
    print("Distortion coefficients:\n", dist.astype(np.float32))

    # 计算重投影误差
    nimg = len(pointsWorld)
    img_error = np.zeros(nimg)
    for i in range(nimg):
        imgpoints2, _ = cv.projectPoints(pointsWorld[i], rvecs[i], tvecs[i], mtx, dist)
        error = cv.norm(pointsPixel[i], imgpoints2, cv.NORM_L2) / len(imgpoints2)
        img_error[i] = error

    good_img = np.where(img_error < 0.5)[0]
    mean_error = np.mean(img_error[good_img])
    print("Reprojection error: ", mean_error)

    # 挑选出重投影误差小于1.0的图片，重新标定相机
    if len(good_img) == 0:
        print("No images with error < 0.5")
    elif len(good_img) == nimg:
        print("All images have error < 0.5")
        pass
    else:
        pointsWorld2 = [pointsWorld[i] for i in good_img]
        pointsPixel2 = [pointsPixel[i] for i in good_img]
        ret, mtx, dist, rvecs, tvecs = cv.calibrateCamera(pointsWorld2, pointsPixel2, gray.shape[::-1], None, None)
        print("Intrinsic matrix:\n", mtx.astype(np.float32))
        print("Distortion coefficients:\n", dist.astype(np.float32))

    return mtx, dist


def undistort_image(img, mtx, dist):
    h, w = img.shape[:2]
    newcameramtx, roi = cv.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
    dst = cv.undistort(img, mtx, dist, None, newcameramtx)
    x, y, w, h = roi
    dst = dst[y:y + h, x:x + w]
    return dst


def undistort_images(imgFolder, mtx, dist):
    imgPaths = []
    for extension in ["jpg", "png", "jpeg"]:
        imgPaths += glob.glob(osp.join(imgFolder, "*.{}".format(extension)))
    if len(imgPaths) == 0:
        print("No images found!")
        return

    outputFolder = osp.join(imgFolder, "undistorted_images")
    if not osp.exists(outputFolder):
        os.makedirs(outputFolder)

    for imgPath in imgPaths:
        img = cv.imread(imgPath)
        dst = undistort_image(img, mtx, dist)
        cv.imwrite(osp.join(outputFolder, osp.basename(imgPath)), dst)

    print("Undistorted images saved to: ", outputFolder)


def calibrate_cameras(imgFolder, chessboardSize, squareSize):
    mtxs = []
    dists = []
    for folder in glob.glob(osp.join(imgFolder, "*")):
        if not osp.isdir(folder):
            continue
        mtx, dist = calibrate_camera(folder, chessboardSize, squareSize)
        mtxs.append(mtx)
        dists.append(dist)
    return mtxs, dists


def write_json_data(mtx, dist, outputFolder):
    data = {
        "intrinsic_matrix": mtx.tolist(),
        "distortion_coefficients": dist.tolist()
    }
    with open(osp.join(outputFolder, "calibration.json"), "w") as f:
        json.dump(data, f, indent=4)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("imgFolder", help="Folder containing images for calibration")
    parser.add_argument("--chessboardSize", help="Size of chessboard (rows, cols)", default="11,8")
    parser.add_argument("--squareSize", help="Size of square in chessboard", default=60)
    args = parser.parse_args()

    chessboardSize = tuple(map(int, args.chessboardSize.split(",")))
    squareSize = float(args.squareSize)