import os
from glob import glob
from os.path import join
import numpy as np
import cv2 as cv
import json

def read_json(input):
    with open(input, "r") as f:
        data = json.load(f)
    return data

def solvePnP(k3d, k2d, K, dist, flag, tryextri=False):
    k2d = np.ascontiguousarray(k2d[:, :2]) # 保留前两列
    # try different initial values:
    if tryextri: # 尝试不同的初始化外参
        def closure(rvec, tvec):
            ret, rvec, tvec = cv.solvePnP(k3d, k2d, K, dist, rvec, tvec, True, flags=flag)
            points2d_repro, xxx = cv.projectPoints(k3d, rvec, tvec, K, dist)
            kpts_repro = points2d_repro.squeeze()
            err = np.linalg.norm(points2d_repro.squeeze() - k2d, axis=1).mean()
            return err, rvec, tvec, kpts_repro

        # create a series of extrinsic parameters looking at the origin
        height_guess = 2.7 # 相机的初始高度猜测
        radius_guess = 4. # 相机的初始水平距离猜测，圆的半径，需要根据自己的实际情况调整
        infos = []
        for theta in np.linspace(0, 2 * np.pi, 180):
            st = np.sin(theta)
            ct = np.cos(theta)
            center = np.array([radius_guess * ct, radius_guess * st, height_guess]).reshape(3, 1)
            R = np.array([
                [-st, ct, 0],
                [0, 0, -1],
                [-ct, -st, 0]
            ])
            tvec = - R @ center
            rvec = cv.Rodrigues(R)[0]
            err, rvec, tvec, kpts_repro = closure(rvec, tvec)
            infos.append({
                'err': err,
                'repro': kpts_repro,
                'rvec': rvec,
                'tvec': tvec
            })
        infos.sort(key=lambda x: x['err'])
        err, rvec, tvec, kpts_repro = infos[0]['err'], infos[0]['rvec'], infos[0]['tvec'], infos[0]['repro']
    else:
        # 直接求解的初值是零向量
        ret, rvec, tvec = cv.solvePnP(k3d, k2d, K, dist, flags=flag)
        points2d_repro, xxx = cv.projectPoints(k3d, rvec, tvec, K, dist)
        kpts_repro = points2d_repro.squeeze()
        err = np.linalg.norm(points2d_repro.squeeze() - k2d, axis=1).mean()
    # print(err)
    return err, rvec, tvec, kpts_repro

# 对单个相机进行外参标定
def _calibrate_extri(k3d, k2d, K, dist, flag, tryfocal=False):
    extri = {}
    methods = [cv.SOLVEPNP_ITERATIVE]
    # 检查关键点数据的数量是否匹配
    if k3d.shape[0] != k2d.shape[0]:
        print('k3d {} doesnot match k2d {}'.format(k3d.shape, k2d.shape))
        length = min(k3d.shape[0], k2d.shape[0])
        k3d = k3d[:length]
        k2d = k2d[:length]
    valididx = k2d[:, 2] > 0 # k2d第三列是置信度，检查是否大于0
    if valididx.sum() < 4: # 筛选出有效的2D和3D关键点，数量大于4
        rvec = np.zeros((1, 3)) # 初始话旋转和平移为0并标记为失败
        tvec = np.zeros((3, 1))
        extri['Rvec'] = rvec
        extri['R'] = cv.Rodrigues(rvec)[0]
        extri['T'] = tvec
        print('[ERROR] Failed to initialize the extrinsic parameters')
        return extri
    k3d = k3d[valididx]
    k2d = k2d[valididx]
    # 优化相机焦距
    # 如果启用焦距优化
    if tryfocal:
        infos = []
        for focal in range(500, 5000, 10):  # 遍历焦距范围
            # 设置焦距值
            K[0, 0] = focal  # 更新 K 的 fx
            K[1, 1] = focal  # 更新 K 的 fy
            for method in methods:
                # 调用 solvePnP
                err, rvec, tvec, kpts_repro = solvePnP(k3d, k2d, K, dist, method)
                # 保存结果
                infos.append({
                    'focal': focal,
                    'err': err,
                    'rvec': rvec,
                    'tvec': tvec,
                    'repro': kpts_repro
                })
        # 根据重投影误差选择最佳焦距
        infos.sort(key=lambda x: x['err'])
        best_result = infos[0]
        focal = best_result['focal']
        err, rvec, tvec, kpts_repro = best_result['err'], best_result['rvec'], best_result['tvec'], best_result['repro']
        # 更新内参中的焦距
        K[0, 0] = focal
        K[1, 1] = focal
        print(f'[INFO] Optimal focal length found: {focal}, reprojection error: {err:.3f}')
    else:
        # 如果不优化焦距，直接调用 solvePnP
        err, rvec, tvec, kpts_repro = solvePnP(k3d, k2d, K, dist, flag)

    # 保存外参结果
    extri['Rvec'] = rvec
    extri['R'] = cv.Rodrigues(rvec)[0]
    extri['T'] = tvec
    center = - extri['R'].T @ tvec
    print(f'[INFO] Camera center: {center.squeeze()}, reprojection error: {err:.3f}')
    return extri

def calibrate_extri(kpts_path, intri_path, flag, tryfocal=False, tryextri=False):
    extri = {}
    intri_data = read_json(intri_path)
    kpts_data = read_json(kpts_path)
    # 获取内参
    camnames = list(intri_data.keys())
    for cam in camnames:
        print(f'[INFO] Processing camera: {cam}')
        K = np.array(intri_data[cam]['K'])
        dist = np.array(intri_data[cam]['dist'])
        k3d = np.array(kpts_data[cam]['keypoints3d'])
        k2d = np.array(kpts_data[cam]['keypoints2d'])

        extri[cam] = _calibrate_extri(k3d, k2d, K, dist, flag, tryfocal=tryfocal)

    return extri


if __name__ == "__main__":
    pass