import numpy as np
import cv2
import matplotlib.pyplot as plt

# ego_cv: x+ is to right, y+ downward, z+ forward
# ego_original: x+ forward, y+ to the left, z+ up
R_correct = np.array([
    [0, 0, 1],
    [-1, 0, 0],
    [0, -1, 0]
], np.float64)
t_correct = np.array([
    [0],
    [0],
    [0]
], np.float64)
    
correction_matrix = np.block([
    [R_correct, t_correct],
    [np.zeros((1, 3)), 1]
])

def skew_matrix(v):
    """
    Generates a 3x3 skew-symmetric matrix from a 3D vector.
    
    Args:
    v (np.array): A 3D vector (numpy array of shape (3,))
    
    Returns:
    np.array: A 3x3 skew-symmetric matrix
    """
    return np.array([
        [0, -v[2], v[1]],
        [v[2], 0, -v[0]],
        [-v[1], v[0], 0]
    ])
    
def pose_matrix(angle, translation):
    def rotation_matrix(angle):
        """
        Generates a 3x3 rotation matrix for a given angle (in radians) around the Z-axis.
        """
        R = np.array([
            [np.cos(angle), -np.sin(angle), 0],
            [np.sin(angle), np.cos(angle), 0],
            [0, 0, 1]
        ])
        return R
    """
    Generates a 4x4 pose matrix given an angle and translation vector.
    Pose matrix is in the form [R, T; 0^T, 1].
    """
    R = rotation_matrix(angle)
    T = np.array(translation).reshape(3, 1)
    pose = np.block([
        [R, T],
        [np.zeros((1, 3)), np.array([[1]])]
    ])


    return pose @ correction_matrix

def apply_intrinsics(K, point_3d):
    """
    Applies the intrinsic matrix to a 3D point to project it onto the image plane.
    """
    point_2d_homogeneous = K @ point_3d[:3]  # Apply intrinsic matrix (ignoring the last element, which is 1)
    point_2d = point_2d_homogeneous / point_2d_homogeneous[2]  # Convert from homogeneous to Cartesian coordinates
    return point_2d

def draw_epipolar_lines(img1, img2, pts1, pts2, F):
    """
    Draws the epipolar lines on both images and marks the points.
    img1, img2: Images on which to draw the lines
    pts1, pts2: Points in the two images
    F: Fundamental matrix
    """
    # Convert points to homogeneous coordinates
    pts1 = np.int32(pts1)
    pts2 = np.int32(pts2)
    
    # Draw the epipolar lines for points in the first image
    lines1 = cv2.computeCorrespondEpilines(pts2.reshape(-1, 1, 2), 2, F)
    lines1 = lines1.reshape(-1, 3)
    img1_with_lines = img1.copy()
    
    for r in lines1:
        color = tuple(np.random.randint(0, 255, 3).tolist())
        x0, y0 = map(int, [0, -r[2] / r[1]])
        x1, y1 = map(int, [img1.shape[1], -(r[2] + r[0] * img1.shape[1]) / r[1]])
        img1_with_lines = cv2.line(img1_with_lines, (x0, y0), (x1, y1), color, 1)
    
    # Draw the epipolar lines for points in the second image
    lines2 = cv2.computeCorrespondEpilines(pts1.reshape(-1, 1, 2), 1, F)
    lines2 = lines2.reshape(-1, 3)
    img2_with_lines = img2.copy()
    
    for r in lines2:
        color = tuple(np.random.randint(0, 255, 3).tolist())
        x0, y0 = map(int, [0, -r[2] / r[1]])
        x1, y1 = map(int, [img2.shape[1], -(r[2] + r[0] * img2.shape[1]) / r[1]])
        img2_with_lines = cv2.line(img2_with_lines, (x0, y0), (x1, y1), color, 1)
    
    return img1_with_lines, img2_with_lines

def find_epipoles(F):
    """
    Finds the epipoles in the image given the Fundamental matrix.
    """
    # Compute the left epipole
    _, _, Vt = np.linalg.svd(F.T)
    e1 = Vt[-1, :]
    e1 = e1 / e1[2]  # Normalize

    # Compute the right epipole
    _, _, Vt = np.linalg.svd(F)
    e2 = Vt[-1, :]
    e2 = e2 / e2[2]  # Normalize

    return e1, e2 

def main():
    # Define pose matrices using rotation around Z-axis and translation
    verge = 30
    P1 = pose_matrix(np.radians(-verge), [0, 1, 0])  # P1 rotated -30 degrees and translated (0, 1, 0)
    P2 = pose_matrix(np.radians(+verge), [0, -1, 0])  # P2 rotated 30 degrees and translated (0, -1, 0)
    D = np.linalg.inv(P1) @ P2

    # Example 3D landmark point in homogeneous coordinates
    X = np.array([10, 0, 0, 1])

    # Project the 3D point into the two camera images
    x1_ego = np.linalg.inv(P1) @ X  # 3D point in camera 1's coordinate system
    x2_ego = np.linalg.inv(P2) @ X  # 3D point in camera 2's coordinate system
    print("X: ", X)
    print("x1_ego: ", x1_ego)
    print("x2_ego: ", x2_ego)

    # Intrinsic matrix (example)
    K = np.array([
        [800, 0, 250],
        [0, 800, 250],
        [0, 0, 1]
    ])
    Kinv = np.linalg.inv(K)


    # Apply the intrinsic matrix to get 2D points
    x1_2d = apply_intrinsics(K, x1_ego)
    x2_2d = apply_intrinsics(K, x2_ego)
    
    print("x1_2d: ", x1_2d)
    print("x2_2d: ", x2_2d)

    # Fundamental matrix (example, normally computed from camera matrices)
    E = skew_matrix(D[:3, 3]) @ D[:3, :3]
    F = Kinv.T @ E @ Kinv

    # Example images
    img1 = np.ones((500, 500, 3), dtype=np.uint8) * 255
    img2 = np.ones((500, 500, 3), dtype=np.uint8) * 255

    # Draw epipolar lines
    pts1 = np.array([x1_2d[:2]])
    pts2 = np.array([x2_2d[:2]])

    # Find epipoles
    e1, e2 = find_epipoles(F)
    print("e1: " , e1)
    print("e2: " , e2)

    # if True:
    # this is wrong
    #     t = D[:3, 3]
    #     t_2d = apply_intrinsics(K, t)
    #     # can we find e1's epipol by projecting p2 into p1 (which is effectively projecting t into p1)
    #     print("t: ", t)
    #     print("t_2d: ", t_2d)
    #     print("e1: ", e1)
    #     print("F.T @ t_2d: ", F.T @ t_2d)
    #     return

    img1_with_lines, img2_with_lines = draw_epipolar_lines(img1, img2, pts1, pts2, F)

    # Plot the results
    plt.figure(figsize=(10, 5))
    plt.subplot(1, 2, 1)
    plt.imshow(img1_with_lines)
    plt.scatter([e1[0]], [e1[1]], color='red', marker='x', label='Epipole')  # Epipole in the first image
    plt.scatter([x1_2d[0]], [x1_2d[1]], color='blue', marker='o', label='Landmark')  # Landmark projection
    plt.title("Image 1 with Epipolar Lines")
    plt.legend()

    plt.subplot(1, 2, 2)
    plt.imshow(img2_with_lines)
    plt.scatter([e2[0]], [e2[1]], color='red', marker='x', label='Epipole')  # Epipole in the second image
    plt.scatter([x2_2d[0]], [x2_2d[1]], color='blue', marker='o', label='Landmark')  # Landmark projection
    plt.title("Image 2 with Epipolar Lines")
    plt.legend()

    plt.show()

if __name__ == "__main__":
    main()