# - * - coding: utf - 8 - *-

import numpy as np
import math as m
from data import feed_date, get_ref_param


def projection_matrix_estimation(image, world):
    'Finding P Matrix using DLT via svd'
    img_pts = np.asarray(image)
    world_pts = np.asarray(world)
    n = world_pts.shape[0]
    if n < 7:
        print("lack of points")
        return
    A = np.zeros((2 * n, 12))
    for i in range(n):
        A[i * 2, 0:4] = -1 * world_pts[i, :]
        A[i * 2, 8:12] = img_pts[i, 0] * world_pts[i, :]
        A[i * 2 + 1, 4:8] = -1 * world_pts[i, :]
        A[i * 2 + 1, 8:12] = img_pts[i, 1] * world_pts[i, :]
    U, D, V = np.linalg.svd(A)
    # print("U={}, D={}, V={}".format(V.shape, D.shape, V.shape))
    P = V[11, :]
    P = np.reshape(P, (3, 4))
    if P[2, 3] != 0:
        P = P / P[2, 3]
    else:
        P = P / 0.001
    return P


def KRT_from_projection(P):
    'QR Decomposition get KRT from P'
    temp = np.linalg.inv(P[0:3, 0:3])
    R, K = np.linalg.qr(temp)
    R = np.linalg.inv(R)
    K = np.linalg.inv(K)
    K = K / K[2, 2]
    # K = abs(K)  # need abs
    T = -1 * np.matmul(temp, P[:, 3])
    return K, R, T


def r_mat(f, w, k):
    '计算R矩阵'
    Rf = np.mat([[m.cos(f), 0, -m.sin(f)],
                 [0, 1, 0],
                 [m.sin(f), 0, m.cos(f)]])
    Rw = np.mat([[1, 0, 0],
                 [0, m.cos(w), -m.sin(w)],
                 [0, m.sin(w), m.cos(w)]])
    Rk = np.mat([[m.cos(k), -m.sin(k), 0],
                 [m.sin(k), m.cos(k), 0],
                 [0, 0, 1]])
    # z y x
    R = Rk * Rf * Rw
    return R


def calc_xy(world, outer, inner, error, R):
    X = [wo[0] for wo in world]
    Y = [wo[1] for wo in world]
    Z = [wo[2] for wo in world]
    Xs, Ys, Zs = outer[0], outer[1], outer[2]
    fx, fy, xo, yo = inner[0], inner[1], inner[2], inner[3]
    k1, k2, p1, p2 = error[0], error[1], error[2], error[3]

    calc_x = [0] * len(X)
    calc_y = [0] * len(Y)

    Ws = np.mat([Xs, Ys, Zs])
    K = np.mat([[fx, 0, xo], [0, fy, yo], [0, 0, 1]])
    for i in range(len(X)):
        # 世界左边 current
        current = np.mat([X[i], Y[i], Z[i]])
        # camera = R * current + t
        xy = R * current.T + Ws.T
        # 归一化
        tmp_x, tmp_y = xy[0]/xy[2], xy[1]/xy[2]
        r2 = tmp_x ** 2 + tmp_y ** 2

        # distorted_x = tmp_x * (1 + k1 * r2 + k2 * r2 ** 2) + 2 * p1 * tmp_x * tmp_y + p2 * (r2 + 2 * tmp_x ** 2)
        # distorted_y = tmp_y * (1 + k1 * r2 + k2 * r2 ** 2) + 2 * p2 * tmp_x * tmp_y + p1 * (r2 + 2 * tmp_y ** 2)
        distorted_x, distorted_y = 0, 0

        calc_x[i] = xo + fx * distorted_x - fx * ((R[0, 0] * (X[i] - Xs) + R[1, 0] * (Y[i] - Ys) + R[2, 0] * (Z[i] - Zs))
                               / (R[0, 2] * (X[i] - Xs) + R[1, 2] * (Y[i] - Ys) + R[2, 2] * (Z[i] - Zs)))
        calc_y[i] = yo + fy * distorted_y - fy * ((R[0, 1] * (X[i] - Xs) + R[1, 1] * (Y[i] - Ys) + R[2, 1] * (Z[i] - Zs))
                               / (R[0, 2] * (X[i] - Xs) + R[1, 2] * (Y[i] - Ys) + R[2, 2] * (Z[i] - Zs)))
    # 返回的是像素
    return calc_x, calc_y


def a_parameter(image_pts, world_pts, outer, inner, error, R, param_num):
    Xs, Ys, Zs = outer[0], outer[1], outer[2]
    fx, fy, xo, yo = inner[0], inner[1], inner[2], inner[3]
    k1, k2, p1, p2 = error[0], error[1], error[2], error[3]
    w, k = outer[4], outer[5]
    x, y = image_pts[0], image_pts[1]

    Ws = np.mat([Xs, Ys, Zs])

    parameter = np.zeros((2, param_num))

    minus = np.array([[world_pts[0] - outer[0]],
                      [world_pts[1] - outer[1]],
                      [world_pts[2] - outer[2]]])
    mean = R.T * np.mat(minus)
    # 6个外方位元素
    parameter[0][0] = (R[0, 0] * fx + R[0, 2] * (x - xo)) / mean[2]
    parameter[0][1] = (R[1, 0] * fx + R[1, 2] * (x - xo)) / mean[2]
    parameter[0][2] = (R[2, 0] * fx + R[2, 2] * (x - xo)) / mean[2]
    parameter[1][0] = (R[0, 1] * fy + R[0, 2] * (y - yo)) / mean[2]
    parameter[1][1] = (R[1, 1] * fy + R[1, 2] * (y - yo)) / mean[2]
    parameter[1][2] = (R[2, 1] * fy + R[2, 2] * (y - yo)) / mean[2]

    parameter[0][3] = (y - yo) * m.sin(w) - (
            ((x - xo) / fx) * ((x - xo) * m.cos(k) - (y - yo) * m.sin(k)) + fx * m.cos(k)) * m.cos(w)
    parameter[0][4] = -fx * m.sin(k) - ((x - xo) / fx) * ((x - xo) * m.sin(k) + (y - yo) * m.cos(k))
    parameter[0][5] = y - yo
    parameter[1][3] = -(x - xo) * m.sin(w) - (
            ((y - yo) / fy) * ((x - xo) * m.cos(k) - (y - yo) * m.sin(k)) - fy * m.cos(k)) * m.cos(w)
    parameter[1][4] = -fy * m.cos(k) - ((y - yo) / fy) * ((x - xo) * m.sin(k) + (y - yo) * m.cos(k))
    parameter[1][5] = -(x - xo)

    current = np.mat([world_pts[0], world_pts[1], world_pts[2]])
    xy = R * current.T + Ws.T
    tmp_x, tmp_y = xy[0] / xy[2], xy[1] / xy[2]
    r2 = tmp_x ** 2 + tmp_y ** 2
    # distorted_x = tmp_x * (1 + k1 * r2 + k2 * r2 ** 2) + 2 * p1 * tmp_x * tmp_y + p2 * (r2 + 2 * tmp_x ** 2)
    # distorted_y = tmp_y * (1 + k1 * r2 + k2 * r2 ** 2) + 2 * p2 * tmp_x * tmp_y + p1 * (r2 + 2 * tmp_y ** 2)
    distorted_x, distorted_y = 0, 0
    # fx, fy, xo, yo  6 + 4
    if param_num > 6:
        parameter[0][6] = (image_pts[0] - xo) / fx + distorted_x
        parameter[0][7] = 0
        parameter[0][8] = 1
        parameter[0][9] = 0

        parameter[1][6] = 0
        parameter[1][7] = (image_pts[1] - yo) / fy + distorted_y
        parameter[1][8] = 0
        parameter[1][9] = 1
    # k1, k2, p1, p2 10 + 4
    if param_num > 10:
        parameter[0][10] = tmp_x * r2
        parameter[0][11] = tmp_x * r2 * r2
        parameter[0][12] = 2 * tmp_x * tmp_y
        parameter[0][13] = r2 + 2 * (tmp_x ** 2)
        parameter[1][10] = tmp_y * r2
        parameter[1][11] = tmp_y * r2 * r2
        parameter[1][12] = r2 + 2 * (tmp_y ** 2)
        parameter[1][13] = 2 * tmp_x * tmp_y
    return parameter


def run(image_pts, world_pts, times=10, t1=1, t2=0.01):
    # 6: 迭代内外元素
    # 10: 内参 14: 误差
    param_num = 14
    n = len(image_pts)

    image, world = [], []
    # 齐次坐标
    for i in range(n):
        tmp1 = image_pts[i]
        tmp2 = world_pts[i]
        tmp1.append(1)
        tmp2.append(1)
        image.append(tmp1)
        world.append(tmp2)

    P = projection_matrix_estimation(image, world)
    K, R, T = KRT_from_projection(P)
    # 参数赋初值
    omiga = kaba = fai = 0
    Xos, Yos, Zos = T[0], T[1], T[2]
    fx, fy, xo, yo = K[0][0], K[1][1], K[0][2], K[1][2]
    print("初始内参: fx=%f, fy=%f, xo=%f, yo=%f" % (fx, fy, xo, yo))
    k1 = k2 = p1 = p2 = 0.00001

    # 误差 and A矩阵
    L = np.mat(np.zeros((2 * n, 1)))
    A = np.mat(np.zeros((2 * n, param_num)))

    d_fx = d_fy = d_xo = d_yo = 100
    dx = dy = dz = do = dk = df = 100
    epoch = 0
    while (abs(dx) > t1) | (abs(dy) > t1) | (abs(dz) > t1) | (abs(do) > t2) | (abs(dk) > t2) | (abs(df) > t2):
        error = [k1, k2, p1, p2]
        outer, inner = [Xos, Yos, Zos, fai, omiga, kaba], [fx, fy, xo, yo]
        R = r_mat(fai, omiga, kaba)
        calc_x, calc_y = calc_xy(world_pts, outer, inner, error, R)
        for i in range(n):
            # 像素点误差  用像素误差太大 一定要用mm
            L[2 * i] = image_pts[i][0] - calc_x[i]
            L[2 * i + 1] = image_pts[i][1] - calc_y[i]

        for i in range(n):
            A[2 * i:2 * i + 2, :] = a_parameter(image_pts[i], world_pts[i], outer, inner, error, R, param_num)

        # 计算改正数
        X = (A.T * A).I * A.T * L

        Xos, Yos, Zos, fai, omiga, kaba = Xos + X[0, 0], Yos + X[1, 0], Zos + X[2, 0], \
                                          fai + X[3, 0], omiga + X[4, 0], kaba + X[5, 0]

        if param_num > 6:
            fx, fy, xo, yo = fx + X[6, 0], fy + X[7, 0], xo + X[8, 0], yo + X[9, 0]
            d_fx, d_fy, d_xo, d_yo = X[6, 0], X[7, 0], X[8, 0], X[9, 0]

        if param_num > 10:
            k1, k2, p1, p2 = k1 + X[10, 0], k2 + X[11, 0], p1 + X[12, 0], p2 + X[13, 0]

        # 变成角度
        for i in range(3, 6):
            X[i, 0] = X[i, 0] * 180 / m.pi
        dx, dy, dz, do, dk, df = X[0, 0], X[1, 0], X[2, 0], X[3, 0], X[4, 0], X[5, 0]

        epoch += 1
        if epoch <= times:
            # print(X[10:14, :])
            print("第 %d 次迭代:Xs=%f,Ys=%f,Zs=%f,fai=%f,omiga=%f,kaba=%f,fx=%f,fy=%f,xo=%f,yo=%f" %
                  (epoch, Xos, Yos, Zos, fai * 180 / m.pi, omiga*180/m.pi, kaba*180/m.pi, fx, fy, xo, yo))
        else:
            print("不收敛")
            break
    if epoch < times:
        print("迭代收敛")
        print("结果为:Xs=%f,Ys=%f,Zs=%f,fai=%f,omiga=%f,kaba=%f,fx=%f,fy=%f,xo=%f,yo=%f" %
              (Xos, Yos, Zos, fai, omiga, kaba, fx, fy, xo, yo))
        print("结果为:k1=%f,k2=%f,p1=%f,p2=%f" % (k1, k2, p1, p2))
        # return [Xos, Yos, Zos, fai * 180 / m.pi, omiga*180/m.pi, kaba*180/m.pi], [fx, fy, xo, yo]


if __name__ == '__main__':
    # 不使用科学记数法
    np.set_printoptions(suppress=True)
    image_pts, world_pts = feed_date()
    run(image_pts, world_pts, times=10, t1=1, t2=0.05)
    print("-------")
    # get_ref_param()