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

import numpy as np
import math as m


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)

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

    for i in range(len(X)):
        world_pts = np.mat([X[i], Y[i], Z[i]])
        # camera = R * world_pts + t mm
        xy = R * world_pts.T + W.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
        # aloha, bate <=> dx, dy  查相机参数把
        # 单位 pixel / mm
        alpha, bate = 3000 / 50, 3000 / 50

        calc_x[i] = xo + alpha * 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 + bate * 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, R, param_num):
    # 只计算外方位元素的偏导
    fx, fy, xo, yo = inner[0], inner[1], inner[2], inner[3]
    w, k = outer[4], outer[5]
    x, y = image_pts[0], image_pts[1]

    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)

    return parameter


def intersection(image_pts, world_pts, inner_element, distort_element, init_value, times=10, thresh=0.1):
    # 6: 迭代内外元素
    param_num = 6
    n = len(image_pts)
    fai, omiga, kaba = init_value[3], init_value[4], init_value[5]
    Xos, Yos, Zos = init_value[0], init_value[1], init_value[2]
    fx, fy, xo, yo = inner_element[0], inner_element[1], inner_element[2], inner_element[3]
    k1, k2, p1, p2 = distort_element[0], distort_element[1], distort_element[2], distort_element[3]

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

    dx = dy = dz = do = dk = df = 100
    epoch = 0
    while (abs(do) > thresh) | (abs(dk) > thresh) | (abs(df) > thresh) | (abs(dx) > thresh) | (abs(dy) > thresh) | (abs(dz) > thresh):
        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, 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]

        # 变成角度
        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("第 %d 次迭代:Xs=%f,Ys=%f,Zs=%f,fai=%f,omiga=%f,kaba=%f" %
                  (epoch, Xos, Yos, Zos, fai * 180 / m.pi, omiga*180/m.pi, kaba*180/m.pi))
        else:
            print("后方交汇不收敛")
            break
    if epoch < times:
        print("后方交汇迭代收敛")
    return [Xos, Yos, Zos, fai * 180 / m.pi, omiga*180/m.pi, kaba*180/m.pi]


def usage():
    print("Usage:")
    print("from posture import intersection")
    print("intersection()")


if __name__ == '__main__':
    # 不使用科学记数法
    np.set_printoptions(suppress=True)
    print("-------")
    usage()
    print("-------")