import numpy as np

def inv_form(R, t):
    """
    计算逆变换矩阵和平移向量
    """
    R1 = np.linalg.inv(R)
    t1 = -(R1 @ t.reshape(-1,1)).flatten()
    return R1, t1

def warp_point(point, R, t):
    """
    点坐标变换
    """
    point = np.append(point, 1)
    tmp = R @ point
    tmp = tmp / tmp[2]
    new_point = tmp + t
    return new_point[:2]

def get_M_G2I(dH, fYaw, fPitch, fRoll, fAz, fPt, imgW, imgH, fd):
    """
    获取从地面坐标系到图像坐标系的变换矩阵
    """
    # 深度矩阵
    M_het = np.array([
        [1, 0, 0],
        [0, 1, 0],
        [0, 0, dH]
    ])

    # 姿态矩阵
    M_yaw = np.array([
        [np.cos(np.deg2rad(fYaw)), -np.sin(np.deg2rad(fYaw)), 0],
        [np.sin(np.deg2rad(fYaw)), np.cos(np.deg2rad(fYaw)), 0],
        [0, 0, 1]
    ])

    M_pitch = np.array([
        [1, 0, 0],
        [0, np.cos(np.deg2rad(fPitch)), -np.sin(np.deg2rad(fPitch))],
        [0, np.sin(np.deg2rad(fPitch)), np.cos(np.deg2rad(fPitch))]
    ])

    M_roll = np.array([
        [np.cos(np.deg2rad(fRoll)), 0, np.sin(np.deg2rad(fRoll))],
        [0, 1, 0],
        [-np.sin(np.deg2rad(fRoll)), 0, np.cos(np.deg2rad(fRoll))]
    ])

    # 伺服矩阵
    M_beta = np.array([
        [np.cos(np.deg2rad(fAz)), -np.sin(np.deg2rad(fAz)), 0],
        [np.sin(np.deg2rad(fAz)), np.cos(np.deg2rad(fAz)), 0],
        [0, 0, 1]
    ])

    M_alaph = np.array([
        [1, 0, 0],
        [0, np.cos(np.deg2rad(fPt)), -np.sin(np.deg2rad(fPt))],
        [0, np.sin(np.deg2rad(fPt)), np.cos(np.deg2rad(fPt))]
    ])

    # 内参矩阵
    M_cam = np.array([
        [fd, 0, imgW/2],
        [0, -fd, imgH/2],
        [0, 0, 1]
    ])

    # 计算最终变换矩阵
    M = M_cam @ M_alaph @ M_beta @ M_roll @ M_pitch @ M_yaw @ M_het
    return M

def main():
    # 相机信息
    imgW = 1280
    imgH = 1024
    f = 48
    dsize = 7.5
    fd = f/dsize*1000

    # 吊舱参数
    fAz = 44.1876869
    fPt = 18.3241043

    # 姿态参数
    fYaw = 165.08250427246094
    fPitch = 4.2472858428955078
    fRoll = -0.37968909740447998

    # 深度
    dH = 2920

    # 计算变换矩阵
    M_G2I = get_M_G2I(dH, fYaw, fPitch, fRoll, fAz, fPt, imgW, imgH, fd)
    M_I2G = np.linalg.inv(M_G2I)

    t = np.array([200, 1000, 1])
    point = np.array([640, 512])
    
    # 计算变换点
    npoint = warp_point(point, M_I2G, t)
    R1, t1 = inv_form(M_I2G, t)
    pxpoint = warp_point(npoint, R1, -t)

    # 绘制视场范围
    import matplotlib.pyplot as plt
    plt.figure()
    
    for fAz in range(44, 61, 5):
        for fPt in range(18, 46, 10):
            M_G2I = get_M_G2I(dH, fYaw, fPitch, fRoll, fAz, fPt, imgW, imgH, fd)
            M_I2G = np.linalg.inv(M_G2I)

            # 计算四个角点
            points = np.array([[0,0,1], [640,0,1], [640,512,1], [0,512,1]])
            transformed_points = []
            
            for p in points:
                p_transformed = M_I2G @ p
                p_transformed = p_transformed / p_transformed[2]
                transformed_points.append(p_transformed)

            # 添加第一个点以闭合多边形
            transformed_points.append(transformed_points[0])
            
            # 转换为numpy数组以便绘图
            points_array = np.array(transformed_points)
            plt.plot(points_array[:,0], points_array[:,1])
            
    plt.gca().set_aspect('equal')
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    main()