import numpy as np
import time
import cv2
import cv2.aruco as aruco

# with np.load('webcam_calibration_output.npz') as X:
#    mtx, dist, _, _ = [X[i] for i in ('mtx','dist','rvecs','tvecs')]

# mtx =
# 2946.48    0    1980.53
# 0    2945.41    1129.25
# 0    0    1

mtx = np.array([
    [2946.48, 0, 1980.53],
    [0, 2945.41, 1129.25],
    [0, 0, 1],
])
# 我的手机拍棋盘的时候图片大小是 4000 x 2250
# ip摄像头拍视频的时候设置的是 1920 x 1080，长宽比是一样的，
# ip摄像头设置分辨率的时候注意一下


dist = np.array([0.226317, -1.21478, 0.00170689, -0.000334551, 1.9892])

#video = "http://admin:admin@192.168.1.2:8081/"  # 手机ip摄像头
# 根据ip摄像头在你手机上生成的ip地址更改，右上角可修改图像分辨率

cap = cv2.VideoCapture(0)
cap.set(3, 320)  # 设置摄像头输出宽
cap.set(4, 240)  # 设置摄像头输出高

font = cv2.FONT_HERSHEY_SIMPLEX  # font for displaying text (below)

# num = 0
while True:
    ret, frame = cap.read()
    # operations on the frame come here

    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    aruco_dict = aruco.Dictionary_get(aruco.DICT_6X6_250)
    parameters = aruco.DetectorParameters_create()

    '''
 detectMarkers(...) 
     detectMarkers(image, dictionary[, corners[, ids[, parameters[, rejectedI 
     mgPoints]]]]) -> corners, ids, rejectedImgPoints 
 '''

    # lists of ids and the corners beloning to each id
    corners, ids, rejectedImgPoints = aruco.detectMarkers(gray,
                                                          aruco_dict,
                                                          parameters=parameters)

    #    if ids != None:
    if ids is not None:

        rvec, tvec, _ = aruco.estimatePoseSingleMarkers(corners, 0.05, mtx, dist)
        # Estimate pose of each marker and return the values rvet and tvec---different
        # from camera coeficcients
        (rvec - tvec).any()  # get rid of that nasty numpy value array error
        numbers = ids.shape[0]
        send = np.zeros([numbers, 3], dtype=float)
        for t in range(numbers):
            a = corners[t]
            x1 = a[0][1][0]
            y1 = a[0][1][1]
            x2 = a[0][2][0]
            y2 = a[0][2][1]
            x3 = a[0][3][0]
            y3 = a[0][3][1]
            x4 = a[0][0][0]
            y4 = a[0][0][1]
            try:
                k1 = (y3 - y1) / (x3 - x1)
                b1 = (x3 * y1 - x1 * y3) / (x3 - x1)
                k2 = (y4 - y2) / (x4 - x2)
                b2 = (x4 * y2 - x2 * y4) / (x4 - x2)
                x5 = -(b1 - b2) / (k1 - k2)
                y5 = k1 * x5 + b1
                send[t][0] = ids[t][0]
                send[t][1] = int(x5)
                send[t][2] = int(y5)
            except:
                print("error=1")
            print("ids[" + str(t) + "][0]=" + str(ids[t][0]))
            print(send[t][1])
            print(send[t][2])


        #        aruco.drawAxis(frame, mtx, dist, rvec, tvec, 0.1) #Draw Axis
        #        aruco.drawDetectedMarkers(frame, corners) #Draw A square around the markers

        for i in range(rvec.shape[0]):
            aruco.drawAxis(frame, mtx, dist, rvec[i, :, :], tvec[i, :, :], 0.03)
            aruco.drawDetectedMarkers(frame, corners)
        ###### DRAW ID #####
    #        cv2.putText(frame, "Id: " + str(ids), (0,64), font, 1, (0,255,0),2,cv2.LINE_AA)

    else:
        ##### DRAW "NO IDS" #####
        cv2.putText(frame, "No Ids", (0, 64), font, 1, (0, 255, 0), 2, cv2.LINE_AA)

        # Display the resulting frame
    cv2.imshow("frame", frame)

    key = cv2.waitKey(1)

    if key == 27:  # 按esc键退出
        print('esc break...')
        cap.release()
        cv2.destroyAllWindows()
        break

    if key == ord(' '):  # 按空格键保存
        #        num = num + 1
        #        filename = "frames_%s.jpg" % num  # 保存一张图像
        filename = str(time.time())[:10] + ".jpg"
        cv2.imwrite(filename, frame)